257 lines
9.1 KiB
Rust
257 lines
9.1 KiB
Rust
use crate::commands::{fs, search, shell};
|
|
use crate::llm::ollama::OllamaProvider;
|
|
use crate::llm::types::{
|
|
Message, ModelProvider, Role, ToolCall, ToolDefinition, ToolFunctionDefinition,
|
|
};
|
|
use crate::state::SessionState;
|
|
use serde::Deserialize;
|
|
use serde_json::json;
|
|
use tauri::State;
|
|
|
|
#[derive(Deserialize)]
|
|
pub struct ProviderConfig {
|
|
pub provider: String, // "ollama"
|
|
pub model: String,
|
|
pub base_url: Option<String>,
|
|
pub enable_tools: Option<bool>,
|
|
}
|
|
|
|
const MAX_TURNS: usize = 10;
|
|
|
|
#[tauri::command]
|
|
pub async fn get_ollama_models(base_url: Option<String>) -> Result<Vec<String>, String> {
|
|
let url = base_url.unwrap_or_else(|| "http://localhost:11434".to_string());
|
|
OllamaProvider::get_models(&url).await
|
|
}
|
|
|
|
#[tauri::command]
|
|
pub async fn chat(
|
|
messages: Vec<Message>,
|
|
config: ProviderConfig,
|
|
state: State<'_, SessionState>,
|
|
) -> Result<Vec<Message>, String> {
|
|
// 1. Setup Provider
|
|
let provider: Box<dyn ModelProvider> = match config.provider.as_str() {
|
|
"ollama" => Box::new(OllamaProvider::new(
|
|
config
|
|
.base_url
|
|
.unwrap_or_else(|| "http://localhost:11434".to_string()),
|
|
)),
|
|
_ => return Err(format!("Unsupported provider: {}", config.provider)),
|
|
};
|
|
|
|
// 2. Define Tools
|
|
let tool_defs = get_tool_definitions();
|
|
let tools = if config.enable_tools.unwrap_or(true) {
|
|
tool_defs.as_slice()
|
|
} else {
|
|
&[]
|
|
};
|
|
|
|
// 3. Agent Loop
|
|
let mut current_history = messages.clone();
|
|
let mut new_messages: Vec<Message> = Vec::new();
|
|
let mut turn_count = 0;
|
|
|
|
loop {
|
|
if turn_count >= MAX_TURNS {
|
|
return Err("Max conversation turns reached.".to_string());
|
|
}
|
|
turn_count += 1;
|
|
|
|
// Call LLM
|
|
let response = provider
|
|
.chat(&config.model, ¤t_history, tools)
|
|
.await
|
|
.map_err(|e| format!("LLM Error: {}", e))?;
|
|
|
|
// Process Response
|
|
if let Some(tool_calls) = response.tool_calls {
|
|
// The Assistant wants to run tools
|
|
let assistant_msg = Message {
|
|
role: Role::Assistant,
|
|
content: response.content.unwrap_or_default(),
|
|
tool_calls: Some(tool_calls.clone()),
|
|
tool_call_id: None,
|
|
};
|
|
|
|
current_history.push(assistant_msg.clone());
|
|
new_messages.push(assistant_msg);
|
|
|
|
// Execute Tools
|
|
for call in tool_calls {
|
|
let output = execute_tool(&call, &state).await;
|
|
|
|
let tool_msg = Message {
|
|
role: Role::Tool,
|
|
content: output,
|
|
tool_calls: None,
|
|
// For Ollama/Simple flow, we just append.
|
|
// For OpenAI strict, this needs to match call.id.
|
|
tool_call_id: call.id,
|
|
};
|
|
|
|
current_history.push(tool_msg.clone());
|
|
new_messages.push(tool_msg);
|
|
}
|
|
} else {
|
|
// Final text response
|
|
let assistant_msg = Message {
|
|
role: Role::Assistant,
|
|
content: response.content.unwrap_or_default(),
|
|
tool_calls: None,
|
|
tool_call_id: None,
|
|
};
|
|
|
|
// We don't push to current_history needed for next loop, because we are done.
|
|
new_messages.push(assistant_msg);
|
|
break;
|
|
}
|
|
}
|
|
|
|
Ok(new_messages)
|
|
}
|
|
|
|
async fn execute_tool(call: &ToolCall, state: &State<'_, SessionState>) -> String {
|
|
let name = call.function.name.as_str();
|
|
// Parse arguments. They come as a JSON string from the LLM abstraction.
|
|
let args: serde_json::Value = match serde_json::from_str(&call.function.arguments) {
|
|
Ok(v) => v,
|
|
Err(e) => return format!("Error parsing arguments: {}", e),
|
|
};
|
|
|
|
match name {
|
|
"read_file" => {
|
|
let path = args["path"].as_str().unwrap_or("").to_string();
|
|
match fs::read_file(path, state.clone()).await {
|
|
Ok(content) => content,
|
|
Err(e) => format!("Error: {}", e),
|
|
}
|
|
}
|
|
"write_file" => {
|
|
let path = args["path"].as_str().unwrap_or("").to_string();
|
|
let content = args["content"].as_str().unwrap_or("").to_string();
|
|
match fs::write_file(path, content, state.clone()).await {
|
|
Ok(_) => "File written successfully.".to_string(),
|
|
Err(e) => format!("Error: {}", e),
|
|
}
|
|
}
|
|
"list_directory" => {
|
|
let path = args["path"].as_str().unwrap_or("").to_string();
|
|
match fs::list_directory(path, state.clone()).await {
|
|
Ok(entries) => serde_json::to_string(&entries).unwrap_or_default(),
|
|
Err(e) => format!("Error: {}", e),
|
|
}
|
|
}
|
|
"search_files" => {
|
|
let query = args["query"].as_str().unwrap_or("").to_string();
|
|
match search::search_files(query, state.clone()).await {
|
|
Ok(results) => serde_json::to_string(&results).unwrap_or_default(),
|
|
Err(e) => format!("Error: {}", e),
|
|
}
|
|
}
|
|
"exec_shell" => {
|
|
let command = args["command"].as_str().unwrap_or("").to_string();
|
|
let args_vec: Vec<String> = args["args"]
|
|
.as_array()
|
|
.map(|arr| {
|
|
arr.iter()
|
|
.map(|v| v.as_str().unwrap_or("").to_string())
|
|
.collect()
|
|
})
|
|
.unwrap_or_default();
|
|
|
|
match shell::exec_shell(command, args_vec, state.clone()).await {
|
|
Ok(output) => serde_json::to_string(&output).unwrap_or_default(),
|
|
Err(e) => format!("Error: {}", e),
|
|
}
|
|
}
|
|
_ => format!("Unknown tool: {}", name),
|
|
}
|
|
}
|
|
|
|
fn get_tool_definitions() -> Vec<ToolDefinition> {
|
|
vec![
|
|
ToolDefinition {
|
|
kind: "function".to_string(),
|
|
function: ToolFunctionDefinition {
|
|
name: "read_file".to_string(),
|
|
description: "Reads the content of a file in the project.".to_string(),
|
|
parameters: json!({
|
|
"type": "object",
|
|
"properties": {
|
|
"path": { "type": "string", "description": "Relative path to the file" }
|
|
},
|
|
"required": ["path"]
|
|
}),
|
|
},
|
|
},
|
|
ToolDefinition {
|
|
kind: "function".to_string(),
|
|
function: ToolFunctionDefinition {
|
|
name: "write_file".to_string(),
|
|
description: "Writes content to a file. Overwrites if exists.".to_string(),
|
|
parameters: json!({
|
|
"type": "object",
|
|
"properties": {
|
|
"path": { "type": "string", "description": "Relative path to the file" },
|
|
"content": { "type": "string", "description": "The full content to write" }
|
|
},
|
|
"required": ["path", "content"]
|
|
}),
|
|
},
|
|
},
|
|
ToolDefinition {
|
|
kind: "function".to_string(),
|
|
function: ToolFunctionDefinition {
|
|
name: "list_directory".to_string(),
|
|
description: "Lists files and directories at a path.".to_string(),
|
|
parameters: json!({
|
|
"type": "object",
|
|
"properties": {
|
|
"path": { "type": "string", "description": "Relative path to list (use '.' for root)" }
|
|
},
|
|
"required": ["path"]
|
|
}),
|
|
},
|
|
},
|
|
ToolDefinition {
|
|
kind: "function".to_string(),
|
|
function: ToolFunctionDefinition {
|
|
name: "search_files".to_string(),
|
|
description: "Searches for text content across all files in the project."
|
|
.to_string(),
|
|
parameters: json!({
|
|
"type": "object",
|
|
"properties": {
|
|
"query": { "type": "string", "description": "The string to search for" }
|
|
},
|
|
"required": ["query"]
|
|
}),
|
|
},
|
|
},
|
|
ToolDefinition {
|
|
kind: "function".to_string(),
|
|
function: ToolFunctionDefinition {
|
|
name: "exec_shell".to_string(),
|
|
description: "Executes a shell command in the project root.".to_string(),
|
|
parameters: json!({
|
|
"type": "object",
|
|
"properties": {
|
|
"command": {
|
|
"type": "string",
|
|
"description": "The command to run (e.g., 'git', 'cargo', 'ls')"
|
|
},
|
|
"args": {
|
|
"type": "array",
|
|
"items": { "type": "string" },
|
|
"description": "Arguments for the command"
|
|
}
|
|
},
|
|
"required": ["command", "args"]
|
|
}),
|
|
},
|
|
},
|
|
]
|
|
}
|