razeri-plugin-api 0.22.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,
}

/// Which part of a request the caller will send again, so a provider that
/// keeps a prompt cache marks only that part.
///
/// A cache write costs more than fresh input, and it pays back only when the
/// same text is asked again. The caller knows whether it will be; the provider
/// cannot tell. A provider with no prompt cache ignores the hint.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum PromptReuse {
    /// The conversation grows by a turn and is sent again: mark all of it.
    /// This is what a request that says nothing gets.
    #[default]
    Conversation,
    /// Only the system prompt comes back. Each call adds turns that are never
    /// asked again, like one file's profile under a prompt shared by every file.
    SystemPrompt,
    /// Nothing comes back: mark nothing.
    Nothing,
}

/// 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>,
    /// Input tokens written to the provider's prompt cache, when the provider
    /// reports them. Written input is priced apart from fresh input and from
    /// cache reads, so a spend figure needs it. These tokens are part of
    /// `prompt_tokens`.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub cache_write_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,
    /// What a second model, offered to the answering one as an advisor, spent
    /// on this request. It is billed at the advisor's rates, so it is not in
    /// the counts above. `None` where no advisor was offered or the provider
    /// has none.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub advisor: Option<AdvisorTokenUsage>,
}

/// What an advisor model spent while it helped answer one request.
///
/// The four token counts are kept apart because they are priced apart, and
/// they do not overlap: `input_tokens` is input read fresh, neither written to
/// nor read from the cache. (The answering model's `prompt_tokens` is the sum
/// of all three input kinds; the advisor's is not, so no count is added twice
/// when the two are priced.)
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct AdvisorTokenUsage {
    /// The advisor's model. Empty when it was offered but never consulted and
    /// the provider names none.
    pub model: String,
    /// How many times the answering model consulted it.
    pub calls: u32,
    /// Input tokens read fresh.
    pub input_tokens: u32,
    /// Input tokens read from the prompt cache.
    pub cache_read_tokens: u32,
    /// Input tokens written to the prompt cache.
    pub cache_write_tokens: u32,
    /// Tokens the advisor generated.
    pub output_tokens: u32,
}

/// 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 streaming chat request that says which part of it will be sent
    /// again (see [`PromptReuse`]).
    ///
    /// The default ignores the hint and delegates to
    /// [`chat_stream`](Self::chat_stream); only a provider with a prompt
    /// cache overrides it.
    async fn chat_stream_with_reuse(
        &self,
        messages: &[Message],
        model: Option<&str>,
        event_tx: Option<mpsc::UnboundedSender<StreamEvent>>,
        tools: Option<Vec<ToolDefinition>>,
        _reuse: PromptReuse,
    ) -> 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::*;

    #[test]
    fn token_usage_round_trips_cache_writes_and_the_advisor() {
        let usage = TokenUsage {
            prompt_tokens: 4_020,
            cached_input_tokens: Some(3_000),
            cache_write_tokens: Some(900),
            completion_tokens: 45,
            total_tokens: 4_065,
            tool_call_count: None,
            prompt_eval_nanos: 0,
            eval_nanos: 0,
            total_nanos: 7,
            load_nanos: 0,
            advisor: Some(AdvisorTokenUsage {
                model: "claude-opus-5-5".to_string(),
                calls: 2,
                input_tokens: 900,
                cache_read_tokens: 300,
                cache_write_tokens: 50,
                output_tokens: 700,
            }),
        };
        let json = serde_json::to_string(&usage).expect("serializes");
        let back: TokenUsage = serde_json::from_str(&json).expect("reads back");
        assert_eq!(back.cache_write_tokens, Some(900));
        assert_eq!(back.advisor, usage.advisor);
    }

    #[test]
    fn a_provider_that_reports_neither_leaves_them_out_of_the_json() {
        let usage: TokenUsage =
            serde_json::from_str(r#"{"prompt_tokens":1,"completion_tokens":2,"total_tokens":3}"#)
                .expect("the older shape still reads");
        assert_eq!(usage.cache_write_tokens, None);
        assert!(usage.advisor.is_none());
        let json = serde_json::to_value(&usage).expect("serializes");
        assert!(json.get("cache_write_tokens").is_none(), "{json}");
        assert!(json.get("advisor").is_none(), "{json}");
    }

    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")
            }
        }
    }
}