Просмотр исходного кода

feat: add openai and anthropic request adapters

main
orangels 1 неделю назад
Родитель
Сommit
1e0926cfba
5 измененных файлов: 358 добавлений и 1 удалений
  1. +25
    -1
      src/cli.rs
  2. +135
    -0
      src/protocols/anthropic.rs
  3. +17
    -0
      src/protocols/mod.rs
  4. +138
    -0
      src/protocols/openai.rs
  5. +43
    -0
      src/runner.rs

+ 25
- 1
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 },
} => {


+ 135
- 0
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<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"));
}
}

+ 17
- 0
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<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}"))
}

+ 138
- 0
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<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"));
}
}

+ 43
- 0
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<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)
}

Загрузка…
Отмена
Сохранить