llm/
assistant_reply.rs

1use schemars::JsonSchema;
2use serde::{Deserialize, Serialize};
3
4use crate::catalog::LlmModel;
5use crate::{LlmResponse, ToolCallError, ToolCallRequest, ToolCallResult, tool_result_parts, tool_result_request};
6
7#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
8#[serde(transparent)]
9pub struct AssistantReply(Vec<AssistantContent>);
10
11#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
12#[serde(tag = "type", rename_all = "camelCase")]
13pub enum AssistantContent {
14    Text {
15        text: String,
16    },
17    Reasoning {
18        #[serde(default, skip_serializing_if = "String::is_empty")]
19        summary: String,
20        #[serde(default, skip_serializing_if = "Option::is_none")]
21        encrypted: Option<EncryptedReasoningContent>,
22    },
23    ToolCall(ToolCallRequest),
24}
25
26#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
27pub struct EncryptedReasoningContent {
28    pub id: String,
29    pub content: String,
30}
31
32impl AssistantReply {
33    pub fn blocks(&self) -> &[AssistantContent] {
34        &self.0
35    }
36
37    pub fn is_empty(&self) -> bool {
38        self.0.is_empty()
39    }
40
41    pub fn text(&self) -> String {
42        self.0
43            .iter()
44            .filter_map(|block| match block {
45                AssistantContent::Text { text } => Some(text.as_str()),
46                AssistantContent::Reasoning { .. } | AssistantContent::ToolCall(_) => None,
47            })
48            .collect()
49    }
50
51    pub fn reasoning(&self) -> String {
52        self.0
53            .iter()
54            .filter_map(|block| match block {
55                AssistantContent::Reasoning { summary, .. } => Some(summary.as_str()),
56                AssistantContent::Text { .. } | AssistantContent::ToolCall(_) => None,
57            })
58            .collect()
59    }
60
61    pub fn tool_calls(&self) -> impl Iterator<Item = &ToolCallRequest> {
62        self.0.iter().filter_map(|block| match block {
63            AssistantContent::ToolCall(call) => Some(call),
64            AssistantContent::Text { .. } | AssistantContent::Reasoning { .. } => None,
65        })
66    }
67
68    pub fn push(&mut self, response: &LlmResponse) {
69        match response {
70            LlmResponse::Text { chunk } if !chunk.is_empty() => match self.0.last_mut() {
71                Some(AssistantContent::Text { text }) => text.push_str(chunk),
72                _ => self.0.push(AssistantContent::text(chunk.as_str())),
73            },
74            LlmResponse::Reasoning { chunk } if !chunk.is_empty() => match self.0.last_mut() {
75                Some(AssistantContent::Reasoning { summary, encrypted: None }) => summary.push_str(chunk),
76                _ => self.0.push(AssistantContent::Reasoning { summary: chunk.clone(), encrypted: None }),
77            },
78            LlmResponse::EncryptedReasoning { id, content } => {
79                let encrypted = EncryptedReasoningContent { id: id.clone(), content: content.clone() };
80                // Encrypted reasoning closes the reasoning block streamed just before it, or
81                // stands alone when the provider streamed no summary for it.
82                match self.0.last_mut() {
83                    Some(AssistantContent::Reasoning { encrypted: open @ None, .. }) => *open = Some(encrypted),
84                    _ => {
85                        self.0.push(AssistantContent::Reasoning { summary: String::new(), encrypted: Some(encrypted) });
86                    }
87                }
88            }
89            LlmResponse::ToolRequestComplete { tool_call } => {
90                self.0.push(AssistantContent::ToolCall(tool_call.clone()));
91            }
92            _ => {}
93        }
94    }
95
96    pub fn estimated_bytes(&self) -> usize {
97        self.0.iter().map(AssistantContent::estimated_bytes).sum()
98    }
99
100    pub(crate) fn pair_tool_calls(&mut self, results: &[Result<ToolCallResult, ToolCallError>]) {
101        let answered = |id: &str| results.iter().any(|result| tool_result_parts(result).0 == id);
102        self.0.retain(|block| !matches!(block, AssistantContent::ToolCall(call) if !answered(&call.id)));
103        for result in results {
104            let id = tool_result_parts(result).0;
105            if !self.tool_calls().any(|call| call.id == id) {
106                self.0.push(AssistantContent::ToolCall(tool_result_request(result)));
107            }
108        }
109    }
110}
111
112impl From<Vec<AssistantContent>> for AssistantReply {
113    fn from(blocks: Vec<AssistantContent>) -> Self {
114        Self(blocks)
115    }
116}
117
118impl FromIterator<AssistantContent> for AssistantReply {
119    fn from_iter<T: IntoIterator<Item = AssistantContent>>(blocks: T) -> Self {
120        Self(blocks.into_iter().collect())
121    }
122}
123
124impl AssistantContent {
125    pub fn text(text: impl Into<String>) -> Self {
126        Self::Text { text: text.into() }
127    }
128
129    fn estimated_bytes(&self) -> usize {
130        match self {
131            Self::Text { text } => text.len(),
132            Self::Reasoning { summary, encrypted } => {
133                encrypted.as_ref().map_or(summary.len(), |encrypted| encrypted.content.len())
134            }
135            Self::ToolCall(call) => call.name.len() + call.arguments.len(),
136        }
137    }
138}
139
140pub(crate) fn replays_encrypted_reasoning(author: Option<&LlmModel>, target: Option<&LlmModel>) -> bool {
141    target.is_some() && author == target
142}
143
144#[cfg(test)]
145mod tests {
146    use super::*;
147    use crate::testing::llm_response;
148
149    #[test]
150    fn push_keeps_each_reasoning_block_where_the_model_wrote_it() {
151        let responses = llm_response()
152            .reasoning(&["Pl", "an"])
153            .encrypted_reasoning("r_1", "first")
154            .text(&["Add", "ing"])
155            .reasoning(&["Check"])
156            .encrypted_reasoning("r_2", "second")
157            .encrypted_reasoning("r_3", "redacted")
158            .tool_call("call_1", "add", &["{}"])
159            .build();
160
161        let mut reply = AssistantReply::default();
162        for response in &responses {
163            reply.push(response);
164        }
165
166        let reasoning = |summary: &str, id: &str, content: &str| AssistantContent::Reasoning {
167            summary: summary.to_string(),
168            encrypted: Some(EncryptedReasoningContent { id: id.to_string(), content: content.to_string() }),
169        };
170        assert_eq!(
171            reply.blocks(),
172            [
173                reasoning("Plan", "r_1", "first"),
174                AssistantContent::text("Adding"),
175                reasoning("Check", "r_2", "second"),
176                reasoning("", "r_3", "redacted"),
177                AssistantContent::ToolCall(ToolCallRequest {
178                    id: "call_1".to_string(),
179                    name: "add".to_string(),
180                    arguments: "{}".to_string(),
181                }),
182            ]
183        );
184    }
185
186    #[test]
187    fn reply_roundtrips_through_json() {
188        let reply = AssistantReply::from(vec![
189            AssistantContent::Reasoning {
190                summary: "plan".to_string(),
191                encrypted: Some(EncryptedReasoningContent { id: "r_1".to_string(), content: "blob".to_string() }),
192            },
193            AssistantContent::text("answer"),
194            AssistantContent::ToolCall(ToolCallRequest {
195                id: "call_1".to_string(),
196                name: "search".to_string(),
197                arguments: "{}".to_string(),
198            }),
199        ]);
200
201        let json = serde_json::to_string(&reply).unwrap();
202
203        assert_eq!(serde_json::from_str::<AssistantReply>(&json).unwrap(), reply);
204    }
205}
llm/assistant_reply.rs

llm/
assistant_reply.rs
