Создание инструмента генерации кода с помощью 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]. Контрибутинг приветствуются!