From ff621d759cca66103a79c872950d61d474105045 Mon Sep 17 00:00:00 2001 From: J / Jacob Babich Date: Thu, 8 Oct 2026 22:58:07 -0400 Subject: [PATCH] chore: groundwork for embeddings --- .gitignore | 2 + src/openai/chat/completion.rs | 4 +- src/openai/embeddings.rs | 97 ++ src/openai/mod.rs | 1 + src/{server.rs => server/chat/completions.rs} | 1096 ++++++++--------- src/server/chat/mod.rs | 1 + src/server/embeddings.rs | 36 + src/server/mod.rs | 58 + src/server/models.rs | 36 + 9 files changed, 742 insertions(+), 589 deletions(-) create mode 100644 src/openai/embeddings.rs rename src/{server.rs => server/chat/completions.rs} (87%) create mode 100644 src/server/chat/mod.rs create mode 100644 src/server/embeddings.rs create mode 100644 src/server/mod.rs create mode 100644 src/server/models.rs diff --git a/.gitignore b/.gitignore index ea8c4bf..68cdb13 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1,3 @@ /target + +.history/ diff --git a/src/openai/chat/completion.rs b/src/openai/chat/completion.rs index 01c1ec2..eaa6669 100644 --- a/src/openai/chat/completion.rs +++ b/src/openai/chat/completion.rs @@ -1092,7 +1092,7 @@ pub struct ApproximateLocation { /// Represents a chat completion response returned by model, based on the provided input. #[derive(Debug, Clone, Deserialize, Serialize)] -pub struct ChatCompletion { +pub struct ChatCompletion { /// A unique identifier for the chat completion. pub id: String, @@ -1103,7 +1103,7 @@ pub struct ChatCompletion { pub created: i64, /// The model used for the chat completion. - pub model: String, + pub model: Model, /// The object type, which is always `chat.completion`. pub object: String, diff --git a/src/openai/embeddings.rs b/src/openai/embeddings.rs new file mode 100644 index 0000000..3e811cc --- /dev/null +++ b/src/openai/embeddings.rs @@ -0,0 +1,97 @@ +use serde::{Deserialize, Serialize}; + +/// Represents the possible input types for embedding requests +#[derive(Debug, Clone, Deserialize, Serialize)] +#[serde(untagged)] +pub enum EmbeddingInput { + /// A single string input + String(String), + /// An array of string inputs + StringArray(Vec), + /// An array of token IDs (numbers) + NumberArray(Vec), + /// An array of arrays of token IDs (numbers) + NumberMatrix(Vec>), +} + +/// Represents the encoding format for embeddings +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)] +#[serde(rename_all = "lowercase")] +pub enum EncodingFormat { + /// Return embeddings as an array of floats + Float, + /// Return embeddings as a base64-encoded string + Base64, +} + +/// Request body for creating embeddings +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct CreateEmbeddingRequest { + /// Input text to embed + pub input: EmbeddingInput, + + /// ID of the model to use for generating embeddings + pub model: Model, + + /// Optional number of dimensions for the output embeddings + #[serde(skip_serializing_if = "Option::is_none")] + pub dimensions: Option, + + /// Format to return the embeddings in + #[serde(skip_serializing_if = "Option::is_none")] + pub encoding_format: Option, + + /// Optional unique identifier for tracking end-user + #[serde(skip_serializing_if = "Option::is_none")] + pub user: Option, +} + +/// Represents a single embedding in the response +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct Embedding { + /// The embedding vector, returned as an array of floats or a base64 string + /// depending on the encoding_format parameter + pub embedding: EmbeddingOutput, + + /// The index of this embedding in the list of embeddings + pub index: u32, + + /// The object type, which is always "embedding" + pub object: String, +} + +/// Represents the possible output formats for embeddings +#[derive(Debug, Clone, Deserialize, Serialize)] +#[serde(untagged)] +pub enum EmbeddingOutput { + /// Array of floats when encoding_format is "float" + FloatArray(Vec), + /// Base64-encoded string when encoding_format is "base64" + Base64String(String), +} + +/// Usage information for an embeddings request +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct Usage { + /// Number of tokens used by the prompt + pub prompt_tokens: u32, + + /// Total number of tokens used by the request + pub total_tokens: u32, +} + +/// Response from the embeddings API +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct CreateEmbeddingResponse { + /// List of generated embeddings + pub data: Vec, + + /// The name of the model used to generate the embeddings + pub model: String, + + /// The object type, which is always "list" + pub object: String, + + /// Usage information for the request + pub usage: Usage, +} diff --git a/src/openai/mod.rs b/src/openai/mod.rs index 2e55b7c..58726e4 100644 --- a/src/openai/mod.rs +++ b/src/openai/mod.rs @@ -1,2 +1,3 @@ pub mod chat; +pub mod embeddings; pub mod models; diff --git a/src/server.rs b/src/server/chat/completions.rs similarity index 87% rename from src/server.rs rename to src/server/chat/completions.rs index 8b8f2ae..78c7197 100644 --- a/src/server.rs +++ b/src/server/chat/completions.rs @@ -1,587 +1,509 @@ -use crate::AppArgs; -use crate::CloudflareCredentials; -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, TypedPath}; -use ecow::EcoString; -use futures::StreamExt; -use headers::Header; -use headers::authorization::Bearer; -use itertools::iproduct; -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_with::DeserializeFromStr; -use serde_with::SerializeDisplay; -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; - -#[derive(Debug, Clone, TypedBuilder)] -pub struct AppState { - app_args: Arc, -} - -#[derive(Deserialize, TypedPath)] -#[typed_path("/v1/chat/completions")] -struct V1ChatCompletionsPath; - -#[derive(Debug, Clone)] -enum Provider { - CloudflareWorkersAi, - GoogleAiStudio, - Groq, - HuggingFace, - Mistral, - OpenRouter, -} - -#[derive( - Clone, Debug, EnumString, DeserializeFromStr, strum::Display, SerializeDisplay, VariantArray, -)] -#[strum(serialize_all = "lowercase")] -enum Shortcut { - XS, - S, - M, - L, - XL, - - Experimental, -} - -#[derive(Debug, Clone, DeserializeFromStr, SerializeDisplay)] -struct ShortcutAndMaybeReasoning(Shortcut, Option); - -impl FromStr for ShortcutAndMaybeReasoning { - type Err = Infallible; // TODO - - fn from_str(s: &str) -> Result { - if let Some((model, reasoning)) = s.split_once(":") { - let model = model.parse().expect("TODO"); - let reasoning = reasoning.parse().expect("TODO"); - - Ok(Self(model, Some(reasoning))) - } else { - let model = s.parse().expect("TODO"); - let reasoning = None; - - Ok(Self(model, reasoning)) - } - } -} - -impl Display for ShortcutAndMaybeReasoning { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - let Self(model, reasoning) = self; - - write!(f, "{model}")?; - if let Some(reasoning) = reasoning { - write!(f, ":")?; - write!(f, "{reasoning}")?; - } - - Ok(()) - } -} - -impl Shortcut { - fn suggestions(&self) -> Vec<(Provider, EcoString)> { - use Provider as P; - use Shortcut as S; - - match self { - S::XS => vec![ - (P::GoogleAiStudio, "gemini-3.5-flash-lite".into()), - (P::OpenRouter, "apodex/apodex-1.1-mini:free".into()), - (P::OpenRouter, "poolside/laguna-xs-2.1:free".into()), - (P::OpenRouter, "cohere/north-mini-code:free".into()), - (P::OpenRouter, "nvidia/nemotron-3.5-lightning:free".into()), - (P::OpenRouter, "google/gemma-4-26b-a4b-it:free".into()), - (P::HuggingFace, "google/gemma-4-26B-A4B-it".into()), - (P::OpenRouter, "poolside/laguna-xs.2: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()), - (P::HuggingFace, "openai/gpt-oss-20b".into()), - (P::OpenRouter, "liquid/lfm-2.5-2.6b:free".into()), - (P::Groq, "groq/compound-mini".into()), - (P::OpenRouter, "openrouter/free".into()), - ], - - S::S => vec![ - (P::CloudflareWorkersAi, "@cf/qwen/qwen3.8-27b".into()), - (P::OpenRouter, "qwen/qwen3.8-27b:free".into()), - (P::Groq, "qwen/qwen3.8-27b".into()), - (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()), - ], - - S::M => vec![ - (P::GoogleAiStudio, "gemini-3.8-flash".into()), - (P::GoogleAiStudio, "gemini-3.7-flash".into()), - (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(), - ), - (P::CloudflareWorkersAi, "@cf/openai/gpt-oss-120b".into()), - (P::OpenRouter, "openai/gpt-oss-120b:free".into()), - (P::Groq, "openai/gpt-oss-120b".into()), - (P::HuggingFace, "openai/gpt-oss-120b".into()), - ], - - 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::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::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::Mistral, "mistral-large-4".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()), - ], - - S::Experimental => vec![(P::OpenRouter, "space-bunny-alpha".into())], - } - } -} - -fn take_random_by_swapping(mut vec1: Vec1, mut rng: impl Rng) -> (T, Option>) { - let end = vec1.len().get() - 1; - 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(reqwest::Response), -} - -impl IntoResponse for ChatCompletionResponse { - fn into_response(self) -> Response { - match self { - ChatCompletionResponse::Immediate(chat_completion) => { - Json(chat_completion).into_response() - } - ChatCompletionResponse::Stream(response) => { - let (sink, stream) = async_channel::bounded(256); - - let handle = Handle::current(); - - 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() - } - } - } -} - -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( - _: V1ChatCompletionsPath, - State(state): State, - 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: Vec = - match provider { - Provider::CloudflareWorkersAi => state - .app_args - .cloudflare_workers_ai_credentials - .iter() - .cloned() - .map( - |CloudflareCredentials { - account_id, - api_key, - }| { - Upstream::builder() - .base_url( - format!("https://api.cloudflare.com/client/v4/accounts/{account_id}/ai") - .parse() - .expect("TODO"), - ) - .api_key(api_key) - .build() - }, - ) - .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)); - } - } - - 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)); - } - - tracing::info!(?model, ?upstream); - - 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 reasoning_effort = reasoning_effort.or(reasoning_in_id.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::builder() - .tls_certs_only( - webpki_root_certs::TLS_SERVER_ROOT_CERTS - .iter() - .filter_map(|cert| Certificate::from_der(cert).ok()), - ) - .build() - .expect("TODO"); - - let mut upstream_request = client - .post({ - let mut url = upstream.base_url; - - { - let mut path_segments = url.path_segments_mut().expect("TODO"); - - let s = V1ChatCompletionsPath.to_string(); - let s = s.strip_prefix("/").unwrap_or(&s); - let s = s.strip_suffix("/").unwrap_or(&s); - - path_segments.extend(s.split("/")); - } - - tracing::info!(?url); - - url - }) - .json(&request); - - 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; - } - }; - - if let Err(upstream_error) = upstream_response.error_for_status_ref() { - tracing::error!(?upstream_error); - match upstream_response.text().await { - Ok(body) => tracing::error!(?body), - Err(body_error) => tracing::error!(?body_error), - }; - continue; - }; - - if request.stream.unwrap_or(false) { - return Ok(ChatCompletionResponse::Stream(upstream_response)); - } 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) -} - -#[derive(Deserialize, TypedPath)] -#[typed_path("/v1/models")] -struct V1ModelsPath; - -#[tracing::instrument] -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).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(), - object: openai::models::ListObject::List, - }) -} - -pub fn create_router() -> Router { - Router::new() - .typed_post(create_chat_completion) - .typed_get(list_models) -} - -pub fn create_app(state: AppState) -> Router<()> { - create_router().with_state(state) -} +use crate::CloudflareCredentials; +use crate::openai; +use crate::server::AppState; +use crate::server::Provider; +use crate::server::Upstream; +use crate::server::take_random_by_swapping; +use axum::Json; +use axum::extract::State; +use axum::http::StatusCode; +use axum::response::IntoResponse; +use axum::response::Response; +use axum::response::Sse; +use axum_extra::routing::TypedPath; +use ecow::EcoString; +use futures::StreamExt; +use headers::Header; +use headers::authorization::Bearer; +use mitsein::vec1::Vec1; +use reqwest::Certificate; +use reqwest_sse::EventSource as _; +use secrecy::ExposeSecret; +use serde::Deserialize; +use serde_with::DeserializeFromStr; +use serde_with::SerializeDisplay; +use std::collections::VecDeque; +use std::convert::Infallible; +use std::fmt::Display; +use std::str::FromStr; +use strum::EnumString; +use strum::VariantArray; +use tokio::runtime::Handle; + +#[derive(Deserialize, TypedPath)] +#[typed_path("/v1/chat/completions")] +pub struct V1ChatCompletionsPath; + +#[derive( + Clone, Debug, EnumString, DeserializeFromStr, strum::Display, SerializeDisplay, VariantArray, +)] +#[strum(serialize_all = "lowercase")] +pub enum Shortcut { + XS, + S, + M, + L, + XL, + + Experimental, +} + +#[derive(Debug, Clone, DeserializeFromStr, SerializeDisplay)] +pub struct ShortcutAndMaybeReasoning( + pub Shortcut, + pub Option, +); + +impl FromStr for ShortcutAndMaybeReasoning { + type Err = Infallible; // TODO + + fn from_str(s: &str) -> Result { + if let Some((model, reasoning)) = s.split_once(":") { + let model = model.parse().expect("TODO"); + let reasoning = reasoning.parse().expect("TODO"); + + Ok(Self(model, Some(reasoning))) + } else { + let model = s.parse().expect("TODO"); + let reasoning = None; + + Ok(Self(model, reasoning)) + } + } +} + +impl Display for ShortcutAndMaybeReasoning { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let Self(model, reasoning) = self; + + write!(f, "{model}")?; + if let Some(reasoning) = reasoning { + write!(f, ":")?; + write!(f, "{reasoning}")?; + } + + Ok(()) + } +} + +impl Shortcut { + fn suggestions(&self) -> Vec<(Provider, EcoString)> { + use Provider as P; + use Shortcut as S; + + match self { + S::XS => vec![ + (P::GoogleAiStudio, "gemini-3.5-flash-lite".into()), + (P::OpenRouter, "apodex/apodex-1.1-mini:free".into()), + (P::OpenRouter, "poolside/laguna-xs-2.1:free".into()), + (P::OpenRouter, "cohere/north-mini-code:free".into()), + (P::OpenRouter, "nvidia/nemotron-3.5-lightning:free".into()), + (P::OpenRouter, "google/gemma-4-26b-a4b-it:free".into()), + (P::HuggingFace, "google/gemma-4-26B-A4B-it".into()), + (P::OpenRouter, "poolside/laguna-xs.2: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()), + (P::HuggingFace, "openai/gpt-oss-20b".into()), + (P::OpenRouter, "liquid/lfm-2.5-2.6b:free".into()), + (P::Groq, "groq/compound-mini".into()), + (P::OpenRouter, "openrouter/free".into()), + ], + + S::S => vec![ + (P::CloudflareWorkersAi, "@cf/qwen/qwen3.8-27b".into()), + (P::OpenRouter, "qwen/qwen3.8-27b:free".into()), + (P::Groq, "qwen/qwen3.8-27b".into()), + (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()), + ], + + S::M => vec![ + (P::GoogleAiStudio, "gemini-3.8-flash".into()), + (P::GoogleAiStudio, "gemini-3.7-flash".into()), + (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(), + ), + (P::CloudflareWorkersAi, "@cf/openai/gpt-oss-120b".into()), + (P::OpenRouter, "openai/gpt-oss-120b:free".into()), + (P::Groq, "openai/gpt-oss-120b".into()), + (P::HuggingFace, "openai/gpt-oss-120b".into()), + ], + + 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::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::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::Mistral, "mistral-large-4".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()), + ], + + S::Experimental => vec![(P::OpenRouter, "space-bunny-alpha".into())], + } + } +} + +pub enum ChatCompletionResponse { + Immediate(openai::chat::completion::ChatCompletion), + Stream(reqwest::Response), +} + +impl IntoResponse for ChatCompletionResponse { + fn into_response(self) -> Response { + match self { + ChatCompletionResponse::Immediate(chat_completion) => { + Json(chat_completion).into_response() + } + ChatCompletionResponse::Stream(response) => { + let (sink, stream) = async_channel::bounded(256); + + let handle = Handle::current(); + + 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() + } + } + } +} + +pub 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))] +pub async fn create_chat_completion( + _: V1ChatCompletionsPath, + State(state): State, + 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: Vec = + match provider { + Provider::CloudflareWorkersAi => state + .app_args + .cloudflare_workers_ai_credentials + .iter() + .cloned() + .map( + |CloudflareCredentials { + account_id, + api_key, + }| { + Upstream::builder() + .base_url( + format!("https://api.cloudflare.com/client/v4/accounts/{account_id}/ai") + .parse() + .expect("TODO"), + ) + .api_key(api_key) + .build() + }, + ) + .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)); + } + } + + 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)); + } + + tracing::info!(?model, ?upstream); + + 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 reasoning_effort = reasoning_effort.or(reasoning_in_id.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::builder() + .tls_certs_only( + webpki_root_certs::TLS_SERVER_ROOT_CERTS + .iter() + .filter_map(|cert| Certificate::from_der(cert).ok()), + ) + .build() + .expect("TODO"); + + let mut upstream_request = client + .post({ + let mut url = upstream.base_url; + + { + let mut path_segments = url.path_segments_mut().expect("TODO"); + + let s = V1ChatCompletionsPath.to_string(); + let s = s.strip_prefix("/").unwrap_or(&s); + let s = s.strip_suffix("/").unwrap_or(&s); + + path_segments.extend(s.split("/")); + } + + tracing::info!(?url); + + url + }) + .json(&request); + + 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; + } + }; + + if let Err(upstream_error) = upstream_response.error_for_status_ref() { + tracing::error!(?upstream_error); + match upstream_response.text().await { + Ok(body) => tracing::error!(?body), + Err(body_error) => tracing::error!(?body_error), + }; + continue; + }; + + if request.stream.unwrap_or(false) { + return Ok(ChatCompletionResponse::Stream(upstream_response)); + } 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) +} diff --git a/src/server/chat/mod.rs b/src/server/chat/mod.rs new file mode 100644 index 0000000..82c6873 --- /dev/null +++ b/src/server/chat/mod.rs @@ -0,0 +1 @@ +pub mod completions; diff --git a/src/server/embeddings.rs b/src/server/embeddings.rs new file mode 100644 index 0000000..5d5f885 --- /dev/null +++ b/src/server/embeddings.rs @@ -0,0 +1,36 @@ +#[derive(Deserialize, TypedPath)] +#[typed_path("/v1/chat/completions")] +pub struct V1ChatCompletionsPath; + +#[derive( + Clone, Debug, EnumString, DeserializeFromStr, strum::Display, SerializeDisplay, VariantArray, +)] +#[strum(serialize_all = "lowercase")] +pub enum Shortcut { + XS, + S, + M, + L, + XL, +} + +impl Shortcut { + fn suggestions(&self) -> Vec<(Provider, EcoString)> { + use Provider as P; + use Shortcut as S; + + match self { + S::XS => vec![], + + S::S => vec![], + + S::M => vec![], + + S::L => vec![], + + S::XL => vec![], + + S::Experimental => vec![], + } + } +} diff --git a/src/server/mod.rs b/src/server/mod.rs new file mode 100644 index 0000000..c00bead --- /dev/null +++ b/src/server/mod.rs @@ -0,0 +1,58 @@ +use crate::AppArgs; +use axum::Router; +use axum_extra::routing::RouterExt; +use mitsein::vec1::Vec1; +use rand::Rng; +use rand::RngExt; +use secrecy::SecretString; +use std::sync::Arc; +use typed_builder::TypedBuilder; +use url::Url; + +mod chat; +mod models; + +#[derive(Debug, Clone, TypedBuilder)] +pub struct AppState { + app_args: Arc, +} + +#[derive(Debug, Clone)] +enum Provider { + CloudflareWorkersAi, + GoogleAiStudio, + Groq, + HuggingFace, + Mistral, + OpenRouter, +} + +fn take_random_by_swapping(mut vec1: Vec1, mut rng: impl Rng) -> (T, Option>) { + let end = vec1.len().get() - 1; + 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, +} + +pub fn create_router() -> Router { + Router::new() + .typed_post(chat::completions::create_chat_completion) + .typed_get(models::list_models) +} + +pub fn create_app(state: AppState) -> Router<()> { + create_router().with_state(state) +} diff --git a/src/server/models.rs b/src/server/models.rs new file mode 100644 index 0000000..3681267 --- /dev/null +++ b/src/server/models.rs @@ -0,0 +1,36 @@ +use crate::openai; +use axum::Json; +use axum_extra::routing::TypedPath; +use itertools::iproduct; +use serde::Deserialize; +use strum::VariantArray; + +#[derive(Deserialize, TypedPath)] +#[typed_path("/v1/models")] +pub struct V1ModelsPath; + +#[tracing::instrument] +pub async fn list_models(_: V1ModelsPath) -> Json> { + let shortcuts = super::chat::completions::Shortcut::VARIANTS.iter().cloned(); + let reasoning_efforts = openai::chat::completion::ReasoningEffort::VARIANTS + .iter() + .cloned() + .map(Some) + .chain(None); + + let models = iproduct!(shortcuts, reasoning_efforts).map(|(shortcut, reasoning)| { + openai::models::Model { + id: super::chat::completions::ShortcutAndMaybeReasoning(shortcut, reasoning) + .to_string(), + created: 0, + object: openai::models::ModelObject::Model, + owned_by: "nobody".into(), + shutdown_date: None, + } + }); + + Json(openai::models::ListModelsResponse { + data: models.collect(), + object: openai::models::ListObject::List, + }) +}