Skip to main content

lellm_core/
response.rs

1//! 响应类型。
2
3use serde::{Deserialize, Serialize};
4
5use super::{ContentBlock, ToolCall};
6
7/// 统一的聊天响应。
8#[derive(Debug, Clone, Serialize, Deserialize)]
9pub struct ChatResponse {
10    /// 响应内容块列表,与 `Message::Assistant` 的 content 类型对齐。
11    pub content: Vec<ContentBlock>,
12    pub usage: TokenUsage,
13    pub raw: serde_json::Value,
14}
15
16impl ChatResponse {
17    /// 构造函数
18    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    /// 借用视图 — 零分配、零拷贝的 tool_call 迭代器
27    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    /// 是否存在 tool_calls
35    pub fn has_tool_calls(&self) -> bool {
36        self.tool_calls().next().is_some()
37    }
38}
39
40/// Token 消耗统计。
41#[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}