flashmind-types 0.2.0

Shared types for the Flashmind AI agent framework
Documentation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
//! LLM wire types — request / response structures and the [`LlmProvider`] trait.
//!
//! This module defines everything needed to communicate with an LLM backend:
//! the provider trait, completion requests and responses, streaming events,
//! model capabilities, and multimodal support (audio, images, video).
//!
//! # Key types
//!
//! | Type | Role |
//! |------|------|
//! | [`LlmProvider`] | Trait implemented by each LLM backend (OpenRouter, Anthropic, Ollama…) |
//! | [`CompletionRequest`] | Single completion request sent to a provider |
//! | [`CompletionResponse`] | Non-streaming response wrapper |
//! | [`StreamEvent`] | Incremental events from the provider's stream |
//! | [`CompletionStream`] | Alias for `Pin<Box<dyn Stream<Item = Result<StreamEvent>>>>` |
//! | [`FinishReason`] | Why generation stopped (stop, length, tool_calls, content_filter) |
//! | [`ToolDefinition`] | OpenAI-compatible JSON schema for tool calling |
//! | [`ModelCapabilities`] | Feature flags (tool_calling, images, reasoning, audio…) |
//! | [`TokenUsage`] | Token consumption reported by the provider |
//! | [`Modality`] | Output modality: text, audio, or image |
//! | [`TtsRequest`], [`SttRequest`], [`VideoGenRequest`] | Multimodal operation requests |

use std::collections::{BTreeMap, HashMap};
use std::pin::Pin;
use std::sync::Arc;

use async_trait::async_trait;
use chrono::{DateTime, Utc};
use futures_core::Stream;
use serde::{Deserialize, Serialize};

use rust_decimal::Decimal;

use crate::event::TurnUsage;
use crate::message::Message;
use crate::model::{
    AgentLlmConfig, Model, Provider, ProviderPreferences, ReasoningLevel, SamplingParams,
};
use crate::session::SessionId;

/// OpenAI-compatible JSON schema for a tool exposed to the LLM.
///
/// Mirrors the structure sent in the `tools` array of chat completion requests.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolDefinition {
    /// Tool name — must match the [`Tool::name`](crate::tool::Tool::name) implementation.
    pub name: String,
    /// Description shown to the model to help it decide when to call this tool.
    pub description: String,
    /// JSON Schema object describing the parameters this tool accepts.
    pub parameters: serde_json::Value,
}

/// A single completion request sent to an LLM provider.
///
/// This is the primary input type for [`LlmProvider::complete`]. The agent runtime
/// assembles this from the conversation history, active tools, and LLM config.
#[derive(Debug)]
pub struct CompletionRequest {
    /// Model to use (includes provider prefix).
    pub model: Model,
    /// Message history in wire format.
    pub messages: Vec<Message>,
    /// Tools available to the model in this turn.
    pub tools: Vec<ToolDefinition>,
    /// Maximum output tokens. If `None`, the provider decides.
    pub max_tokens: Option<u32>,
    /// Whether reasoning/thinking mode is enabled.
    pub reasoning: ReasoningLevel,
    /// Sampling parameters (temperature, top_p, top_k, min_p, penalties).
    pub sampling: SamplingParams,
    /// Output modalities (e.g. text+audio, image). Empty means text-only.
    pub modalities: Vec<Modality>,
    /// Audio output configuration. Only used when modalities includes [`Modality::Audio`].
    pub audio_config: Option<AudioOutputConfig>,
    /// Image generation configuration. Only used when modalities includes [`Modality::Image`].
    pub image_config: Option<ImageGenConfig>,
    /// Opaque user identifier for per-user tracking in provider dashboards.
    pub user: Option<String>,
    /// Provider routing preferences (used by OpenRouter).
    pub provider_preferences: Option<ProviderPreferences>,
    /// Stable conversation identifier for provider cache routing.
    pub session_id: Option<SessionId>,
}

