1use serde::{Deserialize, Serialize};
2use std::fmt;
3
4use crate::message::Message;
5
6#[derive(Clone, Debug, Serialize, Deserialize)]
7pub struct Choice {
8 index: u64,
9 pub message: Message,
10 pub finish_reason: String,
11}
12
13#[derive(Debug, Serialize, Deserialize)]
14pub struct Usage {
15 prompt_tokens: u32,
16 completion_tokens: u32,
17 total_tokens: u32,
18}
19
20#[derive(Debug, Serialize, Deserialize)]
21pub struct ChatResponse {
22 id: String,
23 object: String,
24 created: u64,
25 pub choices: Vec<Choice>,
26 usage: Usage,
27}
28
29impl ChatResponse {
30 pub fn content(&self) -> Option<String> {
31 match self.choices.first() {
32 Some(choice) => {
33 if let Some(c) = choice.message.content.clone() {
34 Some(c)
35 } else {
36 None
37 }
38 }
39 None => None,
40 }
41 }
42
43 pub fn function_call(&self) -> Option<(String, String)> {
44 match self.choices.first() {
45 Some(choice) => {
46 if let Some(f) = choice.message.function_call.clone() {
47 Some((f.name, f.arguments))
48 } else {
49 None
50 }
51 }
52 None => None,
53 }
54 }
55
56 pub fn message(&self) -> Option<Message> {
59 match self.choices.first() {
60 Some(choice) => Some(choice.message.clone()),
61 None => None,
62 }
63 }
64}
65
66impl fmt::Display for Choice {
67 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
68 write!(
69 f,
70 "{{\"index\":{},\"message\":{},\"finish_reason\":\"{}\"}}",
71 self.index, self.message, self.finish_reason
72 )
73 }
74}
75
76impl fmt::Display for ChatResponse {
77 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
78 write!(
79 f,
80 "{{\"id\":\"{}\",\"object\":\"{}\",\"created\":{},\"choices\":[",
81 self.id, self.object, self.created
82 )?;
83 for (i, choice) in self.choices.iter().enumerate() {
84 write!(
85 f,
86 "{}{}",
87 choice,
88 if i == self.choices.len() - 1 { "" } else { "," }
89 )?;
90 }
91 write!(f, "],\"usage\":{}}}", self.usage)
92 }
93}
94
95impl fmt::Display for Usage {
96 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
97 write!(
98 f,
99 "{{\"prompt_tokens\":{},\"completion_tokens\":{},\"total_tokens\":{}}}",
100 self.prompt_tokens, self.completion_tokens, self.total_tokens
101 )
102 }
103}
104
105#[cfg(test)]
106mod tests {
107 use super::*;
108 use crate::message::{FunctionCall, MessageBuilder};
109
110 #[test]
111 fn test_last_content() {
112 let message = MessageBuilder::new()
113 .content("content".to_string())
114 .build()
115 .unwrap();
116 let chat_response = ChatResponse {
117 id: "id".to_string(),
118 object: "object".to_string(),
119 created: 0,
120 choices: vec![Choice {
121 index: 0,
122 message: message.clone(),
123 finish_reason: "finish_reason".to_string(),
124 }],
125 usage: Usage {
126 prompt_tokens: 0,
127 completion_tokens: 0,
128 total_tokens: 0,
129 },
130 };
131 assert_eq!(chat_response.content(), Some("content".to_string()));
132 }
133
134 #[test]
135 fn test_last_function_call() {
136 let message = MessageBuilder::new()
137 .role("role".to_string())
138 .content("content".to_string())
139 .name("name".to_string())
140 .function_call(FunctionCall {
141 name: "name".to_string(),
142 arguments: "{\"example\":\"this\"}".to_string(),
143 })
144 .build()
145 .expect("Failed to build message");
146
147 let chat_response = ChatResponse {
148 id: "id".to_string(),
149 object: "object".to_string(),
150 created: 0,
151 choices: vec![Choice {
152 index: 0,
153 message: message.clone(),
154 finish_reason: "finish_reason".to_string(),
155 }],
156 usage: Usage {
157 prompt_tokens: 0,
158 completion_tokens: 0,
159 total_tokens: 0,
160 },
161 };
162 assert_eq!(
163 chat_response.function_call(),
164 Some(("name".to_string(), "{\"example\":\"this\"}".to_string()))
165 );
166 }
167
168 #[test]
169 fn test_message() {
170 let message = MessageBuilder::new()
171 .role("role".to_string())
172 .content("content".to_string())
173 .name("name".to_string())
174 .function_call(FunctionCall {
175 name: "name".to_string(),
176 arguments: "{\"example\":\"this\"}".to_string(),
177 })
178 .build()
179 .expect("Failed to build message");
180
181 let chat_response = ChatResponse {
182 id: "id".to_string(),
183 object: "object".to_string(),
184 created: 0,
185 choices: vec![Choice {
186 index: 0,
187 message: message.clone(),
188 finish_reason: "finish_reason".to_string(),
189 }],
190 usage: Usage {
191 prompt_tokens: 0,
192 completion_tokens: 0,
193 total_tokens: 0,
194 },
195 };
196 assert_eq!(chat_response.message(), Some(message),);
197 }
198
199 #[test]
200 fn test_display_for_choice() {
201 let message = MessageBuilder::new()
202 .role("role".to_string())
203 .content("content".to_string())
204 .name("name".to_string())
205 .function_call(FunctionCall {
206 name: "name".to_string(),
207 arguments: "{\"example\":\"this\"}".to_string(),
208 })
209 .build()
210 .expect("Failed to build message");
211 let choice = Choice {
212 index: 0,
213 message: message.clone(),
214 finish_reason: "finish_reason".to_string(),
215 };
216 assert_eq!(
217 format!("{}", choice),
218 "{\"index\":0,\"message\":{\"role\":\"role\",\"content\":\"content\",\"name\":\"name\",\"function_call\":{\"name\":\"name\",\"arguments\":\"{\\\"example\\\":\\\"this\\\"}\"}},\"finish_reason\":\"finish_reason\"}"
219 );
220 }
221
222 #[test]
223 fn test_display_chat_response() {
224 let message = MessageBuilder::new()
225 .role("role".to_string())
226 .content("content".to_string())
227 .name("name".to_string())
228 .function_call(FunctionCall {
229 name: "name".to_string(),
230 arguments: "{\"example\":\"this\"}".to_string(),
231 })
232 .build()
233 .expect("Failed to build message");
234 let chat_response = ChatResponse {
235 id: "id".to_string(),
236 object: "object".to_string(),
237 created: 0,
238 choices: vec![Choice {
239 index: 0,
240 message: message.clone(),
241 finish_reason: "finish_reason".to_string(),
242 }],
243 usage: Usage {
244 prompt_tokens: 0,
245 completion_tokens: 0,
246 total_tokens: 0,
247 },
248 };
249 assert_eq!(
250 format!("{}", chat_response),
251 "{\"id\":\"id\",\"object\":\"object\",\"created\":0,\"choices\":[{\"index\":0,\"message\":{\"role\":\"role\",\"content\":\"content\",\"name\":\"name\",\"function_call\":{\"name\":\"name\",\"arguments\":\"{\\\"example\\\":\\\"this\\\"}\"}},\"finish_reason\":\"finish_reason\"}],\"usage\":{\"prompt_tokens\":0,\"completion_tokens\":0,\"total_tokens\":0}}"
252 );
253 }
254
255 #[test]
256 fn test_display_usage() {
257 let usage = Usage {
258 prompt_tokens: 0,
259 completion_tokens: 0,
260 total_tokens: 0,
261 };
262 assert_eq!(
263 format!("{}", usage),
264 "{\"prompt_tokens\":0,\"completion_tokens\":0,\"total_tokens\":0}"
265 );
266 }
267}