razeri_plugin_api/
llm.rs

1// Copyright (C) 2026 Gabriel Shea <gshea@ideapivot.com> - All Rights Reserved
2
3//! LLM client factory trait and registry.
4//!
5//! [`LlmClientFactory`] is the extension point for custom LLM providers.
6//! [`LlmClient`] defines the chat interface that all providers implement.
7//! [`LlmRegistry`] provides string-keyed lookup for provider factories.
8
9use crate::error::LlmError;
10use async_trait::async_trait;
11use serde::{Deserialize, Serialize};
12use std::collections::HashMap;
13use std::sync::Arc;
14use tokio::sync::mpsc;
15
16// ── Message types ──────────────────────────────────────────────────────
17
18/// A single message in a chat conversation.
19#[derive(Debug, Clone, Serialize, Deserialize)]
20pub struct Message {
21    /// Role of the message sender (e.g. `"user"`, `"assistant"`, `"system"`).
22    pub role: String,
23    /// Text content of the message.
24    pub content: String,
25}
26
27/// Response from a chat completion request.
28#[derive(Debug, Clone, Serialize, Deserialize)]
29pub struct ChatResponse {
30    /// Generated text from the model.
31    pub content: String,
32    /// Model identifier used to produce this response.
33    pub model: String,
34    /// Token usage statistics, if reported by the provider.
35    pub usage: Option<TokenUsage>,
36}
37
38/// Token usage statistics from the LLM.
39#[derive(Debug, Clone, Serialize, Deserialize)]
40pub struct TokenUsage {
41    /// Number of tokens in the prompt (input).
42    pub prompt_tokens: u32,
43    /// Input tokens served from cache, when the provider reports them.
44    #[serde(default, skip_serializing_if = "Option::is_none")]
45    pub cached_input_tokens: Option<u32>,
46    /// Number of tokens generated by the model (output).
47    pub completion_tokens: u32,
48    /// Total tokens consumed (prompt + completion).
49    pub total_tokens: u32,
50    /// Tool/function calls consumed during the request, when reported.
51    #[serde(default, skip_serializing_if = "Option::is_none")]
52    pub tool_call_count: Option<u32>,
53    /// Time to evaluate the prompt (nanoseconds). Ollama: `prompt_eval_duration`.
54    #[serde(default)]
55    pub prompt_eval_nanos: u64,
56    /// Time to generate completion tokens (nanoseconds). Ollama: `eval_duration`.
57    #[serde(default)]
58    pub eval_nanos: u64,
59    /// Total wall-clock time for the request (nanoseconds). Ollama: `total_duration`.
60    #[serde(default)]
61    pub total_nanos: u64,
62    /// Time spent loading the model (nanoseconds). Ollama: `load_duration`.
63    #[serde(default)]
64    pub load_nanos: u64,
65}
66
67/// Information about a model available on a provider.
68#[derive(Debug, Clone, Serialize, Deserialize)]
69pub struct ModelInfo {
70    /// Model identifier (e.g. `"llama3:8b"`).
71    pub name: String,
72    /// Model size in bytes, if reported.
73    #[serde(default)]
74    pub size: u64,
75    /// ISO 8601 timestamp of the most recent modification, if reported.
76    #[serde(default)]
77    pub modified_at: String,
78    /// Content digest string, if reported.
79    #[serde(default)]
80    pub digest: String,
81}
82
83// ── Tool types ─────────────────────────────────────────────────────────
84
85/// A tool (function) exposed to the LLM for action execution.
86///
87/// Tools are passed in [`LlmClient::chat`] / [`LlmClient::chat_stream`]
88/// and let the LLM invoke typed workspace commands rather than embedding
89/// JSON payloads in reply text.
90#[derive(Debug, Clone, Serialize, Deserialize)]
91pub struct ToolDefinition {
92    /// Tool name used by the LLM when invoking the tool.
93    pub name: String,
94    /// Human/model readable description of what this tool does.
95    pub description: String,
96    /// JSON Schema object describing the tool's parameters.
97    pub parameters: serde_json::Value,
98}
99
100// ── Streaming types ────────────────────────────────────────────────────
101
102/// Events emitted during a streaming LLM response.
103///
104/// Mirrors the web-archive's `ChatMessage` model which separated
105/// `thinking` content from regular `content`, with an `isStreaming` flag.
106#[derive(Debug, Clone)]
107pub enum StreamEvent {
108    /// A token/chunk of content from the LLM.
109    Token(String),
110    /// Thinking/reasoning content (for thinking-capable models like
111    /// DeepSeek-R1, `QwQ`, Qwen3, etc.).
112    Thinking(String),
113    /// The LLM invoked a workspace tool (function call).
114    ToolCall {
115        /// Provider-assigned call identifier.
116        id: String,
117        /// Tool name.
118        name: String,
119        /// Tool arguments as a JSON string.
120        arguments_json: String,
121    },
122}
123
124// ── LLM client trait ───────────────────────────────────────────────────
125
126/// Unified interface for sending chat messages to any LLM provider.
127///
128/// This trait is implemented by each LLM provider (Ollama, `OpenAI`, Anthropic, etc.).
129#[async_trait]
130pub trait LlmClient: Send + Sync {
131    /// Send a chat completion request and get a text response.
132    ///
133    /// Parameters:
134    /// - `messages`: The conversation messages.
135    /// - `model`: Optional model override; `None` uses the provider default.
136    /// - `tools`: Optional list of tools the LLM can invoke. When non-empty,
137    ///   the LLM may return tool calls instead of (or alongside) text.
138    async fn chat(
139        &self,
140        messages: &[Message],
141        model: Option<&str>,
142        tools: Option<Vec<ToolDefinition>>,
143    ) -> Result<ChatResponse, LlmError>;
144
145    /// Send a streaming chat request, emitting [`StreamEvent`]s via the channel.
146    ///
147    /// Each token/chunk is sent through `event_tx` as it arrives from the
148    /// provider, enabling real-time progress visibility. The full assembled
149    /// response is still returned as a [`ChatResponse`] when the stream
150    /// completes.
151    ///
152    /// Callers pass `None` for `event_tx` to skip streaming events (equivalent
153    /// to calling [`chat`](Self::chat)).
154    ///
155    /// Default implementation falls back to non-streaming `chat()` and sends
156    /// the complete content as a single `Token` event.
157    async fn chat_stream(
158        &self,
159        messages: &[Message],
160        model: Option<&str>,
161        event_tx: Option<mpsc::UnboundedSender<StreamEvent>>,
162        tools: Option<Vec<ToolDefinition>>,
163    ) -> Result<ChatResponse, LlmError> {
164        let response = self.chat(messages, model, tools).await?;
165        if let Some(tx) = event_tx {
166            let _ = tx.send(StreamEvent::Token(response.content.clone()));
167        }
168        Ok(response)
169    }
170
171    /// Send a streaming chat request honoring the caller's generation params.
172    ///
173    /// `max_completion_tokens` bounds the response budget; `temperature` sets
174    /// determinism. The default IGNORES them and delegates to
175    /// [`chat_stream`](Self::chat_stream) so no provider is forced to change —
176    /// only providers that honor generation params (the OpenAI-compatible
177    /// client) override this. The workspace chat/rail path calls this so the
178    /// `LlmRequest`'s `max_tokens`/`temperature` actually reach the wire instead
179    /// of being silently dropped at the gateway bridge.
180    async fn chat_stream_with_params(
181        &self,
182        messages: &[Message],
183        model: Option<&str>,
184        event_tx: Option<mpsc::UnboundedSender<StreamEvent>>,
185        tools: Option<Vec<ToolDefinition>>,
186        _max_completion_tokens: Option<u32>,
187        _temperature: Option<f32>,
188    ) -> Result<ChatResponse, LlmError> {
189        self.chat_stream(messages, model, event_tx, tools).await
190    }
191
192    /// Send a chat completion request and get a JSON-structured response.
193    async fn chat_json(
194        &self,
195        messages: &[Message],
196        model: Option<&str>,
197        tools: Option<Vec<ToolDefinition>>,
198    ) -> Result<serde_json::Value, LlmError>;
199
200    /// List available models on this provider.
201    ///
202    /// # Errors
203    ///
204    /// Returns an error if the operation fails.
205    async fn list_models(&self) -> Result<Vec<ModelInfo>, LlmError>;
206
207    /// Health check — is the provider reachable?
208    ///
209    /// # Errors
210    ///
211    /// Returns an error if the operation fails.
212    async fn health_check(&self) -> Result<bool, LlmError>;
213}
214
215// ── LLM client factory ────────────────────────────────────────────────
216
217/// Factory that creates LLM client instances from configuration.
218///
219/// Third-party LLM providers (Anthropic, Bedrock, vLLM, etc.)
220/// implement this trait to integrate with Razeri's AI profiler.
221pub trait LlmClientFactory: Send + Sync {
222    /// Machine-readable provider name (e.g., "ollama", "openai", "anthropic").
223    fn name(&self) -> &str;
224
225    /// Human-readable display name.
226    fn display_name(&self) -> &str;
227
228    /// Create an LLM client from configuration (`base_url`, `api_key`, etc.).
229    ///
230    /// # Errors
231    ///
232    /// Returns an error if the operation fails.
233    fn create(&self, config: &HashMap<String, String>) -> Result<Arc<dyn LlmClient>, LlmError>;
234
235    /// Validate configuration before creating a client.
236    fn validate_config(&self, config: &HashMap<String, String>) -> Vec<String>;
237
238    /// List required configuration keys (e.g., "`base_url`", "`api_key`").
239    fn required_config_keys(&self) -> Vec<&str>;
240}
241
242// ── LLM registry ──────────────────────────────────────────────────────
243
244/// Registry for LLM providers (string-keyed).
245pub struct LlmRegistry {
246    providers: HashMap<String, Arc<dyn LlmClientFactory>>,
247}
248
249impl LlmRegistry {
250    /// Create an empty registry.
251    #[must_use]
252    pub fn new() -> Self {
253        Self { providers: HashMap::new() }
254    }
255
256    /// Register an LLM provider factory.
257    pub fn register(&mut self, name: String, factory: Arc<dyn LlmClientFactory>) {
258        self.providers.insert(name, factory);
259    }
260
261    /// Look up a provider factory by name.
262    #[must_use]
263    pub fn get(&self, name: &str) -> Option<&Arc<dyn LlmClientFactory>> {
264        self.providers.get(name)
265    }
266
267    /// List all registered providers as (name, `display_name`) pairs.
268    #[must_use]
269    pub fn list(&self) -> Vec<(&str, &str)> {
270        self.providers.values().map(|f| (f.name(), f.display_name())).collect()
271    }
272
273    /// Create a client by looking up the provider and delegating to its factory.
274    ///
275    /// # Errors
276    ///
277    /// Returns an error if the operation fails.
278    pub fn create_client(
279        &self,
280        provider_name: &str,
281        config: &HashMap<String, String>,
282    ) -> Result<Arc<dyn LlmClient>, LlmError> {
283        let factory = self
284            .providers
285            .get(provider_name)
286            .ok_or_else(|| LlmError::UnsupportedProvider(provider_name.to_string()))?;
287        factory.create(config)
288    }
289}
290
291impl Default for LlmRegistry {
292    fn default() -> Self {
293        Self::new()
294    }
295}
296
297#[cfg(test)]
298mod tests {
299    use super::*;
300
301    struct MockLlmClient;
302
303    #[async_trait]
304    impl LlmClient for MockLlmClient {
305        async fn chat(
306            &self,
307            _messages: &[Message],
308            _model: Option<&str>,
309            _tools: Option<Vec<ToolDefinition>>,
310        ) -> Result<ChatResponse, LlmError> {
311            Ok(ChatResponse {
312                content: "test response".to_string(),
313                model: "test-model".to_string(),
314                usage: None,
315            })
316        }
317
318        async fn chat_json(
319            &self,
320            _messages: &[Message],
321            _model: Option<&str>,
322            _tools: Option<Vec<ToolDefinition>>,
323        ) -> Result<serde_json::Value, LlmError> {
324            Ok(serde_json::json!({"response": "test"}))
325        }
326
327        async fn list_models(&self) -> Result<Vec<ModelInfo>, LlmError> {
328            Ok(vec![ModelInfo {
329                name: "test-model".to_string(),
330                size: 0,
331                modified_at: String::new(),
332                digest: String::new(),
333            }])
334        }
335
336        async fn health_check(&self) -> Result<bool, LlmError> {
337            Ok(true)
338        }
339    }
340
341    struct MockLlmFactory;
342
343    impl LlmClientFactory for MockLlmFactory {
344        fn name(&self) -> &'static str {
345            "mock"
346        }
347
348        fn display_name(&self) -> &'static str {
349            "Mock Provider"
350        }
351
352        fn create(
353            &self,
354            _config: &HashMap<String, String>,
355        ) -> Result<Arc<dyn LlmClient>, LlmError> {
356            Ok(Arc::new(MockLlmClient))
357        }
358
359        fn validate_config(&self, _config: &HashMap<String, String>) -> Vec<String> {
360            Vec::new()
361        }
362
363        fn required_config_keys(&self) -> Vec<&str> {
364            vec!["base_url"]
365        }
366    }
367
368    #[test]
369    fn llm_registry_resolves_registered_provider() {
370        let mut registry = LlmRegistry::new();
371        registry.register("mock".to_string(), Arc::new(MockLlmFactory));
372        assert!(registry.get("mock").is_some());
373    }
374
375    #[test]
376    fn llm_registry_create_client_delegates_to_factory() {
377        let mut registry = LlmRegistry::new();
378        registry.register("mock".to_string(), Arc::new(MockLlmFactory));
379        let config = HashMap::from([("base_url".to_string(), "http://localhost".to_string())]);
380        let client = registry.create_client("mock", &config);
381        assert!(client.is_ok());
382    }
383
384    #[test]
385    fn llm_registry_unknown_provider_returns_error() {
386        let registry = LlmRegistry::new();
387        let config = HashMap::new();
388        let result = registry.create_client("nonexistent", &config);
389        assert!(result.is_err());
390        let err = result.err().unwrap();
391        assert!(matches!(err, LlmError::UnsupportedProvider(_)));
392    }
393
394    #[test]
395    fn llm_registry_list_shows_all_providers() {
396        let mut registry = LlmRegistry::new();
397        registry.register("mock1".to_string(), Arc::new(MockLlmFactory));
398        registry.register("mock2".to_string(), Arc::new(MockLlmFactory));
399        assert_eq!(registry.list().len(), 2);
400    }
401
402    #[tokio::test]
403    async fn chat_stream_default_sends_full_content_as_single_token() {
404        let client = MockLlmClient;
405        let (tx, mut rx) = mpsc::unbounded_channel();
406        let messages = vec![Message { role: "user".to_string(), content: "hello".to_string() }];
407
408        let response = client.chat_stream(&messages, None, Some(tx), None).await.unwrap();
409
410        assert_eq!(response.content, "test response");
411        // Default implementation sends the entire content as one Token event
412        let event = rx.recv().await.unwrap();
413        assert!(matches!(event, StreamEvent::Token(ref s) if s == "test response"));
414        // Channel closed after that
415        assert!(rx.recv().await.is_none());
416    }
417
418    #[tokio::test]
419    async fn chat_stream_with_none_tx_returns_response_without_events() {
420        let client = MockLlmClient;
421        let messages = vec![Message { role: "user".to_string(), content: "hello".to_string() }];
422
423        let response = client.chat_stream(&messages, None, None, None).await.unwrap();
424
425        assert_eq!(response.content, "test response");
426    }
427
428    #[test]
429    fn stream_event_token_carries_content() {
430        let event = StreamEvent::Token("hello world".to_string());
431        match event {
432            StreamEvent::Token(s) => assert_eq!(s, "hello world"),
433            StreamEvent::Thinking(_) | StreamEvent::ToolCall { .. } => {
434                panic!("expected Token variant")
435            }
436        }
437    }
438
439    #[test]
440    fn stream_event_thinking_carries_reasoning() {
441        let event = StreamEvent::Thinking("let me analyze...".to_string());
442        match event {
443            StreamEvent::Thinking(s) => assert_eq!(s, "let me analyze..."),
444            StreamEvent::Token(_) | StreamEvent::ToolCall { .. } => {
445                panic!("expected Thinking variant")
446            }
447        }
448    }
449}
razeri_plugin_api/llm.rs

razeri_plugin_api/
llm.rs
