mentra_provider/
model.rs

1use serde::{Deserialize, Serialize};
2use serde_json::Value;
3use std::fmt::Display;
4use time::OffsetDateTime;
5
6/// Metadata describing a model available from a provider.
7#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
8pub struct ModelInfo {
9    pub id: String,
10    pub provider: crate::ProviderId,
11    pub display_name: Option<String>,
12    pub description: Option<String>,
13    pub created_at: Option<OffsetDateTime>,
14    /// How many tokens the model accepts in one request, when the provider
15    /// says so.
16    ///
17    /// Most providers do not: neither Anthropic's nor OpenAI's model listing
18    /// reports a limit, so this is `None` for them and a host that knows the
19    /// number can set it. Gemini reports it as `inputTokenLimit`. It is the
20    /// only thing that lets anything downstream — a compaction threshold, a
21    /// context-usage report — be expressed relative to the model rather than
22    /// as a constant that is wrong for most of them.
23    #[serde(default, skip_serializing_if = "Option::is_none")]
24    pub context_window: Option<usize>,
25}
26
27impl ModelInfo {
28    pub fn new(id: impl Into<String>, provider: impl Into<crate::ProviderId>) -> Self {
29        Self {
30            id: id.into(),
31            provider: provider.into(),
32            display_name: None,
33            description: None,
34            created_at: None,
35            context_window: None,
36        }
37    }
38
39    /// Returns this model with its context window set.
40    pub fn with_context_window(mut self, context_window: usize) -> Self {
41        self.context_window = Some(context_window);
42        self
43    }
44}
45
46/// Selection strategy used when resolving a model from a provider.
47#[derive(Debug, Clone, PartialEq, Eq)]
48pub enum ModelSelector {
49    Id(String),
50    NewestAvailable,
51}
52
53/// Provider-neutral token usage metadata for a completed or in-progress response.
54#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
55pub struct TokenUsage {
56    pub input_tokens: Option<u64>,
57    pub output_tokens: Option<u64>,
58    pub total_tokens: Option<u64>,
59    pub cache_read_input_tokens: Option<u64>,
60    pub cache_creation_input_tokens: Option<u64>,
61    pub reasoning_tokens: Option<u64>,
62    pub thoughts_tokens: Option<u64>,
63    pub tool_input_tokens: Option<u64>,
64}
65
66impl TokenUsage {
67    pub fn is_empty(&self) -> bool {
68        self.input_tokens.is_none()
69            && self.output_tokens.is_none()
70            && self.total_tokens.is_none()
71            && self.cache_read_input_tokens.is_none()
72            && self.cache_creation_input_tokens.is_none()
73            && self.reasoning_tokens.is_none()
74            && self.thoughts_tokens.is_none()
75            && self.tool_input_tokens.is_none()
76    }
77}
78
79/// Provider-neutral chat role labels.
80#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
81#[serde(rename_all = "snake_case")]
82pub enum Role {
83    User,
84    Assistant,
85    Unknown(String),
86}
87
88impl Display for Role {
89    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
90        let value = match self {
91            Self::User => "user",
92            Self::Assistant => "assistant",
93            Self::Unknown(role) => role.as_str(),
94        };
95        f.write_str(value)
96    }
97}
98
99/// Image payload supported by model providers.
100#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
101pub enum ImageSource {
102    Bytes { media_type: String, data: Vec<u8> },
103    Url { url: String },
104}
105
106impl ImageSource {
107    pub fn bytes(media_type: impl Into<String>, data: impl Into<Vec<u8>>) -> Self {
108        Self::Bytes {
109            media_type: media_type.into(),
110            data: data.into(),
111        }
112    }
113
114    pub fn url(url: impl Into<String>) -> Self {
115        Self::Url { url: url.into() }
116    }
117}
118
119/// Tool result payloads supported by provider streams and history replay.
120#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
121#[serde(untagged)]
122pub enum ToolResultContent {
123    Text(String),
124    Structured(Value),
125}
126
127impl ToolResultContent {
128    pub fn text(value: impl Into<String>) -> Self {
129        Self::Text(value.into())
130    }
131
132    pub fn len(&self) -> usize {
133        match self {
134            Self::Text(text) => text.len(),
135            Self::Structured(value) => value.to_string().len(),
136        }
137    }
138
139    pub fn is_empty(&self) -> bool {
140        self.len() == 0
141    }
142
143    pub fn clear(&mut self) {
144        *self = Self::Text(String::new());
145    }
146
147    pub fn as_str(&self) -> &str {
148        match self {
149            Self::Text(text) => text.as_str(),
150            Self::Structured(_) => panic!("ToolResultContent::as_str requires text content"),
151        }
152    }
153
154    pub fn contains(&self, pattern: &str) -> bool {
155        match self {
156            Self::Text(text) => text.contains(pattern),
157            Self::Structured(value) => value.to_string().contains(pattern),
158        }
159    }
160
161    pub fn starts_with(&self, pattern: &str) -> bool {
162        match self {
163            Self::Text(text) => text.starts_with(pattern),
164            Self::Structured(value) => value.to_string().starts_with(pattern),
165        }
166    }
167
168    pub fn push_str(&mut self, value: &str) {
169        match self {
170            Self::Text(text) => text.push_str(value),
171            Self::Structured(existing) => {
172                let mut text = existing.to_string();
173                text.push_str(value);
174                *self = Self::Text(text);
175            }
176        }
177    }
178
179    pub fn to_display_string(&self) -> String {
180        match self {
181            Self::Text(text) => text.clone(),
182            Self::Structured(value) => value.to_string(),
183        }
184    }
185}
186
187impl Default for ToolResultContent {
188    fn default() -> Self {
189        Self::Text(String::new())
190    }
191}
192
193impl From<String> for ToolResultContent {
194    fn from(value: String) -> Self {
195        Self::Text(value)
196    }
197}
198
199impl Display for ToolResultContent {
200    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
201        f.write_str(&self.to_display_string())
202    }
203}
204
205impl PartialEq<&str> for ToolResultContent {
206    fn eq(&self, other: &&str) -> bool {
207        self.to_display_string() == *other
208    }
209}
210
211impl PartialEq<str> for ToolResultContent {
212    fn eq(&self, other: &str) -> bool {
213        self.to_display_string() == other
214    }
215}
216
217impl PartialEq<ToolResultContent> for &str {
218    fn eq(&self, other: &ToolResultContent) -> bool {
219        *self == other.to_display_string()
220    }
221}
222
223impl PartialEq<ToolResultContent> for str {
224    fn eq(&self, other: &ToolResultContent) -> bool {
225        self == other.to_display_string()
226    }
227}
228
229impl From<&str> for ToolResultContent {
230    fn from(value: &str) -> Self {
231        Self::Text(value.to_string())
232    }
233}
234
235/// Provider-neutral hosted tool search action.
236#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
237pub struct HostedToolSearchCall {
238    pub id: String,
239    #[serde(default, skip_serializing_if = "Option::is_none")]
240    pub status: Option<String>,
241    #[serde(default, skip_serializing_if = "Option::is_none")]
242    pub query: Option<String>,
243}
244
245/// Provider-neutral hosted web search actions.
246#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
247#[serde(tag = "type", rename_all = "snake_case")]
248pub enum WebSearchAction {
249    Search {
250        #[serde(default, skip_serializing_if = "Option::is_none")]
251        query: Option<String>,
252        #[serde(default, skip_serializing_if = "Option::is_none")]
253        queries: Option<Vec<String>>,
254    },
255    OpenPage {
256        #[serde(default, skip_serializing_if = "Option::is_none")]
257        url: Option<String>,
258    },
259    FindInPage {
260        #[serde(default, skip_serializing_if = "Option::is_none")]
261        url: Option<String>,
262        #[serde(default, skip_serializing_if = "Option::is_none")]
263        pattern: Option<String>,
264    },
265}
266
267/// Provider-neutral hosted web search call.
268#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
269pub struct HostedWebSearchCall {
270    pub id: String,
271    #[serde(default, skip_serializing_if = "Option::is_none")]
272    pub status: Option<String>,
273    #[serde(default, skip_serializing_if = "Option::is_none")]
274    pub action: Option<WebSearchAction>,
275}
276
277/// Provider-neutral image generation result.
278#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
279#[serde(tag = "type", rename_all = "snake_case")]
280pub enum ImageGenerationResult {
281    Image { source: ImageSource },
282    ArtifactRef { artifact_id: String },
283}
284
285/// Provider-neutral image generation call.
286#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
287pub struct ImageGenerationCall {
288    pub id: String,
289    pub status: String,
290    #[serde(default, skip_serializing_if = "Option::is_none")]
291    pub revised_prompt: Option<String>,
292    #[serde(default, skip_serializing_if = "Option::is_none")]
293    pub result: Option<ImageGenerationResult>,
294}
295
296/// Provider-specific format carried by a provider-neutral reasoning block.
297#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
298#[serde(rename_all = "snake_case")]
299pub enum ReasoningFormat {
300    AnthropicSigned,
301    OpenAiEncrypted,
302    GeminiThought,
303}
304
305/// Origin required to decide whether opaque reasoning metadata is safe to replay.
306#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
307pub struct ReasoningProvenance {
308    pub provider: crate::ProviderId,
309    pub model: String,
310    pub format: ReasoningFormat,
311}
312
313/// A provider-neutral content block exchanged with models.
314#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
315pub enum ContentBlock {
316    Text {
317        text: String,
318    },
319    Thinking {
320        #[serde(default, skip_serializing_if = "String::is_empty")]
321        thinking: String,
322        #[serde(default, skip_serializing_if = "Option::is_none")]
323        signature: Option<String>,
324        #[serde(default, skip_serializing_if = "Option::is_none")]
325        encrypted_content: Option<String>,
326        #[serde(default, skip_serializing_if = "Option::is_none")]
327        id: Option<String>,
328        #[serde(default, skip_serializing_if = "Option::is_none")]
329        provenance: Option<ReasoningProvenance>,
330        #[serde(default, skip_serializing_if = "std::ops::Not::not")]
331        redacted: bool,
332    },
333    Image {
334        source: ImageSource,
335    },
336    ToolUse {
337        id: String,
338        name: String,
339        input: Value,
340    },
341    ToolResult {
342        tool_use_id: String,
343        content: ToolResultContent,
344        is_error: bool,
345    },
346    HostedToolSearch {
347        call: HostedToolSearchCall,
348    },
349    HostedWebSearch {
350        call: HostedWebSearchCall,
351    },
352    ImageGeneration {
353        call: ImageGenerationCall,
354    },
355}
356
357impl ContentBlock {
358    pub fn text(text: impl Into<String>) -> Self {
359        Self::Text { text: text.into() }
360    }
361
362    pub fn thinking(thinking: impl Into<String>) -> Self {
363        Self::Thinking {
364            thinking: thinking.into(),
365            signature: None,
366            encrypted_content: None,
367            id: None,
368            provenance: None,
369            redacted: false,
370        }
371    }
372
373    pub(crate) fn thinking_fallback_text(&self) -> Option<String> {
374        let Self::Thinking {
375            thinking, redacted, ..
376        } = self
377        else {
378            return None;
379        };
380
381        if !thinking.is_empty() {
382            Some(thinking.clone())
383        } else if *redacted {
384            Some("[redacted reasoning]".to_string())
385        } else {
386            Some("[reasoning unavailable]".to_string())
387        }
388    }
389
390    pub fn image_bytes(media_type: impl Into<String>, data: impl Into<Vec<u8>>) -> Self {
391        Self::Image {
392            source: ImageSource::bytes(media_type, data),
393        }
394    }
395
396    pub fn image_url(url: impl Into<String>) -> Self {
397        Self::Image {
398            source: ImageSource::url(url),
399        }
400    }
401}
402
403#[cfg(test)]
404mod tests {
405    use super::*;
406
407    #[test]
408    fn thinking_serde_is_externally_tagged_and_omits_empty_optional_fields() {
409        let block = ContentBlock::Thinking {
410            thinking: "private chain".to_string(),
411            signature: Some("opaque-signature".to_string()),
412            encrypted_content: None,
413            id: None,
414            provenance: Some(ReasoningProvenance {
415                provider: crate::ProviderId::new("anthropic-edge"),
416                model: "claude-test".to_string(),
417                format: ReasoningFormat::AnthropicSigned,
418            }),
419            redacted: false,
420        };
421
422        let json = serde_json::to_value(&block).expect("thinking block should serialize");
423        assert_eq!(json["Thinking"]["thinking"], "private chain");
424        assert_eq!(json["Thinking"]["signature"], "opaque-signature");
425        assert_eq!(json["Thinking"]["provenance"]["provider"], "anthropic-edge");
426        assert_eq!(json["Thinking"]["provenance"]["format"], "anthropic_signed");
427        assert!(json["Thinking"].get("encrypted_content").is_none());
428        assert!(json["Thinking"].get("id").is_none());
429        assert!(json["Thinking"].get("redacted").is_none());
430        assert_eq!(
431            serde_json::from_value::<ContentBlock>(json).expect("thinking block should load"),
432            block
433        );
434    }
435
436    #[test]
437    fn thinking_serde_defaults_omitted_payload_fields() {
438        let block: ContentBlock = serde_json::from_value(serde_json::json!({
439            "Thinking": {
440                "provenance": {
441                    "provider": "anthropic",
442                    "model": "claude-test",
443                    "format": "anthropic_signed"
444                }
445            }
446        }))
447        .expect("omitted thinking payload fields should default");
448
449        assert_eq!(
450            block,
451            ContentBlock::Thinking {
452                thinking: String::new(),
453                signature: None,
454                encrypted_content: None,
455                id: None,
456                provenance: Some(ReasoningProvenance {
457                    provider: crate::ProviderId::new("anthropic"),
458                    model: "claude-test".to_string(),
459                    format: ReasoningFormat::AnthropicSigned,
460                }),
461                redacted: false,
462            }
463        );
464    }
465
466    #[test]
467    fn pre_thinking_content_block_json_still_deserializes_unchanged() {
468        let json = serde_json::json!({"Text":{"text":"legacy"}});
469
470        assert_eq!(
471            serde_json::from_value::<ContentBlock>(json).expect("legacy block should load"),
472            ContentBlock::text("legacy")
473        );
474    }
475}
476
477/// Provider-neutral chat message content.
478#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
479pub struct Message {
480    pub role: Role,
481    pub content: Vec<ContentBlock>,
482}
483
484impl Message {
485    pub fn user(content: ContentBlock) -> Self {
486        Self {
487            role: Role::User,
488            content: vec![content],
489        }
490    }
491
492    pub fn assistant(content: ContentBlock) -> Self {
493        Self {
494            role: Role::Assistant,
495            content: vec![content],
496        }
497    }
498
499    pub fn unknown(role: impl Into<String>, content: ContentBlock) -> Self {
500        Self {
501            role: Role::Unknown(role.into()),
502            content: vec![content],
503        }
504    }
505
506    pub fn text(&self) -> String {
507        self.content
508            .iter()
509            .filter_map(|block| match block {
510                ContentBlock::Text { text } => Some(text.as_str()),
511                _ => None,
512            })
513            .collect::<Vec<_>>()
514            .join("")
515    }
516}
517
518/// Provider-neutral tool choice hint passed to model APIs.
519#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
520pub enum ToolChoice {
521    #[default]
522    Auto,
523    Any,
524    Tool {
525        name: String,
526    },
527}
mentra_provider/model.rs

mentra_provider/
model.rs