impl CompletionRequest {
    /// Build a request with default settings and the supplied messages.
    pub fn new(model: Model, messages: Vec<Message>) -> Self {
        Self::from_config(&AgentLlmConfig::new(model), messages)
    }

    /// Keep the agent's session ID, routing, user and generation settings.
    ///
    /// ```
    /// use flashmind_types::{AgentLlmConfig, CompletionRequest, Message, SessionId};
    ///
    /// let session = SessionId::new();
    /// let llm = AgentLlmConfig::new("openrouter:example/model".parse()?)
    ///     .with_session_id(session);
    /// let request = CompletionRequest::from_config(&llm, vec![Message::user("Hello")]);
    /// assert_eq!(request.session_id, llm.session_id);
    /// # Ok::<(), anyhow::Error>(())
    /// ```
    pub fn from_config(config: &AgentLlmConfig, messages: Vec<Message>) -> Self {
        Self {
            model: config.model.clone(),
            messages,
            tools: Vec::new(),
            max_tokens: config.max_tokens,
            reasoning: config.reasoning,
            sampling: config.sampling.clone(),
            modalities: Vec::new(),
            audio_config: None,
            image_config: None,
            user: config.user.clone(),
            provider_preferences: config.provider_preferences.clone(),
            session_id: config.session_id.clone(),
        }
    }

    /// Use an existing session identity for this request.
    pub fn with_session_id(mut self, session_id: impl Into<SessionId>) -> Self {
        self.session_id = Some(session_id.into());
        self
    }
}

/// Why the LLM stopped generating (used in stream and non-stream responses).
#[derive(Debug, Clone, Copy, PartialEq, Eq, strum::Display, strum::EnumString)]
#[strum(serialize_all = "snake_case")]
pub enum FinishReason {
    /// Natural stop token or end of response.
    Stop,
    /// Model requested one or more tool calls.
    ToolCalls,
    /// Hit `max_tokens` limit without natural stopping point.
    Length,
    /// Output was filtered by content policy.
    ContentFilter,
}

/// High-level category for filtering models by use-case.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, strum::Display, strum::EnumString)]
#[strum(serialize_all = "snake_case")]
pub enum ModelCategory {
    /// General text chat / instruction-following.
    Chat,
    /// Advanced reasoning / chain-of-thought.
    Reasoning,
    /// Can accept images as input.
    Vision,
    /// Generates images from text.
    ImageGeneration,
    /// Edits existing images.
    ImageEditing,
    /// Generates audio / text-to-speech.
    Tts,
    /// Accepts audio input / speech-to-text.
    Stt,
    /// Generates video.
    VideoGeneration,
    /// Accepts video as input.
    VideoInput,
}

/// Pricing information for a model, including input/output token costs and caching discounts.
///
/// All costs are per-token in USD. Use [`TokenUsage::cost`] to compute the total
/// cost of a request from token counts and pricing.
#[derive(Debug, Clone, Default)]
pub struct ModelPricing {
    /// Cost per input token (USD).
    pub prompt: Option<Decimal>,
    /// Cost per output token (USD).
    pub completion: Option<Decimal>,
    /// Cost per image generated (USD).
    pub image: Option<Decimal>,
    /// Cost per cached input token read (USD).
    pub cache_read: Option<Decimal>,
    /// Cost per cache creation (write) input token (USD).
    pub cache_write: Option<Decimal>,
    /// Published reference rates, which may differ from the provider's charges.
    pub reference: Option<ReferencePricing>,
}

/// Reference token rates from a separately identified source.
#[derive(Debug, Clone)]
pub struct ReferencePricing {
    /// Source of the reference rates.
    pub source: ReferencePricingSource,
    /// Reference input price per token in USD.
    pub prompt: Option<Decimal>,
    /// Reference output price per token in USD.
    pub completion: Option<Decimal>,
}

/// Source of a model's reference rates.
#[derive(Debug, Clone, Copy, PartialEq, Eq, strum::Display)]
pub enum ReferencePricingSource {
    /// Prices published by OpenRouter.
    OpenRouter,
    /// Prices published by models.dev.
    #[strum(to_string = "models.dev")]
    ModelsDev,
}

