razeri-plugin-api 0.12.0

// Copyright (C) 2026 Gabriel Shea
// Licensed under either of the Apache License, Version 2.0 or the MIT
// license, at your option.

//! LLM client factory trait and registry.
//!
//! [`LlmClientFactory`] is the extension point for custom LLM providers.
//! [`LlmClient`] defines the chat interface that all providers implement.
//! [`LlmRegistry`] provides string-keyed lookup for provider factories.

use crate::error::LlmError;
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::mpsc;

// ── Message types ──────────────────────────────────────────────────────

/// A single message in a chat conversation.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Message {
    /// Role of the message sender (e.g. `"user"`, `"assistant"`, `"system"`).
    pub role: String,
    /// Text content of the message.
    pub content: String,
}

/// Response from a chat completion request.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatResponse {
    /// Generated text from the model.
    pub content: String,
    /// Model identifier used to produce this response.
    pub model: String,
    /// Token usage statistics, if reported by the provider.
    pub usage: Option<TokenUsage>,
}

/// Token usage statistics from the LLM.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TokenUsage {
    /// Number of tokens in the prompt (input).
    pub prompt_tokens: u32,
    /// Input tokens served from cache, when the provider reports them.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub cached_input_tokens: Option<u32>,
    /// Number of tokens generated by the model (output).
    pub completion_tokens: u32,
    /// Total tokens consumed (prompt + completion).
    pub total_tokens: u32,
    /// Tool/function calls consumed during the request, when reported.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub tool_call_count: Option<u32>,
    /// Time to evaluate the prompt (nanoseconds). Ollama: `prompt_eval_duration`.
    #[serde(default)]
    pub prompt_eval_nanos: u64,
    /// Time to generate completion tokens (nanoseconds). Ollama: `eval_duration`.
    #[serde(default)]
    pub eval_nanos: u64,
    /// Total wall-clock time for the request (nanoseconds). Ollama: `total_duration`.
    #[serde(default)]
    pub total_nanos: u64,
    /// Time spent loading the model (nanoseconds). Ollama: `load_duration`.
    #[serde(default)]
    pub load_nanos: u64,
}

/// Information about a model available on a provider.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelInfo {
    /// Model identifier (e.g. `"llama3:8b"`).
    pub name: String,
    /// Model size in bytes, if reported.
    #[serde(default)]
    pub size: u64,
    /// ISO 8601 timestamp of the most recent modification, if reported.
    #[serde(default)]
    pub modified_at: String,
    /// Content digest string, if reported.
    #[serde(default)]
    pub digest: String,
}

// ── Tool types ─────────────────────────────────────────────────────────

/// A tool (function) exposed to the LLM for action execution.
///
/// Tools are passed in [`LlmClient::chat`] / [`LlmClient::chat_stream`]
/// and let the LLM invoke typed workspace commands rather than embedding
/// JSON payloads in reply text.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolDefinition {
    /// Tool name used by the LLM when invoking the tool.
    pub name: String,
    /// Human/model readable description of what this tool does.
    pub description: String,
    /// JSON Schema object describing the tool's parameters.
    pub parameters: serde_json::Value,
}

// ── Streaming types ────────────────────────────────────────────────────

/// Events emitted during a streaming LLM response.
///
/// Separates `thinking` content from regular content, with a streaming
/// flag, so a caller can render reasoning and output as distinct channels.
#[derive(Debug, Clone)]
pub enum StreamEvent {
    /// A token/chunk of content from the LLM.
    Token(String),
    /// Thinking/reasoning content (for thinking-capable models like
    /// DeepSeek-R1, `QwQ`, Qwen3, etc.).
    Thinking(String),
    /// The LLM invoked a workspace tool (function call).
    ToolCall {
        /// Provider-assigned call identifier.
        id: String,
        /// Tool name.
        name: String,
        /// Tool arguments as a JSON string.
        arguments_json: String,
    },
}

// ── LLM client trait ───────────────────────────────────────────────────

/// Unified interface for sending chat messages to any LLM provider.
///
/// This trait is implemented by each LLM provider (Ollama, `OpenAI`, Anthropic, etc.).
#[async_trait]
pub trait LlmClient: Send + Sync {
    /// Send a chat completion request and get a text response.
    ///
    /// Parameters:
    /// - `messages`: The conversation messages.
    /// - `model`: Optional model override; `None` uses the provider default.
    /// - `tools`: Optional list of tools the LLM can invoke. When non-empty,
    ///   the LLM may return tool calls instead of (or alongside) text.
    async fn chat(
        &self,
        messages: &[Message],
        model: Option<&str>,
        tools: Option<Vec<ToolDefinition>>,
    ) -> Result<ChatResponse, LlmError>;

