razeri_plugin_api/
llm.rs

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

razeri_plugin_api/
llm.rs