/// Metadata about an available LLM model, including capabilities, pricing, and context window.
#[derive(Debug, Clone)]
pub struct ModelInfo {
    /// Model ID used in API requests.
    pub id: String,
    /// Human-readable display name.
    pub name: Option<String>,
    /// Model creation time, when supplied by the provider.
    pub created_at: Option<DateTime<Utc>>,
    /// Context window size in tokens.
    pub context_length: Option<u32>,
    /// Maximum output tokens allowed.
    pub max_completion_tokens: Option<u32>,
    /// Capability flags.
    pub capabilities: ModelCapabilities,
    /// High-level categories (chat, vision, reasoning, etc.).
    pub categories: Vec<ModelCategory>,
    /// Per-token pricing.
    pub pricing: ModelPricing,
    /// Image-specific capability metadata from the provider.
    pub image: ImageModelCapabilities,
}

/// Image capability metadata, when provided in a model catalog.
#[derive(Debug, Clone, Default)]
pub struct ImageModelCapabilities {
    /// Supports image editing.
    pub editing: bool,
    /// Whether style references are supported, if reported.
    pub style_references: Option<bool>,
    /// Maximum number of images for an edit, if reported.
    pub max_edit_images: Option<u32>,
    /// Provider-reported image options and their allowed values or ranges.
    pub parameters: BTreeMap<String, serde_json::Value>,
}

/// Feature flags describing what a model supports.
///
/// Defaults to all `false` via `#[derive(Default)]` — nothing is assumed.
/// Providers override these after inspecting the model's actual capabilities.
/// Use [`all()`](Self::all()) for a curated "full-featured text model" preset.
#[derive(Debug, Clone, Copy, Default)]
pub struct ModelCapabilities {
    /// Can invoke tools/function calling.
    pub tool_calling: bool,
    /// Can accept images as input (vision).
    pub images: bool,
    /// Can accept documents (PDFs, etc.) as input.
    pub documents: bool,
    /// Can accept video as input.
    pub video: bool,
    /// Can accept audio as input (speech-to-text).
    pub audio: bool,
    /// Supports extended chain-of-thought / reasoning mode.
    pub reasoning: bool,
    // Output capabilities
    /// Can generate audio output (text-to-speech inline).
    pub audio_output: bool,
    /// Can generate images as output.
    pub image_generation: bool,
    /// Can generate video as output.
    pub video_generation: bool,
}

impl ModelCapabilities {
    /// Derive high-level categories from capability flags.
    pub fn categories(&self) -> Vec<ModelCategory> {
        let mut cats = vec![ModelCategory::Chat];
        if self.reasoning {
            cats.push(ModelCategory::Reasoning);
        }
        if self.images {
            cats.push(ModelCategory::Vision);
        }
        if self.image_generation {
            cats.push(ModelCategory::ImageGeneration);
        }
        if self.audio_output {
            cats.push(ModelCategory::Tts);
        }
        if self.audio {
            cats.push(ModelCategory::Stt);
        }
        if self.video_generation {
            cats.push(ModelCategory::VideoGeneration);
        }
        if self.video {
            cats.push(ModelCategory::VideoInput);
        }
        cats
    }

    /// Curated "full-featured text model" capability preset.
    ///
    /// Enables common features (tool calling, vision, reasoning, audio input)
    /// but keeps output-generation capabilities (`audio_output`, `image_generation`,
    /// `video_generation`) disabled, as those require dedicated models.
    pub fn all() -> Self {
        Self {
            tool_calling: true,
            images: true,
            documents: true,
            video: false,
            audio: true,
            reasoning: true,
            audio_output: false,
            image_generation: false,
            video_generation: false,
        }
    }
}

/// Token usage counts for a single completion request or turn.
#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
pub struct TokenUsage {
    /// Tokens consumed by the prompt (input messages).
    pub prompt_tokens: u32,
    /// Tokens generated by the model (output/response).
    pub completion_tokens: u32,
    /// Total tokens (prompt + completion).
    pub total_tokens: u32,
    /// Prompt tokens served from the provider's cache.
    pub cache_read_tokens: u32,
    /// Prompt tokens written to the provider's cache.
    pub cache_creation_tokens: u32,
}

