rig/completion/
message.rs

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

rig/completion/
message.rs
