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

mentra_provider/
model.rs