impl TokenUsage {
    /// Compute the cost of this usage given per-token pricing.
    ///
    /// Non-cached input tokens are charged at the `prompt` rate, cached reads at
    /// `cache_read` (falling back to `prompt`), and cache writes at `cache_write`
    /// (falling back to `prompt`). Returns `None` if pricing has no prompt or
    /// completion prices.
    pub fn cost(&self, pricing: &ModelPricing) -> Option<Decimal> {
        let prompt_price = pricing.prompt?;
        let completion_price = pricing.completion?;
        let cache_read_price = pricing.cache_read.unwrap_or(prompt_price);
        let cache_write_price = pricing.cache_write.unwrap_or(prompt_price);

        let non_cached = self
            .prompt_tokens
            .saturating_sub(self.cache_read_tokens + self.cache_creation_tokens);

        Some(
            Decimal::from(non_cached) * prompt_price
                + Decimal::from(self.completion_tokens) * completion_price
                + Decimal::from(self.cache_read_tokens) * cache_read_price
                + Decimal::from(self.cache_creation_tokens) * cache_write_price,
        )
    }
}

impl From<TurnUsage> for TokenUsage {
    fn from(u: TurnUsage) -> Self {
        Self {
            prompt_tokens: u.prompt_tokens,
            completion_tokens: u.completion_tokens,
            total_tokens: u.prompt_tokens + u.completion_tokens,
            cache_read_tokens: u.cache_read_tokens,
            cache_creation_tokens: u.cache_creation_tokens,
        }
    }
}

/// Non-streaming completion response wrapper.
#[derive(Debug, Clone)]
pub struct CompletionResponse {
    /// The assistant's response message.
    pub message: Message,
    /// Token usage for this completion.
    pub usage: TokenUsage,
    /// Why generation stopped.
    pub finish_reason: FinishReason,
}

/// Events emitted incrementally by the provider's stream.
///
/// Each variant corresponds to a server-side SSE/data-line event type.
/// Consumers assemble these into higher-level [`AgentEvent`](crate::event::AgentEvent)s.
#[derive(Debug, Clone)]
pub enum StreamEvent {
    /// Incremental reasoning/chain-of-thought token.
    ReasoningDelta(String),
    /// Incremental assistant text token.
    ContentDelta(String),
    /// The LLM started requesting a tool call at `index`.
    ToolCallStart {
        index: usize,
        id: String,
        name: String,
    },
    /// Incremental JSON argument string for the tool at `index`.
    ToolCallDelta { index: usize, arguments: String },
    /// Token usage snapshot mid-stream (some providers emit this).
    Usage(TokenUsage),
    /// Incremental audio output chunk (raw PCM/MP3 bytes).
    AudioDelta { data: Vec<u8>, format: String },
    /// Out-of-band file delivered by the provider (e.g. generated image).
    FileAttachment {
        filename: String,
        media_type: String,
        data: Vec<u8>,
    },
    /// Stream has ended; payload is the final finish reason.
    Finished(FinishReason),
}

/// Type alias for a stream of completion events from an LLM provider.
///
/// The stream yields `StreamEvent`s (content deltas, tool calls, finish reasons) and is consumed by the agent loop.
pub type CompletionStream =
    Pin<Box<dyn Stream<Item = anyhow::Result<StreamEvent>> + Send + 'static>>;

/// A voice available for text-to-speech generation.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Voice {
    /// Unique voice identifier (e.g., `"alloy"`).
    pub id: String,
    /// Human-readable display name.
    pub name: String,
}

/// Parameters for a text-to-speech request.
pub struct TtsRequest {
    /// TTS model identifier (e.g., `"tts-1"`).
    pub model: String,
    /// Text to synthesize into speech.
    pub input: String,
    /// Voice preset ID (e.g., `"alloy"`, `"echo"`).
    pub voice: String,
    /// Desired output audio format.
    pub response_format: AudioFormat,
}

