use crate::runner::{ModelRequest, ModelResponse}; use anyhow::{Context, Result, bail}; use futures::StreamExt; use reqwest::Client; use serde::Deserialize; use serde_json::json; use std::time::Instant; pub async fn send(client: &Client, request: &ModelRequest) -> Result { let url = super::endpoint_url(&request.base_url, "/v1/messages")?; let response = client .post(url) .header("x-api-key", &request.api_token) .header("anthropic-version", "2023-06-01") .json(&json!({ "model": request.model, "messages": [{"role": "user", "content": request.prompt}], "temperature": request.temperature, "max_tokens": request.max_tokens })) .send() .await .context("failed to send Anthropic messages request")?; let status = response.status(); let status_code = status.as_u16(); let body = response .text() .await .context("failed to read Anthropic response body")?; if !status.is_success() { bail!( "{}", super::upstream_error_message("Anthropic", status_code, &body) ); } let parsed: AnthropicResponse = serde_json::from_str(&body).context("failed to parse Anthropic response JSON")?; let text = parsed .content .into_iter() .find_map(|block| match block { AnthropicContentBlock::Text { text } if !text.is_empty() => Some(text), _ => None, }) .context("Anthropic response missing first text content block")?; Ok(ModelResponse { text, status: status_code, elapsed_ms: 0, first_token_ms: None, }) } pub async fn send_stream(client: &Client, request: &ModelRequest) -> Result { let url = super::endpoint_url(&request.base_url, "/v1/messages")?; let started = Instant::now(); let response = client .post(url) .header("x-api-key", &request.api_token) .header("anthropic-version", "2023-06-01") .json(&json!({ "model": request.model, "messages": [{"role": "user", "content": request.prompt}], "temperature": request.temperature, "max_tokens": request.max_tokens, "stream": true })) .send() .await .context("failed to send Anthropic streaming request")?; let status = response.status(); let status_code = status.as_u16(); if !status.is_success() { let body = response .text() .await .context("failed to read Anthropic error response body")?; bail!( "{}", super::upstream_error_message("Anthropic", status_code, &body) ); } let mut stream = response.bytes_stream(); let mut buffer = super::SseLineBuffer::new(); let mut text = String::new(); let mut first_token_ms: Option = None; let mut current_event = String::new(); let mut done = false; while let Some(chunk) = stream.next().await { let chunk = chunk.context("Anthropic stream interrupted")?; for line in buffer.feed(&chunk) { if let Some(event_type) = line.strip_prefix("event: ") { current_event = event_type.to_string(); continue; } if let Some(data) = line.strip_prefix("data: ") { if current_event == "content_block_delta" { if let Some(content) = serde_json::from_str::(data) .ok() .and_then(|p| p.delta.text) .filter(|c| !c.is_empty()) { if first_token_ms.is_none() { first_token_ms = Some(started.elapsed().as_millis()); } text.push_str(&content); } } else if current_event == "message_stop" { done = true; break; } } if line.is_empty() { current_event.clear(); } } if done { break; } } if text.is_empty() { bail!("Anthropic stream completed without producing any content"); } Ok(ModelResponse { text, status: status_code, elapsed_ms: 0, first_token_ms, }) } #[derive(Debug, Deserialize)] struct AnthropicResponse { content: Vec, } #[derive(Debug, Deserialize)] #[serde(tag = "type")] enum AnthropicContentBlock { #[serde(rename = "text")] Text { text: String }, #[serde(other)] Other, } #[derive(Debug, Deserialize)] struct AnthropicStreamDelta { delta: AnthropicDeltaContent, } #[derive(Debug, Deserialize)] struct AnthropicDeltaContent { text: Option, } #[cfg(test)] mod tests { use crate::runner::ModelRequest; use reqwest::Client; use wiremock::matchers::{body_json, header, method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; #[tokio::test] async fn extracts_first_text_block() { let server = MockServer::start().await; Mock::given(method("POST")) .and(path("/v1/messages")) .and(header("x-api-key", "test-token")) .and(header("anthropic-version", "2023-06-01")) .and(body_json(serde_json::json!({ "model": "claude-test", "messages": [{"role": "user", "content": "hello"}], "temperature": 0.0, "max_tokens": 1024 }))) .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ "content": [ {"type": "text", "text": "hi from claude"} ] }))) .mount(&server) .await; let request = ModelRequest { base_url: server.uri(), api_token: "test-token".to_string(), model: "claude-test".to_string(), prompt: "hello".to_string(), temperature: 0.0, max_tokens: 1024, stream: false }; let response = super::send(&Client::new(), &request) .await .expect("response"); assert_eq!(response.status, 200); assert_eq!(response.text, "hi from claude"); } #[tokio::test] async fn errors_on_non_success_status() { let server = MockServer::start().await; Mock::given(method("POST")) .and(path("/v1/messages")) .respond_with(ResponseTemplate::new(401).set_body_json(serde_json::json!({ "error": { "message": "unauthorized" } }))) .mount(&server) .await; let request = ModelRequest { base_url: server.uri(), api_token: "test-token".to_string(), model: "claude-test".to_string(), prompt: "hello".to_string(), temperature: 0.0, max_tokens: 1024, stream: false }; let error = super::send(&Client::new(), &request) .await .expect_err("non-success should fail"); assert!(error.to_string().contains("401")); assert!(error.to_string().contains("unauthorized")); } #[tokio::test] async fn non_success_error_does_not_leak_sensitive_response_body() { let server = MockServer::start().await; Mock::given(method("POST")) .and(path("/v1/messages")) .respond_with(ResponseTemplate::new(429).set_body_json(serde_json::json!({ "error": { "message": "rate limited secret prompt sk-leaked-token" } }))) .mount(&server) .await; let request = ModelRequest { base_url: server.uri(), api_token: "test-token".to_string(), model: "claude-test".to_string(), prompt: "hello".to_string(), temperature: 0.0, max_tokens: 1024, stream: false }; let message = super::send(&Client::new(), &request) .await .expect_err("non-success should fail") .to_string(); assert!(message.contains("Anthropic")); assert!(message.contains("429")); assert!(message.contains("rate limited")); assert!(!message.contains("secret prompt")); assert!(!message.contains("sk-leaked-token")); } #[tokio::test] async fn base_url_with_v1_prefix_does_not_duplicate_messages_path() { let server = MockServer::start().await; Mock::given(method("POST")) .and(path("/v1/messages")) .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ "content": [ {"type": "text", "text": "prefixed"} ] }))) .mount(&server) .await; let request = ModelRequest { base_url: format!("{}/v1", server.uri()), api_token: "test-token".to_string(), model: "claude-test".to_string(), prompt: "hello".to_string(), temperature: 0.0, max_tokens: 1024, stream: false }; let response = super::send(&Client::new(), &request) .await .expect("response"); assert_eq!(response.text, "prefixed"); } }