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

rig_core/completion/
message.rs
