feat: implement OpenRouter HTTP client with SSE streaming
This commit is contained in:
parent
dfd037acac
commit
463c6a487c
2 changed files with 362 additions and 0 deletions
360
src/client.rs
Normal file
360
src/client.rs
Normal file
|
|
@ -0,0 +1,360 @@
|
|||
use async_trait::async_trait;
|
||||
use futures_util::StreamExt;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use std::fmt;
|
||||
use std::time::Duration;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
#[derive(thiserror::Error, Debug)]
|
||||
pub enum ClientError {
|
||||
#[error("Authentication failed: invalid or revoked OpenRouter API key")]
|
||||
Unauthorized,
|
||||
#[error("Payment required: account balance exhausted")]
|
||||
PaymentRequired,
|
||||
#[error("Rate limited, retry after {0}s")]
|
||||
RateLimited(u64),
|
||||
#[error("Model unavailable: {0}")]
|
||||
ModelUnavailable(String),
|
||||
#[error("Server error ({status}): {message}")]
|
||||
ServerError { status: u16, message: String },
|
||||
#[error("Network error: {0}")]
|
||||
Network(#[from] reqwest::Error),
|
||||
#[error("Stream decoding error: {0}")]
|
||||
StreamError(String),
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait ModelClient: Send + Sync {
|
||||
async fn complete(&self, model: &str, system: &str, user: &str) -> Result<String, ClientError>;
|
||||
|
||||
async fn complete_stream(
|
||||
&self,
|
||||
model: &str,
|
||||
system: &str,
|
||||
user: &str,
|
||||
chunk_sender: mpsc::Sender<String>,
|
||||
) -> Result<String, ClientError>;
|
||||
}
|
||||
|
||||
pub struct OpenRouterClient {
|
||||
api_key: String,
|
||||
base_url: String,
|
||||
http: reqwest::Client,
|
||||
}
|
||||
|
||||
impl fmt::Debug for OpenRouterClient {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.debug_struct("OpenRouterClient")
|
||||
.field("api_key", &"[REDACTED]")
|
||||
.field("base_url", &self.base_url)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct ChatMessage<'a> {
|
||||
role: &'a str,
|
||||
content: &'a str,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct ChatCompletionRequest<'a> {
|
||||
model: &'a str,
|
||||
messages: Vec<ChatMessage<'a>>,
|
||||
stream: bool,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
max_tokens: Option<u32>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
temperature: Option<f32>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct ChatCompletionChoice {
|
||||
message: Option<ChatCompletionMessage>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct ChatCompletionMessage {
|
||||
content: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct ChatCompletionResponse {
|
||||
choices: Vec<ChatCompletionChoice>,
|
||||
}
|
||||
|
||||
impl OpenRouterClient {
|
||||
pub fn new(api_key: String) -> Self {
|
||||
Self::with_base_url(api_key, "https://openrouter.ai/api/v1".to_string())
|
||||
}
|
||||
|
||||
pub fn with_base_url(api_key: String, base_url: String) -> Self {
|
||||
let http = reqwest::Client::builder()
|
||||
.timeout(Duration::from_secs(60))
|
||||
.connect_timeout(Duration::from_secs(10))
|
||||
.build()
|
||||
.unwrap_or_else(|_| reqwest::Client::new());
|
||||
|
||||
Self {
|
||||
api_key,
|
||||
base_url,
|
||||
http,
|
||||
}
|
||||
}
|
||||
|
||||
/// Parses an SSE formatted line into a token delta if available
|
||||
pub fn parse_sse_line(line: &str) -> Option<String> {
|
||||
let trimmed = line.trim();
|
||||
if !trimmed.starts_with("data:") {
|
||||
return None;
|
||||
}
|
||||
|
||||
let payload = trimmed.strip_prefix("data:")?.trim();
|
||||
if payload == "[DONE]" || payload.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
if let Ok(json) = serde_json::from_str::<Value>(payload) {
|
||||
if let Some(choices) = json.get("choices").and_then(|c| c.as_array()) {
|
||||
if let Some(first) = choices.first() {
|
||||
// Try delta.content (streaming)
|
||||
if let Some(content) = first
|
||||
.get("delta")
|
||||
.and_then(|d| d.get("content"))
|
||||
.and_then(|c| c.as_str())
|
||||
{
|
||||
return Some(content.to_string());
|
||||
}
|
||||
// Try text fallback
|
||||
if let Some(text) = first.get("text").and_then(|t| t.as_str()) {
|
||||
return Some(text.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn check_status(status: reqwest::StatusCode, body: &str, headers: &reqwest::header::HeaderMap) -> Result<(), ClientError> {
|
||||
match status.as_u16() {
|
||||
200..=299 => Ok(()),
|
||||
401 => Err(ClientError::Unauthorized),
|
||||
402 => Err(ClientError::PaymentRequired),
|
||||
429 => {
|
||||
let retry_after = headers
|
||||
.get("retry-after")
|
||||
.and_then(|h| h.to_str().ok())
|
||||
.and_then(|s| s.parse::<u64>().ok())
|
||||
.unwrap_or(5);
|
||||
Err(ClientError::RateLimited(retry_after))
|
||||
}
|
||||
404 => Err(ClientError::ModelUnavailable(format!("Model not found (404): {body}"))),
|
||||
s if s >= 500 => Err(ClientError::ServerError {
|
||||
status: s,
|
||||
message: body.to_string(),
|
||||
}),
|
||||
s => Err(ClientError::ServerError {
|
||||
status: s,
|
||||
message: body.to_string(),
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn complete_with_params(
|
||||
&self,
|
||||
model: &str,
|
||||
system: &str,
|
||||
user: &str,
|
||||
max_tokens: Option<u32>,
|
||||
temperature: Option<f32>,
|
||||
) -> Result<String, ClientError> {
|
||||
let endpoint = format!("{}/chat/completions", self.base_url.trim_end_matches('/'));
|
||||
|
||||
let body = ChatCompletionRequest {
|
||||
model,
|
||||
messages: vec![
|
||||
ChatMessage {
|
||||
role: "system",
|
||||
content: system,
|
||||
},
|
||||
ChatMessage {
|
||||
role: "user",
|
||||
content: user,
|
||||
},
|
||||
],
|
||||
stream: false,
|
||||
max_tokens,
|
||||
temperature,
|
||||
};
|
||||
|
||||
let res = self
|
||||
.http
|
||||
.post(&endpoint)
|
||||
.header("Authorization", format!("Bearer {}", self.api_key))
|
||||
.header("HTTP-Referer", "https://github.com/brokecrew/brokecrew")
|
||||
.header("X-Title", "BrokeCrew")
|
||||
.header("Content-Type", "application/json")
|
||||
.json(&body)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
let status = res.status();
|
||||
let headers = res.headers().clone();
|
||||
|
||||
if !status.is_success() {
|
||||
let error_text = res.text().await.unwrap_or_default();
|
||||
return Err(Self::check_status(status, &error_text, &headers).unwrap_err());
|
||||
}
|
||||
|
||||
let resp_json: ChatCompletionResponse = res.json().await?;
|
||||
let reply = resp_json
|
||||
.choices
|
||||
.into_iter()
|
||||
.next()
|
||||
.and_then(|c| c.message)
|
||||
.and_then(|m| m.content)
|
||||
.unwrap_or_default();
|
||||
|
||||
Ok(reply)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ModelClient for OpenRouterClient {
|
||||
async fn complete(&self, model: &str, system: &str, user: &str) -> Result<String, ClientError> {
|
||||
self.complete_with_params(model, system, user, None, None).await
|
||||
}
|
||||
|
||||
async fn complete_stream(
|
||||
&self,
|
||||
model: &str,
|
||||
system: &str,
|
||||
user: &str,
|
||||
chunk_sender: mpsc::Sender<String>,
|
||||
) -> Result<String, ClientError> {
|
||||
let endpoint = format!("{}/chat/completions", self.base_url.trim_end_matches('/'));
|
||||
|
||||
let body = ChatCompletionRequest {
|
||||
model,
|
||||
messages: vec![
|
||||
ChatMessage {
|
||||
role: "system",
|
||||
content: system,
|
||||
},
|
||||
ChatMessage {
|
||||
role: "user",
|
||||
content: user,
|
||||
},
|
||||
],
|
||||
stream: true,
|
||||
max_tokens: None,
|
||||
temperature: None,
|
||||
};
|
||||
|
||||
let res = self
|
||||
.http
|
||||
.post(&endpoint)
|
||||
.header("Authorization", format!("Bearer {}", self.api_key))
|
||||
.header("HTTP-Referer", "https://github.com/brokecrew/brokecrew")
|
||||
.header("X-Title", "BrokeCrew")
|
||||
.header("Content-Type", "application/json")
|
||||
.json(&body)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
let status = res.status();
|
||||
let headers = res.headers().clone();
|
||||
|
||||
if !status.is_success() {
|
||||
let error_text = res.text().await.unwrap_or_default();
|
||||
return Err(Self::check_status(status, &error_text, &headers).unwrap_err());
|
||||
}
|
||||
|
||||
let mut stream = res.bytes_stream();
|
||||
let mut full_text = String::new();
|
||||
let mut buffer = String::new();
|
||||
|
||||
while let Some(item) = stream.next().await {
|
||||
let bytes = item.map_err(|e| ClientError::Network(e))?;
|
||||
let chunk_str = String::from_utf8_lossy(&bytes);
|
||||
buffer.push_str(&chunk_str);
|
||||
|
||||
while let Some(pos) = buffer.find('\n') {
|
||||
let line = buffer[..pos].trim().to_string();
|
||||
buffer.drain(..=pos);
|
||||
|
||||
if let Some(delta) = Self::parse_sse_line(&line) {
|
||||
full_text.push_str(&delta);
|
||||
let _ = chunk_sender.send(delta).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Parse remaining buffer if any
|
||||
if !buffer.trim().is_empty() {
|
||||
if let Some(delta) = Self::parse_sse_line(&buffer) {
|
||||
full_text.push_str(&delta);
|
||||
let _ = chunk_sender.send(delta).await;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(full_text)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_parse_sse_line_delta() {
|
||||
let line = r#"data: {"choices":[{"delta":{"content":"Hello world"}}]}"#;
|
||||
let delta = OpenRouterClient::parse_sse_line(line);
|
||||
assert_eq!(delta, Some("Hello world".to_string()));
|
||||
|
||||
let done_line = "data: [DONE]";
|
||||
assert_eq!(OpenRouterClient::parse_sse_line(done_line), None);
|
||||
|
||||
let ping_line = ": ping";
|
||||
assert_eq!(OpenRouterClient::parse_sse_line(ping_line), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_api_key_redaction_in_debug() {
|
||||
let client = OpenRouterClient::new("sk-or-v1-secret-key-123456789".to_string());
|
||||
let debug_str = format!("{:?}", client);
|
||||
assert!(!debug_str.contains("sk-or-v1-secret-key-123456789"));
|
||||
assert!(debug_str.contains("[REDACTED]"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_check_status_mappings() {
|
||||
let mut headers = reqwest::header::HeaderMap::new();
|
||||
headers.insert("retry-after", reqwest::header::HeaderValue::from_static("12"));
|
||||
|
||||
let err_401 = OpenRouterClient::check_status(reqwest::StatusCode::UNAUTHORIZED, "bad key", &headers).unwrap_err();
|
||||
match err_401 {
|
||||
ClientError::Unauthorized => (),
|
||||
_ => panic!("Expected Unauthorized, got {:?}", err_401),
|
||||
}
|
||||
|
||||
let err_402 = OpenRouterClient::check_status(reqwest::StatusCode::PAYMENT_REQUIRED, "low balance", &headers).unwrap_err();
|
||||
match err_402 {
|
||||
ClientError::PaymentRequired => (),
|
||||
_ => panic!("Expected PaymentRequired, got {:?}", err_402),
|
||||
}
|
||||
|
||||
let err_429 = OpenRouterClient::check_status(reqwest::StatusCode::TOO_MANY_REQUESTS, "rate limited", &headers).unwrap_err();
|
||||
match err_429 {
|
||||
ClientError::RateLimited(secs) => assert_eq!(secs, 12),
|
||||
_ => panic!("Expected RateLimited(12), got {:?}", err_429),
|
||||
}
|
||||
|
||||
let err_503 = OpenRouterClient::check_status(reqwest::StatusCode::SERVICE_UNAVAILABLE, "outage", &headers).unwrap_err();
|
||||
match err_503 {
|
||||
ClientError::ServerError { status, .. } => assert_eq!(status, 503),
|
||||
_ => panic!("Expected ServerError 503, got {:?}", err_503),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -1,7 +1,9 @@
|
|||
//! BrokeCrew: A Rust-Native Multi-Model AI Team Library
|
||||
|
||||
pub mod budget;
|
||||
pub mod client;
|
||||
pub mod persona;
|
||||
|
||||
pub use budget::{BudgetExceeded, BudgetStatus, BudgetTracker, CreditTier};
|
||||
pub use client::{ClientError, ModelClient, OpenRouterClient};
|
||||
pub use persona::{DispatchResult, ModelAttempt, ModelCandidate, Persona, PersonaRegistry, Tier};
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue