reasonkit-core 0.1.8

//! Multi-Backend LLM Orchestration
//!
//! This module provides unified LLM orchestration using the Graniet `llm` crate,
//! supporting multiple backends (OpenAI, Claude, xAI, etc.) with tool use and memory sharing.
//!
//! # Features
//! - Unified API across all LLM providers
//! - Tool/function calling support
//! - Memory sharing between conversations
//! - Streaming responses
//!
//! Enable with: `cargo build --features llm-orchestration-multi`

use serde::{Deserialize, Serialize};
use std::collections::HashMap;

// Re-export the llm crate for direct access
pub use llm;

/// Supported LLM backends
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum LlmBackend {
    /// OpenAI (GPT-4, GPT-4o, etc.)
    OpenAI,
    /// Anthropic (Claude 3/4)
    Anthropic,
    /// xAI (Grok)
    XAI,
    /// Google (Gemini)
    Google,
    /// Local Ollama
    Ollama,
    /// OpenRouter (multi-provider)
    OpenRouter,
}

impl LlmBackend {
    /// Get the environment variable name for the API key
    pub fn api_key_env_var(&self) -> &'static str {
        match self {
            Self::OpenAI => "OPENAI_API_KEY",
            Self::Anthropic => "ANTHROPIC_API_KEY",
            Self::XAI => "XAI_API_KEY",
            Self::Google => "GOOGLE_API_KEY",
            Self::Ollama => "OLLAMA_HOST",
            Self::OpenRouter => "OPENROUTER_API_KEY",
        }
    }

    /// Get the default model for this backend
    pub fn default_model(&self) -> &'static str {
        match self {
            Self::OpenAI => "gpt-4o",
            Self::Anthropic => "claude-sonnet-4-20250514",
            Self::XAI => "grok-2",
            Self::Google => "gemini-2.0-flash",
            Self::Ollama => "llama3.2",
            Self::OpenRouter => "anthropic/claude-sonnet-4",
        }
    }
}

/// Configuration for multi-backend LLM client
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MultiLlmConfig {
    /// Primary backend to use
    pub primary_backend: LlmBackend,
    /// Fallback backends in order of preference
    pub fallback_backends: Vec<LlmBackend>,
    /// Model overrides per backend
    pub model_overrides: HashMap<String, String>,
    /// Maximum tokens for response
    pub max_tokens: Option<u32>,
    /// Temperature (0.0 - 2.0)
    pub temperature: Option<f32>,
    /// Enable streaming responses
    pub streaming: bool,
    /// Timeout in seconds
    pub timeout_secs: u64,
}

impl Default for MultiLlmConfig {
    fn default() -> Self {
        Self {
            primary_backend: LlmBackend::Anthropic,
            fallback_backends: vec![LlmBackend::OpenAI, LlmBackend::OpenRouter],
            model_overrides: HashMap::new(),
            max_tokens: Some(4096),
            temperature: Some(0.7),
            streaming: true,
            timeout_secs: 120,
        }
    }
}

/// A message in the conversation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Message {
    /// Role: "user", "assistant", or "system"
    pub role: String,
    /// Message content
    pub content: String,
    /// Optional name for the message sender
    pub name: Option<String>,
}

impl Message {
    /// Create a new user message
    pub fn user(content: impl Into<String>) -> Self {
        Self {
            role: "user".to_string(),
            content: content.into(),
            name: None,
        }
    }

    /// Create a new assistant message
    pub fn assistant(content: impl Into<String>) -> Self {
        Self {
            role: "assistant".to_string(),
            content: content.into(),
            name: None,
        }
    }

    /// Create a new system message
    pub fn system(content: impl Into<String>) -> Self {
        Self {
            role: "system".to_string(),
            content: content.into(),
            name: None,
        }
    }
}

/// Tool definition for function calling
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolDefinition {
    /// Tool name
    pub name: String,
    /// Tool description for the LLM
    pub description: String,
    /// JSON Schema for parameters
    pub parameters: serde_json::Value,
}

/// Tool call request from the LLM
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolCall {
    /// Tool call ID
    pub id: String,
    /// Tool name
    pub name: String,
    /// Arguments as JSON
    pub arguments: serde_json::Value,
}

/// Response from LLM completion
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CompletionResponse {
    /// Generated content
    pub content: String,
    /// Tool calls requested (if any)
    pub tool_calls: Vec<ToolCall>,
    /// Finish reason
    pub finish_reason: String,
    /// Token usage
    pub usage: TokenUsage,
    /// Backend that processed the request
    pub backend: LlmBackend,
    /// Model used
    pub model: String,
}

/// Token usage statistics
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct TokenUsage {
    /// Input tokens
    pub prompt_tokens: u32,
    /// Output tokens
    pub completion_tokens: u32,
    /// Total tokens
    pub total_tokens: u32,
}

/// Conversation memory for context sharing
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct ConversationMemory {
    /// Conversation ID
    pub id: String,
    /// Messages in the conversation
    pub messages: Vec<Message>,
    /// Metadata
    pub metadata: HashMap<String, serde_json::Value>,
}

impl ConversationMemory {
    /// Create a new conversation memory
    pub fn new(id: impl Into<String>) -> Self {
        Self {
            id: id.into(),
            messages: Vec::new(),
            metadata: HashMap::new(),
        }
    }

    /// Add a message to the conversation
    pub fn add_message(&mut self, message: Message) {
        self.messages.push(message);
    }

    /// Get the last N messages
    pub fn last_messages(&self, n: usize) -> &[Message] {
        let start = self.messages.len().saturating_sub(n);
        &self.messages[start..]
    }