/// Parameters for a speech-to-text (transcription) request.
pub struct SttRequest {
    /// STT model identifier (e.g., `"whisper-1"`).
    pub model: String,
    /// Raw audio bytes.
    pub audio: Vec<u8>,
    /// MIME type of the audio (e.g. "audio/mp3", "audio/wav").
    pub media_type: String,
    /// Optional language hint (ISO 639-1, e.g. "en").
    pub language: Option<String>,
}

/// Output audio format for TTS responses.
///
/// Serialises to lowercase strings matching common audio file extensions.
#[derive(
    Debug, Clone, Copy, Default, strum::EnumString, strum::Display, Serialize, Deserialize,
)]
#[strum(serialize_all = "lowercase")]
#[serde(rename_all = "lowercase")]
pub enum AudioFormat {
    /// MP3 — most widely supported, good compression.
    #[default]
    Mp3,
    /// WAV — uncompressed PCM, highest quality.
    Wav,
    /// Opus — efficient codec optimized for speech.
    Opus,
    /// AAC — Advanced Audio Coding, common in mobile ecosystems.
    Aac,
    /// FLAC — lossless compression.
    Flac,
}

/// Output modality requested in a completion.
///
/// Multiple modalities can be requested simultaneously (e.g., text + audio).
/// An empty list means text-only output.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, strum::Display)]
#[serde(rename_all = "lowercase")]
#[strum(serialize_all = "lowercase")]
pub enum Modality {
    /// Plain text output (default when no modalities specified).
    Text,
    /// Inline audio output (real-time voice).
    Audio,
    /// Image output (generated during completion).
    Image,
}

/// Configuration for audio output in multimodal completions.
#[derive(Debug, Clone)]
pub struct AudioOutputConfig {
    /// Voice preset ID (e.g., `"alloy"`).
    pub voice: String,
    /// Output audio format.
    pub format: AudioFormat,
}

/// Configuration for image generation requests.
#[derive(Debug, Clone, Default)]
pub struct ImageGenConfig {
    /// Aspect ratio for the generated image (e.g., `"16:9"`).
    pub aspect_ratio: Option<String>,
    /// Absolute dimensions (e.g., `"1024x1024"`).
    pub size: Option<String>,
}

impl ImageGenConfig {
    /// Read an image file and return it as a `data:<mime>;base64,...` URI.
    ///
    /// # Arguments
    ///
    /// * `path` - Path to the image file on disk.
    ///
    /// # Returns
    ///
    /// A data URI string with auto-detected MIME type, or an I/O error if the file cannot be read.
    pub fn data_uri_from_path(path: &std::path::Path) -> std::io::Result<String> {
        let bytes = std::fs::read(path)?;
        let mime = mime_from_extension(path);
        Ok(Self::data_uri_from_bytes(&bytes, mime))
    }

    /// Encode raw bytes as a `data:<mime>;base64,...` URI.
    ///
    /// # Arguments
    ///
    /// * `bytes` - Raw binary data (e.g., image file contents).
    /// * `mime` - MIME type string (e.g., `"image/png"`).
    ///
    /// # Returns
    ///
    /// A data URI string with the provided MIME type and base64-encoded payload.
    pub fn data_uri_from_bytes(bytes: &[u8], mime: &str) -> String {
        use base64::Engine;
        let b64 = base64::engine::general_purpose::STANDARD.encode(bytes);
        format!("data:{mime};base64,{b64}")
    }
}

/// Infer MIME type from a file path's extension.
///
/// # Arguments
///
/// * `path` - File path whose extension is inspected.
///
/// # Returns
///
/// A static MIME type string corresponding to the extension, or `"application/octet-stream"`
/// if the extension is unrecognized or missing.
pub fn mime_from_extension(path: &std::path::Path) -> &'static str {
    match path
        .extension()
        .and_then(|e| e.to_str())
        .map(|e| e.to_ascii_lowercase())
        .as_deref()
    {
        Some("png") => "image/png",
        Some("jpg" | "jpeg") => "image/jpeg",
        Some("gif") => "image/gif",
        Some("webp") => "image/webp",
        Some("bmp") => "image/bmp",
        Some("svg") => "image/svg+xml",
        _ => "application/octet-stream",
    }
}

