diff --git a/src/client.rs b/src/client.rs new file mode 100644 index 0000000..ef504ae --- /dev/null +++ b/src/client.rs @@ -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; + + async fn complete_stream( + &self, + model: &str, + system: &str, + user: &str, + chunk_sender: mpsc::Sender, + ) -> Result; +} + +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>, + stream: bool, + #[serde(skip_serializing_if = "Option::is_none")] + max_tokens: Option, + #[serde(skip_serializing_if = "Option::is_none")] + temperature: Option, +} + +#[derive(Deserialize)] +struct ChatCompletionChoice { + message: Option, +} + +#[derive(Deserialize)] +struct ChatCompletionMessage { + content: Option, +} + +#[derive(Deserialize)] +struct ChatCompletionResponse { + choices: Vec, +} + +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 { + 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::(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::().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, + temperature: Option, + ) -> Result { + 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 { + 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, + ) -> Result { + 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), + } + } +} diff --git a/src/lib.rs b/src/lib.rs index 59f87f3..e62adcd 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -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};