|
- use crate::runner::{ModelRequest, ModelResponse};
- use anyhow::{Context, Result, bail};
- use reqwest::Client;
- use serde::Deserialize;
- use serde_json::json;
-
- 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!("Anthropic request failed with status {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,
- })
- }
-
- #[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,
- }
-
- #[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,
- };
-
- 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_string("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,
- };
-
- 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"));
- }
- }
|