    /// Send a streaming chat request, emitting [`StreamEvent`]s via the channel.
    ///
    /// Each token/chunk is sent through `event_tx` as it arrives from the
    /// provider, enabling real-time progress visibility. The full assembled
    /// response is still returned as a [`ChatResponse`] when the stream
    /// completes.
    ///
    /// Callers pass `None` for `event_tx` to skip streaming events (equivalent
    /// to calling [`chat`](Self::chat)).
    ///
    /// Default implementation falls back to non-streaming `chat()` and sends
    /// the complete content as a single `Token` event.
    async fn chat_stream(
        &self,
        messages: &[Message],
        model: Option<&str>,
        event_tx: Option<mpsc::UnboundedSender<StreamEvent>>,
        tools: Option<Vec<ToolDefinition>>,
    ) -> Result<ChatResponse, LlmError> {
        let response = self.chat(messages, model, tools).await?;
        if let Some(tx) = event_tx {
            let _ = tx.send(StreamEvent::Token(response.content.clone()));
        }
        Ok(response)
    }

    /// Send a streaming chat request honoring the caller's generation params.
    ///
    /// `max_completion_tokens` bounds the response budget; `temperature` sets
    /// determinism. The default IGNORES them and delegates to
    /// [`chat_stream`](Self::chat_stream) so no provider is forced to change —
    /// only providers that honor generation params (the OpenAI-compatible
    /// client) override this. The workspace chat/rail path calls this so the
    /// `LlmRequest`'s `max_tokens`/`temperature` actually reach the wire instead
    /// of being silently dropped at the gateway bridge.
    async fn chat_stream_with_params(
        &self,
        messages: &[Message],
        model: Option<&str>,
        event_tx: Option<mpsc::UnboundedSender<StreamEvent>>,
        tools: Option<Vec<ToolDefinition>>,
        _max_completion_tokens: Option<u32>,
        _temperature: Option<f32>,
    ) -> Result<ChatResponse, LlmError> {
        self.chat_stream(messages, model, event_tx, tools).await
    }

    /// Send a chat completion request and get a JSON-structured response.
    async fn chat_json(
        &self,
        messages: &[Message],
        model: Option<&str>,
        tools: Option<Vec<ToolDefinition>>,
    ) -> Result<serde_json::Value, LlmError>;

    /// List available models on this provider.
    ///
    /// # Errors
    ///
    /// Returns an error if the operation fails.
    async fn list_models(&self) -> Result<Vec<ModelInfo>, LlmError>;

    /// Health check — is the provider reachable?
    ///
    /// # Errors
    ///
    /// Returns an error if the operation fails.
    async fn health_check(&self) -> Result<bool, LlmError>;
}

// ── LLM client factory ────────────────────────────────────────────────

/// Factory that creates LLM client instances from configuration.
///
/// Third-party LLM providers (Anthropic, Bedrock, vLLM, etc.)
/// implement this trait to integrate with Razeri's AI profiler.
pub trait LlmClientFactory: Send + Sync {
    /// Machine-readable provider name (e.g., "ollama", "openai", "anthropic").
    fn name(&self) -> &str;

    /// Human-readable display name.
    fn display_name(&self) -> &str;

    /// Create an LLM client from configuration (`base_url`, `api_key`, etc.).
    ///
    /// # Errors
    ///
    /// Returns an error if the operation fails.
    fn create(&self, config: &HashMap<String, String>) -> Result<Arc<dyn LlmClient>, LlmError>;

    /// Validate configuration before creating a client.
    fn validate_config(&self, config: &HashMap<String, String>) -> Vec<String>;

    /// List required configuration keys (e.g., "`base_url`", "`api_key`").
    fn required_config_keys(&self) -> Vec<&str>;
}

// ── LLM registry ──────────────────────────────────────────────────────

/// Registry for LLM providers (string-keyed).
pub struct LlmRegistry {
    providers: HashMap<String, Arc<dyn LlmClientFactory>>,
}

impl LlmRegistry {
    /// Create an empty registry.
    #[must_use]
    pub fn new() -> Self {
        Self { providers: HashMap::new() }
    }

    /// Register an LLM provider factory.
    pub fn register(&mut self, name: String, factory: Arc<dyn LlmClientFactory>) {
        self.providers.insert(name, factory);
    }

    /// Look up a provider factory by name.
    #[must_use]
    pub fn get(&self, name: &str) -> Option<&Arc<dyn LlmClientFactory>> {
        self.providers.get(name)
    }

    /// List all registered providers as (name, `display_name`) pairs.
    #[must_use]
    pub fn list(&self) -> Vec<(&str, &str)> {
        self.providers.values().map(|f| (f.name(), f.display_name())).collect()
    }

    /// Create a client by looking up the provider and delegating to its factory.
    ///
    /// # Errors
    ///
    /// Returns an error if the operation fails.
    pub fn create_client(
        &self,
        provider_name: &str,
        config: &HashMap<String, String>,
    ) -> Result<Arc<dyn LlmClient>, LlmError> {
        let factory = self
            .providers
            .get(provider_name)
            .ok_or_else(|| LlmError::UnsupportedProvider(provider_name.to_string()))?;
        factory.create(config)
    }
}

