| @@ -1,5 +1,6 @@ | |||
| use crate::benchmarks; | |||
| use crate::config::AppConfig; | |||
| use crate::runner::{ModelRequest, run_model_request}; | |||
| use anyhow::Result; | |||
| use clap::{Parser, Subcommand}; | |||
| use std::path::{Path, PathBuf}; | |||
| @@ -86,7 +87,30 @@ pub enum BenchCommand { | |||
| pub async fn dispatch(cli: Cli) -> Result<()> { | |||
| match cli.command { | |||
| Command::Check { .. } => anyhow::bail!("check is not implemented yet"), | |||
| Command::Check { | |||
| config, | |||
| provider, | |||
| model, | |||
| prompt, | |||
| } => { | |||
| let config = AppConfig::load(&config)?; | |||
| let provider = config.provider(provider.as_deref())?; | |||
| let request = ModelRequest { | |||
| base_url: provider.base_url.clone(), | |||
| api_token: provider.api_token.clone(), | |||
| model: model.unwrap_or_else(|| provider.default_model.clone()), | |||
| prompt, | |||
| temperature: 0.0, | |||
| max_tokens: 1024, | |||
| }; | |||
| let response = run_model_request(provider.protocol, request).await?; | |||
| println!("status: {}", response.status); | |||
| println!("elapsed_ms: {}", response.elapsed_ms); | |||
| println!("{}", response.text); | |||
| Ok(()) | |||
| } | |||
| Command::Dataset { | |||
| command: DatasetCommand::Fetch { dataset }, | |||
| } => { | |||
| @@ -1 +1,136 @@ | |||
| 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")); | |||
| } | |||
| } | |||
| @@ -1,2 +1,19 @@ | |||
| pub mod anthropic; | |||
| pub mod openai; | |||
| use anyhow::{Context, Result}; | |||
| use reqwest::Url; | |||
| fn endpoint_url(base_url: &str, path: &str) -> Result<Url> { | |||
| let base = if base_url.ends_with('/') { | |||
| base_url.to_string() | |||
| } else { | |||
| format!("{base_url}/") | |||
| }; | |||
| let path = path.trim_start_matches('/'); | |||
| Url::parse(&base) | |||
| .with_context(|| format!("invalid base_url: {base_url}"))? | |||
| .join(path) | |||
| .with_context(|| format!("failed to join endpoint path: /{path}")) | |||
| } | |||
| @@ -1 +1,139 @@ | |||
| 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, "/chat/completions")?; | |||
| let response = client | |||
| .post(url) | |||
| .bearer_auth(&request.api_token) | |||
| .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 OpenAI chat completion request")?; | |||
| let status = response.status(); | |||
| let status_code = status.as_u16(); | |||
| let body = response | |||
| .text() | |||
| .await | |||
| .context("failed to read OpenAI response body")?; | |||
| if !status.is_success() { | |||
| bail!("OpenAI request failed with status {status_code}: {body}"); | |||
| } | |||
| let parsed: OpenAiResponse = | |||
| serde_json::from_str(&body).context("failed to parse OpenAI response JSON")?; | |||
| let text = parsed | |||
| .choices | |||
| .into_iter() | |||
| .next() | |||
| .and_then(|choice| choice.message.content) | |||
| .filter(|content| !content.is_empty()) | |||
| .context("OpenAI response missing choices[0].message.content")?; | |||
| Ok(ModelResponse { | |||
| text, | |||
| status: status_code, | |||
| elapsed_ms: 0, | |||
| }) | |||
| } | |||
| #[derive(Debug, Deserialize)] | |||
| struct OpenAiResponse { | |||
| choices: Vec<OpenAiChoice>, | |||
| } | |||
| #[derive(Debug, Deserialize)] | |||
| struct OpenAiChoice { | |||
| message: OpenAiMessage, | |||
| } | |||
| #[derive(Debug, Deserialize)] | |||
| struct OpenAiMessage { | |||
| content: 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_chat_completion_text() { | |||
| let server = MockServer::start().await; | |||
| Mock::given(method("POST")) | |||
| .and(path("/chat/completions")) | |||
| .and(header("authorization", "Bearer test-token")) | |||
| .and(body_json(serde_json::json!({ | |||
| "model": "gpt-test", | |||
| "messages": [{"role": "user", "content": "hello"}], | |||
| "temperature": 0.0, | |||
| "max_tokens": 1024 | |||
| }))) | |||
| .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ | |||
| "choices": [{ | |||
| "message": { | |||
| "content": "hi back" | |||
| } | |||
| }] | |||
| }))) | |||
| .mount(&server) | |||
| .await; | |||
| let request = ModelRequest { | |||
| base_url: server.uri(), | |||
| api_token: "test-token".to_string(), | |||
| model: "gpt-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 back"); | |||
| } | |||
| #[tokio::test] | |||
| async fn errors_when_chat_completion_content_is_missing() { | |||
| let server = MockServer::start().await; | |||
| Mock::given(method("POST")) | |||
| .and(path("/chat/completions")) | |||
| .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ | |||
| "choices": [{ | |||
| "message": {} | |||
| }] | |||
| }))) | |||
| .mount(&server) | |||
| .await; | |||
| let request = ModelRequest { | |||
| base_url: format!("{}/", server.uri()), | |||
| api_token: "test-token".to_string(), | |||
| model: "gpt-test".to_string(), | |||
| prompt: "hello".to_string(), | |||
| temperature: 0.0, | |||
| max_tokens: 1024, | |||
| }; | |||
| let error = super::send(&Client::new(), &request) | |||
| .await | |||
| .expect_err("missing content should fail"); | |||
| assert!(error.to_string().contains("missing")); | |||
| } | |||
| } | |||
| @@ -1 +1,44 @@ | |||
| use crate::config::ProtocolKind; | |||
| use crate::protocols; | |||
| use anyhow::Result; | |||
| use reqwest::Client; | |||
| use std::time::Instant; | |||
| #[derive(Debug, Clone)] | |||
| pub struct ModelRequest { | |||
| pub base_url: String, | |||
| pub api_token: String, | |||
| pub model: String, | |||
| pub prompt: String, | |||
| pub temperature: f32, | |||
| pub max_tokens: u32, | |||
| } | |||
| #[derive(Debug, Clone, PartialEq, Eq)] | |||
| pub struct ModelResponse { | |||
| pub text: String, | |||
| pub status: u16, | |||
| pub elapsed_ms: u128, | |||
| } | |||
| pub async fn run_model_request( | |||
| protocol: ProtocolKind, | |||
| request: ModelRequest, | |||
| ) -> Result<ModelResponse> { | |||
| let client = Client::new(); | |||
| run_model_request_with_client(&client, protocol, &request).await | |||
| } | |||
| pub async fn run_model_request_with_client( | |||
| client: &Client, | |||
| protocol: ProtocolKind, | |||
| request: &ModelRequest, | |||
| ) -> Result<ModelResponse> { | |||
| let started = Instant::now(); | |||
| let mut response = match protocol { | |||
| ProtocolKind::Openai => protocols::openai::send(client, request).await?, | |||
| ProtocolKind::Anthropic => protocols::anthropic::send(client, request).await?, | |||
| }; | |||
| response.elapsed_ms = started.elapsed().as_millis(); | |||
| Ok(response) | |||
| } | |||