1use std::{convert::Infallible, str::FromStr};
2
3use crate::OneOrMany;
4use serde::{Deserialize, Serialize};
5use thiserror::Error;
6
7use super::CompletionError;
8
9#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
24#[serde(tag = "role", rename_all = "lowercase")]
25pub enum Message {
26 User { content: OneOrMany<UserContent> },
28
29 Assistant {
31 id: Option<String>,
32 content: OneOrMany<AssistantContent>,
33 },
34}
35
36#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
40#[serde(tag = "type", rename_all = "lowercase")]
41pub enum UserContent {
42 Text(Text),
43 ToolResult(ToolResult),
44 Image(Image),
45 Audio(Audio),
46 Video(Video),
47 Document(Document),
48}
49
50#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
52#[serde(untagged)]
53pub enum AssistantContent {
54 Text(Text),
55 ToolCall(ToolCall),
56 Reasoning(Reasoning),
57}
58
59#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
60pub struct Reasoning {
61 pub id: Option<String>,
62 pub reasoning: Vec<String>,
63}
64
65impl Reasoning {
66 pub fn new(input: &str) -> Self {
68 Self {
69 id: None,
70 reasoning: vec![input.to_string()],
71 }
72 }
73
74 pub fn multi(input: Vec<String>) -> Self {
75 Self {
76 id: None,
77 reasoning: input,
78 }
79 }
80
81 pub fn optional_id(mut self, id: Option<String>) -> Self {
82 self.id = id;
83 self
84 }
85
86 pub fn with_id(mut self, id: String) -> Self {
87 self.id = Some(id);
88 self
89 }
90}
91
92#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
94pub struct ToolResult {
95 pub id: String,
96 #[serde(skip_serializing_if = "Option::is_none")]
97 pub call_id: Option<String>,
98 pub content: OneOrMany<ToolResultContent>,
99}
100
101#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
103#[serde(tag = "type", rename_all = "lowercase")]
104pub enum ToolResultContent {
105 Text(Text),
106 Image(Image),
107}
108
109#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
111pub struct ToolCall {
112 pub id: String,
113 pub call_id: Option<String>,
114 pub function: ToolFunction,
115}
116
117#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
119pub struct ToolFunction {
120 pub name: String,
121 pub arguments: serde_json::Value,
122}
123
124#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
130pub struct Text {
131 pub text: String,
132}
133
134#[derive(Default, Clone, Debug, Deserialize, Serialize, PartialEq)]
136pub struct Image {
137 pub data: String,
138 #[serde(skip_serializing_if = "Option::is_none")]
139 pub format: Option<ContentFormat>,
140 #[serde(skip_serializing_if = "Option::is_none")]
141 pub media_type: Option<ImageMediaType>,
142 #[serde(skip_serializing_if = "Option::is_none")]
143 pub detail: Option<ImageDetail>,
144 #[serde(flatten, skip_serializing_if = "Option::is_none")]
145 pub additional_params: Option<serde_json::Value>,
146}
147
148#[derive(Default, Clone, Debug, Deserialize, Serialize, PartialEq)]
150pub struct Audio {
151 pub data: String,
152 #[serde(skip_serializing_if = "Option::is_none")]
153 pub format: Option<ContentFormat>,
154 #[serde(skip_serializing_if = "Option::is_none")]
155 pub media_type: Option<AudioMediaType>,
156 #[serde(flatten, skip_serializing_if = "Option::is_none")]
157 pub additional_params: Option<serde_json::Value>,
158}
159
160#[derive(Default, Clone, Debug, Deserialize, Serialize, PartialEq)]
162pub struct Video {
163 pub data: String,
164 #[serde(skip_serializing_if = "Option::is_none")]
165 pub format: Option<ContentFormat>,
166 #[serde(skip_serializing_if = "Option::is_none")]
167 pub media_type: Option<VideoMediaType>,
168 #[serde(flatten, skip_serializing_if = "Option::is_none")]
169 pub additional_params: Option<serde_json::Value>,
170}
171
172#[derive(Default, Clone, Debug, Deserialize, Serialize, PartialEq)]
174pub struct Document {
175 pub data: String,
176 #[serde(skip_serializing_if = "Option::is_none")]
177 pub format: Option<ContentFormat>,
178 #[serde(skip_serializing_if = "Option::is_none")]
179 pub media_type: Option<DocumentMediaType>,
180 #[serde(flatten, skip_serializing_if = "Option::is_none")]
181 pub additional_params: Option<serde_json::Value>,
182}
183
184#[derive(Default, Clone, Debug, Deserialize, Serialize, PartialEq)]
186#[serde(rename_all = "lowercase")]
187pub enum ContentFormat {
188 #[default]
189 Base64,
190 String,
191}
192
193#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
195pub enum MediaType {
196 Image(ImageMediaType),
197 Audio(AudioMediaType),
198 Document(DocumentMediaType),
199 Video(VideoMediaType),
200}
201
202#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
205#[serde(rename_all = "lowercase")]
206pub enum ImageMediaType {
207 JPEG,
208 PNG,
209 GIF,
210 WEBP,
211 HEIC,
212 HEIF,
213 SVG,
214}
215
216#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
220#[serde(rename_all = "lowercase")]
221pub enum DocumentMediaType {
222 PDF,
223 TXT,
224 RTF,
225 HTML,
226 CSS,
227 MARKDOWN,
228 CSV,
229 XML,
230 Javascript,
231 Python,
232}
233
234#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
237#[serde(rename_all = "lowercase")]
238pub enum AudioMediaType {
239 WAV,
240 MP3,
241 AIFF,
242 AAC,
243 OGG,
244 FLAC,
245}
246
247#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
250#[serde(rename_all = "lowercase")]
251pub enum VideoMediaType {
252 AVI,
253 MP4,
254 MPEG,
255}
256
257#[derive(Default, Clone, Debug, Deserialize, Serialize, PartialEq)]
259#[serde(rename_all = "lowercase")]
260pub enum ImageDetail {
261 Low,
262 High,
263 #[default]
264 Auto,
265}
266
267impl Message {
272 pub(crate) fn rag_text(&self) -> Option<String> {
275 match self {
276 Message::User { content } => {
277 for item in content.iter() {
278 if let UserContent::Text(Text { text }) = item {
279 return Some(text.clone());
280 }
281 }
282 None
283 }
284 _ => None,
285 }
286 }
287
288 pub fn user(text: impl Into<String>) -> Self {
290 Message::User {
291 content: OneOrMany::one(UserContent::text(text)),
292 }
293 }
294
295 pub fn assistant(text: impl Into<String>) -> Self {
297 Message::Assistant {
298 id: None,
299 content: OneOrMany::one(AssistantContent::text(text)),
300 }
301 }
302
303 pub fn assistant_with_id(id: String, text: impl Into<String>) -> Self {
305 Message::Assistant {
306 id: Some(id),
307 content: OneOrMany::one(AssistantContent::text(text)),
308 }
309 }
310
311 pub fn tool_result(id: impl Into<String>, content: impl Into<String>) -> Self {
313 Message::User {
314 content: OneOrMany::one(UserContent::ToolResult(ToolResult {
315 id: id.into(),
316 call_id: None,
317 content: OneOrMany::one(ToolResultContent::text(content)),
318 })),
319 }
320 }
321
322 pub fn tool_result_with_call_id(
323 id: impl Into<String>,
324 call_id: Option<String>,
325 content: impl Into<String>,
326 ) -> Self {
327 Message::User {
328 content: OneOrMany::one(UserContent::ToolResult(ToolResult {
329 id: id.into(),
330 call_id,
331 content: OneOrMany::one(ToolResultContent::text(content)),
332 })),
333 }
334 }
335}
336
337impl UserContent {
338 pub fn text(text: impl Into<String>) -> Self {
340 UserContent::Text(text.into().into())
341 }
342
343 pub fn image(
345 data: impl Into<String>,
346 format: Option<ContentFormat>,
347 media_type: Option<ImageMediaType>,
348 detail: Option<ImageDetail>,
349 ) -> Self {
350 UserContent::Image(Image {
351 data: data.into(),
352 format,
353 media_type,
354 detail,
355 additional_params: None,
356 })
357 }
358
359 pub fn audio(
361 data: impl Into<String>,
362 format: Option<ContentFormat>,
363 media_type: Option<AudioMediaType>,
364 ) -> Self {
365 UserContent::Audio(Audio {
366 data: data.into(),
367 format,
368 media_type,
369 additional_params: None,
370 })
371 }
372
373 pub fn document(
375 data: impl Into<String>,
376 format: Option<ContentFormat>,
377 media_type: Option<DocumentMediaType>,
378 ) -> Self {
379 UserContent::Document(Document {
380 data: data.into(),
381 format,
382 media_type,
383 additional_params: None,
384 })
385 }
386
387 pub fn tool_result(id: impl Into<String>, content: OneOrMany<ToolResultContent>) -> Self {
389 UserContent::ToolResult(ToolResult {
390 id: id.into(),
391 call_id: None,
392 content,
393 })
394 }
395
396 pub fn tool_result_with_call_id(
398 id: impl Into<String>,
399 call_id: String,
400 content: OneOrMany<ToolResultContent>,
401 ) -> Self {
402 UserContent::ToolResult(ToolResult {
403 id: id.into(),
404 call_id: Some(call_id),
405 content,
406 })
407 }
408}
409
410impl AssistantContent {
411 pub fn text(text: impl Into<String>) -> Self {
413 AssistantContent::Text(text.into().into())
414 }
415
416 pub fn tool_call(
418 id: impl Into<String>,
419 name: impl Into<String>,
420 arguments: serde_json::Value,
421 ) -> Self {
422 AssistantContent::ToolCall(ToolCall {
423 id: id.into(),
424 call_id: None,
425 function: ToolFunction {
426 name: name.into(),
427 arguments,
428 },
429 })
430 }
431
432 pub fn tool_call_with_call_id(
433 id: impl Into<String>,
434 call_id: String,
435 name: impl Into<String>,
436 arguments: serde_json::Value,
437 ) -> Self {
438 AssistantContent::ToolCall(ToolCall {
439 id: id.into(),
440 call_id: Some(call_id),
441 function: ToolFunction {
442 name: name.into(),
443 arguments,
444 },
445 })
446 }
447}
448
449impl ToolResultContent {
450 pub fn text(text: impl Into<String>) -> Self {
452 ToolResultContent::Text(text.into().into())
453 }
454
455 pub fn image(
457 data: impl Into<String>,
458 format: Option<ContentFormat>,
459 media_type: Option<ImageMediaType>,
460 detail: Option<ImageDetail>,
461 ) -> Self {
462 ToolResultContent::Image(Image {
463 data: data.into(),
464 format,
465 media_type,
466 detail,
467 additional_params: None,
468 })
469 }
470}
471
472pub trait MimeType {
474 fn from_mime_type(mime_type: &str) -> Option<Self>
475 where
476 Self: Sized;
477 fn to_mime_type(&self) -> &'static str;
478}
479
480impl MimeType for MediaType {
481 fn from_mime_type(mime_type: &str) -> Option<Self> {
482 ImageMediaType::from_mime_type(mime_type)
483 .map(MediaType::Image)
484 .or_else(|| {
485 DocumentMediaType::from_mime_type(mime_type)
486 .map(MediaType::Document)
487 .or_else(|| AudioMediaType::from_mime_type(mime_type).map(MediaType::Audio))
488 })
489 }
490
491 fn to_mime_type(&self) -> &'static str {
492 match self {
493 MediaType::Image(media_type) => media_type.to_mime_type(),
494 MediaType::Audio(media_type) => media_type.to_mime_type(),
495 MediaType::Document(media_type) => media_type.to_mime_type(),
496 MediaType::Video(media_type) => media_type.to_mime_type(),
497 }
498 }
499}
500
501impl MimeType for ImageMediaType {
502 fn from_mime_type(mime_type: &str) -> Option<Self> {
503 match mime_type {
504 "image/jpeg" => Some(ImageMediaType::JPEG),
505 "image/png" => Some(ImageMediaType::PNG),
506 "image/gif" => Some(ImageMediaType::GIF),
507 "image/webp" => Some(ImageMediaType::WEBP),
508 "image/heic" => Some(ImageMediaType::HEIC),
509 "image/heif" => Some(ImageMediaType::HEIF),
510 "image/svg+xml" => Some(ImageMediaType::SVG),
511 _ => None,
512 }
513 }
514
515 fn to_mime_type(&self) -> &'static str {
516 match self {
517 ImageMediaType::JPEG => "image/jpeg",
518 ImageMediaType::PNG => "image/png",
519 ImageMediaType::GIF => "image/gif",
520 ImageMediaType::WEBP => "image/webp",
521 ImageMediaType::HEIC => "image/heic",
522 ImageMediaType::HEIF => "image/heif",
523 ImageMediaType::SVG => "image/svg+xml",
524 }
525 }
526}
527
528impl MimeType for DocumentMediaType {
529 fn from_mime_type(mime_type: &str) -> Option<Self> {
530 match mime_type {
531 "application/pdf" => Some(DocumentMediaType::PDF),
532 "text/plain" => Some(DocumentMediaType::TXT),
533 "text/rtf" => Some(DocumentMediaType::RTF),
534 "text/html" => Some(DocumentMediaType::HTML),
535 "text/css" => Some(DocumentMediaType::CSS),
536 "text/md" | "text/markdown" => Some(DocumentMediaType::MARKDOWN),
537 "text/csv" => Some(DocumentMediaType::CSV),
538 "text/xml" => Some(DocumentMediaType::XML),
539 "application/x-javascript" | "text/x-javascript" => Some(DocumentMediaType::Javascript),
540 "application/x-python" | "text/x-python" => Some(DocumentMediaType::Python),
541 _ => None,
542 }
543 }
544
545 fn to_mime_type(&self) -> &'static str {
546 match self {
547 DocumentMediaType::PDF => "application/pdf",
548 DocumentMediaType::TXT => "text/plain",
549 DocumentMediaType::RTF => "text/rtf",
550 DocumentMediaType::HTML => "text/html",
551 DocumentMediaType::CSS => "text/css",
552 DocumentMediaType::MARKDOWN => "text/markdown",
553 DocumentMediaType::CSV => "text/csv",
554 DocumentMediaType::XML => "text/xml",
555 DocumentMediaType::Javascript => "application/x-javascript",
556 DocumentMediaType::Python => "application/x-python",
557 }
558 }
559}
560
561impl MimeType for AudioMediaType {
562 fn from_mime_type(mime_type: &str) -> Option<Self> {
563 match mime_type {
564 "audio/wav" => Some(AudioMediaType::WAV),
565 "audio/mp3" => Some(AudioMediaType::MP3),
566 "audio/aiff" => Some(AudioMediaType::AIFF),
567 "audio/aac" => Some(AudioMediaType::AAC),
568 "audio/ogg" => Some(AudioMediaType::OGG),
569 "audio/flac" => Some(AudioMediaType::FLAC),
570 _ => None,
571 }
572 }
573
574 fn to_mime_type(&self) -> &'static str {
575 match self {
576 AudioMediaType::WAV => "audio/wav",
577 AudioMediaType::MP3 => "audio/mp3",
578 AudioMediaType::AIFF => "audio/aiff",
579 AudioMediaType::AAC => "audio/aac",
580 AudioMediaType::OGG => "audio/ogg",
581 AudioMediaType::FLAC => "audio/flac",
582 }
583 }
584}
585
586impl MimeType for VideoMediaType {
587 fn from_mime_type(mime_type: &str) -> Option<Self>
588 where
589 Self: Sized,
590 {
591 match mime_type {
592 "video/avi" => Some(VideoMediaType::AVI),
593 "video/mp4" => Some(VideoMediaType::MP4),
594 "video/mpeg" => Some(VideoMediaType::MPEG),
595 &_ => None,
596 }
597 }
598
599 fn to_mime_type(&self) -> &'static str {
600 match self {
601 VideoMediaType::AVI => "video/avi",
602 VideoMediaType::MP4 => "video/mp4",
603 VideoMediaType::MPEG => "video/mpeg",
604 }
605 }
606}
607
608impl std::str::FromStr for ImageDetail {
609 type Err = ();
610
611 fn from_str(s: &str) -> Result<Self, Self::Err> {
612 match s.to_lowercase().as_str() {
613 "low" => Ok(ImageDetail::Low),
614 "high" => Ok(ImageDetail::High),
615 "auto" => Ok(ImageDetail::Auto),
616 _ => Err(()),
617 }
618 }
619}
620
621impl From<String> for Text {
626 fn from(text: String) -> Self {
627 Text { text }
628 }
629}
630
631impl From<&String> for Text {
632 fn from(text: &String) -> Self {
633 text.to_owned().into()
634 }
635}
636
637impl From<&str> for Text {
638 fn from(text: &str) -> Self {
639 text.to_owned().into()
640 }
641}
642
643impl FromStr for Text {
644 type Err = Infallible;
645
646 fn from_str(s: &str) -> Result<Self, Self::Err> {
647 Ok(s.into())
648 }
649}
650
651impl From<String> for Message {
652 fn from(text: String) -> Self {
653 Message::User {
654 content: OneOrMany::one(UserContent::Text(text.into())),
655 }
656 }
657}
658
659impl From<&str> for Message {
660 fn from(text: &str) -> Self {
661 Message::User {
662 content: OneOrMany::one(UserContent::Text(text.into())),
663 }
664 }
665}
666
667impl From<&String> for Message {
668 fn from(text: &String) -> Self {
669 Message::User {
670 content: OneOrMany::one(UserContent::Text(text.into())),
671 }
672 }
673}
674
675impl From<Text> for Message {
676 fn from(text: Text) -> Self {
677 Message::User {
678 content: OneOrMany::one(UserContent::Text(text)),
679 }
680 }
681}
682
683impl From<Image> for Message {
684 fn from(image: Image) -> Self {
685 Message::User {
686 content: OneOrMany::one(UserContent::Image(image)),
687 }
688 }
689}
690
691impl From<Audio> for Message {
692 fn from(audio: Audio) -> Self {
693 Message::User {
694 content: OneOrMany::one(UserContent::Audio(audio)),
695 }
696 }
697}
698
699impl From<Document> for Message {
700 fn from(document: Document) -> Self {
701 Message::User {
702 content: OneOrMany::one(UserContent::Document(document)),
703 }
704 }
705}
706
707impl From<String> for ToolResultContent {
708 fn from(text: String) -> Self {
709 ToolResultContent::text(text)
710 }
711}
712
713impl From<String> for AssistantContent {
714 fn from(text: String) -> Self {
715 AssistantContent::text(text)
716 }
717}
718
719impl From<String> for UserContent {
720 fn from(text: String) -> Self {
721 UserContent::text(text)
722 }
723}
724
725impl From<AssistantContent> for Message {
726 fn from(content: AssistantContent) -> Self {
727 Message::Assistant {
728 id: None,
729 content: OneOrMany::one(content),
730 }
731 }
732}
733
734impl From<UserContent> for Message {
735 fn from(content: UserContent) -> Self {
736 Message::User {
737 content: OneOrMany::one(content),
738 }
739 }
740}
741
742impl From<OneOrMany<AssistantContent>> for Message {
743 fn from(content: OneOrMany<AssistantContent>) -> Self {
744 Message::Assistant { id: None, content }
745 }
746}
747
748impl From<OneOrMany<UserContent>> for Message {
749 fn from(content: OneOrMany<UserContent>) -> Self {
750 Message::User { content }
751 }
752}
753
754impl From<ToolCall> for Message {
755 fn from(tool_call: ToolCall) -> Self {
756 Message::Assistant {
757 id: None,
758 content: OneOrMany::one(AssistantContent::ToolCall(tool_call)),
759 }
760 }
761}
762
763impl From<ToolResult> for Message {
764 fn from(tool_result: ToolResult) -> Self {
765 Message::User {
766 content: OneOrMany::one(UserContent::ToolResult(tool_result)),
767 }
768 }
769}
770
771impl From<ToolResultContent> for Message {
772 fn from(tool_result_content: ToolResultContent) -> Self {
773 Message::User {
774 content: OneOrMany::one(UserContent::ToolResult(ToolResult {
775 id: String::new(),
776 call_id: None,
777 content: OneOrMany::one(tool_result_content),
778 })),
779 }
780 }
781}
782
783#[derive(Debug, Error)]
789pub enum MessageError {
790 #[error("Message conversion error: {0}")]
791 ConversionError(String),
792}
793
794impl From<MessageError> for CompletionError {
795 fn from(error: MessageError) -> Self {
796 CompletionError::RequestError(error.into())
797 }
798}