impl Default for LlmRegistry {
    fn default() -> Self {
        Self::new()
    }
}

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

    struct MockLlmClient;

    #[async_trait]
    impl LlmClient for MockLlmClient {
        async fn chat(
            &self,
            _messages: &[Message],
            _model: Option<&str>,
            _tools: Option<Vec<ToolDefinition>>,
        ) -> Result<ChatResponse, LlmError> {
            Ok(ChatResponse {
                content: "test response".to_string(),
                model: "test-model".to_string(),
                usage: None,
            })
        }

        async fn chat_json(
            &self,
            _messages: &[Message],
            _model: Option<&str>,
            _tools: Option<Vec<ToolDefinition>>,
        ) -> Result<serde_json::Value, LlmError> {
            Ok(serde_json::json!({"response": "test"}))
        }

        async fn list_models(&self) -> Result<Vec<ModelInfo>, LlmError> {
            Ok(vec![ModelInfo {
                name: "test-model".to_string(),
                size: 0,
                modified_at: String::new(),
                digest: String::new(),
            }])
        }

        async fn health_check(&self) -> Result<bool, LlmError> {
            Ok(true)
        }
    }

    struct MockLlmFactory;

    impl LlmClientFactory for MockLlmFactory {
        fn name(&self) -> &'static str {
            "mock"
        }

        fn display_name(&self) -> &'static str {
            "Mock Provider"
        }

        fn create(
            &self,
            _config: &HashMap<String, String>,
        ) -> Result<Arc<dyn LlmClient>, LlmError> {
            Ok(Arc::new(MockLlmClient))
        }

        fn validate_config(&self, _config: &HashMap<String, String>) -> Vec<String> {
            Vec::new()
        }

        fn required_config_keys(&self) -> Vec<&str> {
            vec!["base_url"]
        }
    }

    #[test]
    fn llm_registry_resolves_registered_provider() {
        let mut registry = LlmRegistry::new();
        registry.register("mock".to_string(), Arc::new(MockLlmFactory));
        assert!(registry.get("mock").is_some());
    }

    #[test]
    fn llm_registry_create_client_delegates_to_factory() {
        let mut registry = LlmRegistry::new();
        registry.register("mock".to_string(), Arc::new(MockLlmFactory));
        let config = HashMap::from([("base_url".to_string(), "http://localhost".to_string())]);
        let client = registry.create_client("mock", &config);
        assert!(client.is_ok());
    }

    #[test]
    fn llm_registry_unknown_provider_returns_error() {
        let registry = LlmRegistry::new();
        let config = HashMap::new();
        let result = registry.create_client("nonexistent", &config);
        assert!(result.is_err());
        let err = result.err().unwrap();
        assert!(matches!(err, LlmError::UnsupportedProvider(_)));
    }

    #[test]
    fn llm_registry_list_shows_all_providers() {
        let mut registry = LlmRegistry::new();
        registry.register("mock1".to_string(), Arc::new(MockLlmFactory));
        registry.register("mock2".to_string(), Arc::new(MockLlmFactory));
        assert_eq!(registry.list().len(), 2);
    }

    #[tokio::test]
    async fn chat_stream_default_sends_full_content_as_single_token() {
        let client = MockLlmClient;
        let (tx, mut rx) = mpsc::unbounded_channel();
        let messages = vec![Message { role: "user".to_string(), content: "hello".to_string() }];

        let response = client.chat_stream(&messages, None, Some(tx), None).await.unwrap();

        assert_eq!(response.content, "test response");
        // Default implementation sends the entire content as one Token event
        let event = rx.recv().await.unwrap();
        assert!(matches!(event, StreamEvent::Token(ref s) if s == "test response"));
        // Channel closed after that
        assert!(rx.recv().await.is_none());
    }

    #[tokio::test]
    async fn chat_stream_with_none_tx_returns_response_without_events() {
        let client = MockLlmClient;
        let messages = vec![Message { role: "user".to_string(), content: "hello".to_string() }];

        let response = client.chat_stream(&messages, None, None, None).await.unwrap();

        assert_eq!(response.content, "test response");
    }

    #[test]
    fn stream_event_token_carries_content() {
        let event = StreamEvent::Token("hello world".to_string());
        match event {
            StreamEvent::Token(s) => assert_eq!(s, "hello world"),
            StreamEvent::Thinking(_) | StreamEvent::ToolCall { .. } => {
                panic!("expected Token variant")
            }
        }
    }

    #[test]
    fn stream_event_thinking_carries_reasoning() {
        let event = StreamEvent::Thinking("let me analyze...".to_string());
        match event {
            StreamEvent::Thinking(s) => assert_eq!(s, "let me analyze..."),
            StreamEvent::Token(_) | StreamEvent::ToolCall { .. } => {
                panic!("expected Thinking variant")
            }
        }
    }
}