groupchat/server/src/services/openrouter.rs

use serde::{Deserialize, Serialize};

const OPENROUTER_API_URL: &str = "https://openrouter.ai/api/v1/chat/completions";

// ── Request types ──

#[derive(Debug, Serialize)]
struct ChatRequest {
    model: String,
    messages: Vec<ChatMessage>,
    max_tokens: Option<u32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    tools: Option<Vec<Tool>>,
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct ChatMessage {
    pub role: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub content: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_calls: Option<Vec<ToolCall>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_call_id: Option<String>,
}

// ── Tool definition types ──

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct Tool {
    pub r#type: String,
    pub function: ToolFunction,
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct ToolFunction {
    pub name: String,
    pub description: String,
    pub parameters: serde_json::Value,
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct ToolCall {
    pub id: String,
    pub r#type: String,
    pub function: ToolCallFunction,
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct ToolCallFunction {
    pub name: String,
    pub arguments: String,
}

// ── Response types ──

#[derive(Debug, Deserialize)]
struct ChatResponse {
    choices: Vec<Choice>,
    model: Option<String>,
    usage: Option<Usage>,
}

#[derive(Debug, Deserialize)]
struct Choice {
    message: ChoiceMessage,
}

#[derive(Debug, Deserialize)]
struct ChoiceMessage {
    content: Option<String>,
    #[serde(default)]
    tool_calls: Option<Vec<ToolCall>>,
}

#[derive(Debug, Deserialize)]
struct Usage {
    prompt_tokens: Option<u32>,
    completion_tokens: Option<u32>,
    total_tokens: Option<u32>,
}

/// Stats returned alongside an AI completion.
#[derive(Debug, Clone, Serialize)]
pub struct CompletionStats {
    pub model: String,
    pub prompt_tokens: u32,
    pub completion_tokens: u32,
    pub total_tokens: u32,
    pub response_ms: u64,
}

/// Result from a chat completion — either a final text response or tool calls.
pub enum ChatCompletionResult {
    /// AI responded with text content.
    Response(String, CompletionStats),
    /// AI wants to call tools. Contains the assistant message (with tool_calls) and stats.
    ToolCalls(ChatMessage, CompletionStats),
}

/// Build the tool definitions for brave_search and web_fetch.
pub fn build_tools() -> Vec<Tool> {
    vec![
        Tool {
            r#type: "function".into(),
            function: ToolFunction {
                name: "brave_search".into(),
                description: "Search the web for current information. Use this when users ask about recent events, need factual data you're unsure about, or want up-to-date information.".into(),
                parameters: serde_json::json!({
                    "type": "object",
                    "properties": {
                        "query": {
                            "type": "string",
                            "description": "The search query"
                        },
                        "count": {
                            "type": "integer",
                            "description": "Number of results (1-10, default 5)"
                        }
                    },
                    "required": ["query"]
                }),
            },
        },
        Tool {
            r#type: "function".into(),
            function: ToolFunction {
                name: "web_fetch".into(),
                description: "Fetch and read the content of a web page. Use this to read articles, documentation, or any URL shared by users or found in search results.".into(),
                parameters: serde_json::json!({
                    "type": "object",
                    "properties": {
                        "url": {
                            "type": "string",
                            "description": "The URL to fetch"
                        }
                    },
                    "required": ["url"]
                }),
            },
        },
    ]
}

/// Send a chat completion request to OpenRouter.
/// Returns either a text response or tool call requests.
pub async fn chat_completion(
    history: Vec<ChatMessage>,
    model_id: &str,
    api_key: &str,
    tools: Option<Vec<Tool>>,
) -> Result<ChatCompletionResult, String> {
    let client = reqwest::Client::new();

    let request_body = ChatRequest {
        model: model_id.to_string(),
        messages: history,
        max_tokens: Some(2048),
        tools,
    };

    let start = std::time::Instant::now();

    let response = client
        .post(OPENROUTER_API_URL)
        .header("Authorization", format!("Bearer {}", api_key))
        .header("Content-Type", "application/json")
        .header("HTTP-Referer", "http://localhost:3001")
        .header("X-Title", "GroupChat")
        .json(&request_body)
        .send()
        .await
        .map_err(|e| format!("OpenRouter request failed: {}", e))?;

    let elapsed_ms = start.elapsed().as_millis() as u64;

    if !response.status().is_success() {
        let status = response.status();
        let body = response.text().await.unwrap_or_default();
        return Err(format!("OpenRouter error {}: {}", status, body));
    }

    let chat_response: ChatResponse = response
        .json()
        .await
        .map_err(|e| format!("Failed to parse OpenRouter response: {}", e))?;

    let choice = chat_response
        .choices
        .first()
        .ok_or_else(|| "No response from OpenRouter".to_string())?;

    let usage = chat_response.usage.unwrap_or(Usage {
        prompt_tokens: None,
        completion_tokens: None,
        total_tokens: None,
    });

    let stats = CompletionStats {
        model: chat_response.model.unwrap_or_else(|| model_id.to_string()),
        prompt_tokens: usage.prompt_tokens.unwrap_or(0),
        completion_tokens: usage.completion_tokens.unwrap_or(0),
        total_tokens: usage.total_tokens.unwrap_or(0),
        response_ms: elapsed_ms,
    };

    // Check if the AI wants to call tools
    if let Some(tool_calls) = &choice.message.tool_calls {
        if !tool_calls.is_empty() {
            // Return the assistant message with tool calls so it can be added to history
            let assistant_msg = ChatMessage {
                role: "assistant".into(),
                content: choice.message.content.clone(),
                tool_calls: Some(tool_calls.clone()),
                tool_call_id: None,
            };
            return Ok(ChatCompletionResult::ToolCalls(assistant_msg, stats));
        }
    }

    // Regular text response
    let content = choice.message.content.clone().unwrap_or_default();
    Ok(ChatCompletionResult::Response(content, stats))
}

/// Build the message history for OpenRouter from stored messages.
/// Includes the system prompt as the first message.
pub fn build_chat_history(
    system_prompt: &str,
    messages: &[(String, String, bool)], // (sender_name, content, is_ai)
) -> Vec<ChatMessage> {
    let mut history = vec![ChatMessage {
        role: "system".to_string(),
        content: Some(system_prompt.to_string()),
        tool_calls: None,
        tool_call_id: None,
    }];

    for (sender_name, content, is_ai) in messages {
        if *is_ai {
            history.push(ChatMessage {
                role: "assistant".to_string(),
                content: Some(content.clone()),
                tool_calls: None,
                tool_call_id: None,
            });
        } else {
            history.push(ChatMessage {
                role: "user".to_string(),
                content: Some(format!("[{}]: {}", sender_name, content)),
                tool_calls: None,
                tool_call_id: None,
            });
        }
    }

    history
}