From e799e80eb300e040901b067c3e871cbc1b6a1abc Mon Sep 17 00:00:00 2001 From: Salman Mohammed Date: Fri, 28 Feb 2025 13:38:40 -0500 Subject: [PATCH] feat: [anthropic] Claude 3.7 Sonnet with extended thinking (#1370) - only works with anthropic provider (wont work with databricks) --- crates/goose-cli/src/session/output.rs | 13 +- crates/goose/src/message.rs | 54 ++++++ crates/goose/src/providers/anthropic.rs | 24 ++- .../goose/src/providers/formats/anthropic.rs | 183 +++++++++++++++++- crates/goose/src/providers/formats/bedrock.rs | 8 + crates/goose/src/providers/formats/openai.rs | 8 + 6 files changed, 282 insertions(+), 8 deletions(-) diff --git a/crates/goose-cli/src/session/output.rs b/crates/goose-cli/src/session/output.rs index 66021695aa..dc1c7a05b5 100644 --- a/crates/goose-cli/src/session/output.rs +++ b/crates/goose-cli/src/session/output.rs @@ -107,8 +107,19 @@ pub fn render_message(message: &Message) { MessageContent::Image(image) => { println!("Image: [data: {}, type: {}]", image.data, image.mime_type); } + MessageContent::Thinking(thinking) => { + if std::env::var("GOOSE_CLI_SHOW_THINKING").is_ok() { + println!("\n{}", style("Thinking:").dim().italic()); + print_markdown(&thinking.thinking, theme); + } + } + MessageContent::RedactedThinking(_) => { + // For redacted thinking, print thinking was redacted + println!("\n{}", style("Thinking:").dim().italic()); + print_markdown("Thinking was redacted", theme); + } _ => { - println!("Message type could not be rendered"); + println!("WARNING: Message content type could not be rendered"); } } } diff --git a/crates/goose/src/message.rs b/crates/goose/src/message.rs index 41193a53f1..0e5781e7cc 100644 --- a/crates/goose/src/message.rs +++ b/crates/goose/src/message.rs @@ -43,6 +43,17 @@ pub struct ToolConfirmationRequest { pub prompt: Option, } +#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)] +pub struct ThinkingContent { + pub thinking: String, + pub signature: String, +} + +#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)] +pub struct RedactedThinkingContent { + pub data: String, +} + #[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)] /// Content passed inside a message, which can be both simple content and tool content #[serde(tag = "type", rename_all = "camelCase")] @@ -52,6 +63,8 @@ pub enum MessageContent { ToolRequest(ToolRequest), ToolResponse(ToolResponse), ToolConfirmationRequest(ToolConfirmationRequest), + Thinking(ThinkingContent), + RedactedThinking(RedactedThinkingContent), } impl MessageContent { @@ -97,6 +110,17 @@ impl MessageContent { prompt, }) } + + pub fn thinking, S2: Into>(thinking: S1, signature: S2) -> Self { + MessageContent::Thinking(ThinkingContent { + thinking: thinking.into(), + signature: signature.into(), + }) + } + + pub fn redacted_thinking>(data: S) -> Self { + MessageContent::RedactedThinking(RedactedThinkingContent { data: data.into() }) + } pub fn as_tool_request(&self) -> Option<&ToolRequest> { if let MessageContent::ToolRequest(ref tool_request) = self { Some(tool_request) @@ -143,6 +167,22 @@ impl MessageContent { _ => None, } } + + /// Get the thinking content if this is a ThinkingContent variant + pub fn as_thinking(&self) -> Option<&ThinkingContent> { + match self { + MessageContent::Thinking(thinking) => Some(thinking), + _ => None, + } + } + + /// Get the redacted thinking content if this is a RedactedThinkingContent variant + pub fn as_redacted_thinking(&self) -> Option<&RedactedThinkingContent> { + match self { + MessageContent::RedactedThinking(redacted) => Some(redacted), + _ => None, + } + } } impl From for MessageContent { @@ -264,6 +304,20 @@ impl Message { )) } + /// Add thinking content to the message + pub fn with_thinking, S2: Into>( + self, + thinking: S1, + signature: S2, + ) -> Self { + self.with_content(MessageContent::thinking(thinking, signature)) + } + + /// Add redacted thinking content to the message + pub fn with_redacted_thinking>(self, data: S) -> Self { + self.with_content(MessageContent::redacted_thinking(data)) + } + /// Get the concatenated text content of the message, separated by newlines pub fn as_concat_text(&self) -> String { self.content diff --git a/crates/goose/src/providers/anthropic.rs b/crates/goose/src/providers/anthropic.rs index 3cfd51a26e..44dae3e530 100644 --- a/crates/goose/src/providers/anthropic.rs +++ b/crates/goose/src/providers/anthropic.rs @@ -1,5 +1,6 @@ use anyhow::Result; use async_trait::async_trait; +use axum::http::HeaderMap; use reqwest::{Client, StatusCode}; use serde_json::Value; use std::time::Duration; @@ -17,6 +18,8 @@ pub const ANTHROPIC_KNOWN_MODELS: &[&str] = &[ "claude-3-5-sonnet-latest", "claude-3-5-haiku-latest", "claude-3-opus-latest", + "claude-3-7-sonnet-20250219", + "claude-3-7-sonnet-latest", ]; pub const ANTHROPIC_DOC_URL: &str = "https://docs.anthropic.com/en/docs/about-claude/models"; @@ -57,7 +60,7 @@ impl AnthropicProvider { }) } - async fn post(&self, payload: Value) -> Result { + async fn post(&self, headers: HeaderMap, payload: Value) -> Result { let base_url = url::Url::parse(&self.host) .map_err(|e| ProviderError::RequestFailed(format!("Invalid base URL: {e}")))?; let url = base_url.join("v1/messages").map_err(|e| { @@ -67,8 +70,7 @@ impl AnthropicProvider { let response = self .client .post(url) - .header("x-api-key", &self.api_key) - .header("anthropic-version", "2023-06-01") + .headers(headers) .json(&payload) .send() .await?; @@ -155,8 +157,22 @@ impl Provider for AnthropicProvider { ) -> Result<(Message, ProviderUsage), ProviderError> { let payload = create_request(&self.model, system, messages, tools)?; + let mut headers = reqwest::header::HeaderMap::new(); + headers.insert("x-api-key", self.api_key.parse().unwrap()); + headers.insert("anthropic-version", "2023-06-01".parse().unwrap()); + + let is_thinking_enabled = std::env::var("ANTHROPIC_THINKING_ENABLED").is_ok(); + if self.model.model_name.starts_with("claude-3-7-sonnet-") && is_thinking_enabled { + // https://docs.anthropic.com/en/docs/build-with-claude/tool-use/token-efficient-tool-use + headers.insert( + "anthropic-beta", + "token-efficient-tools-2025-02-19".parse().unwrap(), + ); + // https://docs.anthropic.com/en/docs/build-with-claude/extended-thinking#extended-output-capabilities-beta + headers.insert("anthropic-beta", "output-128k-2025-02-19".parse().unwrap()); + } // Make request - let response = self.post(payload.clone()).await?; + let response = self.post(headers, payload.clone()).await?; // Parse response let message = response_to_message(response.clone())?; diff --git a/crates/goose/src/providers/formats/anthropic.rs b/crates/goose/src/providers/formats/anthropic.rs index 4eadf4bcd4..2b1aec6503 100644 --- a/crates/goose/src/providers/formats/anthropic.rs +++ b/crates/goose/src/providers/formats/anthropic.rs @@ -60,6 +60,19 @@ pub fn format_messages(messages: &[Message]) -> Vec { MessageContent::ToolConfirmationRequest(_tool_confirmation_request) => { // Skip tool confirmation requests } + MessageContent::Thinking(thinking) => { + content.push(json!({ + "type": "thinking", + "thinking": thinking.thinking, + "signature": thinking.signature + })); + } + MessageContent::RedactedThinking(redacted) => { + content.push(json!({ + "type": "redacted_thinking", + "data": redacted.data + })); + } MessageContent::Image(_) => continue, // Anthropic doesn't support image content yet } } @@ -179,6 +192,24 @@ pub fn response_to_message(response: Value) -> Result { let tool_call = ToolCall::new(name, input.clone()); message = message.with_tool_request(id, Ok(tool_call)); } + Some("thinking") => { + let thinking = block + .get("thinking") + .and_then(|t| t.as_str()) + .ok_or_else(|| anyhow!("Missing thinking content"))?; + let signature = block + .get("signature") + .and_then(|s| s.as_str()) + .ok_or_else(|| anyhow!("Missing thinking signature"))?; + message = message.with_thinking(thinking, signature); + } + Some("redacted_thinking") => { + let data = block + .get("data") + .and_then(|d| d.as_str()) + .ok_or_else(|| anyhow!("Missing redacted_thinking data"))?; + message = message.with_redacted_thinking(data); + } _ => continue, } } @@ -243,10 +274,13 @@ pub fn create_request( return Err(anyhow!("No valid messages to send to Anthropic API")); } + // https://docs.anthropic.com/en/docs/about-claude/models/all-models#model-comparison-table + // Claude 3.7 supports max output tokens up to 8192 + let max_tokens = model_config.max_tokens.unwrap_or(8192); let mut payload = json!({ "model": model_config.model_name, "messages": anthropic_messages, - "max_tokens": model_config.max_tokens.unwrap_or(4096) + "max_tokens": max_tokens, }); // Add system message if present @@ -265,12 +299,38 @@ pub fn create_request( .insert("tools".to_string(), json!(tool_specs)); } - // Add temperature if specified + // Add temperature if specified and not using extended thinking model if let Some(temp) = model_config.temperature { + // Claude 3.7 models with thinking enabled don't support temperature + if !model_config.model_name.starts_with("claude-3-7-sonnet-") { + payload + .as_object_mut() + .unwrap() + .insert("temperature".to_string(), json!(temp)); + } + } + + // Add thinking parameters for claude-3-7-sonnet model + let is_thinking_enabled = std::env::var("ANTHROPIC_THINKING_ENABLED").is_ok(); + if model_config.model_name.starts_with("claude-3-7-sonnet-") && is_thinking_enabled { + // Minimum budget_tokens is 1024 + let budget_tokens = std::env::var("ANTHROPIC_THINKING_BUDGET") + .unwrap_or_else(|_| "16000".to_string()) + .parse() + .unwrap_or(16000); + payload .as_object_mut() .unwrap() - .insert("temperature".to_string(), json!(temp)); + .insert("max_tokens".to_string(), json!(max_tokens + budget_tokens)); + + payload.as_object_mut().unwrap().insert( + "thinking".to_string(), + json!({ + "type": "enabled", + "budget_tokens": budget_tokens + }), + ); } Ok(payload) @@ -361,6 +421,80 @@ mod tests { Ok(()) } + #[test] + fn test_parse_thinking_response() -> Result<()> { + let response = json!({ + "id": "msg_456", + "type": "message", + "role": "assistant", + "content": [ + { + "type": "thinking", + "thinking": "This is a step-by-step thought process...", + "signature": "EuYBCkQYAiJAVbJNBoH7HQiDcMwwAMhWqNyoe4G2xHRprK8ICM8gZzu16i7Se4EiEbmlKqNH1GtwcX1BMK6iLu8bxWn5wPVIFBIMnptdlVal7ZX5iNPFGgwWjX+BntcEOHky4HciMFVef7FpQeqnuiL1Xt7J4OLHZSyu4tcr809AxAbclcJ5dm1xE5gZrUO+/v60cnJM2ipQp4B8/3eHI03KSV6bZR/vMrBSYCV+aa/f5KHX2cRtLGp/Ba+3Tk/efbsg01WSduwAIbR4coVrZLnGJXNyVTFW/Be2kLy/ECZnx8cqvU3oQOg=" + }, + { + "type": "redacted_thinking", + "data": "EmwKAhgBEgy3va3pzix/LafPsn4aDFIT2Xlxh0L5L8rLVyIwxtE3rAFBa8cr3qpP" + }, + { + "type": "text", + "text": "I've analyzed the problem and here's the solution." + } + ], + "model": "claude-3-7-sonnet-20250219", + "stop_reason": "end_turn", + "stop_sequence": null, + "usage": { + "input_tokens": 10, + "output_tokens": 45, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + } + }); + + let message = response_to_message(response.clone())?; + let usage = get_usage(&response)?; + + assert_eq!(message.content.len(), 3); + + if let MessageContent::Thinking(thinking) = &message.content[0] { + assert_eq!( + thinking.thinking, + "This is a step-by-step thought process..." + ); + assert!(thinking + .signature + .starts_with("EuYBCkQYAiJAVbJNBoH7HQiDcMwwAMhWqNyoe4G2xHRprK8ICM8g")); + } else { + panic!("Expected Thinking content at index 0"); + } + + if let MessageContent::RedactedThinking(redacted) = &message.content[1] { + assert_eq!( + redacted.data, + "EmwKAhgBEgy3va3pzix/LafPsn4aDFIT2Xlxh0L5L8rLVyIwxtE3rAFBa8cr3qpP" + ); + } else { + panic!("Expected RedactedThinking content at index 1"); + } + + if let MessageContent::Text(text) = &message.content[2] { + assert_eq!( + text.text, + "I've analyzed the problem and here's the solution." + ); + } else { + panic!("Expected Text content at index 2"); + } + + assert_eq!(usage.input_tokens, Some(10)); + assert_eq!(usage.output_tokens, Some(45)); + assert_eq!(usage.total_tokens, Some(55)); + + Ok(()) + } + #[test] fn test_message_to_anthropic_spec() { let messages = vec![ @@ -436,4 +570,47 @@ mod tests { assert_eq!(spec_array[0]["text"], system); assert!(spec_array[0].get("cache_control").is_some()); } + + #[test] + fn test_create_request_with_thinking() -> Result<()> { + // Save the original env var value if it exists + let original_value = std::env::var("ANTHROPIC_THINKING_ENABLED").ok(); + + // Set the env var for this test + std::env::set_var("ANTHROPIC_THINKING_ENABLED", "true"); + + // Execute the test + let result = (|| { + let model_config = ModelConfig::new("claude-3-7-sonnet-20250219".to_string()); + let system = "You are a helpful assistant."; + let messages = vec![Message::user().with_text("Hello")]; + let tools = vec![]; + + let payload = create_request(&model_config, system, &messages, &tools)?; + + // Verify basic structure + assert_eq!(payload["model"], "claude-3-7-sonnet-20250219"); + assert_eq!(payload["messages"][0]["role"], "user"); + assert_eq!(payload["messages"][0]["content"][0]["text"], "Hello"); + + // Verify thinking parameters + assert!(payload.get("thinking").is_some()); + assert_eq!(payload["thinking"]["type"], "enabled"); + assert!(payload["thinking"]["budget_tokens"].as_i64().unwrap() >= 1024); + + // Temperature should not be present for 3.7 models with thinking + assert!(payload.get("temperature").is_none()); + + Ok(()) + })(); + + // Restore the original env var state + match original_value { + Some(val) => std::env::set_var("ANTHROPIC_THINKING_ENABLED", val), + None => std::env::remove_var("ANTHROPIC_THINKING_ENABLED"), + } + + // Return the test result + result + } } diff --git a/crates/goose/src/providers/formats/bedrock.rs b/crates/goose/src/providers/formats/bedrock.rs index dcedf31bad..dd32b0a415 100644 --- a/crates/goose/src/providers/formats/bedrock.rs +++ b/crates/goose/src/providers/formats/bedrock.rs @@ -34,6 +34,14 @@ pub fn to_bedrock_message_content(content: &MessageContent) -> Result { bail!("Image content is not supported by Bedrock provider yet") } + MessageContent::Thinking(_) => { + // Thinking blocks are not supported in Bedrock - skip + bedrock::ContentBlock::Text("".to_string()) + } + MessageContent::RedactedThinking(_) => { + // Redacted thinking blocks are not supported in Bedrock - skip + bedrock::ContentBlock::Text("".to_string()) + } MessageContent::ToolRequest(tool_req) => { let tool_use_id = tool_req.id.to_string(); let tool_use = if let Ok(call) = tool_req.tool_call.as_ref() { diff --git a/crates/goose/src/providers/formats/openai.rs b/crates/goose/src/providers/formats/openai.rs index f783041633..7273d6fea4 100644 --- a/crates/goose/src/providers/formats/openai.rs +++ b/crates/goose/src/providers/formats/openai.rs @@ -44,6 +44,14 @@ pub fn format_messages(messages: &[Message], image_format: &ImageFormat) -> Vec< } } } + MessageContent::Thinking(_) => { + // Thinking blocks are not directly used in OpenAI format + continue; + } + MessageContent::RedactedThinking(_) => { + // Redacted thinking blocks are not directly used in OpenAI format + continue; + } MessageContent::ToolRequest(request) => match &request.tool_call { Ok(tool_call) => { let sanitized_name = sanitize_function_name(&tool_call.name);