|
- 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<ModelResponse> {
- 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<ModelResponse> {
- 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<u128> = 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::<AnthropicStreamDelta>(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<AnthropicContentBlock>,
- }
-
- #[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<String>,
- }
-
- #[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");
- }
- }
|