rig_core/completion/
message.rs

1use std::{convert::Infallible, str::FromStr};
2
3use crate::OneOrMany;
4use serde::{Deserialize, Serialize};
5use thiserror::Error;
6
7use super::CompletionError;
8
9// ================================================================
10// Message models
11// ================================================================
12
13/// A provider-agnostic chat message.
14///
15/// Messages are role-tagged and may contain one or many content items, including
16/// text, images, audio, documents, tool calls, and tool results. Provider modules
17/// are responsible for translating these generic messages into provider-native
18/// request bodies. That conversion may be lossy when a provider does not support
19/// a particular content type.
20#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
21#[serde(tag = "role", rename_all = "lowercase")]
22pub enum Message {
23    /// System message containing instruction text.
24    System { content: String },
25
26    /// User message containing one or more content types defined by `UserContent`.
27    User { content: OneOrMany<UserContent> },
28
29    /// Assistant message containing one or more content types defined by `AssistantContent`.
30    Assistant {
31        /// Provider-assigned assistant message ID, when available.
32        id: Option<String>,
33        content: OneOrMany<AssistantContent>,
34    },
35}
36
37/// Describes the content of a message, which can be text, a tool result, an image, audio, or
38///  a document. Dependent on provider supporting the content type. Multimedia content is generally
39///  base64 (defined by it's format) encoded but additionally supports urls (for some providers).
40#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
41#[serde(tag = "type", rename_all = "lowercase")]
42pub enum UserContent {
43    /// Plain text user content.
44    Text(Text),
45    /// Result of a tool call returned as user-visible context to the model.
46    ToolResult(ToolResult),
47    /// Image content.
48    Image(Image),
49    /// Audio content.
50    Audio(Audio),
51    /// Video content.
52    Video(Video),
53    /// Document content.
54    Document(Document),
55}
56
57/// Describes responses from a provider which is either text or a tool call.
58#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
59#[serde(untagged)]
60pub enum AssistantContent {
61    /// Plain assistant text.
62    Text(Text),
63    /// Tool call requested by the assistant.
64    ToolCall(ToolCall),
65    /// Structured reasoning emitted by the assistant.
66    Reasoning(Reasoning),
67    /// Image content emitted by the assistant.
68    Image(Image),
69}
70
71#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
72#[serde(tag = "type", content = "content", rename_all = "snake_case")]
73#[non_exhaustive]
74/// A typed reasoning block used by providers that emit structured thinking data.
75pub enum ReasoningContent {
76    /// Plain reasoning text with an optional provider signature.
77    Text {
78        text: String,
79        #[serde(skip_serializing_if = "Option::is_none")]
80        signature: Option<String>,
81    },
82    /// Provider-encrypted reasoning payload.
83    Encrypted(String),
84    /// Redacted reasoning payload preserved as opaque data.
85    Redacted { data: String },
86    /// Provider-generated reasoning summary text.
87    Summary(String),
88}
89
90#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
91#[non_exhaustive]
92/// Assistant reasoning payload with an optional provider-supplied identifier.
93pub struct Reasoning {
94    /// Provider reasoning identifier, when supplied by the upstream API.
95    pub id: Option<String>,
96    /// Ordered reasoning content blocks.
97    pub content: Vec<ReasoningContent>,
98}
99
100impl Reasoning {
101    /// Create a new reasoning item from a single item
102    pub fn new(input: &str) -> Self {
103        Self::new_with_signature(input, None)
104    }
105
106    /// Create a new reasoning item from a single text item and optional signature.
107    pub fn new_with_signature(input: &str, signature: Option<String>) -> Self {
108        Self {
109            id: None,
110            content: vec![ReasoningContent::Text {
111                text: input.to_string(),
112                signature,
113            }],
114        }
115    }
116
117    /// Set or clear the provider reasoning ID.
118    pub fn optional_id(mut self, id: Option<String>) -> Self {
119        self.id = id;
120        self
121    }
122
123    /// Set a provider reasoning ID.
124    pub fn with_id(mut self, id: String) -> Self {
125        self.id = Some(id);
126        self
127    }
128
129    /// Create reasoning content from multiple text blocks.
130    pub fn multi(input: Vec<String>) -> Self {
131        Self {
132            id: None,
133            content: input
134                .into_iter()
135                .map(|text| ReasoningContent::Text {
136                    text,
137                    signature: None,
138                })
139                .collect(),
140        }
141    }
142
143    /// Create a redacted reasoning block.
144    pub fn redacted(data: impl Into<String>) -> Self {
145        Self {
146            id: None,
147            content: vec![ReasoningContent::Redacted { data: data.into() }],
148        }
149    }
150
151    /// Create an encrypted reasoning block.
152    pub fn encrypted(data: impl Into<String>) -> Self {
153        Self {
154            id: None,
155            content: vec![ReasoningContent::Encrypted(data.into())],
156        }
157    }
158
159    /// Create one reasoning block containing summary items.
160    pub fn summaries(input: Vec<String>) -> Self {
161        Self {
162            id: None,
163            content: input.into_iter().map(ReasoningContent::Summary).collect(),
164        }
165    }
166
167    /// Render reasoning as displayable text by joining text-like blocks with newlines.
168    pub fn display_text(&self) -> String {
169        self.content
170            .iter()
171            .filter_map(|content| match content {
172                ReasoningContent::Text { text, .. } => Some(text.as_str()),
173                ReasoningContent::Summary(summary) => Some(summary.as_str()),
174                ReasoningContent::Redacted { data } => Some(data.as_str()),
175                ReasoningContent::Encrypted(_) => None,
176            })
177            .collect::<Vec<_>>()
178            .join("\n")
179    }
180
181    /// Return the first text reasoning block, if present.
182    pub fn first_text(&self) -> Option<&str> {
183        self.content.iter().find_map(|content| match content {
184            ReasoningContent::Text { text, .. } => Some(text.as_str()),
185            _ => None,
186        })
187    }
188
189    /// Return the first signature from text reasoning, if present.
190    pub fn first_signature(&self) -> Option<&str> {
191        self.content.iter().find_map(|content| match content {
192            ReasoningContent::Text {
193                signature: Some(signature),
194                ..
195            } => Some(signature.as_str()),
196            _ => None,
197        })
198    }
199
200    /// Return the first encrypted reasoning payload, if present.
201    pub fn encrypted_content(&self) -> Option<&str> {
202        self.content.iter().find_map(|content| match content {
203            ReasoningContent::Encrypted(data) => Some(data.as_str()),
204            _ => None,
205        })
206    }
207}
208
209/// Tool result content containing information about a tool call and it's resulting content.
210#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
211pub struct ToolResult {
212    /// Tool call ID that this result answers.
213    pub id: String,
214    /// Provider-specific call ID, when distinct from `id`.
215    #[serde(skip_serializing_if = "Option::is_none")]
216    pub call_id: Option<String>,
217    /// One or more content items produced by the tool.
218    pub content: OneOrMany<ToolResultContent>,
219}
220
221/// Describes the content of a tool result, which can be text or an image.
222#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
223#[serde(tag = "type", rename_all = "lowercase")]
224pub enum ToolResultContent {
225    Text(Text),
226    Image(Image),
227}
228
229/// Describes a tool call with an id and function to call, generally produced by a provider.
230#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
231pub struct ToolCall {
232    /// Provider-supplied tool call ID.
233    pub id: String,
234    /// Provider-specific call ID used by some APIs for tool result correlation.
235    pub call_id: Option<String>,
236    /// Function name and JSON arguments requested by the model.
237    pub function: ToolFunction,
238    /// Optional cryptographic signature for the tool call.
239    ///
240    /// This field is used by some providers (e.g., Google) to provide a signature
241    /// that can verify the authenticity and integrity of the tool call. When present,
242    /// it allows verification that the tool call was actually generated by the model
243    /// and has not been tampered with.
244    ///
245    /// This is an optional, provider-specific feature and will be `None` for providers
246    /// that don't support tool call signatures.
247    pub signature: Option<String>,
248    /// Additional provider-specific parameters to be sent to the completion model provider
249    pub additional_params: Option<serde_json::Value>,
250}
251
252impl ToolCall {
253    pub fn new(id: String, function: ToolFunction) -> Self {
254        Self {
255            id,
256            call_id: None,
257            function,
258            signature: None,
259            additional_params: None,
260        }
261    }
262
263    pub fn with_call_id(mut self, call_id: String) -> Self {
264        self.call_id = Some(call_id);
265        self
266    }
267
268    pub fn with_signature(mut self, signature: Option<String>) -> Self {
269        self.signature = signature;
270        self
271    }
272
273    pub fn with_additional_params(mut self, additional_params: Option<serde_json::Value>) -> Self {
274        self.additional_params = additional_params;
275        self
276    }
277}
278
279/// Describes a tool function to call with a name and arguments, generally produced by a provider.
280#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
281pub struct ToolFunction {
282    /// Tool/function name to invoke.
283    pub name: String,
284    /// JSON arguments for the tool/function.
285    pub arguments: serde_json::Value,
286}
287
288impl ToolFunction {
289    /// Create a tool function call payload.
290    pub fn new(name: String, arguments: serde_json::Value) -> Self {
291        Self { name, arguments }
292    }
293}
294
295// ================================================================
296// Base content models
297// ================================================================
298
299/// Basic text content.
300///
301/// `additional_params` carries provider-specific fields that arrive on text
302/// content blocks (e.g. Anthropic returns citation metadata on assistant text
303/// blocks). It is flattened during serialization so the inner JSON keys appear
304/// at the same level as `text`, matching the wire format of provider APIs.
305#[derive(Default, Clone, Debug, Deserialize, Serialize, PartialEq)]
306pub struct Text {
307    /// Text content.
308    pub text: String,
309    /// Provider-specific text fields.
310    #[serde(flatten, skip_serializing_if = "Option::is_none")]
311    pub additional_params: Option<serde_json::Value>,
312}
313
314impl Text {
315    /// Construct a new text block with no provider-specific fields.
316    pub fn new(text: impl Into<String>) -> Self {
317        Self {
318            text: text.into(),
319            additional_params: None,
320        }
321    }
322
323    /// Returns the inner text string.
324    pub fn text(&self) -> &str {
325        &self.text
326    }
327}
328
329impl std::fmt::Display for Text {
330    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
331        let Self { text, .. } = self;
332        write!(f, "{text}")
333    }
334}
335
336/// Image content containing image data and metadata about it.
337#[derive(Default, Clone, Debug, Deserialize, Serialize, PartialEq)]
338pub struct Image {
339    /// Image source data.
340    pub data: DocumentSourceKind,
341    /// Image media type, if known.
342    #[serde(skip_serializing_if = "Option::is_none")]
343    pub media_type: Option<ImageMediaType>,
344    /// Provider-specific image detail preference.
345    #[serde(skip_serializing_if = "Option::is_none")]
346    pub detail: Option<ImageDetail>,
347    /// Provider-specific image fields.
348    #[serde(flatten, skip_serializing_if = "Option::is_none")]
349    pub additional_params: Option<serde_json::Value>,
350}
351
352impl Image {
353    pub fn try_into_url(self) -> Result<String, MessageError> {
354        match self.data {
355            DocumentSourceKind::Url(url) => Ok(url),
356            DocumentSourceKind::Base64(data) => {
357                let Some(media_type) = self.media_type else {
358                    return Err(MessageError::ConversionError(
359                        "A media type is required to create a valid base64-encoded image URL"
360                            .to_string(),
361                    ));
362                };
363
364                Ok(format!(
365                    "data:image/{ty};base64,{data}",
366                    ty = media_type.to_mime_type()
367                ))
368            }
369            unknown => Err(MessageError::ConversionError(format!(
370                "Tried to convert unknown type to a URL: {unknown:?}"
371            ))),
372        }
373    }
374}
375
376/// The kind of image source (to be used).
377#[derive(Debug, Deserialize, Serialize, Clone, PartialEq, Default)]
378#[serde(tag = "type", content = "value", rename_all = "camelCase")]
379#[non_exhaustive]
380pub enum DocumentSourceKind {
381    /// A file URL/URI.
382    Url(String),
383    /// A base-64 encoded string.
384    Base64(String),
385    /// A provider-side uploaded file identifier.
386    FileId(String),
387    /// Raw bytes
388    Raw(Vec<u8>),
389    /// A string (or a string literal).
390    String(String),
391    #[default]
392    /// An unknown file source (there's nothing there).
393    Unknown,
394}
395
396impl DocumentSourceKind {
397    /// Create a URL-backed source.
398    pub fn url(url: &str) -> Self {
399        Self::Url(url.to_string())
400    }
401
402    /// Create a base64-backed source.
403    pub fn base64(base64_string: &str) -> Self {
404        Self::Base64(base64_string.to_string())
405    }
406
407    /// Create a provider file ID-backed source.
408    pub fn file_id(file_id: &str) -> Self {
409        Self::FileId(file_id.to_string())
410    }
411
412    /// Create a raw byte source.
413    pub fn raw(bytes: impl Into<Vec<u8>>) -> Self {
414        Self::Raw(bytes.into())
415    }
416
417    /// Create a string-backed source.
418    pub fn string(input: &str) -> Self {
419        Self::String(input.into())
420    }
421
422    /// Create an unknown source placeholder.
423    pub fn unknown() -> Self {
424        Self::Unknown
425    }
426
427    /// Return the contained URL, base64 string, or file ID, if this source stores one.
428    pub fn try_into_inner(self) -> Option<String> {
429        match self {
430            Self::Url(s) | Self::Base64(s) | Self::FileId(s) => Some(s),
431            _ => None,
432        }
433    }
434}
435
436impl std::fmt::Display for DocumentSourceKind {
437    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
438        match self {
439            Self::Url(string) => write!(f, "{string}"),
440            Self::Base64(string) => write!(f, "{string}"),
441            Self::FileId(string) => write!(f, "{string}"),
442            Self::String(string) => write!(f, "{string}"),
443            Self::Raw(_) => write!(f, "<binary data>"),
444            Self::Unknown => write!(f, "<unknown>"),
445        }
446    }
447}
448
449/// Audio content containing audio data and metadata about it.
450#[derive(Default, Clone, Debug, Deserialize, Serialize, PartialEq)]
451pub struct Audio {
452    /// Audio source data.
453    pub data: DocumentSourceKind,
454    /// Audio media type, if known.
455    #[serde(skip_serializing_if = "Option::is_none")]
456    pub media_type: Option<AudioMediaType>,
457    /// Provider-specific audio fields.
458    #[serde(flatten, skip_serializing_if = "Option::is_none")]
459    pub additional_params: Option<serde_json::Value>,
460}
461
462/// Video content containing video data and metadata about it.
463#[derive(Default, Clone, Debug, Deserialize, Serialize, PartialEq)]
464pub struct Video {
465    /// Video source data.
466    pub data: DocumentSourceKind,
467    /// Video media type, if known.
468    #[serde(skip_serializing_if = "Option::is_none")]
469    pub media_type: Option<VideoMediaType>,
470    /// Provider-specific video fields.
471    #[serde(flatten, skip_serializing_if = "Option::is_none")]
472    pub additional_params: Option<serde_json::Value>,
473}
474
475/// Document content containing document data and metadata about it.
476#[derive(Default, Clone, Debug, Deserialize, Serialize, PartialEq)]
477pub struct Document {
478    /// Document source data.
479    pub data: DocumentSourceKind,
480    /// Document media type, if known.
481    #[serde(skip_serializing_if = "Option::is_none")]
482    pub media_type: Option<DocumentMediaType>,
483    /// Provider-specific document fields.
484    #[serde(flatten, skip_serializing_if = "Option::is_none")]
485    pub additional_params: Option<serde_json::Value>,
486}
487
488/// Describes the format of the content, which can be base64 or string.
489#[derive(Default, Clone, Debug, Deserialize, Serialize, PartialEq)]
490#[serde(rename_all = "lowercase")]
491pub enum ContentFormat {
492    #[default]
493    Base64,
494    String,
495    Url,
496}
497
498/// Helper enum that tracks the media type of the content.
499#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
500pub enum MediaType {
501    Image(ImageMediaType),
502    Audio(AudioMediaType),
503    Document(DocumentMediaType),
504    Video(VideoMediaType),
505}
506
507/// Describes the image media type of the content. Not every provider supports every media type.
508/// Convertible to and from MIME type strings.
509#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
510#[serde(rename_all = "lowercase")]
511pub enum ImageMediaType {
512    JPEG,
513    PNG,
514    GIF,
515    WEBP,
516    HEIC,
517    HEIF,
518    SVG,
519}
520
521/// Describes the document media type of the content. Not every provider supports every media type.
522/// Includes also programming languages as document types for providers who support code running.
523/// Convertible to and from MIME type strings.
524#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
525#[serde(rename_all = "lowercase")]
526pub enum DocumentMediaType {
527    PDF,
528    TXT,
529    RTF,
530    HTML,
531    CSS,
532    MARKDOWN,
533    CSV,
534    XML,
535    Javascript,
536    Python,
537}
538
539impl DocumentMediaType {
540    pub fn is_code(&self) -> bool {
541        matches!(self, Self::Javascript | Self::Python)
542    }
543}
544
545/// Describes the audio media type of the content. Not every provider supports every media type.
546/// Convertible to and from MIME type strings.
547#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
548#[serde(rename_all = "lowercase")]
549pub enum AudioMediaType {
550    WAV,
551    MP3,
552    AIFF,
553    AAC,
554    OGG,
555    FLAC,
556    M4A,
557    PCM16,
558    PCM24,
559}
560
561/// Describes the video media type of the content. Not every provider supports every media type.
562/// Convertible to and from MIME type strings.
563#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
564#[serde(rename_all = "lowercase")]
565pub enum VideoMediaType {
566    AVI,
567    MP4,
568    MPEG,
569    MOV,
570    WEBM,
571}
572
573/// Describes the detail of the image content, which can be low, high, or auto (open-ai specific).
574#[derive(Default, Clone, Debug, Deserialize, Serialize, PartialEq)]
575#[serde(rename_all = "lowercase")]
576pub enum ImageDetail {
577    Low,
578    High,
579    #[default]
580    Auto,
581}
582
583// ================================================================
584// Impl. for message models
585// ================================================================
586
587impl Message {
588    /// This helper method is primarily used to extract the first string prompt from a `Message`.
589    /// Since `Message` might have more than just text content, we need to find the first text.
590    pub(crate) fn rag_text(&self) -> Option<String> {
591        match self {
592            Message::User { content } => {
593                for item in content.iter() {
594                    if let UserContent::Text(Text { text, .. }) = item {
595                        return Some(text.clone());
596                    }
597                }
598                None
599            }
600            Message::System { .. } => None,
601            _ => None,
602        }
603    }
604
605    /// Helper constructor to make creating system messages easier.
606    pub fn system(text: impl Into<String>) -> Self {
607        Message::System {
608            content: text.into(),
609        }
610    }
611
612    /// Helper constructor to make creating user messages easier.
613    pub fn user(text: impl Into<String>) -> Self {
614        Message::User {
615            content: OneOrMany::one(UserContent::text(text)),
616        }
617    }
618
619    /// Helper constructor to make creating assistant messages easier.
620    pub fn assistant(text: impl Into<String>) -> Self {
621        Message::Assistant {
622            id: None,
623            content: OneOrMany::one(AssistantContent::text(text)),
624        }
625    }
626
627    /// Helper constructor to make creating assistant messages easier.
628    pub fn assistant_with_id(id: String, text: impl Into<String>) -> Self {
629        Message::Assistant {
630            id: Some(id),
631            content: OneOrMany::one(AssistantContent::text(text)),
632        }
633    }
634
635    /// Helper constructor to make creating tool result messages easier.
636    pub fn tool_result(id: impl Into<String>, content: impl Into<String>) -> Self {
637        Message::User {
638            content: OneOrMany::one(UserContent::ToolResult(ToolResult {
639                id: id.into(),
640                call_id: None,
641                content: OneOrMany::one(ToolResultContent::text(content)),
642            })),
643        }
644    }
645
646    pub fn tool_result_with_call_id(
647        id: impl Into<String>,
648        call_id: Option<String>,
649        content: impl Into<String>,
650    ) -> Self {
651        Message::User {
652            content: OneOrMany::one(UserContent::ToolResult(ToolResult {
653                id: id.into(),
654                call_id,
655                content: OneOrMany::one(ToolResultContent::text(content)),
656            })),
657        }
658    }
659}
660
661impl UserContent {
662    /// Helper constructor to make creating user text content easier.
663    pub fn text(text: impl Into<String>) -> Self {
664        UserContent::Text(text.into().into())
665    }
666
667    /// Helper constructor to make creating user image content easier.
668    pub fn image_base64(
669        data: impl Into<String>,
670        media_type: Option<ImageMediaType>,
671        detail: Option<ImageDetail>,
672    ) -> Self {
673        UserContent::Image(Image {
674            data: DocumentSourceKind::Base64(data.into()),
675            media_type,
676            detail,
677            additional_params: None,
678        })
679    }
680
681    /// Helper constructor to make creating user image content from raw unencoded bytes easier.
682    pub fn image_raw(
683        data: impl Into<Vec<u8>>,
684        media_type: Option<ImageMediaType>,
685        detail: Option<ImageDetail>,
686    ) -> Self {
687        UserContent::Image(Image {
688            data: DocumentSourceKind::Raw(data.into()),
689            media_type,
690            detail,
691            ..Default::default()
692        })
693    }
694
695    /// Helper constructor to make creating user image content easier.
696    pub fn image_url(
697        url: impl Into<String>,
698        media_type: Option<ImageMediaType>,
699        detail: Option<ImageDetail>,
700    ) -> Self {
701        UserContent::Image(Image {
702            data: DocumentSourceKind::Url(url.into()),
703            media_type,
704            detail,
705            additional_params: None,
706        })
707    }
708
709    /// Helper constructor to make creating user audio content easier.
710    pub fn audio(data: impl Into<String>, media_type: Option<AudioMediaType>) -> Self {
711        UserContent::Audio(Audio {
712            data: DocumentSourceKind::Base64(data.into()),
713            media_type,
714            additional_params: None,
715        })
716    }
717
718    /// Helper constructor to make creating user audio content from raw unencoded bytes easier.
719    pub fn audio_raw(data: impl Into<Vec<u8>>, media_type: Option<AudioMediaType>) -> Self {
720        UserContent::Audio(Audio {
721            data: DocumentSourceKind::Raw(data.into()),
722            media_type,
723            ..Default::default()
724        })
725    }
726
727    /// Helper to create an audio resource from a URL
728    pub fn audio_url(url: impl Into<String>, media_type: Option<AudioMediaType>) -> Self {
729        UserContent::Audio(Audio {
730            data: DocumentSourceKind::Url(url.into()),
731            media_type,
732            ..Default::default()
733        })
734    }
735
736    /// Helper constructor to make creating user video content easier.
737    pub fn video(data: impl Into<String>, media_type: Option<VideoMediaType>) -> Self {
738        UserContent::Video(Video {
739            data: DocumentSourceKind::Base64(data.into()),
740            media_type,
741            additional_params: None,
742        })
743    }
744
745    /// Helper constructor to make creating user video content from raw unencoded bytes easier.
746    pub fn video_raw(data: impl Into<Vec<u8>>, media_type: Option<VideoMediaType>) -> Self {
747        UserContent::Video(Video {
748            data: DocumentSourceKind::Raw(data.into()),
749            media_type,
750            ..Default::default()
751        })
752    }
753
754    /// Helper to create a video resource from a URL
755    pub fn video_url(url: impl Into<String>, media_type: Option<VideoMediaType>) -> Self {
756        UserContent::Video(Video {
757            data: DocumentSourceKind::Url(url.into()),
758            media_type,
759            ..Default::default()
760        })
761    }
762
763    /// Helper constructor to make creating user document content easier.
764    /// This creates a document that assumes the data being passed in is a raw string.
765    pub fn document(data: impl Into<String>, media_type: Option<DocumentMediaType>) -> Self {
766        let data: String = data.into();
767        UserContent::Document(Document {
768            data: DocumentSourceKind::string(&data),
769            media_type,
770            additional_params: None,
771        })
772    }
773
774    /// Helper to create a document from raw unencoded bytes
775    pub fn document_raw(data: impl Into<Vec<u8>>, media_type: Option<DocumentMediaType>) -> Self {
776        UserContent::Document(Document {
777            data: DocumentSourceKind::Raw(data.into()),
778            media_type,
779            ..Default::default()
780        })
781    }
782
783    /// Helper to create a document from a URL
784    pub fn document_url(url: impl Into<String>, media_type: Option<DocumentMediaType>) -> Self {
785        UserContent::Document(Document {
786            data: DocumentSourceKind::Url(url.into()),
787            media_type,
788            ..Default::default()
789        })
790    }
791
792    /// Helper constructor to make creating user tool result content easier.
793    pub fn tool_result(id: impl Into<String>, content: OneOrMany<ToolResultContent>) -> Self {
794        UserContent::ToolResult(ToolResult {
795            id: id.into(),
796            call_id: None,
797            content,
798        })
799    }
800
801    /// Helper constructor to make creating user tool result content easier.
802    pub fn tool_result_with_call_id(
803        id: impl Into<String>,
804        call_id: String,
805        content: OneOrMany<ToolResultContent>,
806    ) -> Self {
807        UserContent::ToolResult(ToolResult {
808            id: id.into(),
809            call_id: Some(call_id),
810            content,
811        })
812    }
813}
814
815impl AssistantContent {
816    /// Helper constructor to make creating assistant text content easier.
817    pub fn text(text: impl Into<String>) -> Self {
818        AssistantContent::Text(text.into().into())
819    }
820
821    /// Helper constructor to make creating assistant image content easier.
822    pub fn image_base64(
823        data: impl Into<String>,
824        media_type: Option<ImageMediaType>,
825        detail: Option<ImageDetail>,
826    ) -> Self {
827        AssistantContent::Image(Image {
828            data: DocumentSourceKind::Base64(data.into()),
829            media_type,
830            detail,
831            additional_params: None,
832        })
833    }
834
835    /// Helper constructor to make creating assistant tool call content easier.
836    pub fn tool_call(
837        id: impl Into<String>,
838        name: impl Into<String>,
839        arguments: serde_json::Value,
840    ) -> Self {
841        AssistantContent::ToolCall(ToolCall::new(
842            id.into(),
843            ToolFunction {
844                name: name.into(),
845                arguments,
846            },
847        ))
848    }
849
850    pub fn tool_call_with_call_id(
851        id: impl Into<String>,
852        call_id: String,
853        name: impl Into<String>,
854        arguments: serde_json::Value,
855    ) -> Self {
856        AssistantContent::ToolCall(
857            ToolCall::new(
858                id.into(),
859                ToolFunction {
860                    name: name.into(),
861                    arguments,
862                },
863            )
864            .with_call_id(call_id),
865        )
866    }
867
868    pub fn reasoning(reasoning: impl AsRef<str>) -> Self {
869        AssistantContent::Reasoning(Reasoning::new(reasoning.as_ref()))
870    }
871}
872
873impl ToolResultContent {
874    /// Helper constructor to make creating tool result text content easier.
875    pub fn text(text: impl Into<String>) -> Self {
876        ToolResultContent::Text(text.into().into())
877    }
878
879    /// Helper constructor to make tool result images from a base64-encoded string.
880    pub fn image_base64(
881        data: impl Into<String>,
882        media_type: Option<ImageMediaType>,
883        detail: Option<ImageDetail>,
884    ) -> Self {
885        ToolResultContent::Image(Image {
886            data: DocumentSourceKind::Base64(data.into()),
887            media_type,
888            detail,
889            additional_params: None,
890        })
891    }
892
893    /// Helper constructor to make tool result images from a base64-encoded string.
894    pub fn image_raw(
895        data: impl Into<Vec<u8>>,
896        media_type: Option<ImageMediaType>,
897        detail: Option<ImageDetail>,
898    ) -> Self {
899        ToolResultContent::Image(Image {
900            data: DocumentSourceKind::Raw(data.into()),
901            media_type,
902            detail,
903            ..Default::default()
904        })
905    }
906
907    /// Helper constructor to make tool result images from a URL.
908    pub fn image_url(
909        url: impl Into<String>,
910        media_type: Option<ImageMediaType>,
911        detail: Option<ImageDetail>,
912    ) -> Self {
913        ToolResultContent::Image(Image {
914            data: DocumentSourceKind::Url(url.into()),
915            media_type,
916            detail,
917            additional_params: None,
918        })
919    }
920
921    /// Parse a tool output string into appropriate ToolResultContent(s).
922    ///
923    /// Supports three formats:
924    /// 1. Simple text: Any string → `OneOrMany::one(Text)`
925    /// 2. Image JSON: `{"type": "image", "data": "...", "mimeType": "..."}` → `OneOrMany::one(Image)`
926    /// 3. Hybrid JSON: `{"response": {...}, "parts": [...]}` → `OneOrMany::many([Text, Image, ...])`
927    ///
928    /// If JSON parsing fails, treats the entire string as text.
929    pub fn from_tool_output(output: impl Into<String>) -> OneOrMany<ToolResultContent> {
930        let output_str = output.into();
931
932        if let Ok(json) = serde_json::from_str::<serde_json::Value>(&output_str) {
933            if json.get("response").is_some() || json.get("parts").is_some() {
934                let mut results: Vec<ToolResultContent> = Vec::new();
935
936                if let Some(response) = json.get("response") {
937                    results.push(ToolResultContent::Text(Text::new(response.to_string())));
938                }
939
940                if let Some(parts) = json.get("parts").and_then(|p| p.as_array()) {
941                    for part in parts {
942                        let is_image = part
943                            .get("type")
944                            .and_then(|t| t.as_str())
945                            .is_some_and(|t| t == "image");
946
947                        if !is_image {
948                            continue;
949                        }
950
951                        if let (Some(data), Some(mime_type)) = (
952                            part.get("data").and_then(|v| v.as_str()),
953                            part.get("mimeType").and_then(|v| v.as_str()),
954                        ) {
955                            let data_kind =
956                                if data.starts_with("http://") || data.starts_with("https://") {
957                                    DocumentSourceKind::Url(data.to_string())
958                                } else {
959                                    DocumentSourceKind::Base64(data.to_string())
960                                };
961
962                            results.push(ToolResultContent::Image(Image {
963                                data: data_kind,
964                                media_type: ImageMediaType::from_mime_type(mime_type),
965                                detail: None,
966                                additional_params: None,
967                            }));
968                        }
969                    }
970                }
971
972                if !results.is_empty() {
973                    return OneOrMany::many(results).unwrap_or_else(|_| {
974                        OneOrMany::one(ToolResultContent::Text(output_str.into()))
975                    });
976                }
977            }
978
979            let is_image = json
980                .get("type")
981                .and_then(|v| v.as_str())
982                .is_some_and(|t| t == "image");
983
984            if is_image
985                && let (Some(data), Some(mime_type)) = (
986                    json.get("data").and_then(|v| v.as_str()),
987                    json.get("mimeType").and_then(|v| v.as_str()),
988                )
989            {
990                let data_kind = if data.starts_with("http://") || data.starts_with("https://") {
991                    DocumentSourceKind::Url(data.to_string())
992                } else {
993                    DocumentSourceKind::Base64(data.to_string())
994                };
995
996                return OneOrMany::one(ToolResultContent::Image(Image {
997                    data: data_kind,
998                    media_type: ImageMediaType::from_mime_type(mime_type),
999                    detail: None,
1000                    additional_params: None,
1001                }));
1002            }
1003        }
1004
1005        OneOrMany::one(ToolResultContent::Text(output_str.into()))
1006    }
1007}
1008
1009/// Trait for converting between MIME types and media types.
1010pub trait MimeType {
1011    fn from_mime_type(mime_type: &str) -> Option<Self>
1012    where
1013        Self: Sized;
1014    fn to_mime_type(&self) -> &'static str;
1015}
1016
1017impl MimeType for MediaType {
1018    fn from_mime_type(mime_type: &str) -> Option<Self> {
1019        ImageMediaType::from_mime_type(mime_type)
1020            .map(MediaType::Image)
1021            .or_else(|| {
1022                DocumentMediaType::from_mime_type(mime_type)
1023                    .map(MediaType::Document)
1024                    .or_else(|| {
1025                        AudioMediaType::from_mime_type(mime_type)
1026                            .map(MediaType::Audio)
1027                            .or_else(|| {
1028                                VideoMediaType::from_mime_type(mime_type).map(MediaType::Video)
1029                            })
1030                    })
1031            })
1032    }
1033
1034    fn to_mime_type(&self) -> &'static str {
1035        match self {
1036            MediaType::Image(media_type) => media_type.to_mime_type(),
1037            MediaType::Audio(media_type) => media_type.to_mime_type(),
1038            MediaType::Document(media_type) => media_type.to_mime_type(),
1039            MediaType::Video(media_type) => media_type.to_mime_type(),
1040        }
1041    }
1042}
1043
1044impl MimeType for ImageMediaType {
1045    fn from_mime_type(mime_type: &str) -> Option<Self> {
1046        match mime_type {
1047            "image/jpeg" => Some(ImageMediaType::JPEG),
1048            "image/png" => Some(ImageMediaType::PNG),
1049            "image/gif" => Some(ImageMediaType::GIF),
1050            "image/webp" => Some(ImageMediaType::WEBP),
1051            "image/heic" => Some(ImageMediaType::HEIC),
1052            "image/heif" => Some(ImageMediaType::HEIF),
1053            "image/svg+xml" => Some(ImageMediaType::SVG),
1054            _ => None,
1055        }
1056    }
1057
1058    fn to_mime_type(&self) -> &'static str {
1059        match self {
1060            ImageMediaType::JPEG => "image/jpeg",
1061            ImageMediaType::PNG => "image/png",
1062            ImageMediaType::GIF => "image/gif",
1063            ImageMediaType::WEBP => "image/webp",
1064            ImageMediaType::HEIC => "image/heic",
1065            ImageMediaType::HEIF => "image/heif",
1066            ImageMediaType::SVG => "image/svg+xml",
1067        }
1068    }
1069}
1070
1071impl MimeType for DocumentMediaType {
1072    fn from_mime_type(mime_type: &str) -> Option<Self> {
1073        match mime_type {
1074            "application/pdf" => Some(DocumentMediaType::PDF),
1075            "text/plain" => Some(DocumentMediaType::TXT),
1076            "text/rtf" => Some(DocumentMediaType::RTF),
1077            "text/html" => Some(DocumentMediaType::HTML),
1078            "text/css" => Some(DocumentMediaType::CSS),
1079            "text/md" | "text/markdown" => Some(DocumentMediaType::MARKDOWN),
1080            "text/csv" => Some(DocumentMediaType::CSV),
1081            "text/xml" => Some(DocumentMediaType::XML),
1082            "application/x-javascript" | "text/x-javascript" => Some(DocumentMediaType::Javascript),
1083            "application/x-python" | "text/x-python" => Some(DocumentMediaType::Python),
1084            _ => None,
1085        }
1086    }
1087
1088    fn to_mime_type(&self) -> &'static str {
1089        match self {
1090            DocumentMediaType::PDF => "application/pdf",
1091            DocumentMediaType::TXT => "text/plain",
1092            DocumentMediaType::RTF => "text/rtf",
1093            DocumentMediaType::HTML => "text/html",
1094            DocumentMediaType::CSS => "text/css",
1095            DocumentMediaType::MARKDOWN => "text/markdown",
1096            DocumentMediaType::CSV => "text/csv",
1097            DocumentMediaType::XML => "text/xml",
1098            DocumentMediaType::Javascript => "application/x-javascript",
1099            DocumentMediaType::Python => "application/x-python",
1100        }
1101    }
1102}
1103
1104impl MimeType for AudioMediaType {
1105    fn from_mime_type(mime_type: &str) -> Option<Self> {
1106        match mime_type {
1107            "audio/wav" => Some(AudioMediaType::WAV),
1108            "audio/mp3" => Some(AudioMediaType::MP3),
1109            "audio/aiff" => Some(AudioMediaType::AIFF),
1110            "audio/aac" => Some(AudioMediaType::AAC),
1111            "audio/ogg" => Some(AudioMediaType::OGG),
1112            "audio/flac" => Some(AudioMediaType::FLAC),
1113            "audio/m4a" => Some(AudioMediaType::M4A),
1114            "audio/pcm16" => Some(AudioMediaType::PCM16),
1115            "audio/pcm24" => Some(AudioMediaType::PCM24),
1116            _ => None,
1117        }
1118    }
1119
1120    fn to_mime_type(&self) -> &'static str {
1121        match self {
1122            AudioMediaType::WAV => "audio/wav",
1123            AudioMediaType::MP3 => "audio/mp3",
1124            AudioMediaType::AIFF => "audio/aiff",
1125            AudioMediaType::AAC => "audio/aac",
1126            AudioMediaType::OGG => "audio/ogg",
1127            AudioMediaType::FLAC => "audio/flac",
1128            AudioMediaType::M4A => "audio/m4a",
1129            AudioMediaType::PCM16 => "audio/pcm16",
1130            AudioMediaType::PCM24 => "audio/pcm24",
1131        }
1132    }
1133}
1134
1135impl MimeType for VideoMediaType {
1136    fn from_mime_type(mime_type: &str) -> Option<Self>
1137    where
1138        Self: Sized,
1139    {
1140        match mime_type {
1141            "video/avi" => Some(VideoMediaType::AVI),
1142            "video/mp4" => Some(VideoMediaType::MP4),
1143            "video/mpeg" => Some(VideoMediaType::MPEG),
1144            "video/mov" => Some(VideoMediaType::MOV),
1145            "video/webm" => Some(VideoMediaType::WEBM),
1146            &_ => None,
1147        }
1148    }
1149
1150    fn to_mime_type(&self) -> &'static str {
1151        match self {
1152            VideoMediaType::AVI => "video/avi",
1153            VideoMediaType::MP4 => "video/mp4",
1154            VideoMediaType::MPEG => "video/mpeg",
1155            VideoMediaType::MOV => "video/mov",
1156            VideoMediaType::WEBM => "video/webm",
1157        }
1158    }
1159}
1160
1161impl std::str::FromStr for ImageDetail {
1162    type Err = ();
1163
1164    fn from_str(s: &str) -> Result<Self, Self::Err> {
1165        match s.to_lowercase().as_str() {
1166            "low" => Ok(ImageDetail::Low),
1167            "high" => Ok(ImageDetail::High),
1168            "auto" => Ok(ImageDetail::Auto),
1169            _ => Err(()),
1170        }
1171    }
1172}
1173
1174// ================================================================
1175// FromStr, From<String>, and From<&str> impls
1176// ================================================================
1177
1178impl From<String> for Text {
1179    fn from(text: String) -> Self {
1180        Text {
1181            text,
1182            additional_params: None,
1183        }
1184    }
1185}
1186
1187impl From<&String> for Text {
1188    fn from(text: &String) -> Self {
1189        text.to_owned().into()
1190    }
1191}
1192
1193impl From<&str> for Text {
1194    fn from(text: &str) -> Self {
1195        text.to_owned().into()
1196    }
1197}
1198
1199impl FromStr for Text {
1200    type Err = Infallible;
1201
1202    fn from_str(s: &str) -> Result<Self, Self::Err> {
1203        Ok(s.into())
1204    }
1205}
1206
1207impl From<&Message> for Message {
1208    fn from(msg: &Message) -> Self {
1209        msg.clone()
1210    }
1211}
1212
1213impl From<String> for Message {
1214    fn from(text: String) -> Self {
1215        Message::User {
1216            content: OneOrMany::one(UserContent::Text(text.into())),
1217        }
1218    }
1219}
1220
1221impl From<&str> for Message {
1222    fn from(text: &str) -> Self {
1223        Message::User {
1224            content: OneOrMany::one(UserContent::Text(text.into())),
1225        }
1226    }
1227}
1228
1229impl From<&String> for Message {
1230    fn from(text: &String) -> Self {
1231        Message::User {
1232            content: OneOrMany::one(UserContent::Text(text.into())),
1233        }
1234    }
1235}
1236
1237impl From<Text> for Message {
1238    fn from(text: Text) -> Self {
1239        Message::User {
1240            content: OneOrMany::one(UserContent::Text(text)),
1241        }
1242    }
1243}
1244
1245impl From<Image> for Message {
1246    fn from(image: Image) -> Self {
1247        Message::User {
1248            content: OneOrMany::one(UserContent::Image(image)),
1249        }
1250    }
1251}
1252
1253impl From<Audio> for Message {
1254    fn from(audio: Audio) -> Self {
1255        Message::User {
1256            content: OneOrMany::one(UserContent::Audio(audio)),
1257        }
1258    }
1259}
1260
1261impl From<Document> for Message {
1262    fn from(document: Document) -> Self {
1263        Message::User {
1264            content: OneOrMany::one(UserContent::Document(document)),
1265        }
1266    }
1267}
1268
1269impl From<String> for ToolResultContent {
1270    fn from(text: String) -> Self {
1271        ToolResultContent::text(text)
1272    }
1273}
1274
1275impl From<String> for AssistantContent {
1276    fn from(text: String) -> Self {
1277        AssistantContent::text(text)
1278    }
1279}
1280
1281impl From<String> for UserContent {
1282    fn from(text: String) -> Self {
1283        UserContent::text(text)
1284    }
1285}
1286
1287impl From<AssistantContent> for Message {
1288    fn from(content: AssistantContent) -> Self {
1289        Message::Assistant {
1290            id: None,
1291            content: OneOrMany::one(content),
1292        }
1293    }
1294}
1295
1296impl From<UserContent> for Message {
1297    fn from(content: UserContent) -> Self {
1298        Message::User {
1299            content: OneOrMany::one(content),
1300        }
1301    }
1302}
1303
1304impl From<OneOrMany<AssistantContent>> for Message {
1305    fn from(content: OneOrMany<AssistantContent>) -> Self {
1306        Message::Assistant { id: None, content }
1307    }
1308}
1309
1310impl From<OneOrMany<UserContent>> for Message {
1311    fn from(content: OneOrMany<UserContent>) -> Self {
1312        Message::User { content }
1313    }
1314}
1315
1316impl From<ToolCall> for Message {
1317    fn from(tool_call: ToolCall) -> Self {
1318        Message::Assistant {
1319            id: None,
1320            content: OneOrMany::one(AssistantContent::ToolCall(tool_call)),
1321        }
1322    }
1323}
1324
1325impl From<ToolResult> for Message {
1326    fn from(tool_result: ToolResult) -> Self {
1327        Message::User {
1328            content: OneOrMany::one(UserContent::ToolResult(tool_result)),
1329        }
1330    }
1331}
1332
1333impl From<ToolResultContent> for Message {
1334    fn from(tool_result_content: ToolResultContent) -> Self {
1335        Message::User {
1336            content: OneOrMany::one(UserContent::ToolResult(ToolResult {
1337                id: String::new(),
1338                call_id: None,
1339                content: OneOrMany::one(tool_result_content),
1340            })),
1341        }
1342    }
1343}
1344
1345#[derive(Default, Clone, Debug, Deserialize, Serialize, PartialEq)]
1346#[serde(rename_all = "snake_case")]
1347pub enum ToolChoice {
1348    #[default]
1349    Auto,
1350    None,
1351    Required,
1352    Specific {
1353        function_names: Vec<String>,
1354    },
1355}
1356
1357// ================================================================
1358// Error types
1359// ================================================================
1360
1361/// Error type to represent issues with converting messages to and from specific provider messages.
1362#[derive(Debug, Error)]
1363pub enum MessageError {
1364    #[error("Message conversion error: {0}")]
1365    ConversionError(String),
1366}
1367
1368impl From<MessageError> for CompletionError {
1369    fn from(error: MessageError) -> Self {
1370        CompletionError::RequestError(error.into())
1371    }
1372}
1373
1374#[cfg(test)]
1375mod tests {
1376    use super::{Message, Reasoning, ReasoningContent};
1377
1378    #[test]
1379    fn reasoning_constructors_and_accessors_work() {
1380        let single = Reasoning::new("think");
1381        assert_eq!(single.first_text(), Some("think"));
1382        assert_eq!(single.first_signature(), None);
1383
1384        let signed = Reasoning::new_with_signature("signed", Some("sig-1".to_string()));
1385        assert_eq!(signed.first_text(), Some("signed"));
1386        assert_eq!(signed.first_signature(), Some("sig-1"));
1387
1388        let multi = Reasoning::multi(vec!["a".to_string(), "b".to_string()]);
1389        assert_eq!(multi.display_text(), "a\nb");
1390        assert_eq!(multi.first_text(), Some("a"));
1391
1392        let redacted = Reasoning::redacted("redacted-value");
1393        assert_eq!(redacted.display_text(), "redacted-value");
1394        assert_eq!(redacted.first_text(), None);
1395
1396        let encrypted = Reasoning::encrypted("enc");
1397        assert_eq!(encrypted.encrypted_content(), Some("enc"));
1398        assert_eq!(encrypted.display_text(), "");
1399
1400        let summaries = Reasoning::summaries(vec!["s1".to_string(), "s2".to_string()]);
1401        assert_eq!(summaries.display_text(), "s1\ns2");
1402        assert_eq!(summaries.encrypted_content(), None);
1403    }
1404
1405    #[test]
1406    fn reasoning_content_serde_roundtrip() {
1407        let variants = vec![
1408            ReasoningContent::Text {
1409                text: "plain".to_string(),
1410                signature: Some("sig".to_string()),
1411            },
1412            ReasoningContent::Encrypted("opaque".to_string()),
1413            ReasoningContent::Redacted {
1414                data: "redacted".to_string(),
1415            },
1416            ReasoningContent::Summary("summary".to_string()),
1417        ];
1418
1419        for variant in variants {
1420            let json = serde_json::to_string(&variant).expect("serialize");
1421            let roundtrip: ReasoningContent = serde_json::from_str(&json).expect("deserialize");
1422            assert_eq!(roundtrip, variant);
1423        }
1424    }
1425
1426    #[test]
1427    fn system_message_constructor_and_serde_roundtrip() {
1428        let message = Message::system("You are concise.");
1429
1430        match &message {
1431            Message::System { content } => assert_eq!(content, "You are concise."),
1432            _ => panic!("Expected system message"),
1433        }
1434
1435        let json = serde_json::to_string(&message).expect("serialize");
1436        let roundtrip: Message = serde_json::from_str(&json).expect("deserialize");
1437        assert_eq!(roundtrip, message);
1438    }
1439}
rig_core/completion/message.rs

rig_core/completion/
message.rs
