feat: Improved support for Anthropic's extended thinking

This commit is contained in:
2026-07-17 17:26:41 -06:00
parent 078e6e3744
commit 8629c1ca15
9 changed files with 121 additions and 24 deletions
+1 -1
View File
@@ -529,7 +529,7 @@ fn extract_chat_completions(data: &Value) -> Result<ChatCompletionsOutput> {
bail!("Invalid response data: {data}"); bail!("Invalid response data: {data}");
} }
let output = ChatCompletionsOutput { text, tool_calls }; let output = ChatCompletionsOutput { text, tool_calls, ..Default::default() };
Ok(output) Ok(output)
} }
+48 -12
View File
@@ -168,12 +168,22 @@ pub async fn claude_chat_completions_streaming(
let mut function_arguments = String::new(); let mut function_arguments = String::new();
let mut function_id = String::new(); let mut function_id = String::new();
let mut reasoning_state = 0; let mut reasoning_state = 0;
let mut thinking_text = String::new();
let mut thinking_signature = String::new();
let handle = |message: SseMessage| -> Result<bool> { let handle = |message: SseMessage| -> Result<bool> {
let data: Value = serde_json::from_str(&message.data)?; let data: Value = serde_json::from_str(&message.data)?;
debug!("stream-data: {data}"); debug!("stream-data: {data}");
if let Some(typ) = data["type"].as_str() { if let Some(typ) = data["type"].as_str() {
match typ { match typ {
"content_block_start" => { "content_block_start" => {
if let (Some("redacted_thinking"), Some(redacted_data)) = (
data["content_block"]["type"].as_str(),
data["content_block"]["data"].as_str(),
) {
handler.thinking_block(ThinkingBlock::RedactedThinking {
data: redacted_data.to_string(),
});
}
if let (Some("tool_use"), Some(name), Some(id)) = ( if let (Some("tool_use"), Some(name), Some(id)) = (
data["content_block"]["type"].as_str(), data["content_block"]["type"].as_str(),
data["content_block"]["name"].as_str(), data["content_block"]["name"].as_str(),
@@ -206,7 +216,10 @@ pub async fn claude_chat_completions_streaming(
handler.text("<think>\n")?; handler.text("<think>\n")?;
reasoning_state = 1; reasoning_state = 1;
} }
thinking_text.push_str(text);
handler.text(text)?; handler.text(text)?;
} else if let Some(signature) = data["delta"]["signature"].as_str() {
thinking_signature.push_str(signature);
} else if let (true, Some(partial_json)) = ( } else if let (true, Some(partial_json)) = (
!function_name.is_empty(), !function_name.is_empty(),
data["delta"]["partial_json"].as_str(), data["delta"]["partial_json"].as_str(),
@@ -218,6 +231,10 @@ pub async fn claude_chat_completions_streaming(
if reasoning_state == 1 { if reasoning_state == 1 {
handler.text("\n</think>\n\n")?; handler.text("\n</think>\n\n")?;
reasoning_state = 0; reasoning_state = 0;
handler.thinking_block(ThinkingBlock::Thinking {
thinking: std::mem::take(&mut thinking_text),
signature: std::mem::take(&mut thinking_signature),
});
} }
if !function_name.is_empty() { if !function_name.is_empty() {
let arguments: Value = if function_arguments.is_empty() { let arguments: Value = if function_arguments.is_empty() {
@@ -313,18 +330,24 @@ pub fn claude_build_chat_completions_body(
}) => { }) => {
let mut assistant_parts = vec![]; let mut assistant_parts = vec![];
let mut user_parts = vec![]; let mut user_parts = vec![];
if !text.is_empty() { for (index, tool_result) in tool_results.iter().enumerate() {
assistant_parts.push(json!({ for block in &tool_result.thinking {
"type": "text", assistant_parts.push(json!(block));
"text": text, }
})) let round_text = if index == 0 && !text.is_empty() {
} Some(text.as_str())
for tool_result in tool_results { } else {
if let Some(round_text) = &tool_result.text { tool_result.text.as_deref()
assistant_parts.push(json!({ };
"type": "text", if let Some(round_text) = round_text {
"text": round_text, let round_text = strip_think_tag(round_text);
})) let round_text = round_text.trim();
if !round_text.is_empty() {
assistant_parts.push(json!({
"type": "text",
"text": round_text,
}))
}
} }
assistant_parts.push(json!({ assistant_parts.push(json!({
"type": "tool_use", "type": "tool_use",
@@ -409,12 +432,24 @@ pub fn claude_extract_chat_completions(data: &Value) -> Result<ChatCompletionsOu
let mut text = String::new(); let mut text = String::new();
let mut reasoning = None; let mut reasoning = None;
let mut tool_calls = vec![]; let mut tool_calls = vec![];
let mut thinking = vec![];
if let Some(list) = data["content"].as_array() { if let Some(list) = data["content"].as_array() {
for item in list { for item in list {
match item["type"].as_str() { match item["type"].as_str() {
Some("thinking") => { Some("thinking") => {
if let Some(v) = item["thinking"].as_str() { if let Some(v) = item["thinking"].as_str() {
reasoning = Some(v.to_string()); reasoning = Some(v.to_string());
thinking.push(ThinkingBlock::Thinking {
thinking: v.to_string(),
signature: item["signature"].as_str().unwrap_or_default().to_string(),
});
}
}
Some("redacted_thinking") => {
if let Some(v) = item["data"].as_str() {
thinking.push(ThinkingBlock::RedactedThinking {
data: v.to_string(),
});
} }
} }
Some("text") => { Some("text") => {
@@ -453,6 +488,7 @@ pub fn claude_extract_chat_completions(data: &Value) -> Result<ChatCompletionsOu
let output = ChatCompletionsOutput { let output = ChatCompletionsOutput {
text: text.to_string(), text: text.to_string(),
tool_calls, tool_calls,
thinking,
}; };
Ok(output) Ok(output)
} }
+1 -1
View File
@@ -244,6 +244,6 @@ fn extract_chat_completions(data: &Value) -> Result<ChatCompletionsOutput> {
if text.is_empty() && tool_calls.is_empty() { if text.is_empty() && tool_calls.is_empty() {
bail!("Invalid response data: {data}"); bail!("Invalid response data: {data}");
} }
let output = ChatCompletionsOutput { text, tool_calls }; let output = ChatCompletionsOutput { text, tool_calls, ..Default::default() };
Ok(output) Ok(output)
} }
+11 -3
View File
@@ -295,6 +295,7 @@ pub struct ChatCompletionsData {
pub struct ChatCompletionsOutput { pub struct ChatCompletionsOutput {
pub text: String, pub text: String,
pub tool_calls: Vec<ToolCall>, pub tool_calls: Vec<ToolCall>,
pub thinking: Vec<ThinkingBlock>,
} }
impl ChatCompletionsOutput { impl ChatCompletionsOutput {
@@ -435,6 +436,7 @@ pub async fn call_chat_completions(
let ChatCompletionsOutput { let ChatCompletionsOutput {
mut text, mut text,
tool_calls, tool_calls,
thinking,
.. ..
} = ret; } = ret;
if !text.is_empty() { if !text.is_empty() {
@@ -445,7 +447,10 @@ pub async fn call_chat_completions(
ctx.app.config.print_markdown(&text)?; ctx.app.config.print_markdown(&text)?;
} }
} }
let tool_results = eval_tool_calls(ctx, tool_calls).await?; let mut tool_results = eval_tool_calls(ctx, tool_calls).await?;
if let Some(first) = tool_results.first_mut() {
first.thinking = thinking;
}
tool_results tool_results
.iter() .iter()
.for_each(|res| ctx.tool_scope.tool_tracker.record_call(res.call.clone())); .for_each(|res| ctx.tool_scope.tool_tracker.record_call(res.call.clone()));
@@ -479,13 +484,16 @@ pub async fn call_chat_completions_streaming(
render_ret?; render_ret?;
let (text, tool_calls) = handler.take(); let (text, tool_calls, thinking) = handler.take();
match send_ret { match send_ret {
Ok(_) => { Ok(_) => {
if !text.is_empty() && !text.ends_with('\n') { if !text.is_empty() && !text.ends_with('\n') {
println!(); println!();
} }
let tool_results = eval_tool_calls(ctx, tool_calls).await?; let mut tool_results = eval_tool_calls(ctx, tool_calls).await?;
if let Some(first) = tool_results.first_mut() {
first.thinking = thinking;
}
tool_results tool_results
.iter() .iter()
.for_each(|res| ctx.tool_scope.tool_tracker.record_call(res.call.clone())); .for_each(|res| ctx.tool_scope.tool_tracker.record_call(res.call.clone()));
+11
View File
@@ -188,6 +188,17 @@ pub struct ImageUrl {
pub url: String, pub url: String,
} }
/// An extended-thinking block returned by Anthropic-protocol models.
/// Serialized to match the API wire format (`type: thinking` / `type: redacted_thinking`)
/// so blocks can be replayed verbatim, signature intact, in subsequent
/// tool-loop rounds as the API requires.
#[derive(Debug, Clone, Deserialize, Serialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ThinkingBlock {
Thinking { thinking: String, signature: String },
RedactedThinking { data: String },
}
#[derive(Debug, Clone, Deserialize, Serialize)] #[derive(Debug, Clone, Deserialize, Serialize)]
pub struct MessageContentToolCalls { pub struct MessageContentToolCalls {
pub tool_results: Vec<ToolResult>, pub tool_results: Vec<ToolResult>,
+2 -2
View File
@@ -533,7 +533,7 @@ pub fn openai_extract_chat_completions(data: &Value) -> Result<ChatCompletionsOu
} else { } else {
text.to_string() text.to_string()
}; };
let output = ChatCompletionsOutput { text, tool_calls }; let output = ChatCompletionsOutput { text, tool_calls, ..Default::default() };
Ok(output) Ok(output)
} }
@@ -696,7 +696,7 @@ pub fn openai_extract_responses(data: &Value) -> Result<ChatCompletionsOutput> {
if text.is_empty() && tool_calls.is_empty() { if text.is_empty() && tool_calls.is_empty() {
bail!("Invalid response data: {data}"); bail!("Invalid response data: {data}");
} }
Ok(ChatCompletionsOutput { text, tool_calls }) Ok(ChatCompletionsOutput { text, tool_calls, ..Default::default() })
} }
pub async fn openai_responses_streaming( pub async fn openai_responses_streaming(
+13 -4
View File
@@ -1,4 +1,4 @@
use super::{ToolCall, catch_error}; use super::{ThinkingBlock, ToolCall, catch_error};
use crate::utils::AbortSignal; use crate::utils::AbortSignal;
use anyhow::{Context, Result, anyhow, bail}; use anyhow::{Context, Result, anyhow, bail};
@@ -13,6 +13,7 @@ pub struct SseHandler {
abort_signal: AbortSignal, abort_signal: AbortSignal,
buffer: String, buffer: String,
tool_calls: Vec<ToolCall>, tool_calls: Vec<ToolCall>,
thinking: Vec<ThinkingBlock>,
last_tool_calls: Vec<ToolCall>, last_tool_calls: Vec<ToolCall>,
max_call_repeats: usize, max_call_repeats: usize,
call_repeat_chain_len: usize, call_repeat_chain_len: usize,
@@ -26,6 +27,7 @@ impl SseHandler {
abort_signal, abort_signal,
buffer: String::new(), buffer: String::new(),
tool_calls: Vec::new(), tool_calls: Vec::new(),
thinking: Vec::new(),
last_tool_calls: Vec::new(), last_tool_calls: Vec::new(),
max_call_repeats: 2, max_call_repeats: 2,
call_repeat_chain_len: 3, call_repeat_chain_len: 3,
@@ -170,6 +172,10 @@ impl SseHandler {
message message
} }
pub fn thinking_block(&mut self, block: ThinkingBlock) {
self.thinking.push(block);
}
pub fn abort(&self) -> AbortSignal { pub fn abort(&self) -> AbortSignal {
self.abort_signal.clone() self.abort_signal.clone()
} }
@@ -179,11 +185,14 @@ impl SseHandler {
&self.last_tool_calls &self.last_tool_calls
} }
pub fn take(self) -> (String, Vec<ToolCall>) { pub fn take(self) -> (String, Vec<ToolCall>, Vec<ThinkingBlock>) {
let Self { let Self {
buffer, tool_calls, .. buffer,
tool_calls,
thinking,
..
} = self; } = self;
(buffer, tool_calls) (buffer, tool_calls, thinking)
} }
} }
+1 -1
View File
@@ -322,7 +322,7 @@ fn gemini_extract_chat_completions_text(data: &Value) -> Result<ChatCompletionsO
bail!("Invalid response data: {data}"); bail!("Invalid response data: {data}");
} }
} }
let output = ChatCompletionsOutput { text, tool_calls }; let output = ChatCompletionsOutput { text, tool_calls, ..Default::default() };
Ok(output) Ok(output)
} }
+33
View File
@@ -5,6 +5,7 @@ pub(crate) mod todo;
pub(crate) mod user_interaction; pub(crate) mod user_interaction;
use crate::{ use crate::{
client::ThinkingBlock,
config::{Agent, RequestContext}, config::{Agent, RequestContext},
graph, graph,
utils::*, utils::*,
@@ -205,6 +206,8 @@ pub struct ToolResult {
pub output: Value, pub output: Value,
#[serde(default, skip_serializing_if = "Option::is_none")] #[serde(default, skip_serializing_if = "Option::is_none")]
pub text: Option<String>, pub text: Option<String>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub thinking: Vec<ThinkingBlock>,
} }
impl ToolResult { impl ToolResult {
@@ -213,6 +216,7 @@ impl ToolResult {
call, call,
output, output,
text: None, text: None,
thinking: vec![],
} }
} }
} }
@@ -1914,4 +1918,33 @@ mod tests {
assert_eq!(result.call.name, "my_tool"); assert_eq!(result.call.name, "my_tool");
assert_eq!(result.output, json!({"result": "ok"})); assert_eq!(result.output, json!({"result": "ok"}));
} }
#[test]
fn thinking_block_matches_anthropic_wire_format() {
let block = ThinkingBlock::Thinking {
thinking: "chain of thought".to_string(),
signature: "sig123".to_string(),
};
assert_eq!(
serde_json::to_value(&block).unwrap(),
json!({"type": "thinking", "thinking": "chain of thought", "signature": "sig123"})
);
let redacted = ThinkingBlock::RedactedThinking {
data: "opaque".to_string(),
};
assert_eq!(
serde_json::to_value(&redacted).unwrap(),
json!({"type": "redacted_thinking", "data": "opaque"})
);
}
#[test]
fn tool_result_deserializes_without_text_and_thinking() {
let yaml = "call:\n name: my_tool\n arguments: {}\noutput: ok\n";
let result: ToolResult = serde_yaml::from_str(yaml).unwrap();
assert_eq!(result.call.name, "my_tool");
assert!(result.text.is_none());
assert!(result.thinking.is_empty());
}
} }