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

rig/completion/
message.rs
