656 lines
20 KiB
Rust
656 lines
20 KiB
Rust
use super::thinking::ThinkingStateManager;
|
|
use super::{LlmProvider, create_http_client};
|
|
use anyhow::{Context, Result, bail};
|
|
use async_trait::async_trait;
|
|
use serde::{Deserialize, Serialize};
|
|
use std::sync::Arc;
|
|
use std::time::Duration;
|
|
|
|
/// Anthropic Claude API client
|
|
pub struct AnthropicClient {
|
|
api_key: String,
|
|
model: String,
|
|
client: reqwest::Client,
|
|
thinking_enabled: bool,
|
|
thinking_budget_tokens: u32,
|
|
max_tokens: u32,
|
|
temperature: f32,
|
|
top_p: Option<f32>,
|
|
thinking_state: Option<Arc<ThinkingStateManager>>,
|
|
}
|
|
|
|
#[derive(Debug, Serialize)]
|
|
struct MessagesRequest {
|
|
model: String,
|
|
max_tokens: u32,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
temperature: Option<f32>,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
top_p: Option<f32>,
|
|
messages: Vec<AnthropicMessage>,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
system: Option<Vec<SystemContent>>,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
thinking: Option<ThinkingConfig>,
|
|
stream: bool,
|
|
}
|
|
|
|
#[derive(Debug, Serialize, Clone)]
|
|
struct SystemContent {
|
|
#[serde(rename = "type")]
|
|
content_type: String,
|
|
text: String,
|
|
}
|
|
|
|
#[derive(Debug, Serialize)]
|
|
struct ThinkingConfig {
|
|
#[serde(rename = "type")]
|
|
thinking_type: String,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
budget_tokens: Option<u32>,
|
|
}
|
|
|
|
#[derive(Debug, Serialize, Deserialize, Clone)]
|
|
struct AnthropicMessage {
|
|
role: String,
|
|
content: AnthropicContent,
|
|
}
|
|
|
|
#[derive(Debug, Serialize, Deserialize, Clone)]
|
|
#[serde(untagged)]
|
|
enum AnthropicContent {
|
|
Text(String),
|
|
Blocks(Vec<ContentBlock>),
|
|
}
|
|
|
|
#[derive(Debug, Serialize, Deserialize, Clone)]
|
|
struct ContentBlock {
|
|
#[serde(rename = "type")]
|
|
content_type: String,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
text: Option<String>,
|
|
}
|
|
|
|
#[derive(Debug, Deserialize)]
|
|
struct MessagesResponse {
|
|
content: Vec<ResponseContentBlock>,
|
|
}
|
|
|
|
#[derive(Debug, Deserialize)]
|
|
struct ResponseContentBlock {
|
|
#[serde(rename = "type")]
|
|
content_type: String,
|
|
text: String,
|
|
}
|
|
|
|
#[derive(Debug, Deserialize)]
|
|
struct ErrorResponse {
|
|
error: AnthropicError,
|
|
}
|
|
|
|
#[derive(Debug, Deserialize)]
|
|
struct AnthropicError {
|
|
#[serde(rename = "type")]
|
|
error_type: String,
|
|
message: String,
|
|
}
|
|
|
|
// --- Streaming SSE event structures ---
|
|
|
|
#[derive(Debug, Deserialize)]
|
|
struct SseEvent {
|
|
#[serde(rename = "type")]
|
|
event_type: String,
|
|
#[serde(default)]
|
|
message: Option<SseMessage>,
|
|
#[serde(default)]
|
|
index: Option<u32>,
|
|
#[serde(default)]
|
|
content_block: Option<SseContentBlock>,
|
|
#[serde(default)]
|
|
delta: Option<SseDelta>,
|
|
#[serde(default)]
|
|
usage: Option<SseUsage>,
|
|
}
|
|
|
|
#[derive(Debug, Deserialize)]
|
|
struct SseMessage {
|
|
#[serde(default)]
|
|
content: Option<Vec<SseContentBlock>>,
|
|
}
|
|
|
|
#[derive(Debug, Deserialize)]
|
|
struct SseContentBlock {
|
|
#[serde(rename = "type")]
|
|
content_type: String,
|
|
#[serde(default)]
|
|
thinking: Option<String>,
|
|
#[serde(default)]
|
|
text: Option<String>,
|
|
}
|
|
|
|
#[derive(Debug, Deserialize)]
|
|
struct SseDelta {
|
|
#[serde(rename = "type")]
|
|
delta_type: Option<String>,
|
|
#[serde(default)]
|
|
thinking: Option<String>,
|
|
#[serde(default)]
|
|
text: Option<String>,
|
|
}
|
|
|
|
#[derive(Debug, Deserialize)]
|
|
struct SseUsage {
|
|
#[serde(default)]
|
|
output_tokens: Option<u32>,
|
|
}
|
|
|
|
impl AnthropicClient {
|
|
pub fn new(api_key: &str, model: &str) -> Result<Self> {
|
|
let client = create_http_client(Duration::from_secs(60))?;
|
|
|
|
Ok(Self {
|
|
api_key: api_key.to_string(),
|
|
model: model.to_string(),
|
|
client,
|
|
thinking_enabled: false,
|
|
thinking_budget_tokens: 1024,
|
|
max_tokens: 500,
|
|
temperature: 0.7,
|
|
top_p: None,
|
|
thinking_state: None,
|
|
})
|
|
}
|
|
|
|
pub fn with_timeout(mut self, timeout: Duration) -> Result<Self> {
|
|
self.client = create_http_client(timeout)?;
|
|
Ok(self)
|
|
}
|
|
|
|
pub fn with_thinking(mut self, enabled: bool) -> Self {
|
|
self.thinking_enabled = enabled;
|
|
self
|
|
}
|
|
|
|
pub fn with_thinking_budget_tokens(mut self, budget_tokens: u32) -> Self {
|
|
self.thinking_budget_tokens = budget_tokens;
|
|
self
|
|
}
|
|
|
|
pub fn with_max_tokens(mut self, max_tokens: u32) -> Self {
|
|
self.max_tokens = max_tokens;
|
|
self
|
|
}
|
|
|
|
pub fn with_temperature(mut self, temperature: f32) -> Self {
|
|
self.temperature = temperature;
|
|
self
|
|
}
|
|
|
|
pub fn with_top_p(mut self, top_p: f32) -> Self {
|
|
self.top_p = Some(top_p);
|
|
self
|
|
}
|
|
|
|
pub fn with_thinking_state(mut self, state: Arc<ThinkingStateManager>) -> Self {
|
|
self.thinking_state = Some(state);
|
|
self
|
|
}
|
|
|
|
pub async fn list_models(&self) -> Result<Vec<String>> {
|
|
Ok(ANTHROPIC_MODELS.iter().map(|&m| m.to_string()).collect())
|
|
}
|
|
|
|
pub async fn validate_key(&self) -> Result<bool> {
|
|
let url = "https://api.anthropic.com/v1/messages";
|
|
|
|
let request = MessagesRequest {
|
|
model: self.model.clone(),
|
|
max_tokens: 5,
|
|
temperature: Some(0.0),
|
|
top_p: None,
|
|
messages: vec![AnthropicMessage {
|
|
role: "user".to_string(),
|
|
content: AnthropicContent::Text("Hi".to_string()),
|
|
}],
|
|
system: None,
|
|
thinking: None,
|
|
stream: false,
|
|
};
|
|
|
|
let response = self
|
|
.client
|
|
.post(url)
|
|
.header("x-api-key", &self.api_key)
|
|
.header("anthropic-version", "2023-06-01")
|
|
.header("Content-Type", "application/json")
|
|
.json(&request)
|
|
.send()
|
|
.await;
|
|
|
|
match response {
|
|
Ok(resp) => {
|
|
if resp.status().is_success() {
|
|
Ok(true)
|
|
} else {
|
|
let status = resp.status();
|
|
if status.as_u16() == 401 {
|
|
Ok(false)
|
|
} else {
|
|
let text = resp.text().await.unwrap_or_default();
|
|
bail!("Anthropic API error: {} - {}", status, text)
|
|
}
|
|
}
|
|
}
|
|
Err(e) => Err(e.into()),
|
|
}
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl LlmProvider for AnthropicClient {
|
|
async fn generate(&self, prompt: &str) -> Result<String> {
|
|
let messages = vec![AnthropicMessage {
|
|
role: "user".to_string(),
|
|
content: AnthropicContent::Text(prompt.to_string()),
|
|
}];
|
|
|
|
self.messages_request_with_retry(messages, None).await
|
|
}
|
|
|
|
async fn generate_with_system(&self, system: &str, user: &str) -> Result<String> {
|
|
let messages = vec![AnthropicMessage {
|
|
role: "user".to_string(),
|
|
content: AnthropicContent::Text(user.to_string()),
|
|
}];
|
|
|
|
let system = if system.is_empty() {
|
|
None
|
|
} else {
|
|
Some(vec![SystemContent {
|
|
content_type: "text".to_string(),
|
|
text: system.to_string(),
|
|
}])
|
|
};
|
|
|
|
self.messages_request_with_retry(messages, system).await
|
|
}
|
|
|
|
async fn is_available(&self) -> bool {
|
|
self.validate_key().await.unwrap_or(false)
|
|
}
|
|
|
|
fn name(&self) -> &str {
|
|
"anthropic"
|
|
}
|
|
}
|
|
|
|
impl AnthropicClient {
|
|
async fn messages_request_with_retry(
|
|
&self,
|
|
messages: Vec<AnthropicMessage>,
|
|
system: Option<Vec<SystemContent>>,
|
|
) -> Result<String> {
|
|
let mut last_error = None;
|
|
|
|
for attempt in 1..=3 {
|
|
match self
|
|
.messages_request(messages.clone(), system.clone())
|
|
.await
|
|
{
|
|
Ok(result) => return Ok(result),
|
|
Err(e) => {
|
|
let err_msg = e.to_string();
|
|
let is_retryable = err_msg.contains("timeout")
|
|
|| err_msg.contains("connection")
|
|
|| err_msg.contains("temporary")
|
|
|| err_msg.contains("5")
|
|
&& (err_msg.contains("500")
|
|
|| err_msg.contains("502")
|
|
|| err_msg.contains("503")
|
|
|| err_msg.contains("504"));
|
|
|
|
if !is_retryable || attempt == 3 {
|
|
last_error = Some(e);
|
|
break;
|
|
}
|
|
|
|
tokio::time::sleep(Duration::from_millis(500 * 2u64.pow(attempt - 1))).await;
|
|
}
|
|
}
|
|
}
|
|
|
|
Err(last_error.unwrap_or_else(|| anyhow::anyhow!("Request failed after retries")))
|
|
}
|
|
|
|
async fn messages_request(
|
|
&self,
|
|
messages: Vec<AnthropicMessage>,
|
|
system: Option<Vec<SystemContent>>,
|
|
) -> Result<String> {
|
|
if self.thinking_enabled {
|
|
self.streaming_messages_request(messages, system).await
|
|
} else {
|
|
self.non_streaming_messages_request(messages, system).await
|
|
}
|
|
}
|
|
|
|
async fn non_streaming_messages_request(
|
|
&self,
|
|
messages: Vec<AnthropicMessage>,
|
|
system: Option<Vec<SystemContent>>,
|
|
) -> Result<String> {
|
|
let url = "https://api.anthropic.com/v1/messages";
|
|
|
|
let temperature = if self.temperature == 0.0 {
|
|
None
|
|
} else {
|
|
Some(self.temperature)
|
|
};
|
|
|
|
let request = MessagesRequest {
|
|
model: self.model.clone(),
|
|
max_tokens: self.max_tokens,
|
|
temperature,
|
|
top_p: self.top_p,
|
|
messages,
|
|
system,
|
|
thinking: Some(ThinkingConfig {
|
|
thinking_type: "disabled".to_string(),
|
|
budget_tokens: None,
|
|
}),
|
|
stream: false,
|
|
};
|
|
|
|
let response = self
|
|
.client
|
|
.post(url)
|
|
.header("x-api-key", &self.api_key)
|
|
.header("anthropic-version", "2023-06-01")
|
|
.header("Content-Type", "application/json")
|
|
.json(&request)
|
|
.send()
|
|
.await
|
|
.context("Failed to send request to Anthropic")?;
|
|
|
|
let status = response.status();
|
|
|
|
if !status.is_success() {
|
|
let text = response.text().await.unwrap_or_default();
|
|
|
|
if let Ok(error) = serde_json::from_str::<ErrorResponse>(&text) {
|
|
bail!(
|
|
"Anthropic API error: {} ({})",
|
|
error.error.message,
|
|
error.error.error_type
|
|
);
|
|
}
|
|
|
|
bail!("Anthropic API error: {} - {}", status, text);
|
|
}
|
|
|
|
let result: MessagesResponse = response
|
|
.json()
|
|
.await
|
|
.context("Failed to parse Anthropic response")?;
|
|
|
|
result
|
|
.content
|
|
.into_iter()
|
|
.find(|c| c.content_type == "text")
|
|
.map(|c| c.text.trim().to_string())
|
|
.filter(|s| !s.is_empty())
|
|
.ok_or_else(|| anyhow::anyhow!("No text response from Anthropic"))
|
|
}
|
|
|
|
/// Streaming request for thinking mode, filters thinking content blocks
|
|
async fn streaming_messages_request(
|
|
&self,
|
|
messages: Vec<AnthropicMessage>,
|
|
system: Option<Vec<SystemContent>>,
|
|
) -> Result<String> {
|
|
let url = "https://api.anthropic.com/v1/messages";
|
|
|
|
let thinking = ThinkingConfig {
|
|
thinking_type: "enabled".to_string(),
|
|
budget_tokens: Some(self.thinking_budget_tokens),
|
|
};
|
|
|
|
// max_tokens must exceed budget_tokens
|
|
let max_tokens = (self.max_tokens).max(self.thinking_budget_tokens + 100);
|
|
|
|
let request = MessagesRequest {
|
|
model: self.model.clone(),
|
|
max_tokens,
|
|
temperature: None, // must be omitted for thinking mode
|
|
top_p: None,
|
|
messages,
|
|
system,
|
|
thinking: Some(thinking),
|
|
stream: true,
|
|
};
|
|
|
|
let response = self
|
|
.client
|
|
.post(url)
|
|
.header("x-api-key", &self.api_key)
|
|
.header("anthropic-version", "2023-06-01")
|
|
.header("Content-Type", "application/json")
|
|
.header("Accept", "text/event-stream")
|
|
.json(&request)
|
|
.send()
|
|
.await
|
|
.context("Failed to send streaming request to Anthropic")?;
|
|
|
|
let status = response.status();
|
|
|
|
if !status.is_success() {
|
|
let text = response.text().await.unwrap_or_default();
|
|
|
|
if let Ok(error) = serde_json::from_str::<ErrorResponse>(&text) {
|
|
bail!(
|
|
"Anthropic API error: {} ({})",
|
|
error.error.message,
|
|
error.error.error_type
|
|
);
|
|
}
|
|
|
|
bail!("Anthropic API error: {} - {}", status, text);
|
|
}
|
|
|
|
let mut content_buffer = String::new();
|
|
let mut in_thinking = false;
|
|
let mut has_reasoning = false;
|
|
let mut has_content = false;
|
|
|
|
let thinking_state = self.thinking_state.as_ref();
|
|
|
|
let mut byte_stream = response.bytes_stream();
|
|
let mut line_buffer = String::new();
|
|
|
|
use futures_util::StreamExt;
|
|
|
|
while let Some(chunk) = byte_stream.next().await {
|
|
let chunk = chunk.context("Failed to read streaming response chunk")?;
|
|
let chunk_str =
|
|
String::from_utf8(chunk.to_vec()).context("Invalid UTF-8 in stream chunk")?;
|
|
|
|
line_buffer.push_str(&chunk_str);
|
|
|
|
while let Some(line_end) = line_buffer.find('\n') {
|
|
let line = line_buffer[..line_end].trim().to_string();
|
|
line_buffer = line_buffer[line_end + 1..].to_string();
|
|
|
|
if line.is_empty() {
|
|
continue;
|
|
}
|
|
|
|
// Parse SSE event line
|
|
if let Some(data) = line.strip_prefix("data: ") {
|
|
if let Ok(event) = serde_json::from_str::<SseEvent>(data) {
|
|
match event.event_type.as_str() {
|
|
"content_block_start" => {
|
|
if let Some(ref block) = event.content_block {
|
|
if block.content_type == "thinking" {
|
|
in_thinking = true;
|
|
if !has_reasoning {
|
|
has_reasoning = true;
|
|
if let Some(state) = thinking_state {
|
|
state.start_thinking();
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
"content_block_delta" => {
|
|
if let Some(ref delta) = event.delta {
|
|
// Thinking delta - ignore content but track state
|
|
if delta.thinking.is_some() {
|
|
continue;
|
|
}
|
|
|
|
// Text delta - collect
|
|
if in_thinking && delta.text.is_some() {
|
|
// Transition from thinking to text
|
|
if let Some(state) = thinking_state {
|
|
state.end_thinking();
|
|
}
|
|
in_thinking = false;
|
|
}
|
|
if let Some(ref text) = delta.text
|
|
&& !text.is_empty()
|
|
{
|
|
has_content = true;
|
|
content_buffer.push_str(text);
|
|
}
|
|
}
|
|
}
|
|
"content_block_stop" => {
|
|
if in_thinking {
|
|
if let Some(state) = thinking_state {
|
|
state.end_thinking();
|
|
}
|
|
in_thinking = false;
|
|
}
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Ensure thinking state is ended
|
|
if let Some(state) = thinking_state {
|
|
state.end_thinking();
|
|
}
|
|
|
|
let result = content_buffer.trim().to_string();
|
|
|
|
if result.is_empty() {
|
|
if has_reasoning && !has_content {
|
|
bail!(
|
|
"Anthropic returned thinking content but no final answer. \
|
|
The model may have entered an incomplete thinking state. \
|
|
Please try again or disable thinking mode."
|
|
);
|
|
}
|
|
bail!(
|
|
"No response from Anthropic. \
|
|
If thinking mode is enabled, try disabling it or ensure the model supports it."
|
|
);
|
|
}
|
|
|
|
Ok(result)
|
|
}
|
|
}
|
|
|
|
/// Available Anthropic models (Claude 4 series with extended thinking)
|
|
pub const ANTHROPIC_MODELS: &[&str] = &[
|
|
"claude-opus-4-7",
|
|
"claude-sonnet-4-6",
|
|
"claude-haiku-4-5",
|
|
// Legacy models
|
|
"claude-3-opus-20240229",
|
|
"claude-3-sonnet-20240229",
|
|
"claude-3-haiku-20240307",
|
|
"claude-2.1",
|
|
"claude-2.0",
|
|
"claude-instant-1.2",
|
|
];
|
|
|
|
pub fn is_valid_model(model: &str) -> bool {
|
|
ANTHROPIC_MODELS.contains(&model)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_model_validation_claude4() {
|
|
assert!(is_valid_model("claude-opus-4-7"));
|
|
assert!(is_valid_model("claude-sonnet-4-6"));
|
|
assert!(is_valid_model("claude-haiku-4-5"));
|
|
assert!(is_valid_model("claude-3-sonnet-20240229"));
|
|
assert!(!is_valid_model("invalid-model"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_thinking_config_serialization() {
|
|
let config = ThinkingConfig {
|
|
thinking_type: "enabled".to_string(),
|
|
budget_tokens: Some(2048),
|
|
};
|
|
let json = serde_json::to_string(&config).unwrap();
|
|
assert!(json.contains(r#""type":"enabled""#));
|
|
assert!(json.contains(r#""budget_tokens":2048"#));
|
|
}
|
|
|
|
#[test]
|
|
fn test_thinking_config_disabled_serialization() {
|
|
let config = ThinkingConfig {
|
|
thinking_type: "disabled".to_string(),
|
|
budget_tokens: None,
|
|
};
|
|
let json = serde_json::to_string(&config).unwrap();
|
|
assert_eq!(json, r#"{"type":"disabled"}"#);
|
|
}
|
|
|
|
#[test]
|
|
fn test_system_content_serialization() {
|
|
let content = SystemContent {
|
|
content_type: "text".to_string(),
|
|
text: "You are helpful.".to_string(),
|
|
};
|
|
let json = serde_json::to_string(&content).unwrap();
|
|
assert!(json.contains(r#""type":"text""#));
|
|
}
|
|
|
|
#[test]
|
|
fn test_sse_event_parsing_content_block_start() {
|
|
let json = r#"{"type":"content_block_start","index":0,"content_block":{"type":"thinking","thinking":""}}"#;
|
|
let event: SseEvent = serde_json::from_str(json).unwrap();
|
|
assert_eq!(event.event_type, "content_block_start");
|
|
assert_eq!(event.content_block.unwrap().content_type, "thinking");
|
|
}
|
|
|
|
#[test]
|
|
fn test_sse_event_parsing_text_delta() {
|
|
let json = r#"{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}}"#;
|
|
let event: SseEvent = serde_json::from_str(json).unwrap();
|
|
assert_eq!(event.event_type, "content_block_delta");
|
|
assert_eq!(event.delta.unwrap().text, Some("Hello".to_string()));
|
|
}
|
|
|
|
#[test]
|
|
fn test_anthropic_content_text() {
|
|
let msg = AnthropicMessage {
|
|
role: "user".to_string(),
|
|
content: AnthropicContent::Text("Hello".to_string()),
|
|
};
|
|
let json = serde_json::to_string(&msg).unwrap();
|
|
assert!(json.contains(r#""content":"Hello""#));
|
|
}
|
|
}
|