/// Image format requested from a provider.
#[derive(
    Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, strum::Display, strum::EnumString,
)]
#[serde(rename_all = "lowercase")]
#[strum(serialize_all = "lowercase")]
pub enum ImageOutputFormat {
    /// PNG image.
    Png,
    /// JPEG image.
    Jpeg,
    /// WebP image.
    Webp,
}

impl ImageOutputFormat {
    /// MIME type corresponding to the format.
    pub fn media_type(self) -> &'static str {
        match self {
            Self::Png => "image/png",
            Self::Jpeg => "image/jpeg",
            Self::Webp => "image/webp",
        }
    }
}

/// Parameters for generating an image via a generative model.
#[derive(Debug, Clone)]
pub struct ImageGenRequest {
    /// Image generation model identifier.
    pub model: String,
    /// Text description of the desired image.
    pub prompt: String,
    /// Image dimensions (e.g., `"1024x1024"`).
    pub size: Option<String>,
    /// Aspect ratio (e.g., `"16:9"`).
    pub aspect_ratio: Option<String>,
    /// Quality preset (e.g., `"standard"`, `"hd"`).
    pub quality: Option<String>,
    /// Style preset (e.g., `"natural"`, `"vivid"`).
    pub style: Option<String>,
    /// Number of images to generate.
    pub n: Option<u32>,
    /// Reference images as URLs or data URIs.
    pub input_references: Vec<String>,
    /// Requested output format, if supported by the provider.
    pub output_format: Option<ImageOutputFormat>,
}

/// Parameters for editing one or more images.
#[derive(Debug, Clone)]
pub struct ImageEditRequest {
    /// Provider model identifier.
    pub model: String,
    /// Instructions describing the edit.
    pub prompt: String,
    /// Input images as URLs or data URIs, beginning with the base image.
    pub images: Vec<String>,
    /// Output resolution tier.
    pub resolution: Option<String>,
    /// Output aspect ratio.
    pub aspect_ratio: Option<String>,
    /// Provider-specific quality tier.
    pub quality: Option<String>,
    /// Mask as a URL or data URI. Transparent pixels indicate editable areas.
    pub mask: Option<String>,
    /// Requested output format, if supported by the provider.
    pub output_format: Option<ImageOutputFormat>,
}

/// Parameters for generating a video from text prompts or image frames.
#[derive(Debug, Clone)]
pub struct VideoGenRequest {
    /// Video generation model identifier.
    pub model: String,
    /// Text description of the desired video.
    pub description: String,
    /// Video resolution (e.g., `"720p"`, `"1080p"`).
    pub resolution: Option<String>,
    /// Aspect ratio (e.g., `"16:9"`).
    pub aspect_ratio: Option<String>,
    /// Duration in seconds.
    pub duration: Option<u32>,
    /// Whether to generate accompanying audio.
    pub generate_audio: Option<bool>,
    /// Image URLs for first/last frame (image-to-video).
    pub frame_images: Vec<FrameImage>,
    /// Reference image URLs for style guidance.
    pub input_references: Vec<String>,
    /// Reference videos as URLs or data URIs.
    pub video_references: Vec<String>,
}

/// An image used as a frame reference in video generation.
#[derive(Debug, Clone)]
pub struct FrameImage {
    /// URL or data URI of the frame image.
    pub url: String,
    /// Frame position — `"first_frame"` or `"last_frame"`.
    pub frame_type: String,
}

