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 _ => { 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}