Building a Code Generation Tool with Rust and Local LLMs by Ollama

This is a reaction to the ChatGPT o1-preview release. An attempt to add logic to an open source LLM that can be run at home on a modest GPU or even a CPU.

I'm currently working on a Rust-based tool that automates code generation, compilation, and testing using large language models (LLMs). The idea is to interact with the LLM to generate code based on user-provided explanations, compile the code, resolve dependencies, and run tests to ensure everything works as expected.

I wanted to optimize the process of coding functions based on natural language descriptions. I wanted to create a system where I could enter an explanation of what a function should do, and have the tool handle the rest from code generation to testing.

The tool starts by asking the user to explain the function they want to create. It then interacts with LLM to generate the function's code, compiles it, and checks for compilation errors. If errors are found, the tool attempts to fix them, perhaps by adding dependencies or rewriting the code. Once the code has been successfully compiled, it generates tests for the function, runs them, and again handles any errors, iteratively improving the code or tests.

The first step is to get an explanation from the user:

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

Using this explanation, the tool creates a request to send LLM:

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

Here is the `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(
"#;

This prompt tells LLM to generate Rust code for the function based on the user's explanation.

After generating the code, the tool tries to compile it:

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

If compilation fails, it checks whether the problem is due to missing dependencies

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

Depending on the LLM response, it may add the necessary dependencies to the `Cargo.toml` file:

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);

After successful compilation of the code, the tool generates tests:

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);

Then it runs the tests:

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

If the tests fail, the tool decides whether to rewrite the code or the tests, depending on where the error is:

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
}

To improve efficiency, the tool has a caching system:

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()
    }
};

This avoids redundant calculations by storing previous results and retrieving them when the same data is re-entered.

Here is a more detailed diagram of the logic of operation:

Before running the tool, please follow these steps:

Make sure you have Rust installed. You can install it [здесь]
Required to interact with LLM. Install with [официального сайта Ollama]
Download the model:

ollama run gemma2:27b

After loading the model, you can say “hello” to the model to check if it works correctly. After that, you can press “Ctrl+D” to exit the model.

cargo run

You will be asked to explain what the function should do:

Explain what the function should do:

Provide a detailed explanation and the tool will do the rest.

Let's say I enter:

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

The tool will generate the corresponding Rust function, handle all dependencies (e.g. add `serde` and `serde_json`), generate tests and run them. The final output will be displayed and the result will be saved in the `sandbox` folder.

Generated result:

[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());
    }
}

The source code is available at [GitHub]. Contributing is welcome!