Создание инструмента генерации кода с помощью Rust и локальных LLM от Ollama

Это реакция на выпуск ChatGPT o1-preview. Попытка добавить логику в LLM с открытым исходным кодом, которые можно запустить дома на скромном GPU или даже на CPU

Сейчас я работаю над инструментом на основе Rust, который автоматизирует генерацию, компиляцию и тестирование кода с использованием больших языковых моделей (LLM). Идея заключается в том, чтобы взаимодействовать с LLM для генерации кода на основе предоставленных пользователем объяснений, компилировать код, разрешать зависимости и запускать тесты, чтобы убедиться, что все работает так, как ожидается.

Я хотел бы оптимизировать процесс кодирования функций на основе описаний на естественном языке. Я хотел создать систему, в которой я мог бы ввести объяснение того, что должна делать функция, и чтобы инструмент обрабатывал все остальное от генерации кода до тестирования.

Работа инструмента начинается с того, что запрашивает у пользователя объяснение функции, которую он хочет создать. Затем он взаимодействует с LLM для генерации кода функции, компилирует его и проверяет на наличие ошибок компиляции. Если ошибки обнаружены, инструмент пытается их устранить, возможно, путем добавления зависимостей или переписывания кода. После успешной компиляции кода он генерирует тесты для функции, запускает их и снова обрабатывает любые ошибки, итеративно улучшая код или тесты.

Первый шаг — получить объяснение от пользователя:

println!("Explain what the function should do:"); let mut explanation = String::new(); std::io::stdin().read_line(&mut explanation).unwrap(

Используя это объяснение, инструмент создает запрос для отправки LLM:

let generate_code_prompt = construct_prompt(     generate_code_prompt_template,     vec![&explanation], );

Вот `generate_code_prompt_template`:

let generate_code_prompt_template = r#" {{{0}}}  Write on Rust language code of this function (without example of usage like main function): ```rust fn solution( "#;

Этот prompt сообщает LLM о необходимости сгенерировать код Rust для функции на основе объяснения пользователя.

После генерации кода инструмент пытается его скомпилировать:

create_rust_project(&code, "", ""); let (mut exit_code, mut output) = cargo("build", &mut cache);

Если компиляция завершается неудачей, проверяется, связана ли проблема с отсутствием зависимостей

let build_dependencies_req_prompt = construct_prompt(     build_dependencies_req_prompt_template,     vec![&explanation, &code, &output], );

В зависимости от ответа LLM он может добавить необходимые зависимости в файл `Cargo.toml`:

let build_dependencies_prompt = construct_prompt(     build_dependencies_prompt_template,     vec![&explanation, &code], ); let build_dependencies_result = llm_request(&build_dependencies_prompt, &mut cache); dependencies = extract_code(&build_dependencies_result);

После успешной компиляции кода инструмент генерирует тесты:

let generate_test_prompt = construct_prompt(     generate_test_prompt_template,     vec![&explanation, &code], ); let generation_test_result = llm_request(&generate_test_prompt, &mut cache); code_test = extract_code(&generation_test_result);

Затем он запускает тесты:

let (exit_code_immut, output_immut) = cargo("test", &mut cache);

Если тесты не пройдены, инструмент решает, следует ли переписать код или тесты, в зависимости от того, где находится ошибка:

let rewrite_code_req_prompt_template_prompt = construct_prompt(     rewrite_code_req_prompt_template,     vec![&explanation, &code, &code_test, &output], ); let rewrite_code_req_result = llm_request(&rewrite_code_req_prompt_template_prompt, &mut cache); if extract_number(&rewrite_code_req_result) == 1 {     // Rewrite code } else {     // Rewrite tests }

Для повышения эффективности в инструменте реализована система кэширования:

let result_str_opt = cache.get(&key); let result_str = match result_str_opt {     None => {         // Run command and cache result     }     Some(result) => {         result.to_string()     } };

Это позволяет избежать избыточных вычислений за счет сохранения предыдущих результатов и их извлечения при повторном вводе тех же данных.

Вот более подробная схема логики работы:

Перед запуском инструмента выполните следующие шаги:

1. Убедитесь, что у вас установлен Rust. Вы можете установить его [здесь]

2. Требуется для взаимодействия с LLM. Установите с [официального сайта Ollama]

3. Загрузите модель:

ollama run gemma2:27b

После загрузки модели вы можете сказать «привет» модели, чтобы проверить, правильно ли она работает. После этого вы можете нажать «Ctrl+D», чтобы выйти из модели.

cargo run

Вам будет предложено объяснить, что должна делать функция:

Explain what the function should do:

Предоставьте подробное объяснение, и инструмент сделает все остальное.

Допустим, я ввожу:

parse json string and return struct User (age, name)

Инструмент сгенерирует соответствующую функцию Rust, обработает все зависимости (например, добавит `serde` и `serde_json`), сгенерирует тесты и запустит их. Окончательный вывод будет отображен, а результат будет сохранен в папке `sandbox`.

Сгенерированный результат:

[dependencies] serde = { version = "1.0", features = ["derive"] } serde_json = "1.0"  use serde::{Deserialize, Serialize};  #[derive(Deserialize, Serialize, Debug)] struct User {     name: String,     age: u32, }  fn solution(json_string: &str) -> Result<User, serde_json::Error> {     let user: User = serde_json::from_str(json_string)?;     Ok(user) }  #[cfg(test)] mod tests {     use super::*;      #[test]     fn test_solution() {         let json_string = r#"{"name": "John Doe", "age": 30}"#;         let user = solution(json_string).unwrap();         assert_eq!(user.name, "John Doe");         assert_eq!(user.age, 30);     }      #[test]     fn test_solution_invalid_json() {         let json_string = r#"{"name": "John Doe", "age": }"#;         assert!(solution(json_string).is_err());     } }

Исходный код доступен на [GitHub]. Контрибутинг приветствуются!