rig_core/completion/
message.rs

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

rig_core/completion/
message.rs
