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

@igumnov
18.09.2024 23:04 UTC
Первоисточник

Комментарии

@MountainGoat
18.09.2024 18:52 UTC
0

Чем лучше плагина для VSCode "Continue", к которому кроме перечисленного ещё и какая-то IDE присобачена с плагинами?

@igumnov
18.09.2024 19:07 UTC
0
  1. Это не плагин автокомплита кода, который выдает LLM на основе окружающего кода где стоит курсор разработчки

  2. Это LLM-генератор кода, зависимостей для сборщика, тестов и тулза, которая компилирует, запускает тесты и вносит правки на основе этого в предлагаемый код

18.09.2024 19:11 UTC
0

Ну я это всё через Continue делаю. Он только новые файлы создавать не умеет, ему надо руками пустой файл создать, и он его заполнит.

18.09.2024 19:25 UTC
0

Те он предлагает код и перед этим скрытно от вас запускает тесты у вас на рабочей машине, что бы подтвердить, что код рабочий?

18.09.2024 19:41 UTC
+2

Да чур меня, чтобы у меня LLM чего-то скрытно запускала. Она назапускает.

Теперь понял в чём смысл.

24.09.2024 14:56 UTC
0

Запланировал добавить запуск того что делает ИИ в изолированом докере. Add `--environment` Parameter to Support Docker Execution · Issue #19 · evgenyigumnov/rustsn (github.com)

@Moog_Prodigy
18.09.2024 22:57 UTC
0

Интересный подход, таким образом можно с локальными моделями куда лучше работать. Прям намного качественней. Без увеличения моделей и переобучения.

@igumnov
24.09.2024 14:57 UTC
0

Так и есть! Вот вышел недавно qwen2.5-coder:1.5b весит всего 900 мб - работает быстро даже на скромном железе. Вот тикет:
Ollama "qwen2.5-coder:1.5b" model support (parse problems in "parse_llm_response" functions) · Issue #27 · evgenyigumnov/rustsn (github.com)

@StrikerXP
20.09.2024 11:02 UTC
0

Он только на Rust-е пишет или на других языках тоже?

@igumnov
20.09.2024 11:03 UTC
0

Планирую Java добавить, еще можно TypeScript, Python, Scala будет добавить

20.09.2024 11:58 UTC
0

Я относительно новичок и пишу на TypeScript и Python. Соотвественно интересуют интеграция данных языков. Ваш проэкт мне весьма интересен:)

24.09.2024 14:58 UTC
+1

Добавил поддержку Typescript и Python