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

rig/completion/
message.rs
