diff --git a/Cargo.lock b/Cargo.lock index 4d0eeaf..12d7ead 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -82,6 +82,18 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "async-channel" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "924ed96dd52d1b75e9c1a3e6275715fd320f5f9439fb5a4a11fa51f4221158d2" +dependencies = [ + "concurrent-queue", + "event-listener-strategy", + "futures-core", + "pin-project-lite", +] + [[package]] name = "async-compression" version = "0.4.46" @@ -94,6 +106,28 @@ dependencies = [ "tokio", ] +[[package]] +name = "async-stream" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b5a71a6f37880a80d1d7f19efd781e4b5de42c88f0722cc13bcb6cc2cfe8476" +dependencies = [ + "async-stream-impl", + "futures-core", + "pin-project-lite", +] + +[[package]] +name = "async-stream-impl" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "async-trait" version = "0.1.92" @@ -147,6 +181,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "31b698c5f9a010f6573133b09e0de5408834d0c82f8d7475a89fc1867a71cd90" dependencies = [ "axum-core", + "axum-macros", "base64 0.22.1", "bytes", "form_urlencoded", @@ -374,14 +409,20 @@ checksum = "1c133bc6a41be0d194c306b5506d15e6feeea7b1d6604bd3f8310dfb2ca96486" name = "cloud-llms" version = "0.1.0" dependencies = [ + "async-channel", "axum", "axum-extra", "clap", + "ecow", "enum-map", "enumset", + "futures", "headers", + "indexmap", + "mitsein", "rand 0.10.3", "reqwest", + "reqwest-sse", "schemars", "secrecy", "serde", @@ -438,6 +479,15 @@ version = "0.4.33" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6e8ccc4ea9f6acc32d102c0f6d471d11d913ad15f20c04de743374861fa1d414" +[[package]] +name = "concurrent-queue" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ca0197aee26d1ae37445ee532fefce43251d24cc7c166799f4d46817f1d3973" +dependencies = [ + "crossbeam-utils", +] + [[package]] name = "cookie" version = "0.18.2" @@ -654,6 +704,12 @@ version = "1.0.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555" +[[package]] +name = "ecow" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2dca6141d10838e13f9b166c190254a3f9b5d8379521f9b7510810061cef17a4" + [[package]] name = "either" version = "1.18.0" @@ -723,6 +779,26 @@ version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" +[[package]] +name = "event-listener" +version = "5.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a23add41df1562121a9393cb065eab5146a1242410f23a644851e90cfd669d2" +dependencies = [ + "parking", + "pin-project-lite", +] + +[[package]] +name = "event-listener-strategy" +version = "0.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8be9f3dfaaffdae2972880079a491a1a8bb7cbed0b8dd7a347f668b4150a3b93" +dependencies = [ + "event-listener", + "pin-project-lite", +] + [[package]] name = "find-msvc-tools" version = "0.1.14" @@ -761,6 +837,21 @@ version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" +[[package]] +name = "futures" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a31d2a3fbaaeb2af2368bbdd904aa8e812d3c04a1ee10d3171f52d556e5d0a3" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + [[package]] name = "futures-channel" version = "0.3.34" @@ -768,6 +859,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b1f9e3d69d39e4862ffed03ed071a76f9a13ba1d9109d355b0f0aa6b15e393c4" dependencies = [ "futures-core", + "futures-sink", ] [[package]] @@ -776,6 +868,17 @@ version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "92d699e522242e69e3003b94ecc1f960f3a5e015aa7c5d7486e65ad01dd94f5e" +[[package]] +name = "futures-executor" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "031b47cf1a3c6cc8bc2fc76cd437f521619387907d469316e7c0bc278f1f5432" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + [[package]] name = "futures-io" version = "0.3.34" @@ -811,6 +914,7 @@ version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc" dependencies = [ + "futures-channel", "futures-core", "futures-io", "futures-macro", @@ -1251,6 +1355,15 @@ version = "1.70.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695" +[[package]] +name = "itertools" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b192c782037fadd9cfa75548310488aabdbf3d2da73885b31bd0abd03351285" +dependencies = [ + "either", +] + [[package]] name = "itoa" version = "1.0.18" @@ -1430,6 +1543,20 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "mitsein" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a49aef75c004fb496a23f175395415ecfe2a630e4c6175ae6f7d9cc1a19332e" +dependencies = [ + "either", + "indexmap", + "itertools", + "schemars", + "serde", + "smallvec", +] + [[package]] name = "moka" version = "0.12.16" @@ -1522,6 +1649,12 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7c87def4c32ab89d880effc9e097653c8da5d6ef28e6b539d313baaacfbafcbe" +[[package]] +name = "parking" +version = "2.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f38d5652c16fde515bb1ecef450ab0f6a219d619a7274976324d5e377f7dceba" + [[package]] name = "parking_lot" version = "0.12.5" @@ -1857,6 +1990,19 @@ dependencies = [ "web-sys", ] +[[package]] +name = "reqwest-sse" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97e36f9508407d80f97d0715f92ffdfe7aca394c7fffbf238ae80775b23330cb" +dependencies = [ + "async-stream", + "reqwest", + "tokio", + "tokio-stream", + "tokio-util", +] + [[package]] name = "resolv-conf" version = "0.7.6" @@ -2004,6 +2150,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "687274d293b6cdc6e73e0fee520bf2049650090d7164f87672d212a3c530cf4a" dependencies = [ "dyn-clone", + "indexmap", "ref-cast", "schemars_derive", "serde", @@ -2431,6 +2578,17 @@ dependencies = [ "tokio", ] +[[package]] +name = "tokio-stream" +version = "0.1.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a3d06f0b082ba57c26b79407372e57cf2a1e28124f78e9479fe80322cf53420b" +dependencies = [ + "futures-core", + "pin-project-lite", + "tokio", +] + [[package]] name = "tokio-tungstenite" version = "0.29.0" diff --git a/Cargo.toml b/Cargo.toml index 2149bd8..3d6826d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,14 +4,20 @@ version = "0.1.0" edition = "2024" [dependencies] -axum = { version = "0.8.9", features = ["http2", "multipart", "ws"] } +async-channel = "2.5.0" +axum = { version = "0.8.9", features = ["http2", "macros", "multipart", "ws"] } axum-extra = { version = "0.12.6", features = ["typed-header", "typed-routing"] } clap = { version = "4.6.7", features = ["derive", "env"] } +ecow = "0.3.1" enum-map = { version = "3.1.0", features = ["serde"] } enumset = { version = "1.1.14", features = ["serde"] } +futures = "0.3.34" headers = "0.4.2" +indexmap = "2.14.2" +mitsein = { version = "0.9.0", features = ["indexmap"] } rand = "0.10.3" reqwest = { version = "0.13.5", features = ["brotli", "cookies", "deflate", "form", "gzip", "hickory-dns", "json", "multipart", "query", "socks", "stream", "zstd"] } +reqwest-sse = "0.2.0" schemars = "1.2.2" secrecy = { version = "0.10.3", features = ["serde"] } serde = { version = "1.0.229", features = ["derive"] } diff --git a/src/main.rs b/src/main.rs index ca85cb1..26c03d4 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,7 +1,10 @@ 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; @@ -60,13 +63,49 @@ async fn main() { let addr = SocketAddr::new(app_args.ip, app_args.port); + let (tx, rx) = async_channel::unbounded(); + let app = server::create_app( server::AppState::builder() .app_args(app_args.into()) + .stream_requester(tx) .build(), ); let listener = tokio::net::TcpListener::bind(addr).await.expect("TODO"); - axum::serve(listener, app).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 { + axum::serve(listener, app).await.expect("TODO"); + }); } diff --git a/src/openai/chat/mod.rs b/src/openai/chat/mod.rs index 6b31b14..c893d10 100644 --- a/src/openai/chat/mod.rs +++ b/src/openai/chat/mod.rs @@ -1 +1 @@ -pub mod completion; \ No newline at end of file +pub mod completion; diff --git a/src/server.rs b/src/server.rs index 1356249..c791a33 100644 --- a/src/server.rs +++ b/src/server.rs @@ -1,25 +1,54 @@ -use std::sync::Arc; - use crate::AppArgs; use crate::openai; use axum::Json; use axum::Router; use axum::extract::State; +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 ecow::EcoString; +use ecow::string::ToEcoString; +use futures::Stream; +use futures::StreamExt as _; +use futures::TryStreamExt as _; +use headers::Header; +use headers::authorization::Bearer; +use indexmap::IndexMap; +use mitsein::vec1::Vec1; +use rand::Rng; +use rand::RngExt; +use reqwest_sse::EventSource; +use secrecy::ExposeSecret; +use secrecy::SecretString; use serde::Deserialize; use serde::Serialize; +use std::collections::VecDeque; +use std::hash::Hash; +use std::pin::Pin; +use std::sync::Arc; +use tokio::sync::oneshot; 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)] #[typed_path("/chat/completions")] struct ChatCompletionsPath; +#[derive(Debug, Clone)] enum Provider { CloudflareWorkersAi, GoogleAiStudio, @@ -29,7 +58,7 @@ enum Provider { OpenRouter, } -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Clone)] enum Shortcut { XS, S, @@ -39,7 +68,7 @@ enum Shortcut { } impl Shortcut { - fn suggestions(&self) -> Vec<(Provider, String)> { + fn suggestions(&self) -> Vec<(Provider, EcoString)> { use Provider as P; use Shortcut as S; @@ -52,7 +81,10 @@ impl Shortcut { (P::OpenRouter, "nvidia/nemotron-3.5-lightning:free".into()), (P::OpenRouter, "google/gemma-4-26b-a4b-it:free".into()), (P::OpenRouter, "poolside/laguna-xs.2:free".into()), - (P::OpenRouter, "nvidia/nemotron-3-nano-omni-30b-a3b-reasoning:free".into()), + ( + P::OpenRouter, + "nvidia/nemotron-3-nano-omni-30b-a3b-reasoning:free".into(), + ), (P::CloudflareWorkersAi, "@cf/openai/gpt-oss-20b".into()), (P::OpenRouter, "openai/gpt-oss-20b:free".into()), (P::Groq, "openai/gpt-oss-20b".into()), @@ -77,7 +109,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, "nvidia/nemotron-3-super-120b-a12b:free".into()), + ( + P::OpenRouter, + "nvidia/nemotron-3-super-120b-a12b:free".into(), + ), (P::CloudflareWorkersAi, "@cf/openai/gpt-oss-120b".into()), (P::OpenRouter, "openai/gpt-oss-120b:free".into()), (P::Groq, "openai/gpt-oss-120b".into()), @@ -87,8 +122,9 @@ impl Shortcut { (P::HuggingFace, "zai-org/GLM-5.3-Flash".into()), (P::HuggingFace, "deepseek-ai/DeepSeek-V4.1-Flash".into()), (P::OpenRouter, "nex-agi/nex-n2.5-pro:free".into()), + (P::HuggingFace, "MiniMaxAI/MiniMax-M3".into()), (P::OpenRouter, "thinkingmachines/inkling-small:free".into()), - (P::HuggingFace, "thinkingmachines/Inkling-Small".into()) + (P::HuggingFace, "thinkingmachines/Inkling-Small".into()), ], S::XL => vec![ (P::CloudflareWorkersAi, "@cf/zai-org/glm-5.3".into()), @@ -96,27 +132,311 @@ impl Shortcut { (P::Mistral, "zai-glm-5-3".into()), (P::HuggingFace, "moonshotai/Kimi-K3".into()), (P::HuggingFace, "Qwen/Qwen3.8-2.4T-A95B".into()), - (P::HuggingFace, "MiniMaxAI/MiniMax-M3".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()), ], } } } -enum ChatCompletionResponse { - Immediate(openai::chat::completion::ChatCompletion), - Stream(()), +fn take_random_by_swapping(mut vec1: Vec1, mut rng: impl Rng) -> (T, Option>) { + let end = vec1.len().get(); + let index = rng.random_range(0..end); + + vec1.swap(index, end); + let (body, element) = vec1.into_rtail_and_head(); + + let body = body.try_into().ok(); + + (element, body) } +#[derive(Debug, TypedBuilder)] +struct Upstream { + base_url: Url, + + #[builder(setter(strip_option))] + api_key: Option, +} + +enum ChatCompletionResponse { + Immediate(openai::chat::completion::ChatCompletion), + Stream { + response: reqwest::Response, + stream_requester: StreamRequester, + }, +} + +impl IntoResponse for ChatCompletionResponse { + fn into_response(self) -> Response { + match self { + ChatCompletionResponse::Immediate(chat_completion) => { + Json(chat_completion).into_response() + } + ChatCompletionResponse::Stream { + response, + stream_requester, + } => { + let (sse_sink, sse_stream) = async_channel::bounded(256); + + stream_requester + .try_send((response, sse_sink)) + .expect("this is an unbounded channel"); + + Sse::new(sse_stream).into_response() + } + } + } +} + +enum ChatCompletionError { + NoneSucceeded, +} + +impl IntoResponse for ChatCompletionError { + fn into_response(self) -> Response { + match self { + ChatCompletionError::NoneSucceeded => StatusCode::SERVICE_UNAVAILABLE.into_response(), + } + } +} + +#[axum::debug_handler] +#[tracing::instrument(skip(state, request))] async fn create_chat_completion( _: ChatCompletionsPath, State(state): State, Json(request): Json>, -) -> Result { - if request.stream.unwrap_or(false) { - } else { +) -> Result { + let suggestions = request.model.suggestions(); + let mut model_upstreams = VecDeque::new(); + + for (provider, model) in 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() + .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)); + } } + + while let Some((model, upstreams)) = model_upstreams.pop_front() { + let (upstream, upstreams) = take_random_by_swapping(upstreams, &mut rand::rng()); + if let Some(upstreams) = upstreams { + model_upstreams.push_back((model.clone(), upstreams)); + } + + let openai::chat::completion::CreateChatCompletionRequest { + model: _shortcut, + messages, + audio, + frequency_penalty, + function_call, + functions, + logit_bias, + logprobs, + max_completion_tokens, + max_tokens, + metadata, + modalities, + moderation, + n, + parallel_tool_calls, + prediction, + presence_penalty, + prompt_cache_key, + prompt_cache_options, + prompt_cache_retention, + reasoning_effort, + response_format, + safety_identifier, + seed, + service_tier, + stop, + store, + stream, + stream_options, + temperature, + tool_choice, + tools, + top_logprobs, + top_p, + user, + verbosity, + web_search_options, + } = request.clone(); + + let request = openai::chat::completion::CreateChatCompletionRequest { + model, + messages, + audio, + frequency_penalty, + function_call, + functions, + logit_bias, + logprobs, + max_completion_tokens, + max_tokens, + metadata, + modalities, + moderation, + n, + parallel_tool_calls, + prediction, + presence_penalty, + prompt_cache_key, + prompt_cache_options, + prompt_cache_retention, + reasoning_effort, + response_format, + safety_identifier, + seed, + service_tier, + stop, + store, + stream, + stream_options, + temperature, + tool_choice, + tools, + top_logprobs, + top_p, + user, + verbosity, + web_search_options, + }; + + let client = reqwest::Client::new(); + + let mut upstream_request = client.post({ + let mut url = upstream.base_url.clone(); + + { + let mut segments = url.path_segments_mut().expect("TODO"); + + segments.push("chat"); + segments.push("completions"); + } + + url + }); + + if let Some(api_key) = upstream.api_key { + let authorization = + headers::Authorization::bearer(api_key.expose_secret()).expect("TODO"); + upstream_request = upstream_request.header(headers::Authorization::::name(), { + let mut vec = vec![]; + + authorization.encode(&mut vec); + + let [one] = vec.try_into().expect("TODO"); + + one + }); + } + + let upstream_response = match upstream_request.send().await { + Ok(upstream_response) => upstream_response, + Err(upstream_error) => { + tracing::error!(?upstream_error); + continue; + } + }; + + let upstream_response = match upstream_response.error_for_status() { + Ok(upstream_response) => upstream_response, + Err(upstream_error) => { + tracing::error!(?upstream_error); + continue; + } + }; + + if request.stream.unwrap_or(false) { + return Ok(ChatCompletionResponse::Stream { + response: upstream_response, + stream_requester: state.stream_requester, + }); + } else { + let response = match upstream_response.json().await { + Ok(response) => response, + Err(json_error) => { + tracing::error!(?json_error); + continue; + } + }; + + return Ok(ChatCompletionResponse::Immediate(response)); + } + } + + Err(ChatCompletionError::NoneSucceeded) } pub fn create_router() -> Router {