| @@ -1,5 +1,6 @@ | |||||
| use crate::benchmarks; | use crate::benchmarks; | ||||
| use crate::config::AppConfig; | use crate::config::AppConfig; | ||||
| use crate::runner::{ModelRequest, run_model_request}; | |||||
| use anyhow::Result; | use anyhow::Result; | ||||
| use clap::{Parser, Subcommand}; | use clap::{Parser, Subcommand}; | ||||
| use std::path::{Path, PathBuf}; | use std::path::{Path, PathBuf}; | ||||
| @@ -86,7 +87,30 @@ pub enum BenchCommand { | |||||
| pub async fn dispatch(cli: Cli) -> Result<()> { | pub async fn dispatch(cli: Cli) -> Result<()> { | ||||
| match cli.command { | 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::Dataset { | ||||
| command: DatasetCommand::Fetch { 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 anthropic; | ||||
| pub mod openai; | 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) | |||||
| } | |||||