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

rig/completion/
message.rs
