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}");
}
let output = ChatCompletionsOutput { text, tool_calls };
let output = ChatCompletionsOutput { text, tool_calls, ..Default::default() };
Ok(output)
}
+43 -7
View File
@@ -168,12 +168,22 @@ pub async fn claude_chat_completions_streaming(
let mut function_arguments = String::new();
let mut function_id = String::new();
let mut reasoning_state = 0;
let mut thinking_text = String::new();
let mut thinking_signature = String::new();
let handle = |message: SseMessage| -> Result<bool> {
let data: Value = serde_json::from_str(&message.data)?;
debug!("stream-data: {data}");
if let Some(typ) = data["type"].as_str() {
match typ {
"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)) = (
data["content_block"]["type"].as_str(),
data["content_block"]["name"].as_str(),
@@ -206,7 +216,10 @@ pub async fn claude_chat_completions_streaming(
handler.text("<think>\n")?;
reasoning_state = 1;
}
thinking_text.push_str(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)) = (
!function_name.is_empty(),
data["delta"]["partial_json"].as_str(),
@@ -218,6 +231,10 @@ pub async fn claude_chat_completions_streaming(
if reasoning_state == 1 {
handler.text("\n</think>\n\n")?;
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() {
let arguments: Value = if function_arguments.is_empty() {
@@ -313,19 +330,25 @@ pub fn claude_build_chat_completions_body(
}) => {
let mut assistant_parts = vec![];
let mut user_parts = vec![];
if !text.is_empty() {
assistant_parts.push(json!({
"type": "text",
"text": text,
}))
for (index, tool_result) in tool_results.iter().enumerate() {
for block in &tool_result.thinking {
assistant_parts.push(json!(block));
}
for tool_result in tool_results {
if let Some(round_text) = &tool_result.text {
let round_text = if index == 0 && !text.is_empty() {
Some(text.as_str())
} else {
tool_result.text.as_deref()
};
if let Some(round_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!({
"type": "tool_use",
"id": tool_result.call.id,
@@ -409,12 +432,24 @@ pub fn claude_extract_chat_completions(data: &Value) -> Result<ChatCompletionsOu
let mut text = String::new();
let mut reasoning = None;
let mut tool_calls = vec![];
let mut thinking = vec![];
if let Some(list) = data["content"].as_array() {
for item in list {
match item["type"].as_str() {
Some("thinking") => {
if let Some(v) = item["thinking"].as_str() {
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") => {
@@ -453,6 +488,7 @@ pub fn claude_extract_chat_completions(data: &Value) -> Result<ChatCompletionsOu
let output = ChatCompletionsOutput {
text: text.to_string(),
tool_calls,
thinking,
};
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() {
bail!("Invalid response data: {data}");
}
let output = ChatCompletionsOutput { text, tool_calls };
let output = ChatCompletionsOutput { text, tool_calls, ..Default::default() };
Ok(output)
}
+11 -3
View File
@@ -295,6 +295,7 @@ pub struct ChatCompletionsData {
pub struct ChatCompletionsOutput {
pub text: String,
pub tool_calls: Vec<ToolCall>,
pub thinking: Vec<ThinkingBlock>,
}
impl ChatCompletionsOutput {
@@ -435,6 +436,7 @@ pub async fn call_chat_completions(
let ChatCompletionsOutput {
mut text,
tool_calls,
thinking,
..
} = ret;
if !text.is_empty() {
@@ -445,7 +447,10 @@ pub async fn call_chat_completions(
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
.iter()
.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?;
let (text, tool_calls) = handler.take();
let (text, tool_calls, thinking) = handler.take();
match send_ret {
Ok(_) => {
if !text.is_empty() && !text.ends_with('\n') {
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
.iter()
.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,
}
/// 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)]
pub struct MessageContentToolCalls {
pub tool_results: Vec<ToolResult>,
+2 -2
View File
@@ -533,7 +533,7 @@ pub fn openai_extract_chat_completions(data: &Value) -> Result<ChatCompletionsOu
} else {
text.to_string()
};
let output = ChatCompletionsOutput { text, tool_calls };
let output = ChatCompletionsOutput { text, tool_calls, ..Default::default() };
Ok(output)
}
@@ -696,7 +696,7 @@ pub fn openai_extract_responses(data: &Value) -> Result<ChatCompletionsOutput> {
if text.is_empty() && tool_calls.is_empty() {
bail!("Invalid response data: {data}");
}
Ok(ChatCompletionsOutput { text, tool_calls })
Ok(ChatCompletionsOutput { text, tool_calls, ..Default::default() })
}
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 anyhow::{Context, Result, anyhow, bail};
@@ -13,6 +13,7 @@ pub struct SseHandler {
abort_signal: AbortSignal,
buffer: String,
tool_calls: Vec<ToolCall>,
thinking: Vec<ThinkingBlock>,
last_tool_calls: Vec<ToolCall>,
max_call_repeats: usize,
call_repeat_chain_len: usize,
@@ -26,6 +27,7 @@ impl SseHandler {
abort_signal,
buffer: String::new(),
tool_calls: Vec::new(),
thinking: Vec::new(),
last_tool_calls: Vec::new(),
max_call_repeats: 2,
call_repeat_chain_len: 3,
@@ -170,6 +172,10 @@ impl SseHandler {
message
}
pub fn thinking_block(&mut self, block: ThinkingBlock) {
self.thinking.push(block);
}
pub fn abort(&self) -> AbortSignal {
self.abort_signal.clone()
}
@@ -179,11 +185,14 @@ impl SseHandler {
&self.last_tool_calls
}
pub fn take(self) -> (String, Vec<ToolCall>) {
pub fn take(self) -> (String, Vec<ToolCall>, Vec<ThinkingBlock>) {
let Self {
buffer, tool_calls, ..
buffer,
tool_calls,
thinking,
..
} = 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}");
}
}
let output = ChatCompletionsOutput { text, tool_calls };
let output = ChatCompletionsOutput { text, tool_calls, ..Default::default() };
Ok(output)
}
+33
View File
@@ -5,6 +5,7 @@ pub(crate) mod todo;
pub(crate) mod user_interaction;
use crate::{
client::ThinkingBlock,
config::{Agent, RequestContext},
graph,
utils::*,
@@ -205,6 +206,8 @@ pub struct ToolResult {
pub output: Value,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub text: Option<String>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub thinking: Vec<ThinkingBlock>,
}
impl ToolResult {
@@ -213,6 +216,7 @@ impl ToolResult {
call,
output,
text: None,
thinking: vec![],
}
}
}
@@ -1914,4 +1918,33 @@ mod tests {
assert_eq!(result.call.name, "my_tool");
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());
}
}