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

razeri_plugin_api/
llm.rs
