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    /// Input tokens written to the provider's prompt cache, when the provider
49    /// reports them. Written input is priced apart from fresh input and from
50    /// cache reads, so a spend figure needs it. These tokens are part of
51    /// `prompt_tokens`.
52    #[serde(default, skip_serializing_if = "Option::is_none")]
53    pub cache_write_tokens: Option<u32>,
54    /// Number of tokens generated by the model (output).
55    pub completion_tokens: u32,
56    /// Total tokens consumed (prompt + completion).
57    pub total_tokens: u32,
58    /// Tool/function calls consumed during the request, when reported.
59    #[serde(default, skip_serializing_if = "Option::is_none")]
60    pub tool_call_count: Option<u32>,
61    /// Time to evaluate the prompt (nanoseconds). Ollama: `prompt_eval_duration`.
62    #[serde(default)]
63    pub prompt_eval_nanos: u64,
64    /// Time to generate completion tokens (nanoseconds). Ollama: `eval_duration`.
65    #[serde(default)]
66    pub eval_nanos: u64,
67    /// Total wall-clock time for the request (nanoseconds). Ollama: `total_duration`.
68    #[serde(default)]
69    pub total_nanos: u64,
70    /// Time spent loading the model (nanoseconds). Ollama: `load_duration`.
71    #[serde(default)]
72    pub load_nanos: u64,
73    /// What a second model, offered to the answering one as an advisor, spent
74    /// on this request. It is billed at the advisor's rates, so it is not in
75    /// the counts above. `None` where no advisor was offered or the provider
76    /// has none.
77    #[serde(default, skip_serializing_if = "Option::is_none")]
78    pub advisor: Option<AdvisorTokenUsage>,
79}
80
81/// What an advisor model spent while it helped answer one request.
82///
83/// The four token counts are kept apart because they are priced apart, and
84/// they do not overlap: `input_tokens` is input read fresh, neither written to
85/// nor read from the cache. (The answering model's `prompt_tokens` is the sum
86/// of all three input kinds; the advisor's is not, so no count is added twice
87/// when the two are priced.)
88#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
89pub struct AdvisorTokenUsage {
90    /// The advisor's model. Empty when it was offered but never consulted and
91    /// the provider names none.
92    pub model: String,
93    /// How many times the answering model consulted it.
94    pub calls: u32,
95    /// Input tokens read fresh.
96    pub input_tokens: u32,
97    /// Input tokens read from the prompt cache.
98    pub cache_read_tokens: u32,
99    /// Input tokens written to the prompt cache.
100    pub cache_write_tokens: u32,
101    /// Tokens the advisor generated.
102    pub output_tokens: u32,
103}
104
105/// Information about a model available on a provider.
106#[derive(Debug, Clone, Serialize, Deserialize)]
107pub struct ModelInfo {
108    /// Model identifier (e.g. `"llama3:8b"`).
109    pub name: String,
110    /// Model size in bytes, if reported.
111    #[serde(default)]
112    pub size: u64,
113    /// ISO 8601 timestamp of the most recent modification, if reported.
114    #[serde(default)]
115    pub modified_at: String,
116    /// Content digest string, if reported.
117    #[serde(default)]
118    pub digest: String,
119}
120
121// ── Tool types ─────────────────────────────────────────────────────────
122
123/// A tool (function) exposed to the LLM for action execution.
124///
125/// Tools are passed in [`LlmClient::chat`] / [`LlmClient::chat_stream`]
126/// and let the LLM invoke typed workspace commands rather than embedding
127/// JSON payloads in reply text.
128#[derive(Debug, Clone, Serialize, Deserialize)]
129pub struct ToolDefinition {
130    /// Tool name used by the LLM when invoking the tool.
131    pub name: String,
132    /// Human/model readable description of what this tool does.
133    pub description: String,
134    /// JSON Schema object describing the tool's parameters.
135    pub parameters: serde_json::Value,
136}
137
138// ── Streaming types ────────────────────────────────────────────────────
139
140/// Events emitted during a streaming LLM response.
141///
142/// Separates `thinking` content from regular content, with a streaming
143/// flag, so a caller can render reasoning and output as distinct channels.
144#[derive(Debug, Clone)]
145pub enum StreamEvent {
146    /// A token/chunk of content from the LLM.
147    Token(String),
148    /// Thinking/reasoning content (for thinking-capable models like
149    /// DeepSeek-R1, `QwQ`, Qwen3, etc.).
150    Thinking(String),
151    /// The LLM invoked a workspace tool (function call).
152    ToolCall {
153        /// Provider-assigned call identifier.
154        id: String,
155        /// Tool name.
156        name: String,
157        /// Tool arguments as a JSON string.
158        arguments_json: String,
159    },
160}
161
162// ── LLM client trait ───────────────────────────────────────────────────
163
164/// Unified interface for sending chat messages to any LLM provider.
165///
166/// This trait is implemented by each LLM provider (Ollama, `OpenAI`, Anthropic, etc.).
167#[async_trait]
168pub trait LlmClient: Send + Sync {
169    /// Send a chat completion request and get a text response.
170    ///
171    /// Parameters:
172    /// - `messages`: The conversation messages.
173    /// - `model`: Optional model override; `None` uses the provider default.
174    /// - `tools`: Optional list of tools the LLM can invoke. When non-empty,
175    ///   the LLM may return tool calls instead of (or alongside) text.
176    async fn chat(
177        &self,
178        messages: &[Message],
179        model: Option<&str>,
180        tools: Option<Vec<ToolDefinition>>,
181    ) -> Result<ChatResponse, LlmError>;
182
183    /// Send a streaming chat request, emitting [`StreamEvent`]s via the channel.
184    ///
185    /// Each token/chunk is sent through `event_tx` as it arrives from the
186    /// provider, enabling real-time progress visibility. The full assembled
187    /// response is still returned as a [`ChatResponse`] when the stream
188    /// completes.
189    ///
190    /// Callers pass `None` for `event_tx` to skip streaming events (equivalent
191    /// to calling [`chat`](Self::chat)).
192    ///
193    /// Default implementation falls back to non-streaming `chat()` and sends
194    /// the complete content as a single `Token` event.
195    async fn chat_stream(
196        &self,
197        messages: &[Message],
198        model: Option<&str>,
199        event_tx: Option<mpsc::UnboundedSender<StreamEvent>>,
200        tools: Option<Vec<ToolDefinition>>,
201    ) -> Result<ChatResponse, LlmError> {
202        let response = self.chat(messages, model, tools).await?;
203        if let Some(tx) = event_tx {
204            let _ = tx.send(StreamEvent::Token(response.content.clone()));
205        }
206        Ok(response)
207    }
208
209    /// Send a streaming chat request honoring the caller's generation params.
210    ///
211    /// `max_completion_tokens` bounds the response budget; `temperature` sets
212    /// determinism. The default IGNORES them and delegates to
213    /// [`chat_stream`](Self::chat_stream) so no provider is forced to change —
214    /// only providers that honor generation params (the OpenAI-compatible
215    /// client) override this. The workspace chat/rail path calls this so the
216    /// `LlmRequest`'s `max_tokens`/`temperature` actually reach the wire instead
217    /// of being silently dropped at the gateway bridge.
218    async fn chat_stream_with_params(
219        &self,
220        messages: &[Message],
221        model: Option<&str>,
222        event_tx: Option<mpsc::UnboundedSender<StreamEvent>>,
223        tools: Option<Vec<ToolDefinition>>,
224        _max_completion_tokens: Option<u32>,
225        _temperature: Option<f32>,
226    ) -> Result<ChatResponse, LlmError> {
227        self.chat_stream(messages, model, event_tx, tools).await
228    }
229
230    /// Send a chat completion request and get a JSON-structured response.
231    async fn chat_json(
232        &self,
233        messages: &[Message],
234        model: Option<&str>,
235        tools: Option<Vec<ToolDefinition>>,
236    ) -> Result<serde_json::Value, LlmError>;
237
238    /// List available models on this provider.
239    ///
240    /// # Errors
241    ///
242    /// Returns an error if the operation fails.
243    async fn list_models(&self) -> Result<Vec<ModelInfo>, LlmError>;
244
245    /// Health check — is the provider reachable?
246    ///
247    /// # Errors
248    ///
249    /// Returns an error if the operation fails.
250    async fn health_check(&self) -> Result<bool, LlmError>;
251}
252
253// ── LLM client factory ────────────────────────────────────────────────
254
255/// Factory that creates LLM client instances from configuration.
256///
257/// Third-party LLM providers (Anthropic, Bedrock, vLLM, etc.)
258/// implement this trait to integrate with Razeri's AI profiler.
259pub trait LlmClientFactory: Send + Sync {
260    /// Machine-readable provider name (e.g., "ollama", "openai", "anthropic").
261    fn name(&self) -> &str;
262
263    /// Human-readable display name.
264    fn display_name(&self) -> &str;
265
266    /// Create an LLM client from configuration (`base_url`, `api_key`, etc.).
267    ///
268    /// # Errors
269    ///
270    /// Returns an error if the operation fails.
271    fn create(&self, config: &HashMap<String, String>) -> Result<Arc<dyn LlmClient>, LlmError>;
272
273    /// Validate configuration before creating a client.
274    fn validate_config(&self, config: &HashMap<String, String>) -> Vec<String>;
275
276    /// List required configuration keys (e.g., "`base_url`", "`api_key`").
277    fn required_config_keys(&self) -> Vec<&str>;
278}
279
280// ── LLM registry ──────────────────────────────────────────────────────
281
282/// Registry for LLM providers (string-keyed).
283pub struct LlmRegistry {
284    providers: HashMap<String, Arc<dyn LlmClientFactory>>,
285}
286
287impl LlmRegistry {
288    /// Create an empty registry.
289    #[must_use]
290    pub fn new() -> Self {
291        Self { providers: HashMap::new() }
292    }
293
294    /// Register an LLM provider factory.
295    pub fn register(&mut self, name: String, factory: Arc<dyn LlmClientFactory>) {
296        self.providers.insert(name, factory);
297    }
298
299    /// Look up a provider factory by name.
300    #[must_use]
301    pub fn get(&self, name: &str) -> Option<&Arc<dyn LlmClientFactory>> {
302        self.providers.get(name)
303    }
304
305    /// List all registered providers as (name, `display_name`) pairs.
306    #[must_use]
307    pub fn list(&self) -> Vec<(&str, &str)> {
308        self.providers.values().map(|f| (f.name(), f.display_name())).collect()
309    }
310
311    /// Create a client by looking up the provider and delegating to its factory.
312    ///
313    /// # Errors
314    ///
315    /// Returns an error if the operation fails.
316    pub fn create_client(
317        &self,
318        provider_name: &str,
319        config: &HashMap<String, String>,
320    ) -> Result<Arc<dyn LlmClient>, LlmError> {
321        let factory = self
322            .providers
323            .get(provider_name)
324            .ok_or_else(|| LlmError::UnsupportedProvider(provider_name.to_string()))?;
325        factory.create(config)
326    }
327}
328
329impl Default for LlmRegistry {
330    fn default() -> Self {
331        Self::new()
332    }
333}
334
335#[cfg(test)]
336mod tests {
337    use super::*;
338
339    #[test]
340    fn token_usage_round_trips_cache_writes_and_the_advisor() {
341        let usage = TokenUsage {
342            prompt_tokens: 4_020,
343            cached_input_tokens: Some(3_000),
344            cache_write_tokens: Some(900),
345            completion_tokens: 45,
346            total_tokens: 4_065,
347            tool_call_count: None,
348            prompt_eval_nanos: 0,
349            eval_nanos: 0,
350            total_nanos: 7,
351            load_nanos: 0,
352            advisor: Some(AdvisorTokenUsage {
353                model: "claude-opus-5-5".to_string(),
354                calls: 2,
355                input_tokens: 900,
356                cache_read_tokens: 300,
357                cache_write_tokens: 50,
358                output_tokens: 700,
359            }),
360        };
361        let json = serde_json::to_string(&usage).expect("serializes");
362        let back: TokenUsage = serde_json::from_str(&json).expect("reads back");
363        assert_eq!(back.cache_write_tokens, Some(900));
364        assert_eq!(back.advisor, usage.advisor);
365    }
366
367    #[test]
368    fn a_provider_that_reports_neither_leaves_them_out_of_the_json() {
369        let usage: TokenUsage =
370            serde_json::from_str(r#"{"prompt_tokens":1,"completion_tokens":2,"total_tokens":3}"#)
371                .expect("the older shape still reads");
372        assert_eq!(usage.cache_write_tokens, None);
373        assert!(usage.advisor.is_none());
374        let json = serde_json::to_value(&usage).expect("serializes");
375        assert!(json.get("cache_write_tokens").is_none(), "{json}");
376        assert!(json.get("advisor").is_none(), "{json}");
377    }
378
379    struct MockLlmClient;
380
381    #[async_trait]
382    impl LlmClient for MockLlmClient {
383        async fn chat(
384            &self,
385            _messages: &[Message],
386            _model: Option<&str>,
387            _tools: Option<Vec<ToolDefinition>>,
388        ) -> Result<ChatResponse, LlmError> {
389            Ok(ChatResponse {
390                content: "test response".to_string(),
391                model: "test-model".to_string(),
392                usage: None,
393            })
394        }
395
396        async fn chat_json(
397            &self,
398            _messages: &[Message],
399            _model: Option<&str>,
400            _tools: Option<Vec<ToolDefinition>>,
401        ) -> Result<serde_json::Value, LlmError> {
402            Ok(serde_json::json!({"response": "test"}))
403        }
404
405        async fn list_models(&self) -> Result<Vec<ModelInfo>, LlmError> {
406            Ok(vec![ModelInfo {
407                name: "test-model".to_string(),
408                size: 0,
409                modified_at: String::new(),
410                digest: String::new(),
411            }])
412        }
413
414        async fn health_check(&self) -> Result<bool, LlmError> {
415            Ok(true)
416        }
417    }
418
419    struct MockLlmFactory;
420
421    impl LlmClientFactory for MockLlmFactory {
422        fn name(&self) -> &'static str {
423            "mock"
424        }
425
426        fn display_name(&self) -> &'static str {
427            "Mock Provider"
428        }
429
430        fn create(
431            &self,
432            _config: &HashMap<String, String>,
433        ) -> Result<Arc<dyn LlmClient>, LlmError> {
434            Ok(Arc::new(MockLlmClient))
435        }
436
437        fn validate_config(&self, _config: &HashMap<String, String>) -> Vec<String> {
438            Vec::new()
439        }
440
441        fn required_config_keys(&self) -> Vec<&str> {
442            vec!["base_url"]
443        }
444    }
445
446    #[test]
447    fn llm_registry_resolves_registered_provider() {
448        let mut registry = LlmRegistry::new();
449        registry.register("mock".to_string(), Arc::new(MockLlmFactory));
450        assert!(registry.get("mock").is_some());
451    }
452
453    #[test]
454    fn llm_registry_create_client_delegates_to_factory() {
455        let mut registry = LlmRegistry::new();
456        registry.register("mock".to_string(), Arc::new(MockLlmFactory));
457        let config = HashMap::from([("base_url".to_string(), "http://localhost".to_string())]);
458        let client = registry.create_client("mock", &config);
459        assert!(client.is_ok());
460    }
461
462    #[test]
463    fn llm_registry_unknown_provider_returns_error() {
464        let registry = LlmRegistry::new();
465        let config = HashMap::new();
466        let result = registry.create_client("nonexistent", &config);
467        assert!(result.is_err());
468        let err = result.err().unwrap();
469        assert!(matches!(err, LlmError::UnsupportedProvider(_)));
470    }
471
472    #[test]
473    fn llm_registry_list_shows_all_providers() {
474        let mut registry = LlmRegistry::new();
475        registry.register("mock1".to_string(), Arc::new(MockLlmFactory));
476        registry.register("mock2".to_string(), Arc::new(MockLlmFactory));
477        assert_eq!(registry.list().len(), 2);
478    }
479
480    #[tokio::test]
481    async fn chat_stream_default_sends_full_content_as_single_token() {
482        let client = MockLlmClient;
483        let (tx, mut rx) = mpsc::unbounded_channel();
484        let messages = vec![Message { role: "user".to_string(), content: "hello".to_string() }];
485
486        let response = client.chat_stream(&messages, None, Some(tx), None).await.unwrap();
487
488        assert_eq!(response.content, "test response");
489        // Default implementation sends the entire content as one Token event
490        let event = rx.recv().await.unwrap();
491        assert!(matches!(event, StreamEvent::Token(ref s) if s == "test response"));
492        // Channel closed after that
493        assert!(rx.recv().await.is_none());
494    }
495
496    #[tokio::test]
497    async fn chat_stream_with_none_tx_returns_response_without_events() {
498        let client = MockLlmClient;
499        let messages = vec![Message { role: "user".to_string(), content: "hello".to_string() }];
500
501        let response = client.chat_stream(&messages, None, None, None).await.unwrap();
502
503        assert_eq!(response.content, "test response");
504    }
505
506    #[test]
507    fn stream_event_token_carries_content() {
508        let event = StreamEvent::Token("hello world".to_string());
509        match event {
510            StreamEvent::Token(s) => assert_eq!(s, "hello world"),
511            StreamEvent::Thinking(_) | StreamEvent::ToolCall { .. } => {
512                panic!("expected Token variant")
513            }
514        }
515    }
516
517    #[test]
518    fn stream_event_thinking_carries_reasoning() {
519        let event = StreamEvent::Thinking("let me analyze...".to_string());
520        match event {
521            StreamEvent::Thinking(s) => assert_eq!(s, "let me analyze..."),
522            StreamEvent::Token(_) | StreamEvent::ToolCall { .. } => {
523                panic!("expected Thinking variant")
524            }
525        }
526    }
527}
razeri_plugin_api/llm.rs

razeri_plugin_api/
llm.rs