    /// Clear all messages
    pub fn clear(&mut self) {
        self.messages.clear();
    }

    /// Get total message count
    pub fn len(&self) -> usize {
        self.messages.len()
    }

    /// Check if empty
    pub fn is_empty(&self) -> bool {
        self.messages.is_empty()
    }
}

/// Builder for tool definitions
pub struct ToolBuilder {
    name: String,
    description: String,
    properties: serde_json::Map<String, serde_json::Value>,
    required: Vec<String>,
}

impl ToolBuilder {
    /// Create a new tool builder
    pub fn new(name: impl Into<String>, description: impl Into<String>) -> Self {
        Self {
            name: name.into(),
            description: description.into(),
            properties: serde_json::Map::new(),
            required: Vec::new(),
        }
    }

    /// Add a string parameter
    pub fn string_param(
        mut self,
        name: impl Into<String>,
        description: impl Into<String>,
        required: bool,
    ) -> Self {
        let name = name.into();
        self.properties.insert(
            name.clone(),
            serde_json::json!({
                "type": "string",
                "description": description.into()
            }),
        );
        if required {
            self.required.push(name);
        }
        self
    }

    /// Add a number parameter
    pub fn number_param(
        mut self,
        name: impl Into<String>,
        description: impl Into<String>,
        required: bool,
    ) -> Self {
        let name = name.into();
        self.properties.insert(
            name.clone(),
            serde_json::json!({
                "type": "number",
                "description": description.into()
            }),
        );
        if required {
            self.required.push(name);
        }
        self
    }

    /// Add a boolean parameter
    pub fn bool_param(
        mut self,
        name: impl Into<String>,
        description: impl Into<String>,
        required: bool,
    ) -> Self {
        let name = name.into();
        self.properties.insert(
            name.clone(),
            serde_json::json!({
                "type": "boolean",
                "description": description.into()
            }),
        );
        if required {
            self.required.push(name);
        }
        self
    }

    /// Add an array parameter
    pub fn array_param(
        mut self,
        name: impl Into<String>,
        description: impl Into<String>,
        item_type: &str,
        required: bool,
    ) -> Self {
        let name = name.into();
        self.properties.insert(
            name.clone(),
            serde_json::json!({
                "type": "array",
                "items": { "type": item_type },
                "description": description.into()
            }),
        );
        if required {
            self.required.push(name);
        }
        self
    }

    /// Add an enum parameter
    pub fn enum_param(
        mut self,
        name: impl Into<String>,
        description: impl Into<String>,
        values: &[&str],
        required: bool,
    ) -> Self {
        let name = name.into();
        self.properties.insert(
            name.clone(),
            serde_json::json!({
                "type": "string",
                "enum": values,
                "description": description.into()
            }),
        );
        if required {
            self.required.push(name);
        }
        self
    }

    /// Build the tool definition
    pub fn build(self) -> ToolDefinition {
        ToolDefinition {
            name: self.name,
            description: self.description,
            parameters: serde_json::json!({
                "type": "object",
                "properties": self.properties,
                "required": self.required
            }),
        }
    }
}

/// Check if a backend is available (has API key configured)
pub fn is_backend_available(backend: LlmBackend) -> bool {
    std::env::var(backend.api_key_env_var()).is_ok()
}

/// Get all available backends
pub fn available_backends() -> Vec<LlmBackend> {
    [
        LlmBackend::OpenAI,
        LlmBackend::Anthropic,
        LlmBackend::XAI,
        LlmBackend::Google,
        LlmBackend::Ollama,
        LlmBackend::OpenRouter,
    ]
    .into_iter()
    .filter(|b| is_backend_available(*b))
    .collect()
}

/// Estimate token count for text (rough approximation)
pub fn estimate_tokens(text: &str) -> u32 {
    // Rough estimate: ~4 chars per token for English
    (text.len() / 4) as u32
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_config_default() {
        let config = MultiLlmConfig::default();
        assert_eq!(config.primary_backend, LlmBackend::Anthropic);
        assert!(config.streaming);
    }

    #[test]
    fn test_message_construction() {
        let msg = Message::user("Hello");
        assert_eq!(msg.role, "user");
        assert_eq!(msg.content, "Hello");

        let msg = Message::system("You are helpful");
        assert_eq!(msg.role, "system");
    }

    #[test]
    fn test_tool_builder() {
        let tool = ToolBuilder::new("search", "Search for information")
            .string_param("query", "Search query", true)
            .number_param("limit", "Max results", false)
            .enum_param("type", "Search type", &["web", "images", "news"], false)
            .build();

        assert_eq!(tool.name, "search");
        assert!(tool.parameters["required"]
            .as_array()
            .unwrap()
            .contains(&serde_json::json!("query")));
    }

    #[test]
    fn test_conversation_memory() {
        let mut memory = ConversationMemory::new("test-conv");
        memory.add_message(Message::user("Hello"));
        memory.add_message(Message::assistant("Hi there!"));

        assert_eq!(memory.len(), 2);
        assert_eq!(memory.last_messages(1).len(), 1);
        assert_eq!(memory.last_messages(1)[0].role, "assistant");
    }

    #[test]
    fn test_backend_env_vars() {
        assert_eq!(LlmBackend::OpenAI.api_key_env_var(), "OPENAI_API_KEY");
        assert_eq!(LlmBackend::Anthropic.api_key_env_var(), "ANTHROPIC_API_KEY");
    }

    #[test]
    fn test_token_estimate() {
        let text = "Hello, world! This is a test.";
        let tokens = estimate_tokens(text);
        assert!(tokens > 0);
        assert!(tokens < 100);
    }
}