Skip to main content

flashmind_types/
llm.rs

1//! LLM wire types — request / response structures and the [`LlmProvider`] trait.
2//!
3//! This module defines everything needed to communicate with an LLM backend:
4//! the provider trait, completion requests and responses, streaming events,
5//! model capabilities, and multimodal support (audio, images, video).
6//!
7//! # Key types
8//!
9//! | Type | Role |
10//! |------|------|
11//! | [`LlmProvider`] | Trait implemented by each LLM backend (OpenRouter, Anthropic, Ollama…) |
12//! | [`CompletionRequest`] | Single completion request sent to a provider |
13//! | [`CompletionResponse`] | Non-streaming response wrapper |
14//! | [`StreamEvent`] | Incremental events from the provider's stream |
15//! | [`CompletionStream`] | Alias for `Pin<Box<dyn Stream<Item = Result<StreamEvent>>>>` |
16//! | [`FinishReason`] | Why generation stopped (stop, length, tool_calls, content_filter) |
17//! | [`ToolDefinition`] | OpenAI-compatible JSON schema for tool calling |
18//! | [`ModelCapabilities`] | Feature flags (tool_calling, images, reasoning, audio…) |
19//! | [`TokenUsage`] | Token consumption reported by the provider |
20//! | [`Modality`] | Output modality: text, audio, or image |
21//! | [`TtsRequest`], [`SttRequest`], [`VideoGenRequest`] | Multimodal operation requests |
22
23use std::collections::{BTreeMap, HashMap};
24use std::pin::Pin;
25use std::sync::Arc;
26
27use async_trait::async_trait;
28use chrono::{DateTime, Utc};
29use futures_core::Stream;
30use serde::{Deserialize, Serialize};
31
32use rust_decimal::Decimal;
33
34use crate::event::TurnUsage;
35use crate::message::Message;
36use crate::model::{
37    AgentLlmConfig, Model, Provider, ProviderPreferences, ReasoningLevel, SamplingParams,
38};
39use crate::session::SessionId;
40
41/// OpenAI-compatible JSON schema for a tool exposed to the LLM.
42///
43/// Mirrors the structure sent in the `tools` array of chat completion requests.
44#[derive(Debug, Clone, Serialize, Deserialize)]
45pub struct ToolDefinition {
46    /// Tool name — must match the [`Tool::name`](crate::tool::Tool::name) implementation.
47    pub name: String,
48    /// Description shown to the model to help it decide when to call this tool.
49    pub description: String,
50    /// JSON Schema object describing the parameters this tool accepts.
51    pub parameters: serde_json::Value,
52}
53
54/// A single completion request sent to an LLM provider.
55///
56/// This is the primary input type for [`LlmProvider::complete`]. The agent runtime
57/// assembles this from the conversation history, active tools, and LLM config.
58#[derive(Debug)]
59pub struct CompletionRequest {
60    /// Model to use (includes provider prefix).
61    pub model: Model,
62    /// Message history in wire format.
63    pub messages: Vec<Message>,
64    /// Tools available to the model in this turn.
65    pub tools: Vec<ToolDefinition>,
66    /// Maximum output tokens. If `None`, the provider decides.
67    pub max_tokens: Option<u32>,
68    /// Whether reasoning/thinking mode is enabled.
69    pub reasoning: ReasoningLevel,
70    /// Sampling parameters (temperature, top_p, top_k, min_p, penalties).
71    pub sampling: SamplingParams,
72    /// Output modalities (e.g. text+audio, image). Empty means text-only.
73    pub modalities: Vec<Modality>,
74    /// Audio output configuration. Only used when modalities includes [`Modality::Audio`].
75    pub audio_config: Option<AudioOutputConfig>,
76    /// Image generation configuration. Only used when modalities includes [`Modality::Image`].
77    pub image_config: Option<ImageGenConfig>,
78    /// Opaque user identifier for per-user tracking in provider dashboards.
79    pub user: Option<String>,
80    /// Provider routing preferences (used by OpenRouter).
81    pub provider_preferences: Option<ProviderPreferences>,
82    /// Stable conversation identifier for provider cache routing.
83    pub session_id: Option<SessionId>,
84}
85
86/// A provider-issued context window to replay without editing its contents.
87#[derive(Debug, Clone)]
88pub struct NativeCompaction {
89    pub context: crate::ReasoningDetails,
90    pub summary: String,
91    pub usage: TokenUsage,
92}
93
94impl CompletionRequest {
95    /// Build a request with default settings and the supplied messages.
96    pub fn new(model: Model, messages: Vec<Message>) -> Self {
97        Self::from_config(&AgentLlmConfig::new(model), messages)
98    }
99
100    /// Keep the agent's session ID, routing, user and generation settings.
101    ///
102    /// ```
103    /// use flashmind_types::{AgentLlmConfig, CompletionRequest, Message, SessionId};
104    ///
105    /// let session = SessionId::new();
106    /// let llm = AgentLlmConfig::new("openrouter:example/model".parse()?)
107    ///     .with_session_id(session);
108    /// let request = CompletionRequest::from_config(&llm, vec![Message::user("Hello")]);
109    /// assert_eq!(request.session_id, llm.session_id);
110    /// # Ok::<(), anyhow::Error>(())
111    /// ```
112    pub fn from_config(config: &AgentLlmConfig, messages: Vec<Message>) -> Self {
113        Self {
114            model: config.model.clone(),
115            messages,
116            tools: Vec::new(),
117            max_tokens: config.max_tokens,
118            reasoning: config.reasoning,
119            sampling: config.sampling.clone(),
120            modalities: Vec::new(),
121            audio_config: None,
122            image_config: None,
123            user: config.user.clone(),
124            provider_preferences: config.provider_preferences.clone(),
125            session_id: config.session_id.clone(),
126        }
127    }
128
129    /// Use an existing session identity for this request.
130    pub fn with_session_id(mut self, session_id: impl Into<SessionId>) -> Self {
131        self.session_id = Some(session_id.into());
132        self
133    }
134}
135
136/// Why the LLM stopped generating (used in stream and non-stream responses).
137#[derive(Debug, Clone, Copy, PartialEq, Eq, strum::Display, strum::EnumString)]
138#[strum(serialize_all = "snake_case")]
139pub enum FinishReason {
140    /// Natural stop token or end of response.
141    Stop,
142    /// Model requested one or more tool calls.
143    ToolCalls,
144    /// Hit `max_tokens` limit without natural stopping point.
145    Length,
146    /// Output was filtered by content policy.
147    ContentFilter,
148}
149
150/// High-level category for filtering models by use-case.
151#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, strum::Display, strum::EnumString)]
152#[strum(serialize_all = "snake_case")]
153pub enum ModelCategory {
154    /// General text chat / instruction-following.
155    Chat,
156    /// Advanced reasoning / chain-of-thought.
157    Reasoning,
158    /// Can accept images as input.
159    Vision,
160    /// Generates images from text.
161    ImageGeneration,
162    /// Edits existing images.
163    ImageEditing,
164    /// Generates audio / text-to-speech.
165    Tts,
166    /// Accepts audio input / speech-to-text.
167    Stt,
168    /// Generates video.
169    VideoGeneration,
170    /// Generates video from an image.
171    ImageToVideo,
172    /// Accepts video as input.
173    VideoInput,
174}
175
176/// Pricing information for a model, including input/output token costs and caching discounts.
177///
178/// All costs are per-token in USD. Use [`TokenUsage::cost`] to compute the total
179/// cost of a request from token counts and pricing.
180#[derive(Debug, Clone, Default)]
181pub struct ModelPricing {
182    /// Cost per input token (USD).
183    pub prompt: Option<Decimal>,
184    /// Cost per output token (USD).
185    pub completion: Option<Decimal>,
186    /// Cost per image generated (USD).
187    pub image: Option<Decimal>,
188    /// Cost per cached input token read (USD).
189    pub cache_read: Option<Decimal>,
190    /// Cost per cache creation (write) input token (USD).
191    pub cache_write: Option<Decimal>,
192    /// Published reference rates, which may differ from the provider's charges.
193    pub reference: Option<ReferencePricing>,
194}
195
196/// Reference token rates from a separately identified source.
197#[derive(Debug, Clone)]
198pub struct ReferencePricing {
199    /// Source of the reference rates.
200    pub source: ReferencePricingSource,
201    /// Reference input price per token in USD.
202    pub prompt: Option<Decimal>,
203    /// Reference output price per token in USD.
204    pub completion: Option<Decimal>,
205}
206
207/// Source of a model's reference rates.
208#[derive(Debug, Clone, Copy, PartialEq, Eq, strum::Display)]
209pub enum ReferencePricingSource {
210    /// Prices published by OpenRouter.
211    OpenRouter,
212    /// Prices published by models.dev.
213    #[strum(to_string = "models.dev")]
214    ModelsDev,
215}
216
217/// Metadata about an available LLM model, including capabilities, pricing, and context window.
218#[derive(Debug, Clone)]
219pub struct ModelInfo {
220    /// Model ID used in API requests.
221    pub id: String,
222    /// Human-readable display name.
223    pub name: Option<String>,
224    /// Model creation time, when supplied by the provider.
225    pub created_at: Option<DateTime<Utc>>,
226    /// Context window size in tokens.
227    pub context_length: Option<u32>,
228    /// Maximum output tokens allowed.
229    pub max_completion_tokens: Option<u32>,
230    /// Capability flags.
231    pub capabilities: ModelCapabilities,
232    /// High-level categories (chat, vision, reasoning, etc.).
233    pub categories: Vec<ModelCategory>,
234    /// Per-token pricing.
235    pub pricing: ModelPricing,
236    /// Image-specific capability metadata from the provider.
237    pub image: ImageModelCapabilities,
238}
239
240/// Image capability metadata, when provided in a model catalog.
241#[derive(Debug, Clone, Default)]
242pub struct ImageModelCapabilities {
243    /// Supports image editing.
244    pub editing: bool,
245    /// Whether style references are supported, if reported.
246    pub style_references: Option<bool>,
247    /// Maximum number of images for an edit, if reported.
248    pub max_edit_images: Option<u32>,
249    /// Provider-reported image options and their allowed values or ranges.
250    pub parameters: BTreeMap<String, serde_json::Value>,
251}
252
253/// Feature flags describing what a model supports.
254///
255/// Defaults to all `false` via `#[derive(Default)]` — nothing is assumed.
256/// Providers override these after inspecting the model's actual capabilities.
257/// Use [`all()`](Self::all()) for a curated "full-featured text model" preset.
258#[derive(Debug, Clone, Copy, Default)]
259pub struct ModelCapabilities {
260    /// Can invoke tools/function calling.
261    pub tool_calling: bool,
262    /// Can accept images as input (vision).
263    pub images: bool,
264    /// Can accept documents (PDFs, etc.) as input.
265    pub documents: bool,
266    /// Can accept video as input.
267    pub video: bool,
268    /// Can accept audio as input (speech-to-text).
269    pub audio: bool,
270    /// Supports extended chain-of-thought / reasoning mode.
271    pub reasoning: bool,
272    // Output capabilities
273    /// Can generate audio output (text-to-speech inline).
274    pub audio_output: bool,
275    /// Can generate images as output.
276    pub image_generation: bool,
277    /// Can generate video as output.
278    pub video_generation: bool,
279}
280
281impl ModelCapabilities {
282    /// Derive high-level categories from capability flags.
283    pub fn categories(&self) -> Vec<ModelCategory> {
284        let mut cats = vec![ModelCategory::Chat];
285        if self.reasoning {
286            cats.push(ModelCategory::Reasoning);
287        }
288        if self.images {
289            cats.push(ModelCategory::Vision);
290        }
291        if self.image_generation {
292            cats.push(ModelCategory::ImageGeneration);
293        }
294        if self.audio_output {
295            cats.push(ModelCategory::Tts);
296        }
297        if self.audio {
298            cats.push(ModelCategory::Stt);
299        }
300        if self.video_generation {
301            cats.push(ModelCategory::VideoGeneration);
302        }
303        if self.video {
304            cats.push(ModelCategory::VideoInput);
305        }
306        cats
307    }
308
309    /// Curated "full-featured text model" capability preset.
310    ///
311    /// Enables common features (tool calling, vision, reasoning, audio input)
312    /// but keeps output-generation capabilities (`audio_output`, `image_generation`,
313    /// `video_generation`) disabled, as those require dedicated models.
314    pub fn all() -> Self {
315        Self {
316            tool_calling: true,
317            images: true,
318            documents: true,
319            video: false,
320            audio: true,
321            reasoning: true,
322            audio_output: false,
323            image_generation: false,
324            video_generation: false,
325        }
326    }
327}
328
329/// Token usage counts for a single completion request or turn.
330#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
331pub struct TokenUsage {
332    /// Tokens consumed by the prompt (input messages).
333    pub prompt_tokens: u32,
334    /// Tokens generated by the model (output/response).
335    pub completion_tokens: u32,
336    /// Total tokens (prompt + completion).
337    pub total_tokens: u32,
338    /// Prompt tokens served from the provider's cache.
339    pub cache_read_tokens: u32,
340    /// Prompt tokens written to the provider's cache.
341    pub cache_creation_tokens: u32,
342}
343
344impl TokenUsage {
345    /// Compute the cost of this usage given per-token pricing.
346    ///
347    /// Non-cached input tokens are charged at the `prompt` rate, cached reads at
348    /// `cache_read` (falling back to `prompt`), and cache writes at `cache_write`
349    /// (falling back to `prompt`). Returns `None` if pricing has no prompt or
350    /// completion prices.
351    pub fn cost(&self, pricing: &ModelPricing) -> Option<Decimal> {
352        let prompt_price = pricing.prompt?;
353        let completion_price = pricing.completion?;
354        let cache_read_price = pricing.cache_read.unwrap_or(prompt_price);
355        let cache_write_price = pricing.cache_write.unwrap_or(prompt_price);
356
357        let non_cached = self
358            .prompt_tokens
359            .saturating_sub(self.cache_read_tokens + self.cache_creation_tokens);
360
361        Some(
362            Decimal::from(non_cached) * prompt_price
363                + Decimal::from(self.completion_tokens) * completion_price
364                + Decimal::from(self.cache_read_tokens) * cache_read_price
365                + Decimal::from(self.cache_creation_tokens) * cache_write_price,
366        )
367    }
368}
369
370impl From<TurnUsage> for TokenUsage {
371    fn from(u: TurnUsage) -> Self {
372        Self {
373            prompt_tokens: u.prompt_tokens,
374            completion_tokens: u.completion_tokens,
375            total_tokens: u.prompt_tokens + u.completion_tokens,
376            cache_read_tokens: u.cache_read_tokens,
377            cache_creation_tokens: u.cache_creation_tokens,
378        }
379    }
380}
381
382/// Non-streaming completion response wrapper.
383#[derive(Debug, Clone)]
384pub struct CompletionResponse {
385    /// The assistant's response message.
386    pub message: Message,
387    /// Token usage for this completion.
388    pub usage: TokenUsage,
389    /// Why generation stopped.
390    pub finish_reason: FinishReason,
391}
392
393/// Events emitted incrementally by the provider's stream.
394///
395/// Each variant corresponds to a server-side SSE/data-line event type.
396/// Consumers assemble these into higher-level [`AgentEvent`](crate::event::AgentEvent)s.
397#[derive(Debug, Clone)]
398pub enum StreamEvent {
399    /// Incremental reasoning/chain-of-thought token.
400    ReasoningDelta(String),
401    /// Provider reasoning fragments, including signatures and encrypted state.
402    ReasoningDetailsDelta(crate::ReasoningDetails),
403    /// Incremental assistant text token.
404    ContentDelta(String),
405    /// The LLM started requesting a tool call at `index`.
406    ToolCallStart {
407        index: usize,
408        id: String,
409        name: String,
410    },
411    /// Incremental JSON argument string for the tool at `index`.
412    ToolCallDelta { index: usize, arguments: String },
413    /// Token usage snapshot mid-stream (some providers emit this).
414    Usage(TokenUsage),
415    /// Incremental audio output chunk (raw PCM/MP3 bytes).
416    AudioDelta { data: Vec<u8>, format: String },
417    /// Out-of-band file delivered by the provider (e.g. generated image).
418    FileAttachment {
419        filename: String,
420        media_type: String,
421        data: Vec<u8>,
422    },
423    /// Stream has ended; payload is the final finish reason.
424    Finished(FinishReason),
425}
426
427/// Type alias for a stream of completion events from an LLM provider.
428///
429/// The stream yields `StreamEvent`s (content deltas, tool calls, finish reasons) and is consumed by the agent loop.
430pub type CompletionStream =
431    Pin<Box<dyn Stream<Item = anyhow::Result<StreamEvent>> + Send + 'static>>;
432
433/// A voice available for text-to-speech generation.
434#[derive(Debug, Clone, Serialize, Deserialize)]
435pub struct Voice {
436    /// Unique voice identifier (e.g., `"alloy"`).
437    pub id: String,
438    /// Human-readable display name.
439    pub name: String,
440}
441
442/// Parameters for a text-to-speech request.
443pub struct TtsRequest {
444    /// TTS model identifier (e.g., `"tts-1"`).
445    pub model: String,
446    /// Text to synthesize into speech.
447    pub input: String,
448    /// Voice preset ID (e.g., `"alloy"`, `"echo"`).
449    pub voice: String,
450    /// Desired output audio format.
451    pub response_format: AudioFormat,
452}
453
454/// Parameters for a speech-to-text (transcription) request.
455pub struct SttRequest {
456    /// STT model identifier (e.g., `"whisper-1"`).
457    pub model: String,
458    /// Raw audio bytes.
459    pub audio: Vec<u8>,
460    /// MIME type of the audio (e.g. "audio/mp3", "audio/wav").
461    pub media_type: String,
462    /// Optional language hint (ISO 639-1, e.g. "en").
463    pub language: Option<String>,
464}
465
466/// Output audio format for TTS responses.
467///
468/// Serialises to lowercase strings matching common audio file extensions.
469#[derive(
470    Debug, Clone, Copy, Default, strum::EnumString, strum::Display, Serialize, Deserialize,
471)]
472#[strum(serialize_all = "lowercase")]
473#[serde(rename_all = "lowercase")]
474pub enum AudioFormat {
475    /// MP3 — most widely supported, good compression.
476    #[default]
477    Mp3,
478    /// WAV — uncompressed PCM, highest quality.
479    Wav,
480    /// Opus — efficient codec optimized for speech.
481    Opus,
482    /// AAC — Advanced Audio Coding, common in mobile ecosystems.
483    Aac,
484    /// FLAC — lossless compression.
485    Flac,
486}
487
488/// Output modality requested in a completion.
489///
490/// Multiple modalities can be requested simultaneously (e.g., text + audio).
491/// An empty list means text-only output.
492#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, strum::Display)]
493#[serde(rename_all = "lowercase")]
494#[strum(serialize_all = "lowercase")]
495pub enum Modality {
496    /// Plain text output (default when no modalities specified).
497    Text,
498    /// Inline audio output (real-time voice).
499    Audio,
500    /// Image output (generated during completion).
501    Image,
502}
503
504/// Configuration for audio output in multimodal completions.
505#[derive(Debug, Clone)]
506pub struct AudioOutputConfig {
507    /// Voice preset ID (e.g., `"alloy"`).
508    pub voice: String,
509    /// Output audio format.
510    pub format: AudioFormat,
511}
512
513/// Configuration for image generation requests.
514#[derive(Debug, Clone, Default)]
515pub struct ImageGenConfig {
516    /// Aspect ratio for the generated image (e.g., `"16:9"`).
517    pub aspect_ratio: Option<String>,
518    /// Absolute dimensions (e.g., `"1024x1024"`).
519    pub size: Option<String>,
520}
521
522impl ImageGenConfig {
523    /// Read an image file and return it as a `data:<mime>;base64,...` URI.
524    ///
525    /// # Arguments
526    ///
527    /// * `path` - Path to the image file on disk.
528    ///
529    /// # Returns
530    ///
531    /// A data URI string with auto-detected MIME type, or an I/O error if the file cannot be read.
532    pub fn data_uri_from_path(path: &std::path::Path) -> std::io::Result<String> {
533        let bytes = std::fs::read(path)?;
534        let mime = mime_from_extension(path);
535        Ok(Self::data_uri_from_bytes(&bytes, mime))
536    }
537
538    /// Encode raw bytes as a `data:<mime>;base64,...` URI.
539    ///
540    /// # Arguments
541    ///
542    /// * `bytes` - Raw binary data (e.g., image file contents).
543    /// * `mime` - MIME type string (e.g., `"image/png"`).
544    ///
545    /// # Returns
546    ///
547    /// A data URI string with the provided MIME type and base64-encoded payload.
548    pub fn data_uri_from_bytes(bytes: &[u8], mime: &str) -> String {
549        use base64::Engine;
550        let b64 = base64::engine::general_purpose::STANDARD.encode(bytes);
551        format!("data:{mime};base64,{b64}")
552    }
553}
554
555/// Infer MIME type from a file path's extension.
556///
557/// # Arguments
558///
559/// * `path` - File path whose extension is inspected.
560///
561/// # Returns
562///
563/// A static MIME type string corresponding to the extension, or `"application/octet-stream"`
564/// if the extension is unrecognized or missing.
565pub fn mime_from_extension(path: &std::path::Path) -> &'static str {
566    match path
567        .extension()
568        .and_then(|e| e.to_str())
569        .map(|e| e.to_ascii_lowercase())
570        .as_deref()
571    {
572        Some("png") => "image/png",
573        Some("jpg" | "jpeg") => "image/jpeg",
574        Some("gif") => "image/gif",
575        Some("webp") => "image/webp",
576        Some("bmp") => "image/bmp",
577        Some("svg") => "image/svg+xml",
578        _ => "application/octet-stream",
579    }
580}
581
582/// Image format requested from a provider.
583#[derive(
584    Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, strum::Display, strum::EnumString,
585)]
586#[serde(rename_all = "lowercase")]
587#[strum(serialize_all = "lowercase")]
588pub enum ImageOutputFormat {
589    /// PNG image.
590    Png,
591    /// JPEG image.
592    Jpeg,
593    /// WebP image.
594    Webp,
595}
596
597impl ImageOutputFormat {
598    /// MIME type corresponding to the format.
599    pub fn media_type(self) -> &'static str {
600        match self {
601            Self::Png => "image/png",
602            Self::Jpeg => "image/jpeg",
603            Self::Webp => "image/webp",
604        }
605    }
606}
607
608/// Parameters for generating an image via a generative model.
609#[derive(Debug, Clone)]
610pub struct ImageGenRequest {
611    /// Image generation model identifier.
612    pub model: String,
613    /// Text description of the desired image.
614    pub prompt: String,
615    /// Image dimensions (e.g., `"1024x1024"`).
616    pub size: Option<String>,
617    /// Aspect ratio (e.g., `"16:9"`).
618    pub aspect_ratio: Option<String>,
619    /// Quality preset (e.g., `"standard"`, `"hd"`).
620    pub quality: Option<String>,
621    /// Style preset (e.g., `"natural"`, `"vivid"`).
622    pub style: Option<String>,
623    /// Number of images to generate.
624    pub n: Option<u32>,
625    /// Reference images as URLs or data URIs.
626    pub input_references: Vec<String>,
627    /// Requested output format, if supported by the provider.
628    pub output_format: Option<ImageOutputFormat>,
629}
630
631/// Parameters for editing one or more images.
632#[derive(Debug, Clone)]
633pub struct ImageEditRequest {
634    /// Provider model identifier.
635    pub model: String,
636    /// Instructions describing the edit.
637    pub prompt: String,
638    /// Input images as URLs or data URIs, beginning with the base image.
639    pub images: Vec<String>,
640    /// Output resolution tier.
641    pub resolution: Option<String>,
642    /// Output aspect ratio.
643    pub aspect_ratio: Option<String>,
644    /// Provider-specific quality tier.
645    pub quality: Option<String>,
646    /// Mask as a URL or data URI. Transparent pixels indicate editable areas.
647    pub mask: Option<String>,
648    /// Requested output format, if supported by the provider.
649    pub output_format: Option<ImageOutputFormat>,
650}
651
652/// Parameters for generating a video from text prompts or image frames.
653#[derive(Debug, Clone)]
654pub struct VideoGenRequest {
655    /// Video generation model identifier.
656    pub model: String,
657    /// Text description of the desired video.
658    pub description: String,
659    /// Video resolution (e.g., `"720p"`, `"1080p"`).
660    pub resolution: Option<String>,
661    /// Aspect ratio (e.g., `"16:9"`).
662    pub aspect_ratio: Option<String>,
663    /// Duration in seconds.
664    pub duration: Option<u32>,
665    /// Whether to generate accompanying audio.
666    pub generate_audio: Option<bool>,
667    /// Image URLs for first/last frame (image-to-video).
668    pub frame_images: Vec<FrameImage>,
669    /// Reference image URLs for style guidance.
670    pub input_references: Vec<String>,
671    /// Reference videos as URLs or data URIs.
672    pub video_references: Vec<String>,
673}
674
675/// An image used as a frame reference in video generation.
676#[derive(Debug, Clone)]
677pub struct FrameImage {
678    /// URL or data URI of the frame image.
679    pub url: String,
680    /// Frame position — `"first_frame"` or `"last_frame"`.
681    pub frame_type: String,
682}
683
684/// Trait implemented by each LLM backend (OpenRouter, Anthropic, Ollama, etc.).
685///
686/// Providers are registered in a [`ProviderRegistry`] at startup and selected
687/// per-turn based on the active model's [`Model::provider`](crate::model::Model::provider) field.
688///
689/// # Required method
690///
691/// Only [`complete`](Self::complete) is required — it returns a stream of [`StreamEvent`]s.
692/// All other methods have default implementations suitable for simple providers.
693///
694/// # Concurrency
695///
696/// Implementations must be `Send + Sync` since they're shared via `Arc` across async tasks.
697#[async_trait]
698pub trait LlmProvider: Send + Sync {
699    /// Human-readable name of this provider (for logs and error messages).
700    fn name(&self) -> &str;
701
702    /// Which [`crate::model::Provider`] variant this instance implements.
703    fn provider(&self) -> Provider;
704
705    /// Query credit balances when supported. Failure must not fail a completed turn.
706    async fn credit_balance(&self) -> anyhow::Result<Option<crate::CreditBalance>> {
707        Ok(None)
708    }
709
710    /// Query the provider for the context window size of `model`. Return `None`
711    /// if the provider doesn't expose this information.
712    async fn context_window(&self, _model: &Model) -> Option<u32> {
713        None
714    }
715
716    /// Probe or return hard-coded capabilities for `model`. Used to detect features like reasoning or multimodal support.
717    async fn capabilities(&self, _model: &Model) -> ModelCapabilities {
718        ModelCapabilities {
719            tool_calling: true,
720            images: true,
721            documents: true,
722            video: false,
723            audio: false,
724            reasoning: false,
725            ..Default::default()
726        }
727    }
728
729    /// List models available from this provider. Return `None` if unsupported.
730    async fn list_models(&self) -> Option<Vec<ModelInfo>> {
731        None
732    }
733
734    /// Issue a completion request and return a stream of [`StreamEvent`]s.
735    ///
736    /// This is the core method that every provider must implement. The returned
737    /// stream should yield incremental events (text deltas, tool calls, usage)
738    /// and terminate with a [`StreamEvent::Finished`] event.
739    fn complete(&self, request: CompletionRequest) -> CompletionStream;
740
741    /// Whether the selected transport can replay this request's native checkpoints.
742    fn supports_native_context(&self, _request: &CompletionRequest) -> bool {
743        true
744    }
745
746    /// Compact using the provider's native API. `None` selects local summarization.
747    async fn compact(
748        &self,
749        _request: CompletionRequest,
750        _instructions: &str,
751    ) -> anyhow::Result<Option<NativeCompaction>> {
752        Ok(None)
753    }
754
755    /// Synthesise speech from text as a stream of [`StreamEvent::AudioDelta`] chunks,
756    /// ending with [`StreamEvent::Finished`]. Consumers can play chunks as they arrive
757    /// for real-time playback.
758    fn text_to_speech(&self, _request: TtsRequest) -> CompletionStream {
759        let name = self.name().to_string();
760        Box::pin(futures_util::stream::once(async move {
761            Err(anyhow::anyhow!(
762                "Provider '{name}' does not support text-to-speech"
763            ))
764        }))
765    }
766
767    /// Transcribe audio to text. Returns a stream of [`StreamEvent::ContentDelta`] chunks,
768    /// ending with [`StreamEvent::Finished`].
769    fn transcribe(&self, _request: SttRequest) -> CompletionStream {
770        let name = self.name().to_string();
771        Box::pin(futures_util::stream::once(async move {
772            Err(anyhow::anyhow!(
773                "Provider '{name}' does not support speech-to-text"
774            ))
775        }))
776    }
777
778    /// List available voices for TTS. Return `None` if unsupported.
779    async fn list_voices(&self, _model: &str) -> Option<Vec<Voice>> {
780        None
781    }
782
783    /// Generate an image via a dedicated image API (e.g. DALL-E).
784    /// Default implementation returns an error.
785    fn generate_image(&self, _request: ImageGenRequest) -> CompletionStream {
786        Box::pin(futures_util::stream::once(async move {
787            Err(anyhow::anyhow!(
788                "Provider does not support image generation"
789            ))
790        }))
791    }
792
793    /// Edit an existing image using a provider's image editing API.
794    fn edit_image(&self, _request: ImageEditRequest) -> CompletionStream {
795        Box::pin(futures_util::stream::once(async move {
796            Err(anyhow::anyhow!("Provider does not support image editing"))
797        }))
798    }
799
800    /// Generate a video. Returns a stream that yields progress updates as
801    /// [`StreamEvent::ContentDelta`] and the final video as [`StreamEvent::FileAttachment`].
802    /// Default implementation returns an error.
803    fn generate_video(&self, _request: VideoGenRequest) -> CompletionStream {
804        Box::pin(futures_util::stream::once(async move {
805            Err(anyhow::anyhow!(
806                "Provider does not support video generation"
807            ))
808        }))
809    }
810
811    /// Update URL routing table (for providers like Ollama with multiple backends).
812    fn update_routing(&self, _routing: &HashMap<String, String>) {}
813
814    /// Wrap `self` in an `Arc<dyn LlmProvider>` for use with [`Agent`](crate) builders.
815    fn into_arc(self) -> Arc<dyn LlmProvider>
816    where
817        Self: Sized + 'static,
818    {
819        Arc::new(self)
820    }
821}
822
823/// Shared registry mapping each [`crate::model::Provider`] to its concrete instance.
824///
825/// Created once at startup via the CLI; cloned into every agent so they can
826/// switch providers without reconnecting.
827pub type ProviderRegistry = Arc<HashMap<Provider, Arc<dyn LlmProvider>>>;
828
829#[cfg(test)]
830mod tests {
831    use super::*;
832
833    #[test]
834    fn requests_keep_the_sessions_config_and_identity() {
835        use uuid::Uuid;
836
837        let uuid = Uuid::new_v4();
838        let mut config = AgentLlmConfig::new("openrouter:example/model".parse().unwrap())
839            .with_session_id(uuid)
840            .with_user("example-user")
841            .with_reasoning(ReasoningLevel::Low)
842            .with_provider_preferences(ProviderPreferences {
843                zdr: Some(true),
844                ..Default::default()
845            });
846        config.max_tokens = Some(100);
847        config.sampling.temperature = Some(Decimal::ONE);
848        let first = CompletionRequest::from_config(&config, vec![Message::user("hello")]);
849        let second = CompletionRequest::from_config(&config, vec![Message::user("continue")]);
850        assert_eq!(first.session_id, second.session_id);
851        assert_eq!(
852            first.session_id.as_ref().map(SessionId::as_str),
853            Some(uuid.to_string().as_str())
854        );
855        assert_eq!(first.user, config.user);
856        assert_eq!(first.reasoning, config.reasoning);
857        assert_eq!(first.max_tokens, config.max_tokens);
858        assert_eq!(first.sampling.temperature, config.sampling.temperature);
859        assert_eq!(first.provider_preferences.unwrap().zdr, Some(true));
860
861        let one_off = CompletionRequest::new(config.model.clone(), vec![Message::user("hello")]);
862        assert!(one_off.session_id.is_none());
863        let one_off = one_off.with_session_id(uuid);
864        assert_eq!(one_off.session_id, second.session_id);
865    }
866    use crate::event::TurnUsage;
867
868    #[test]
869    fn token_usage_from_turn_usage() {
870        let turn = TurnUsage {
871            prompt_tokens: 100,
872            completion_tokens: 50,
873            ..TurnUsage::default()
874        };
875        let token: TokenUsage = turn.into();
876        assert_eq!(token.prompt_tokens, 100);
877        assert_eq!(token.completion_tokens, 50);
878        assert_eq!(token.total_tokens, 150);
879    }
880
881    #[test]
882    fn token_usage_default_is_zero() {
883        let usage = TokenUsage::default();
884        assert_eq!(usage.total_tokens, 0);
885    }
886
887    #[test]
888    fn model_capabilities_all() {
889        let caps = ModelCapabilities::all();
890        assert!(caps.tool_calling);
891        assert!(caps.images);
892        assert!(caps.reasoning);
893    }
894
895    #[test]
896    fn model_capabilities_default_is_conservative() {
897        let caps = ModelCapabilities::default();
898        assert!(!caps.tool_calling);
899        assert!(!caps.reasoning);
900    }
901
902    #[test]
903    fn finish_reason_display() {
904        assert_eq!(FinishReason::Stop.to_string(), "stop");
905        assert_eq!(FinishReason::ToolCalls.to_string(), "tool_calls");
906        assert_eq!(FinishReason::Length.to_string(), "length");
907    }
908
909    #[test]
910    fn finish_reason_from_str() {
911        let stop: FinishReason = "stop".parse().unwrap();
912        assert_eq!(stop, FinishReason::Stop);
913        let tc: FinishReason = "tool_calls".parse().unwrap();
914        assert_eq!(tc, FinishReason::ToolCalls);
915    }
916
917    #[test]
918    fn tool_definition_serde() {
919        let def = ToolDefinition {
920            name: "file_read".into(),
921            description: "Read a file".into(),
922            parameters: serde_json::json!({"type": "object", "properties": {}}),
923        };
924        let json = serde_json::to_string(&def).unwrap();
925        let parsed: ToolDefinition = serde_json::from_str(&json).unwrap();
926        assert_eq!(parsed.name, "file_read");
927    }
928
929    #[test]
930    fn audio_format_default() {
931        assert_eq!(AudioFormat::default().to_string(), "mp3");
932    }
933}