1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
|
use crate::config::Input;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct Message {
pub role: MessageRole,
pub content: MessageContent,
}
impl Message {
pub fn new(input: &Input) -> Self {
Self {
role: MessageRole::User,
content: input.to_message_content(),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum MessageRole {
System,
Assistant,
User,
}
#[allow(dead_code)]
impl MessageRole {
pub fn is_system(&self) -> bool {
matches!(self, MessageRole::System)
}
pub fn is_user(&self) -> bool {
matches!(self, MessageRole::User)
}
pub fn is_assistant(&self) -> bool {
matches!(self, MessageRole::Assistant)
}
}
#[derive(Debug, Clone, Deserialize, Serialize)]
#[serde(untagged)]
pub enum MessageContent {
Text(String),
Array(Vec<MessageContentPart>),
}
impl MessageContent {
pub fn render_input(&self, resolve_url_fn: impl Fn(&str) -> String) -> String {
match self {
MessageContent::Text(text) => text.to_string(),
MessageContent::Array(list) => {
let (mut concated_text, mut files) = (String::new(), vec![]);
for item in list {
match item {
MessageContentPart::Text { text } => {
concated_text = format!("{concated_text} {text}")
}
MessageContentPart::ImageUrl { image_url } => {
files.push(resolve_url_fn(&image_url.url))
}
}
}
if !concated_text.is_empty() {
concated_text = format!(" -- {concated_text}")
}
format!(".file {}{}", files.join(" "), concated_text)
}
}
}
pub fn merge_prompt(&mut self, replace_fn: impl Fn(&str) -> String) {
match self {
MessageContent::Text(text) => *text = replace_fn(text),
MessageContent::Array(list) => {
if list.is_empty() {
list.push(MessageContentPart::Text {
text: replace_fn(""),
})
} else if let Some(MessageContentPart::Text { text }) = list.get_mut(0) {
*text = replace_fn(text)
}
}
}
}
}
#[derive(Debug, Clone, Deserialize, Serialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum MessageContentPart {
Text { text: String },
ImageUrl { image_url: ImageUrl },
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct ImageUrl {
pub url: String,
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::InputContext;
#[test]
fn test_serde() {
assert_eq!(
serde_json::to_string(&Message::new(&Input::from_str(
"Hello World",
InputContext::default()
)))
.unwrap(),
"{\"role\":\"user\",\"content\":\"Hello World\"}"
);
}
}
|
