diff options
| author | sigoden <sigoden@gmail.com> | 2025-02-26 08:08:34 +0800 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2025-02-26 08:08:34 +0800 |
| commit | 6b28b1c4fcf57edc4375703bca34eeafa597c905 (patch) | |
| tree | d6f6e6d1ac51f46ba258d46128a5f31d3c874645 /src/client/bedrock.rs | |
| parent | d8b1aef97d00e2f66d5af32ad075d990c7dfdba1 (diff) | |
| download | aichat-6b28b1c4fcf57edc4375703bca34eeafa597c905.tar.gz | |
feat: support claude-3-7-sonnet (#1201)
Diffstat (limited to 'src/client/bedrock.rs')
| -rw-r--r-- | src/client/bedrock.rs | 33 |
1 files changed, 28 insertions, 5 deletions
diff --git a/src/client/bedrock.rs b/src/client/bedrock.rs index ebc0154..089f1c2 100644 --- a/src/client/bedrock.rs +++ b/src/client/bedrock.rs @@ -198,6 +198,7 @@ async fn chat_completions_streaming( let mut function_name = String::new(); let mut function_arguments = String::new(); let mut function_id = String::new(); + let mut reasoning_state = 0; let mut stream = res.bytes_stream(); let mut buffer = BytesMut::new(); @@ -240,11 +241,21 @@ async fn chat_completions_streaming( "contentBlockDelta" => { if let Some(text) = data["delta"]["text"].as_str() { handler.text(text)?; + } else if let Some(text) = data["delta"]["reasoningContent"]["text"].as_str() { + if reasoning_state == 0 { + handler.text("<think>\n")?; + reasoning_state = 1; + } + handler.text(text)?; } else if let Some(input) = data["delta"]["toolUse"]["input"].as_str() { function_arguments.push_str(input); } } "contentBlockStop" => { + if reasoning_state == 1 { + handler.text("\n</think>\n\n")?; + reasoning_state = 0; + } if !function_name.is_empty() { let arguments: Value = function_arguments.parse().with_context(|| { format!("Tool call '{function_name}' have non-JSON arguments '{function_arguments}'") @@ -448,12 +459,20 @@ fn build_chat_completions_body(data: ChatCompletionsData, model: &Model) -> Resu } fn extract_chat_completions(data: &Value) -> Result<ChatCompletionsOutput> { - let mut texts = vec![]; + let mut text = String::new(); + let mut reasoning = None; let mut tool_calls = vec![]; if let Some(array) = data["output"]["message"]["content"].as_array() { for item in array { - if let Some(text) = item["text"].as_str() { - texts.push(text); + if let Some(v) = item["text"].as_str() { + if !text.is_empty() { + text.push_str("\n\n"); + } + text.push_str(v); + } else if let Some(reasoning_text) = item["reasoningContent"]["reasoningText"].as_object() { + if let Some(text) = json_str_from_map(reasoning_text, "text") { + reasoning = Some(text.to_string()); + } } else if let Some(tool_use) = item["toolUse"].as_object() { if let (Some(id), Some(name), Some(input)) = ( json_str_from_map(tool_use, "toolUseId"), @@ -470,12 +489,16 @@ fn extract_chat_completions(data: &Value) -> Result<ChatCompletionsOutput> { } } - if texts.is_empty() && tool_calls.is_empty() { + if let Some(reasoning) = reasoning { + text = format!("<think>\n{reasoning}\n</think>\n\n{text}") + } + + if text.is_empty() && tool_calls.is_empty() { bail!("Invalid response data: {data}"); } let output = ChatCompletionsOutput { - text: texts.join("\n\n"), + text, tool_calls, id: None, input_tokens: data["usage"]["inputTokens"].as_u64(), |
