storkit/src-tauri/src/commands/chat.rs

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, &current_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"]
                }),
            },
        },
    ]
}