diff --git a/src/llm/rig/mod.rs b/src/llm/rig/mod.rs index 7ab64d5..24c0f50 100644 --- a/src/llm/rig/mod.rs +++ b/src/llm/rig/mod.rs @@ -10,7 +10,7 @@ use rig_core::{ client::{CompletionClient, Nothing, VerifyClient}, completion::{AssistantContent, CompletionError, CompletionModel}, http_client::{HttpClientExt, ReqwestClient}, - providers::{deepseek, moonshot, ollama}, + providers::{anthropic, deepseek, moonshot, ollama}, streaming::StreamedAssistantContent, wasm_compat::{WasmCompatSend, WasmCompatSync}, }; @@ -23,6 +23,8 @@ pub struct LlmClientConfig { pub max_tokens: u64, pub temperature: f64, pub timeout: Duration, + /// Anthropic thinking 预算 token（配置缺省 1024）。 + pub thinking_budget_tokens: Option, } impl Default for LlmClientConfig { @@ -31,6 +33,7 @@ impl Default for LlmClientConfig { max_tokens: 500, temperature: 0.7, timeout: Duration::from_secs(30), + thinking_budget_tokens: None, } } } @@ -44,6 +47,7 @@ pub enum Backend { Ollama(ollama::Client), DeepSeek(deepseek::Client), Kimi(moonshot::Client), + Anthropic(anthropic::Client), } /// 基于 rig 的 LLM 客户端门面。 @@ -85,6 +89,7 @@ impl LlmClient { max_tokens: config.llm.max_tokens as u64, temperature: config.llm.temperature as f64, timeout: Duration::from_secs(config.llm.timeout), + thinking_budget_tokens: config.llm.thinking_budget_tokens, }; let provider = manager.llm_provider().to_string(); let model = manager.llm_model().to_string(); @@ -105,8 +110,12 @@ impl LlmClient { let key = required_api_key(&provider, manager.get_api_key())?; Backend::Kimi(build_kimi_client(&key, &base_url, http)?) } + "anthropic" => { + let key = required_api_key(&provider, manager.get_api_key())?; + Backend::Anthropic(build_anthropic_client(&key, &base_url, http)?) + } // 其余 provider 在后续 ticket 接入，扩张期内新门面尚未被应用使用。 - "openai" | "anthropic" | "openrouter" => { + "openai" | "openrouter" => { bail!("Provider '{}' is not available in the new LLM backend yet", provider) } _ => bail!("Unknown LLM provider: {}", provider), @@ -189,6 +198,30 @@ impl LlmClient { })), stream: self.thinking_enabled, }, + // Anthropic：thinking 时省略温度、max_tokens 不低于预算 + 100，预算透传。 + Backend::Anthropic(_) => { + let budget = self.config.thinking_budget_tokens.unwrap_or(1024); + let max_tokens = if self.thinking_enabled { + self.config.max_tokens.max(u64::from(budget) + 100) + } else { + self.config.max_tokens + }; + let thinking_param = if self.thinking_enabled { + serde_json::json!({ "thinking": { "type": "enabled", "budget_tokens": budget } }) + } else { + serde_json::json!({ "thinking": { "type": "disabled" } }) + }; + RequestOptions { + temperature: if self.thinking_enabled { + None + } else { + Some(self.config.temperature) + }, + max_tokens: Some(max_tokens), + additional_params: Some(thinking_param), + stream: self.thinking_enabled, + } + } } } } @@ -221,6 +254,10 @@ where let model = client.completion_model(self.model.as_str()); self.run(&model, system, user, &options).await } + Backend::Anthropic(client) => { + let model = client.completion_model(self.model.as_str()); + self.run(&model, system, user, &options).await + } } } @@ -230,6 +267,7 @@ where Backend::Ollama(client) => client.verify().await.is_ok(), Backend::DeepSeek(client) => client.verify().await.is_ok(), Backend::Kimi(client) => client.verify().await.is_ok(), + Backend::Anthropic(client) => client.verify().await.is_ok(), } } @@ -346,6 +384,16 @@ fn build_kimi_client(key: &str, base_url: &str, http: ReqwestClient) -> Result Result { + anthropic::Client::builder() + .api_key(key) + .base_url(base_url) + .http_client(http) + .build() + .map_err(|e| anyhow::anyhow!("Failed to build Anthropic client: {}", e)) +} + /// 组装统一的补全请求（preamble/temperature/max_tokens/additional_params）。 fn build_request( model: &M, @@ -454,6 +502,7 @@ mod tests { max_tokens: 123, temperature: 0.5, timeout: Duration::from_secs(30), + thinking_budget_tokens: None, } } @@ -757,4 +806,166 @@ mod tests { assert!(!supports_thinking(provider), "{provider} should not support thinking"); } } + + // ---- Anthropic ---- + + const ANTHROPIC_OK: &str = r#"{ + "id": "msg_01", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "hello from claude"}], + "model": "claude-sonnet-4-6", + "stop_reason": "end_turn", + "stop_sequence": null, + "usage": {"input_tokens": 5, "output_tokens": 3} + }"#; + + fn anthropic_client( + recorder: RecordingHttpClient, + thinking: bool, + budget: Option, + ) -> (LlmClient, RecordingHttpClient) { + let backend = Backend::Anthropic( + anthropic::Client::builder() + .api_key("sk-ant-test") + .base_url("https://api.anthropic.com/v1") + .http_client(recorder.clone()) + .build() + .expect("build anthropic client with mock backend"), + ); + let mut config = test_config(); + config.thinking_budget_tokens = budget; + let client = LlmClient::new(backend, "claude-sonnet-4-6", "anthropic", config, thinking, None); + (client, recorder) + } + + #[tokio::test] + async fn anthropic_normal_request_shape() { + let recorder = RecordingHttpClient::new(ANTHROPIC_OK); + let (client, recorder) = anthropic_client(recorder, false, None); + + let text = client.generate(Some("sys"), "hi").await.unwrap(); + assert_eq!(text, "hello from claude"); + + let captured = recorder.requests(); + assert_eq!(captured[0].uri, "https://api.anthropic.com/v1/messages"); + let body: serde_json::Value = serde_json::from_slice(&captured[0].body).unwrap(); + assert_eq!(body["model"], "claude-sonnet-4-6"); + assert_eq!(body["max_tokens"], 123); + assert_eq!(body["temperature"], 0.5); + assert_eq!(body["system"][0]["text"], "sys"); + assert_eq!(body["messages"][0]["role"], "user"); + assert_eq!(body["messages"][0]["content"][0]["type"], "text"); + assert_eq!(body["messages"][0]["content"][0]["text"], "hi"); + // 与旧实现一致：非 thinking 显式传 disabled + assert_eq!(body["thinking"]["type"], "disabled"); + assert!(body.get("stream").is_none(), "非流式请求不应带 stream 字段"); + } + + #[test] + fn anthropic_thinking_request_options() { + let recorder = RecordingHttpClient::new(ANTHROPIC_OK); + let (client, _) = anthropic_client(recorder, true, Some(2000)); + + let options = client.request_options(); + assert_eq!(options.temperature, None, "thinking 时必须省略温度"); + // max_tokens 不低于预算 + 100 + assert_eq!(options.max_tokens, Some(2100)); + assert_eq!(options.stream, true); + let params = options.additional_params.unwrap(); + assert_eq!(params["thinking"]["type"], "enabled"); + assert_eq!(params["thinking"]["budget_tokens"], 2000); + } + + #[test] + fn anthropic_thinking_default_budget() { + let recorder = RecordingHttpClient::new(ANTHROPIC_OK); + let (client, _) = anthropic_client(recorder, true, None); + + let options = client.request_options(); + let params = options.additional_params.unwrap(); + assert_eq!(params["thinking"]["budget_tokens"], 1024); + // 缺省预算时 max_tokens = max(123, 1124) = 1124 + assert_eq!(options.max_tokens, Some(1124)); + } + + const ANTHROPIC_SSE_THINK_THEN_TEXT: &str = concat!( + "data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"thinking_delta\",\"thinking\":\"let me think\"}}\n\n", + "data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"hello\"}}\n\n", + "data: {\"type\":\"content_block_stop\",\"index\":0}\n\n", + "data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":5}}\n\n", + ); + + const ANTHROPIC_SSE_REASONING_ONLY: &str = concat!( + "data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"thinking_delta\",\"thinking\":\"only thoughts\"}}\n\n", + "data: {\"type\":\"content_block_stop\",\"index\":0}\n\n", + "data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":5}}\n\n", + ); + + fn anthropic_streaming_client( + sse: &str, + ) -> ( + LlmClient, + Arc, + Arc, + ) { + let start_count = Arc::new(AtomicUsize::new(0)); + let end_count = Arc::new(AtomicUsize::new(0)); + let state = Arc::new( + ThinkingStateManager::new() + .on_thinking_start({ + let c = start_count.clone(); + move || { + c.fetch_add(1, Ordering::SeqCst); + } + }) + .on_thinking_end({ + let c = end_count.clone(); + move || { + c.fetch_add(1, Ordering::SeqCst); + } + }), + ); + + let backend = Backend::Anthropic( + anthropic::Client::builder() + .api_key("sk-ant-test") + .base_url("https://api.anthropic.com/v1") + .http_client(MockStreamingClient { sse_bytes: sse.to_string().into() }) + .build() + .expect("build anthropic client with streaming mock"), + ); + let client = LlmClient::new( + backend, + "claude-sonnet-4-6", + "anthropic", + test_config(), + true, + Some(state), + ); + (client, start_count, end_count) + } + + #[tokio::test] + async fn anthropic_streaming_aggregates_text_and_drives_thinking() { + let (client, start_count, end_count) = + anthropic_streaming_client(ANTHROPIC_SSE_THINK_THEN_TEXT); + + let text = client.generate(None, "hi").await.unwrap(); + assert_eq!(text, "hello"); + assert_eq!(start_count.load(Ordering::SeqCst), 1); + assert_eq!(end_count.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn anthropic_streaming_reasoning_only_errors_with_hint() { + let (client, _, _) = anthropic_streaming_client(ANTHROPIC_SSE_REASONING_ONLY); + + let err = client.generate(None, "hi").await.unwrap_err(); + let msg = err.to_string(); + assert!( + msg.contains("Anthropic returned reasoning content but no final answer"), + "got: {msg}" + ); + } } \ No newline at end of file