diff --git a/src/lib.rs b/src/lib.rs new file mode 100644 index 0000000..cce761f --- /dev/null +++ b/src/lib.rs @@ -0,0 +1,69 @@ +use clap::Parser; +use secrecy::SecretString; +use std::{convert::Infallible, net::IpAddr, str::FromStr}; +use tracing_subscriber::EnvFilter; + +pub mod openai; +pub mod server; + +#[derive(Debug, Clone)] +pub struct CloudflareCredentials { + pub account_id: String, + pub api_key: SecretString, +} + +impl FromStr for CloudflareCredentials { + type Err = Infallible; // TODO + + fn from_str(s: &str) -> Result { + let (account_id, api_key) = s.split_once(":").expect("TODO"); + + let account_id = account_id.into(); + let api_key = api_key.into(); + + Ok(Self { + account_id, + api_key, + }) + } +} + +#[derive(Debug, Clone, Parser)] +pub struct AppArgs { + #[arg(long, env, default_value_t = [127, 0, 0, 1].into())] + pub ip: IpAddr, + + #[arg(long, env, default_value_t = 3126)] + pub port: u16, + + #[arg(long)] + pub cloudflare_workers_ai_credentials: Vec, + + #[arg(long)] + pub google_ai_studio_api_key: Vec, + + #[arg(long)] + pub groq_api_key: Vec, + + #[arg(long)] + pub hugging_face_api_key: Vec, + + #[arg(long)] + pub mistral_api_key: Vec, + + #[arg(long)] + pub openrouter_api_key: Vec, +} + +#[derive(Parser)] +pub struct Args { + #[arg( + long = "tracing-directives", + env = "RUST_LOG", + default_value = "warn,katniss_ai=debug" + )] + pub env_filter: EnvFilter, + + #[command(flatten)] + pub app_args: AppArgs, +} diff --git a/src/main.rs b/src/main.rs index 92d2929..78748c8 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,55 +1,11 @@ -use std::net::{IpAddr, SocketAddr}; - use clap::Parser; -use futures::StreamExt as _; -use reqwest_sse::EventSource; -use secrecy::SecretString; -use tokio::{join, pin, task::LocalSet}; -use tracing_subscriber::{EnvFilter, fmt::format::FmtSpan}; - -mod openai; -mod server; - -#[derive(Parser, Debug)] -struct AppArgs { - #[arg(long, env, default_value_t = [127, 0, 0, 1].into())] - ip: IpAddr, - - #[arg(long, env, default_value_t = 3126)] - port: u16, - - #[arg(long)] - cloudflare_workers_ai_api_key: Vec, - - #[arg(long)] - google_ai_studio_api_key: Vec, - - #[arg(long)] - groq_api_key: Vec, - - #[arg(long)] - hugging_face_api_key: Vec, - - #[arg(long)] - mistral_api_key: Vec, - - #[arg(long)] - openrouter_api_key: Vec, -} - -#[derive(Parser)] -struct Args { - #[arg(long = "tracing-directives", env = "RUST_LOG", default_value = "warn,katniss_ai=debug")] - env_filter: EnvFilter, - - #[command(flatten)] - app_args: AppArgs, -} +use std::net::SocketAddr; +use tracing_subscriber::fmt::format::FmtSpan; #[tokio::main] async fn main() { let args = Parser::parse(); - let Args { + let katniss_ai::Args { env_filter, app_args, } = args; @@ -61,54 +17,20 @@ async fn main() { .init(); tracing::debug!(?app_args, "using"); - rustls::crypto::aws_lc_rs::default_provider().install_default().unwrap(); + rustls::crypto::aws_lc_rs::default_provider() + .install_default() + .unwrap(); let addr = SocketAddr::new(app_args.ip, app_args.port); - let (tx, rx) = async_channel::unbounded(); - - let app = server::create_app( - server::AppState::builder() + let app = katniss_ai::server::create_app( + katniss_ai::server::AppState::builder() .app_args(app_args.into()) - .stream_requester(tx) .build(), ); let listener = tokio::net::TcpListener::bind(addr).await.expect("TODO"); - let stream_responses = async move { - let local_set = LocalSet::new(); - - local_set - .run_until(async { - pin!(rx); - - while let Some((response, sse_sink)) = rx.next().await { - local_set.spawn_local(async move { - match response.events().await { - Ok(mut events) => { - while let Some(event_result) = events.next().await { - let message = event_result - .map(|event| { - axum::response::sse::Event::default().data(event.data) - }) - .map_err(Into::into); - - let _ = sse_sink.send(message).await; - } - } - Err(event_source_error) => { - let _ = sse_sink.send(Err(event_source_error.into())).await; - } - }; - }); - } - }) - .await; - }; - - join!(stream_responses, async move { - tracing::info!(?addr, "listening"); - axum::serve(listener, app).await.expect("TODO"); - }); + tracing::info!(?addr, "listening"); + axum::serve(listener, app).await.expect("TODO"); } diff --git a/src/openai/chat/completion.rs b/src/openai/chat/completion.rs index ade639a..01c1ec2 100644 --- a/src/openai/chat/completion.rs +++ b/src/openai/chat/completion.rs @@ -1,7 +1,7 @@ use serde::{Deserialize, Serialize}; use serde_with::{DeserializeFromStr, SerializeDisplay}; -use strum::{EnumString, VariantArray}; use std::collections::HashMap; +use strum::{EnumString, VariantArray}; // ============================================================================ // Request: Create Chat Completion @@ -918,7 +918,9 @@ pub enum PromptCacheRetention { } /// Constrains reasoning effort for reasoning models. -#[derive(Clone, Debug, EnumString, DeserializeFromStr, strum::Display, SerializeDisplay, VariantArray)] +#[derive( + Clone, Debug, EnumString, DeserializeFromStr, strum::Display, SerializeDisplay, VariantArray, +)] #[strum(serialize_all = "lowercase")] pub enum ReasoningEffort { /// Disable reasoning. diff --git a/src/server.rs b/src/server.rs index 042dbb2..f7569d6 100644 --- a/src/server.rs +++ b/src/server.rs @@ -1,4 +1,5 @@ use crate::AppArgs; +use crate::CloudflareCredentials; use crate::openai; use axum::Json; use axum::Router; @@ -7,9 +8,9 @@ use axum::http::StatusCode; use axum::response::IntoResponse; use axum::response::Response; use axum::response::Sse; -use axum_extra::routing::RouterExt; -use axum_extra::routing::TypedPath; +use axum_extra::routing::{RouterExt, TypedPath}; use ecow::EcoString; +use futures::StreamExt; use headers::Header; use headers::authorization::Bearer; use itertools::iproduct; @@ -17,31 +18,26 @@ use mitsein::vec1::Vec1; use rand::Rng; use rand::RngExt; use reqwest::Certificate; +use reqwest_sse::EventSource as _; use secrecy::ExposeSecret; use secrecy::SecretString; use serde::Deserialize; -use serde::Serialize; use serde_with::DeserializeFromStr; use serde_with::SerializeDisplay; -use strum::EnumString; -use strum::VariantArray; use std::collections::VecDeque; use std::convert::Infallible; use std::fmt::Display; use std::str::FromStr; use std::sync::Arc; +use strum::EnumString; +use strum::VariantArray; +use tokio::runtime::Handle; use typed_builder::TypedBuilder; use url::Url; -type SSEEventResult = Result>; -type StreamRequest = (reqwest::Response, async_channel::Sender); -type StreamRequester = async_channel::Sender; - #[derive(Debug, Clone, TypedBuilder)] pub struct AppState { app_args: Arc, - - stream_requester: StreamRequester, } #[derive(Deserialize, TypedPath)] @@ -58,7 +54,9 @@ enum Provider { OpenRouter, } -#[derive(Clone, Debug, EnumString, DeserializeFromStr, strum::Display, SerializeDisplay, VariantArray)] +#[derive( + Clone, Debug, EnumString, DeserializeFromStr, strum::Display, SerializeDisplay, VariantArray, +)] #[strum(serialize_all = "lowercase")] enum Shortcut { XS, @@ -70,7 +68,6 @@ enum Shortcut { Experimental, } - #[derive(Debug, Clone, DeserializeFromStr, SerializeDisplay)] struct ShortcutAndMaybeReasoning(Shortcut, Option); @@ -106,7 +103,6 @@ impl Display for ShortcutAndMaybeReasoning { } } - impl Shortcut { fn suggestions(&self) -> Vec<(Provider, EcoString)> { use Provider as P; @@ -139,6 +135,7 @@ impl Shortcut { (P::HuggingFace, "Qwen/Qwen3.8-27B".into()), (P::HuggingFace, "meta-models/Muse-Glimmer-30B".into()), (P::OpenRouter, "google/gemma-4-31b-it:free".into()), + (P::HuggingFace, "google/gemma-4-31B-it".into()), (P::Mistral, "mistral-small-2603".into()), (P::Groq, "groq/compound".into()), (P::OpenRouter, "openrouter/fusion".into()), @@ -150,6 +147,10 @@ impl Shortcut { (P::Mistral, "mistral-medium-3.5".into()), (P::OpenRouter, "inclusionai/ling-3.0-flash-vl:free".into()), (P::OpenRouter, "poolside/laguna-s-2.1:free".into()), + ( + P::OpenRouter, + "inclusionai/ling-3.0-flash-sante:free".into(), + ), ( P::OpenRouter, "nvidia/nemotron-3-super-120b-a12b:free".into(), @@ -162,28 +163,29 @@ impl Shortcut { S::L => vec![ (P::CloudflareWorkersAi, "@cf/zai-org/glm-5.3-flash".into()), (P::HuggingFace, "zai-org/GLM-5.3-Flash".into()), + (P::HuggingFace, "zai-org/GLM-5.3-Flash-BF16".into()), (P::HuggingFace, "deepseek-ai/DeepSeek-V4.1-Flash".into()), - (P::OpenRouter, "nex-agi/nex-n2.5-pro:free".into()), + (P::HuggingFace, "XiaomiMiMo/MiMo-V2.6-Flash-RL".into()), + (P::OpenRouter, "inclusionai/ling-3.1-flash".into()), + (P::OpenRouter, "dots-studio/dots-3-note-preview:free".into()), (P::HuggingFace, "MiniMaxAI/MiniMax-M3".into()), (P::OpenRouter, "thinkingmachines/inkling-small:free".into()), (P::HuggingFace, "thinkingmachines/Inkling-Small".into()), ], S::XL => vec![ - (P::CloudflareWorkersAi, "@cf/zai-org/glm-5.3".into()), (P::HuggingFace, "zai-org/GLM-5.3".into()), + (P::CloudflareWorkersAi, "@cf/zai-org/glm-5.3".into()), (P::Mistral, "zai-glm-5-3".into()), + (P::HuggingFace, "tencent/Hy4-preview".into()), + (P::HuggingFace, "XiaomiMiMo/MiMo-V2.6-Pro-RL".into()), (P::HuggingFace, "moonshotai/Kimi-K3".into()), (P::HuggingFace, "Qwen/Qwen3.8-2.4T-A95B".into()), (P::OpenRouter, "thinkingmachines/inkling:free".into()), (P::HuggingFace, "thinkingmachines/Inkling".into()), - (P::HuggingFace, "tencent/Hy4-preview".into()), - (P::HuggingFace, "deepseek-ai/DeepSeek-V4-Pro".into()), ], - S::Experimental => vec![ - (P::OpenRouter, "space-bunny-alpha".into()) - ] + S::Experimental => vec![(P::OpenRouter, "space-bunny-alpha".into())], } } @@ -225,10 +227,7 @@ struct Upstream { enum ChatCompletionResponse { Immediate(openai::chat::completion::ChatCompletion), - Stream { - response: reqwest::Response, - stream_requester: StreamRequester, - }, + Stream(reqwest::Response), } impl IntoResponse for ChatCompletionResponse { @@ -237,17 +236,33 @@ impl IntoResponse for ChatCompletionResponse { ChatCompletionResponse::Immediate(chat_completion) => { Json(chat_completion).into_response() } - ChatCompletionResponse::Stream { - response, - stream_requester, - } => { - let (sse_sink, sse_stream) = async_channel::bounded(256); + ChatCompletionResponse::Stream(response) => { + let (sink, stream) = async_channel::bounded(256); - stream_requester - .try_send((response, sse_sink)) - .expect("this is an unbounded channel"); + let handle = Handle::current(); - Sse::new(sse_stream).into_response() + std::thread::spawn(move || { + handle.block_on(async move { + match response.events().await { + Ok(mut events) => { + while let Some(event_result) = events.next().await { + let message = event_result + .map(|event| { + axum::response::sse::Event::default().data(event.data) + }) + .map_err(Box::::from); + + let _ = sink.send(message).await; + } + } + Err(event_source_error) => { + let _ = sink.send(Err(event_source_error.into())).await; + } + }; + }); + }); + + Sse::new(stream).into_response() } } } @@ -270,79 +285,102 @@ impl IntoResponse for ChatCompletionError { async fn create_chat_completion( _: V1ChatCompletionsPath, State(state): State, - Json(request): Json>, + Json(request): Json< + openai::chat::completion::CreateChatCompletionRequest, + >, ) -> Result { let ShortcutAndMaybeReasoning(shortcut, reasoning_in_id) = request.model.clone(); let mut model_upstreams = VecDeque::new(); for (provider, model) in shortcut.suggestions() { - let upstreams = match provider { - Provider::CloudflareWorkersAi => vec![], // TODO - Provider::GoogleAiStudio => state - .app_args - .google_ai_studio_api_key - .iter() - .cloned() - .map(|api_key| { - Upstream::builder() + let upstreams: Vec = + match provider { + Provider::CloudflareWorkersAi => state + .app_args + .cloudflare_workers_ai_credentials + .iter() + .cloned() + .map( + |CloudflareCredentials { + account_id, + api_key, + }| { + Upstream::builder() .base_url( - "https://generativelanguage.googleapis.com/v1beta/openai" + format!("https://api.cloudflare.com/client/v4/accounts/{account_id}/ai") .parse() .expect("TODO"), ) .api_key(api_key) .build() - }) - .collect(), - Provider::Groq => state - .app_args - .groq_api_key - .iter() - .cloned() - .map(|api_key| { - Upstream::builder() - .base_url("https://api.groq.com/openai".parse().expect("TODO")) - .api_key(api_key) - .build() - }) - .collect(), - Provider::HuggingFace => state - .app_args - .hugging_face_api_key - .iter() - .cloned() - .map(|api_key| { - Upstream::builder() - .base_url("https://router.huggingface.co".parse().expect("TODO")) - .api_key(api_key) - .build() - }) - .collect(), - Provider::Mistral => state - .app_args - .mistral_api_key - .iter() - .cloned() - .map(|api_key| { - Upstream::builder() - .base_url("https://api.mistral.ai".parse().expect("TODO")) - .api_key(api_key) - .build() - }) - .collect(), - Provider::OpenRouter => state - .app_args - .openrouter_api_key - .iter() - .cloned() - .map(|api_key| { - Upstream::builder() - .base_url("https://openrouter.ai/api".parse().expect("TODO")) - .api_key(api_key) - .build() - }) - .collect(), - }; + }, + ) + .collect(), + Provider::GoogleAiStudio => state + .app_args + .google_ai_studio_api_key + .iter() + .cloned() + .map(|api_key| { + Upstream::builder() + .base_url( + "https://generativelanguage.googleapis.com/v1beta/openai" + .parse() + .expect("TODO"), + ) + .api_key(api_key) + .build() + }) + .collect(), + Provider::Groq => state + .app_args + .groq_api_key + .iter() + .cloned() + .map(|api_key| { + Upstream::builder() + .base_url("https://api.groq.com/openai".parse().expect("TODO")) + .api_key(api_key) + .build() + }) + .collect(), + Provider::HuggingFace => state + .app_args + .hugging_face_api_key + .iter() + .cloned() + .map(|api_key| { + Upstream::builder() + .base_url("https://router.huggingface.co".parse().expect("TODO")) + .api_key(api_key) + .build() + }) + .collect(), + Provider::Mistral => state + .app_args + .mistral_api_key + .iter() + .cloned() + .map(|api_key| { + Upstream::builder() + .base_url("https://api.mistral.ai".parse().expect("TODO")) + .api_key(api_key) + .build() + }) + .collect(), + Provider::OpenRouter => state + .app_args + .openrouter_api_key + .iter() + .cloned() + .map(|api_key| { + Upstream::builder() + .base_url("https://openrouter.ai/api".parse().expect("TODO")) + .api_key(api_key) + .build() + }) + .collect(), + }; if let Ok(upstreams) = Vec1::try_from(upstreams) { model_upstreams.push_back((model, upstreams)); @@ -397,7 +435,11 @@ async fn create_chat_completion( web_search_options, } = request.clone(); - let reasoning_effort = Some(reasoning_effort.or(reasoning_in_id.clone()).unwrap_or(shortcut.default_reasoning())); + let reasoning_effort = Some( + reasoning_effort + .or(reasoning_in_id.clone()) + .unwrap_or(shortcut.default_reasoning()), + ); let request = openai::chat::completion::CreateChatCompletionRequest { model, @@ -443,7 +485,7 @@ async fn create_chat_completion( .tls_certs_only( webpki_root_certs::TLS_SERVER_ROOT_CERTS .iter() - .filter_map(|cert| Certificate::from_der(&cert).ok()), + .filter_map(|cert| Certificate::from_der(cert).ok()), ) .build() .expect("TODO"); @@ -500,10 +542,7 @@ async fn create_chat_completion( }; if request.stream.unwrap_or(false) { - return Ok(ChatCompletionResponse::Stream { - response: upstream_response, - stream_requester: state.stream_requester, - }); + return Ok(ChatCompletionResponse::Stream(upstream_response)); } else { let response = match upstream_response.json().await { Ok(response) => response, @@ -525,19 +564,25 @@ async fn create_chat_completion( struct V1ModelsPath; #[tracing::instrument] -async fn list_models(_: V1ModelsPath) -> Json> { - let shortcuts = Shortcut::VARIANTS.into_iter().cloned(); - let reasoning_efforts = openai::chat::completion::ReasoningEffort::VARIANTS.into_iter().cloned().map(Some).chain(None); +async fn list_models( + _: V1ModelsPath, +) -> Json> { + let shortcuts = Shortcut::VARIANTS.iter().cloned(); + let reasoning_efforts = openai::chat::completion::ReasoningEffort::VARIANTS + .iter() + .cloned() + .map(Some) + .chain(None); - let models = iproduct!(shortcuts, reasoning_efforts) - .into_iter() - .map(|(shortcut, reasoning)| openai::models::Model { + let models = iproduct!(shortcuts, reasoning_efforts).map(|(shortcut, reasoning)| { + openai::models::Model { id: ShortcutAndMaybeReasoning(shortcut, reasoning), created: 0, object: openai::models::ModelObject::Model, owned_by: "nobody".into(), shutdown_date: None, - }); + } + }); Json(openai::models::ListModelsResponse { data: models.collect(),