feat: Improved support for Anthropic's extended thinking
This commit is contained in:
@@ -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)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
+43
-7
@@ -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,19 +330,25 @@ 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,
|
|
||||||
}))
|
|
||||||
}
|
}
|
||||||
for tool_result in tool_results {
|
let round_text = if index == 0 && !text.is_empty() {
|
||||||
if let Some(round_text) = &tool_result.text {
|
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!({
|
assistant_parts.push(json!({
|
||||||
"type": "text",
|
"type": "text",
|
||||||
"text": round_text,
|
"text": round_text,
|
||||||
}))
|
}))
|
||||||
}
|
}
|
||||||
|
}
|
||||||
assistant_parts.push(json!({
|
assistant_parts.push(json!({
|
||||||
"type": "tool_use",
|
"type": "tool_use",
|
||||||
"id": tool_result.call.id,
|
"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 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)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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
@@ -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()));
|
||||||
|
|||||||
@@ -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>,
|
||||||
|
|||||||
@@ -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
@@ -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)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -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)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -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());
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user