/// Trait implemented by each LLM backend (OpenRouter, Anthropic, Ollama, etc.).
///
/// Providers are registered in a [`ProviderRegistry`] at startup and selected
/// per-turn based on the active model's [`Model::provider`](crate::model::Model::provider) field.
///
/// # Required method
///
/// Only [`complete`](Self::complete) is required — it returns a stream of [`StreamEvent`]s.
/// All other methods have default implementations suitable for simple providers.
///
/// # Concurrency
///
/// Implementations must be `Send + Sync` since they're shared via `Arc` across async tasks.
#[async_trait]
pub trait LlmProvider: Send + Sync {
    /// Human-readable name of this provider (for logs and error messages).
    fn name(&self) -> &str;

    /// Which [`crate::model::Provider`] variant this instance implements.
    fn provider(&self) -> Provider;

    /// Query the provider for the context window size of `model`. Return `None`
    /// if the provider doesn't expose this information.
    async fn context_window(&self, _model: &Model) -> Option<u32> {
        None
    }

    /// Probe or return hard-coded capabilities for `model`. Used to detect features like reasoning or multimodal support.
    async fn capabilities(&self, _model: &Model) -> ModelCapabilities {
        ModelCapabilities {
            tool_calling: true,
            images: true,
            documents: true,
            video: false,
            audio: false,
            reasoning: false,
            ..Default::default()
        }
    }

    /// List models available from this provider. Return `None` if unsupported.
    async fn list_models(&self) -> Option<Vec<ModelInfo>> {
        None
    }

    /// Issue a completion request and return a stream of [`StreamEvent`]s.
    ///
    /// This is the core method that every provider must implement. The returned
    /// stream should yield incremental events (text deltas, tool calls, usage)
    /// and terminate with a [`StreamEvent::Finished`] event.
    fn complete(&self, request: CompletionRequest) -> CompletionStream;

    /// Synthesise speech from text as a stream of [`StreamEvent::AudioDelta`] chunks,
    /// ending with [`StreamEvent::Finished`]. Consumers can play chunks as they arrive
    /// for real-time playback.
    fn text_to_speech(&self, _request: TtsRequest) -> CompletionStream {
        let name = self.name().to_string();
        Box::pin(futures_util::stream::once(async move {
            Err(anyhow::anyhow!(
                "Provider '{name}' does not support text-to-speech"
            ))
        }))
    }

    /// Transcribe audio to text. Returns a stream of [`StreamEvent::ContentDelta`] chunks,
    /// ending with [`StreamEvent::Finished`].
    fn transcribe(&self, _request: SttRequest) -> CompletionStream {
        let name = self.name().to_string();
        Box::pin(futures_util::stream::once(async move {
            Err(anyhow::anyhow!(
                "Provider '{name}' does not support speech-to-text"
            ))
        }))
    }

    /// List available voices for TTS. Return `None` if unsupported.
    async fn list_voices(&self, _model: &str) -> Option<Vec<Voice>> {
        None
    }

    /// Generate an image via a dedicated image API (e.g. DALL-E).
    /// Default implementation returns an error.
    fn generate_image(&self, _request: ImageGenRequest) -> CompletionStream {
        Box::pin(futures_util::stream::once(async move {
            Err(anyhow::anyhow!(
                "Provider does not support image generation"
            ))
        }))
    }

    /// Edit an existing image using a provider's image editing API.
    fn edit_image(&self, _request: ImageEditRequest) -> CompletionStream {
        Box::pin(futures_util::stream::once(async move {
            Err(anyhow::anyhow!("Provider does not support image editing"))
        }))
    }

    /// Generate a video. Returns a stream that yields progress updates as
    /// [`StreamEvent::ContentDelta`] and the final video as [`StreamEvent::FileAttachment`].
    /// Default implementation returns an error.
    fn generate_video(&self, _request: VideoGenRequest) -> CompletionStream {
        Box::pin(futures_util::stream::once(async move {
            Err(anyhow::anyhow!(
                "Provider does not support video generation"
            ))
        }))
    }

    /// Update URL routing table (for providers like Ollama with multiple backends).
    fn update_routing(&self, _routing: &HashMap<String, String>) {}

