1use serde::{Deserialize, Serialize};
4
5use super::{ContentBlock, ToolCall};
6
7#[derive(Debug, Clone, Serialize, Deserialize)]
9pub struct ChatResponse {
10 pub content: Vec<ContentBlock>,
12 pub usage: TokenUsage,
13 pub raw: serde_json::Value,
14}
15
16impl ChatResponse {
17 pub fn new(content: Vec<ContentBlock>, usage: TokenUsage, raw: serde_json::Value) -> Self {
19 Self {
20 content,
21 usage,
22 raw,
23 }
24 }
25
26 pub fn tool_calls(&self) -> impl Iterator<Item = &ToolCall> {
28 self.content.iter().filter_map(|block| match block {
29 ContentBlock::ToolCall(call) => Some(call),
30 _ => None,
31 })
32 }
33
34 pub fn has_tool_calls(&self) -> bool {
36 self.tool_calls().next().is_some()
37 }
38}
39
40#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize)]
42pub struct TokenUsage {
43 pub prompt_tokens: u32,
44 pub completion_tokens: u32,
45 pub total_tokens: u32,
46}
47
48#[cfg(test)]
49mod tests {
50 use super::*;
51 use crate::{ContentBlock, ToolCall};
52
53 #[test]
54 fn test_chat_response_tool_calls_iterator() {
55 let tc = ToolCall {
56 id: "1".into(),
57 name: "test".into(),
58 arguments: serde_json::json!({}),
59 };
60 let content = vec![
61 ContentBlock::text("hello".to_string()),
62 ContentBlock::ToolCall(tc.clone()),
63 ];
64 let resp = ChatResponse::new(content, TokenUsage::default(), serde_json::json!(null));
65 let calls: Vec<_> = resp.tool_calls().collect();
66 assert_eq!(calls.len(), 1);
67 assert_eq!(calls[0].id, "1");
68 }
69
70 #[test]
71 fn test_chat_response_has_tool_calls() {
72 let tc = ToolCall {
73 id: "1".into(),
74 name: "test".into(),
75 arguments: serde_json::json!({}),
76 };
77 let content = vec![
78 ContentBlock::text("hello".to_string()),
79 ContentBlock::ToolCall(tc),
80 ];
81 let resp = ChatResponse::new(content, TokenUsage::default(), serde_json::json!(null));
82 assert!(resp.has_tool_calls());
83 }
84
85 #[test]
86 fn test_chat_response_no_tool_calls() {
87 let content = vec![ContentBlock::text("hello".to_string())];
88 let resp = ChatResponse::new(content, TokenUsage::default(), serde_json::json!(null));
89 assert!(!resp.has_tool_calls());
90 }
91}