diff --git a/src/cli.rs b/src/cli.rs index a58cdb2..8ca2f57 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -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 }, } => { diff --git a/src/protocols/anthropic.rs b/src/protocols/anthropic.rs index 8b13789..31f7db7 100644 --- a/src/protocols/anthropic.rs +++ b/src/protocols/anthropic.rs @@ -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 { + 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, +} + +#[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")); + } +} diff --git a/src/protocols/mod.rs b/src/protocols/mod.rs index 943612b..e5d661d 100644 --- a/src/protocols/mod.rs +++ b/src/protocols/mod.rs @@ -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 { + 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}")) +} diff --git a/src/protocols/openai.rs b/src/protocols/openai.rs index 8b13789..1a91a9f 100644 --- a/src/protocols/openai.rs +++ b/src/protocols/openai.rs @@ -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 { + 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, +} + +#[derive(Debug, Deserialize)] +struct OpenAiChoice { + message: OpenAiMessage, +} + +#[derive(Debug, Deserialize)] +struct OpenAiMessage { + content: 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_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")); + } +} diff --git a/src/runner.rs b/src/runner.rs index 8b13789..14516d1 100644 --- a/src/runner.rs +++ b/src/runner.rs @@ -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 { + 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 { + 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) +}