    /// Wrap `self` in an `Arc<dyn LlmProvider>` for use with [`Agent`](crate) builders.
    fn into_arc(self) -> Arc<dyn LlmProvider>
    where
        Self: Sized + 'static,
    {
        Arc::new(self)
    }
}

/// Shared registry mapping each [`crate::model::Provider`] to its concrete instance.
///
/// Created once at startup via the CLI; cloned into every agent so they can
/// switch providers without reconnecting.
pub type ProviderRegistry = Arc<HashMap<Provider, Arc<dyn LlmProvider>>>;

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn requests_keep_the_sessions_config_and_identity() {
        use uuid::Uuid;

        let uuid = Uuid::new_v4();
        let mut config = AgentLlmConfig::new("openrouter:example/model".parse().unwrap())
            .with_session_id(uuid)
            .with_user("example-user")
            .with_reasoning(ReasoningLevel::Low)
            .with_provider_preferences(ProviderPreferences {
                zdr: Some(true),
                ..Default::default()
            });
        config.max_tokens = Some(100);
        config.sampling.temperature = Some(Decimal::ONE);
        let first = CompletionRequest::from_config(&config, vec![Message::user("hello")]);
        let second = CompletionRequest::from_config(&config, vec![Message::user("continue")]);
        assert_eq!(first.session_id, second.session_id);
        assert_eq!(
            first.session_id.as_ref().map(SessionId::as_str),
            Some(uuid.to_string().as_str())
        );
        assert_eq!(first.user, config.user);
        assert_eq!(first.reasoning, config.reasoning);
        assert_eq!(first.max_tokens, config.max_tokens);
        assert_eq!(first.sampling.temperature, config.sampling.temperature);
        assert_eq!(first.provider_preferences.unwrap().zdr, Some(true));

        let one_off = CompletionRequest::new(config.model.clone(), vec![Message::user("hello")]);
        assert!(one_off.session_id.is_none());
        let one_off = one_off.with_session_id(uuid);
        assert_eq!(one_off.session_id, second.session_id);
    }
    use crate::event::TurnUsage;

    #[test]
    fn token_usage_from_turn_usage() {
        let turn = TurnUsage {
            prompt_tokens: 100,
            completion_tokens: 50,
            ..TurnUsage::default()
        };
        let token: TokenUsage = turn.into();
        assert_eq!(token.prompt_tokens, 100);
        assert_eq!(token.completion_tokens, 50);
        assert_eq!(token.total_tokens, 150);
    }

    #[test]
    fn token_usage_default_is_zero() {
        let usage = TokenUsage::default();
        assert_eq!(usage.total_tokens, 0);
    }

    #[test]
    fn model_capabilities_all() {
        let caps = ModelCapabilities::all();
        assert!(caps.tool_calling);
        assert!(caps.images);
        assert!(caps.reasoning);
    }

    #[test]
    fn model_capabilities_default_is_conservative() {
        let caps = ModelCapabilities::default();
        assert!(!caps.tool_calling);
        assert!(!caps.reasoning);
    }

    #[test]
    fn finish_reason_display() {
        assert_eq!(FinishReason::Stop.to_string(), "stop");
        assert_eq!(FinishReason::ToolCalls.to_string(), "tool_calls");
        assert_eq!(FinishReason::Length.to_string(), "length");
    }

    #[test]
    fn finish_reason_from_str() {
        let stop: FinishReason = "stop".parse().unwrap();
        assert_eq!(stop, FinishReason::Stop);
        let tc: FinishReason = "tool_calls".parse().unwrap();
        assert_eq!(tc, FinishReason::ToolCalls);
    }

    #[test]
    fn tool_definition_serde() {
        let def = ToolDefinition {
            name: "file_read".into(),
            description: "Read a file".into(),
            parameters: serde_json::json!({"type": "object", "properties": {}}),
        };
        let json = serde_json::to_string(&def).unwrap();
        let parsed: ToolDefinition = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed.name, "file_read");
    }

    #[test]
    fn audio_format_default() {
        assert_eq!(AudioFormat::default().to_string(), "mp3");
    }
}