ragit_api/
message.rs

1use crate::api_provider::ApiProvider;
2use ragit_pdl::{
3    Message,
4    MessageContent,
5    encode_base64,
6};
7use serde_json::{Map, Value};
8
9pub fn message_content_to_json(message: &MessageContent, api_provider: &ApiProvider) -> Value {
10    match message {
11        MessageContent::String(s) => {
12            let mut content = Map::new();
13            content.insert(String::from("type"), "text".into());
14            content.insert(String::from("text"), s.to_string().into());
15
16            content.into()
17        },
18        MessageContent::Image { image_type, bytes } => match api_provider {
19            ApiProvider::Anthropic => {
20                let mut content = Map::new();
21                content.insert(String::from("type"), "image".into());
22
23                let mut source = Map::new();
24                source.insert(String::from("type"), "base64".into());
25                source.insert(String::from("media_type"), image_type.get_media_type().into());
26                source.insert(String::from("data"), encode_base64(bytes).into());
27
28                content.insert(String::from("source"), source.into());
29                content.into()
30            },
31            // TODO: cohere?
32            _ => {  // assume the others are all openai-compatible
33                let mut content = Map::new();
34                content.insert(String::from("type"), "image_url".into());
35
36                let mut image_url = Map::new();
37                image_url.insert(String::from("url"), format!("data:{};base64,{}", image_type.get_media_type(), encode_base64(bytes)).into());
38                content.insert(String::from("image_url"), image_url.into());
39                content.into()
40            },
41        },
42    }
43}
44
45pub fn message_contents_to_json_array(contents: &[MessageContent], api_provider: &ApiProvider) -> Value {
46    if contents.len() == 1 && contents[0].is_string() {
47        Value::String(contents[0].unwrap_str().into())
48    }
49
50    else {
51        Value::Array(contents.iter().map(
52            |content| message_content_to_json(content, api_provider)
53        ).collect())
54    }
55}
56
57pub fn message_to_json(message: &Message, api_provider: &ApiProvider) -> Value {
58    let mut result = Map::new();
59    result.insert(String::from("role"), message.role.to_api_string().into());
60
61    match (api_provider, message.content.len()) {
62        (_, 0) => panic!("a message without any content"),
63        (ApiProvider::Anthropic, 1) if matches!(&message.content[0], MessageContent::String(_)) => match &message.content[0] {
64            MessageContent::String(s) => {
65                result.insert(String::from("content"), s.to_string().into());
66            },
67            MessageContent::Image { .. } => unreachable!(),
68        },
69        (ApiProvider::Anthropic | ApiProvider::Cohere | ApiProvider::OpenAi { .. }, _) => {
70            result.insert(String::from("content"), message_contents_to_json_array(&message.content, api_provider));
71        },
72        (ApiProvider::Test(_), _) => unreachable!(),
73    }
74
75    result.into()
76}