storkit/src-tauri/src/llm/ollama.rs

use crate::llm::types::{
    CompletionResponse, FunctionCall, Message, ModelProvider, Role, ToolCall, ToolDefinition,
};
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use serde_json::Value;

pub struct OllamaProvider {
    base_url: String,
}

impl OllamaProvider {
    pub fn new(base_url: String) -> Self {
        Self { base_url }
    }

    pub async fn get_models(base_url: &str) -> Result<Vec<String>, String> {
        let client = reqwest::Client::new();
        let url = format!("{}/api/tags", base_url.trim_end_matches('/'));

        let res = client
            .get(&url)
            .send()
            .await
            .map_err(|e| format!("Request failed: {}", e))?;

        if !res.status().is_success() {
            let status = res.status();
            let text = res.text().await.unwrap_or_default();
            return Err(format!("Ollama API error {}: {}", status, text));
        }

        let body: OllamaTagsResponse = res
            .json()
            .await
            .map_err(|e| format!("Failed to parse response: {}", e))?;

        Ok(body.models.into_iter().map(|m| m.name).collect())
    }
}

#[derive(Deserialize)]
struct OllamaTagsResponse {
    models: Vec<OllamaModelTag>,
}

#[derive(Deserialize)]
struct OllamaModelTag {
    name: String,
}

// --- Request Types ---

#[derive(Serialize)]
struct OllamaRequest<'a> {
    model: &'a str,
    messages: Vec<OllamaRequestMessage>,
    stream: bool,
    #[serde(skip_serializing_if = "is_empty_tools")]
    tools: &'a [ToolDefinition],
}

fn is_empty_tools(tools: &&[ToolDefinition]) -> bool {
    tools.is_empty()
}

#[derive(Serialize)]
struct OllamaRequestMessage {
    role: Role,
    content: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    tool_calls: Option<Vec<OllamaRequestToolCall>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    tool_call_id: Option<String>,
}

#[derive(Serialize)]
struct OllamaRequestToolCall {
    function: OllamaRequestFunctionCall,
    #[serde(rename = "type")]
    kind: String,
}

#[derive(Serialize)]
struct OllamaRequestFunctionCall {
    name: String,
    arguments: Value,
}

// --- Response Types ---

#[derive(Deserialize)]
struct OllamaResponse {
    message: OllamaResponseMessage,
}

#[derive(Deserialize)]
struct OllamaResponseMessage {
    content: String,
    tool_calls: Option<Vec<OllamaResponseToolCall>>,
}

#[derive(Deserialize)]
struct OllamaResponseToolCall {
    function: OllamaResponseFunctionCall,
}

#[derive(Deserialize)]
struct OllamaResponseFunctionCall {
    name: String,
    arguments: Value, // Ollama returns Object, we convert to String for internal storage
}

#[async_trait]
impl ModelProvider for OllamaProvider {
    async fn chat(
        &self,
        model: &str,
        messages: &[Message],
        tools: &[ToolDefinition],
    ) -> Result<CompletionResponse, String> {
        let client = reqwest::Client::new();
        let url = format!("{}/api/chat", self.base_url.trim_end_matches('/'));

        // Convert domain Messages to Ollama Messages (handling String -> Object args mismatch)
        let ollama_messages: Vec<OllamaRequestMessage> = messages
            .iter()
            .map(|m| {
                let tool_calls = m.tool_calls.as_ref().map(|calls| {
                    calls
                        .iter()
                        .map(|tc| {
                            // Try to parse string args as JSON, fallback to string value if fails
                            let args_val: Value = serde_json::from_str(&tc.function.arguments)
                                .unwrap_or(Value::String(tc.function.arguments.clone()));

                            OllamaRequestToolCall {
                                kind: tc.kind.clone(),
                                function: OllamaRequestFunctionCall {
                                    name: tc.function.name.clone(),
                                    arguments: args_val,
                                },
                            }
                        })
                        .collect()
                });

                OllamaRequestMessage {
                    role: m.role.clone(),
                    content: m.content.clone(),
                    tool_calls,
                    tool_call_id: m.tool_call_id.clone(),
                }
            })
            .collect();

        let request_body = OllamaRequest {
            model,
            messages: ollama_messages,
            stream: false,
            tools,
        };

        let res = client
            .post(&url)
            .json(&request_body)
            .send()
            .await
            .map_err(|e| format!("Request failed: {}", e))?;

        if !res.status().is_success() {
            let status = res.status();
            let text = res.text().await.unwrap_or_default();
            return Err(format!("Ollama API error {}: {}", status, text));
        }

        let response_body: OllamaResponse = res
            .json()
            .await
            .map_err(|e| format!("Failed to parse response: {}", e))?;

        // Convert Response back to Domain types
        let content = if response_body.message.content.is_empty() {
            None
        } else {
            Some(response_body.message.content)
        };

        let tool_calls = response_body.message.tool_calls.map(|calls| {
            calls
                .into_iter()
                .map(|tc| ToolCall {
                    id: None, // Ollama doesn't typically send IDs
                    kind: "function".to_string(),
                    function: FunctionCall {
                        name: tc.function.name,
                        arguments: tc.function.arguments.to_string(), // Convert Object -> String
                    },
                })
                .collect()
        });

        Ok(CompletionResponse {
            content,
            tool_calls,
        })
    }
}