From f4f657c219d5cde622498185c7d306a0b2379c32 Mon Sep 17 00:00:00 2001 From: Jacob Date: Tue, 13 Oct 2026 23:37:24 -0400 Subject: [PATCH] feat: novita, nvidia nim, vercel ai gateway --- src/lib.rs | 9 + src/server/chat/completions.rs | 1110 +++++++++++++++++--------------- src/server/decisions.rs | 9 +- src/server/embeddings.rs | 565 ++++++++-------- src/server/mod.rs | 3 + src/server/system_one.rs | 16 + 6 files changed, 918 insertions(+), 794 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index eea25f2..d9c9ed8 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -54,9 +54,18 @@ pub struct AppArgs { #[arg(long)] pub mistral_api_key: Vec, + #[arg(long)] + pub novita_ai_api_key: Vec, + + #[arg(long)] + pub nvidia_nim_api_key: Vec, + #[arg(long)] pub openrouter_api_key: Vec, + #[arg(long)] + pub vercel_ai_gateway_api_key: Vec, + #[arg(long)] pub z_ai_api_key: Vec, } diff --git a/src/server/chat/completions.rs b/src/server/chat/completions.rs index 791d967..670ec98 100644 --- a/src/server/chat/completions.rs +++ b/src/server/chat/completions.rs @@ -1,513 +1,597 @@ -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::HuggingFace, "Qwen/Qwen3.6-35B-A3B".into()), - (P::GoogleAiStudio, "gemma-4-26b-a4b-it".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.1:free".into()), - (P::ZAi, "glm-4.7-flash".into()), - (P::OpenRouter, "cohere/north-mini-code:free".into()), - (P::OpenRouter, "nvidia/nemotron-3.5-lightning:free".into()), - (P::ZAi, "glm-4.6v-flash".into()), - (P::ZAi, "glm-4.5-flash".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::GoogleAiStudio, "gemma-4-31b-it".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()), - (P::GoogleAiStudio, "gemini-3.1-pro-preview".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() - .url( - format!("https://api.cloudflare.com/client/v4/accounts/{account_id}/ai/v1/chat/completions") - .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() - .url( - "https://generativelanguage.googleapis.com/v1beta/openai/v1/chat/completions" - .parse() - .expect("TODO"), - ) - .api_key(api_key) - .build() - }) - .collect(), - Provider::Groq => state - .app_args - .groq_api_key - .iter() - .cloned() - .map(|api_key| { - Upstream::builder() - .url("https://api.groq.com/openai/v1/chat/completions".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() - .url("https://router.huggingface.co/v1/chat/completions".parse().expect("TODO")) - .api_key(api_key) - .build() - }) - .collect(), - Provider::LiquidAi => unreachable!(), - Provider::Mistral => state - .app_args - .mistral_api_key - .iter() - .cloned() - .map(|api_key| { - Upstream::builder() - .url("https://api.mistral.ai/v1/chat/completions".parse().expect("TODO")) - .api_key(api_key) - .build() - }) - .collect(), - Provider::OpenRouter => state - .app_args - .openrouter_api_key - .iter() - .cloned() - .map(|api_key| { - Upstream::builder() - .url("https://openrouter.ai/api/v1/chat/completions".parse().expect("TODO")) - .api_key(api_key) - .build() - }) - .collect(), - Provider::ZAi => state - .app_args - .z_ai_api_key - .iter() - .cloned() - .map(|api_key| { - Upstream::builder() - .url("https://api.z.ai/api/paas/v4/chat/completions".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(upstream.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) -} +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::NovitaAi, "apodex/apodex-1.1-mini".into()), + (P::OpenRouter, "apodex/apodex-1.1-mini:free".into()), + (P::VercelAiGateway, "openai/gpt-5.4-nano".into()), + (P::HuggingFace, "Qwen/Qwen3.6-35B-A3B".into()), + (P::VercelAiGateway, "liquid/lfm2.5-24b-a2b".into()), + (P::GoogleAiStudio, "gemma-4-26b-a4b-it".into()), + (P::OpenRouter, "google/gemma-4-26b-a4b-it:free".into()), + (P::HuggingFace, "google/gemma-4-26B-A4B-it".into()), + (P::NvidiaNim, "poolside/laguna-xs-2.1".into()), + (P::OpenRouter, "poolside/laguna-xs-2.1:free".into()), + (P::ZAi, "glm-4.7-flash".into()), + (P::NvidiaNim, "google/diffusiongemma-26b-a4b-it".into()), + (P::OpenRouter, "cohere/north-mini-code:free".into()), + (P::NvidiaNim, "nvidia/nemotron-3.5-lightning-30b-a3b".into()), + ( + P::VercelAiGateway, + "nvidia/nemotron-3.5-lightning-free".into(), + ), + (P::OpenRouter, "nvidia/nemotron-3.5-lightning:free".into()), + (P::ZAi, "glm-4.6v-flash".into()), + (P::ZAi, "glm-4.5-flash".into()), + (P::OpenRouter, "poolside/laguna-xs.2:free".into()), + ( + P::NvidiaNim, + "nvidia/nemotron-3-nano-omni-30b-a3b-reasoning".into(), + ), + ( + P::OpenRouter, + "nvidia/nemotron-3-nano-omni-30b-a3b-reasoning:free".into(), + ), + (P::NvidiaNim, "openai/gpt-oss-20b".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::NvidiaNim, "meta/llama-3.2-11b-vision-instruct".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::NvidiaNim, "meta/muse-glimmer-30b".into()), + (P::HuggingFace, "meta-models/Muse-Glimmer-30B".into()), + (P::GoogleAiStudio, "gemma-4-31b-it".into()), + (P::NvidiaNim, "google/gemma-4-31b-it".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::VercelAiGateway, "alibaba/qwen3.7-flash".into()), + (P::VercelAiGateway, "alibaba/qwen3.8-omni-flash".into()), + (P::Mistral, "mistral-medium-3.5".into()), + (P::NovitaAi, "inclusionai/ling-3.0-flash-vl".into()), + (P::OpenRouter, "inclusionai/ling-3.0-flash-vl:free".into()), + (P::VercelAiGateway, "inclusionai/ling-3.0-flash-vl".into()), + (P::VercelAiGateway, "poolside/laguna-s-2.1-free".into()), + (P::OpenRouter, "poolside/laguna-s-2.1:free".into()), + (P::VercelAiGateway, "stepfun/step-3.7-flash".into()), + (P::NovitaAi, "inclusionai/ling-3.0-flash-fin".into()), + (P::OpenRouter, "inclusionai/ling-3.0-flash-fin:free".into()), + (P::VercelAiGateway, "inclusionai/ling-3.0-flash-fin".into()), + (P::NovitaAi, "inclusionai/ling-3.0-flash-sante".into()), + ( + P::OpenRouter, + "inclusionai/ling-3.0-flash-sante:free".into(), + ), + ( + P::VercelAiGateway, + "inclusionai/ling-3.0-flash-sante".into(), + ), + (P::VercelAiGateway, "mistral/mistral-large-3".into()), + (P::NvidiaNim, "nvidia/nemotron-3-super-120b-a12b".into()), + ( + P::OpenRouter, + "nvidia/nemotron-3-super-120b-a12b:free".into(), + ), + (P::Groq, "openai/gpt-oss-120b".into()), + (P::CloudflareWorkersAi, "@cf/openai/gpt-oss-120b".into()), + (P::OpenRouter, "openai/gpt-oss-120b:free".into()), + (P::HuggingFace, "openai/gpt-oss-120b".into()), + (P::NvidiaNim, "meta/llama-3.2-90b-vision-instruct".into()), + ], + + S::L => vec![ + (P::CloudflareWorkersAi, "@cf/zai-org/glm-5.3-flash".into()), + (P::NvidiaNim, "z-ai/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::VercelAiGateway, "openai/gpt-5.2".into()), + (P::NvidiaNim, "deepseek-ai/deepseek-v4.1-flash".into()), + (P::HuggingFace, "deepseek-ai/DeepSeek-V4.1-Flash".into()), + (P::VercelAiGateway, "xiaomi/mimo-v2.6-flash".into()), + (P::HuggingFace, "XiaomiMiMo/MiMo-V2.6-Flash-RL".into()), + (P::VercelAiGateway, "inclusionai/ling-3.1-flash".into()), + (P::OpenRouter, "inclusionai/ling-3.1-flash".into()), + (P::VercelAiGateway, "xiaomi/mimo-v2.5".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()), + (P::NvidiaNim, "nvidia/nemotron-3-ultra-550b-a55b".into()), + ], + + S::XL => vec![ + (P::NvidiaNim, "z-ai/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::Mistral, "mistral-large-4".into()), + (P::NvidiaNim, "moonshotai/kimi-k3".into()), + (P::HuggingFace, "moonshotai/Kimi-K3".into()), + (P::HuggingFace, "Qwen/Qwen3.8-2.4T-A95B".into()), + (P::VercelAiGateway, "xiaomi/mimo-v2.5-pro".into()), + (P::VercelAiGateway, "tencent/hy3".into()), + (P::OpenRouter, "thinkingmachines/inkling:free".into()), + (P::HuggingFace, "thinkingmachines/Inkling".into()), + (P::VercelAiGateway, "thinkingmachines/inkling".into()), + (P::GoogleAiStudio, "gemini-3.1-pro-preview".into()), + (P::VercelAiGateway, "moonshotai/kimi-k2.7-code".into()), + ], + + S::Experimental => vec![ + (P::OpenRouter, "space-bunny-alpha".into()), + (P::VercelAiGateway, "stealth/glyph-cluster".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() + .url( + format!("https://api.cloudflare.com/client/v4/accounts/{account_id}/ai/v1/chat/completions") + .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() + .url( + "https://generativelanguage.googleapis.com/v1beta/openai/v1/chat/completions" + .parse() + .expect("TODO"), + ) + .api_key(api_key) + .build() + }) + .collect(), + Provider::Groq => state + .app_args + .groq_api_key + .iter() + .cloned() + .map(|api_key| { + Upstream::builder() + .url("https://api.groq.com/openai/v1/chat/completions".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() + .url("https://router.huggingface.co/v1/chat/completions".parse().expect("TODO")) + .api_key(api_key) + .build() + }) + .collect(), + Provider::LiquidAi => unreachable!(), + Provider::Mistral => state + .app_args + .mistral_api_key + .iter() + .cloned() + .map(|api_key| { + Upstream::builder() + .url("https://api.mistral.ai/v1/chat/completions".parse().expect("TODO")) + .api_key(api_key) + .build() + }) + .collect(), + Provider::NvidiaNim => state + .app_args + .nvidia_nim_api_key + .iter() + .cloned() + .map(|api_key| { + Upstream::builder() + .url("https://integrate.api.nvidia.com/v1/chat/completions".parse().expect("TODO")) + .api_key(api_key) + .build() + }) + .collect(), + Provider::NovitaAi => state + .app_args + .novita_ai_api_key + .iter() + .cloned() + .map(|api_key| { + Upstream::builder() + .url("https://api.novita.ai/openai/v1/chat/completions".parse().expect("TODO")) + .api_key(api_key) + .build() + }) + .collect(), + Provider::OpenRouter => state + .app_args + .openrouter_api_key + .iter() + .cloned() + .map(|api_key| { + Upstream::builder() + .url("https://openrouter.ai/api/v1/chat/completions".parse().expect("TODO")) + .api_key(api_key) + .build() + }) + .collect(), + Provider::VercelAiGateway => state + .app_args + .vercel_ai_gateway_api_key + .iter() + .cloned() + .map(|api_key| { + Upstream::builder() + .url("https://ai-gateway.vercel.sh/v1/chat/completions".parse().expect("TODO")) + .api_key(api_key) + .build() + }) + .collect(), + Provider::ZAi => state + .app_args + .z_ai_api_key + .iter() + .cloned() + .map(|api_key| { + Upstream::builder() + .url("https://api.z.ai/api/paas/v4/chat/completions".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(upstream.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/decisions.rs b/src/server/decisions.rs index 23d48b3..7e3ada8 100644 --- a/src/server/decisions.rs +++ b/src/server/decisions.rs @@ -224,8 +224,13 @@ pub async fn create_decision( .into_iter() .map(|(choice_id, probability)| { openai::decisions::ScoreProbability { - value: choice_id.parse().expect("it should be a number, I think"), - label: legend.get(&choice_id).expect("it should exist, I think").clone(), + value: choice_id + .parse() + .expect("it should be a number, I think"), + label: legend + .get(&choice_id) + .expect("it should exist, I think") + .clone(), probability, } }) diff --git a/src/server/embeddings.rs b/src/server/embeddings.rs index e249bb9..baa0642 100644 --- a/src/server/embeddings.rs +++ b/src/server/embeddings.rs @@ -1,279 +1,286 @@ -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_extra::routing::TypedPath; -use ecow::EcoString; -use headers::Header; -use headers::authorization::Bearer; -use mitsein::vec1::Vec1; -use reqwest::Certificate; -use secrecy::ExposeSecret; -use serde::Deserialize; -use serde_with::DeserializeFromStr; -use serde_with::SerializeDisplay; -use std::collections::VecDeque; -use strum::EnumString; -use strum::VariantArray; - -#[derive(Deserialize, TypedPath)] -#[typed_path("/v1/embeddings")] -pub struct V1EmbeddingsPath; - -#[derive( - Clone, Debug, EnumString, DeserializeFromStr, strum::Display, SerializeDisplay, VariantArray, -)] -#[strum(serialize_all = "lowercase")] -pub enum Shortcut { - #[strum(serialize = "codestral-embed-2505")] - CodestralEmbed2505, - - #[strum(serialize = "lfm2.5-embedding-350m")] - Lfm2_5Embedding350M, - - #[strum(serialize = "mistral-embed-2312")] - MistralEmbed2312, - - #[strum(serialize = "nemotron-3-embed-1b")] - Nemotron3Embed1B, -} - -impl Shortcut { - fn suggestions(&self) -> Vec<(Provider, EcoString)> { - use Provider as P; - use Shortcut as S; - - match self { - S::CodestralEmbed2505 => vec![ - (P::Mistral, "codestral-embed-2505".into()) - ], - S::Lfm2_5Embedding350M => { - vec![(P::OpenRouter, "liquid/lfm-2.5-embedding-350m:free".into())] - }, - S::MistralEmbed2312 => vec![ - (P::Mistral, "mistral-embed-2312".into()) - ], - S::Nemotron3Embed1B => { - vec![(P::OpenRouter, "nvidia/nemotron-3-embed-1b:free".into())] - } - } - } -} - -pub struct EmbeddingsResponse(openai::embeddings::CreateEmbeddingResponse); - -impl IntoResponse for EmbeddingsResponse { - fn into_response(self) -> Response { - let Self(embeddings) = self; - - Json(embeddings).into_response() - } -} - -pub enum EmbeddingsError { - NoneSucceeded, -} - -impl IntoResponse for EmbeddingsError { - fn into_response(self) -> Response { - match self { - EmbeddingsError::NoneSucceeded => StatusCode::SERVICE_UNAVAILABLE.into_response(), - } - } -} - -#[tracing::instrument(skip(state, request))] -pub async fn create_embeddings( - _: V1EmbeddingsPath, - State(state): State, - Json(request): Json>, -) -> Result { - let shortcut = 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() - .url( - format!("https://api.cloudflare.com/client/v4/accounts/{account_id}/ai/v1/embeddings") - .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() - .url( - "https://generativelanguage.googleapis.com/v1beta/openai/v1/embeddings" - .parse() - .expect("TODO"), - ) - .api_key(api_key) - .build() - }) - .collect(), - Provider::Groq => state - .app_args - .groq_api_key - .iter() - .cloned() - .map(|api_key| { - Upstream::builder() - .url("https://api.groq.com/openai/v1/embeddings".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() - .url("https://router.huggingface.co/v1/embeddings".parse().expect("TODO")) - .api_key(api_key) - .build() - }) - .collect(), - Provider::LiquidAi => unreachable!(), - Provider::Mistral => state - .app_args - .mistral_api_key - .iter() - .cloned() - .map(|api_key| { - Upstream::builder() - .url("https://api.mistral.ai/v1/embeddings".parse().expect("TODO")) - .api_key(api_key) - .build() - }) - .collect(), - Provider::OpenRouter => state - .app_args - .openrouter_api_key - .iter() - .cloned() - .map(|api_key| { - Upstream::builder() - .url("https://openrouter.ai/api/v1/embeddings".parse().expect("TODO")) - .api_key(api_key) - .build() - }) - .collect(), - Provider::ZAi => unreachable!(), - }; - - 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::embeddings::CreateEmbeddingRequest { - model: _shortcut, - input, - dimensions, - encoding_format, - user, - } = request.clone(); - - let request = openai::embeddings::CreateEmbeddingRequest { - model, - input, - dimensions, - encoding_format, - user, - }; - - - 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(upstream.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; - }; - - - let response = match upstream_response.json().await { - Ok(response) => response, - Err(json_error) => { - tracing::error!(?json_error); - continue; - } - }; - - return Ok(EmbeddingsResponse(response)); - } - Err(EmbeddingsError::NoneSucceeded) -} +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_extra::routing::TypedPath; +use ecow::EcoString; +use headers::Header; +use headers::authorization::Bearer; +use mitsein::vec1::Vec1; +use reqwest::Certificate; +use secrecy::ExposeSecret; +use serde::Deserialize; +use serde_with::DeserializeFromStr; +use serde_with::SerializeDisplay; +use std::collections::VecDeque; +use strum::EnumString; +use strum::VariantArray; + +#[derive(Deserialize, TypedPath)] +#[typed_path("/v1/embeddings")] +pub struct V1EmbeddingsPath; + +#[derive( + Clone, Debug, EnumString, DeserializeFromStr, strum::Display, SerializeDisplay, VariantArray, +)] +#[strum(serialize_all = "lowercase")] +pub enum Shortcut { + #[strum(serialize = "codestral-embed-2505")] + CodestralEmbed2505, + + #[strum(serialize = "lfm2.5-embedding-350m")] + Lfm2_5Embedding350M, + + #[strum(serialize = "mistral-embed-2312")] + MistralEmbed2312, + + #[strum(serialize = "nemotron-3-embed-1b")] + Nemotron3Embed1B, +} + +impl Shortcut { + fn suggestions(&self) -> Vec<(Provider, EcoString)> { + use Provider as P; + use Shortcut as S; + + match self { + S::CodestralEmbed2505 => vec![(P::Mistral, "codestral-embed-2505".into())], + S::Lfm2_5Embedding350M => { + vec![(P::OpenRouter, "liquid/lfm-2.5-embedding-350m:free".into())] + } + S::MistralEmbed2312 => vec![(P::Mistral, "mistral-embed-2312".into())], + S::Nemotron3Embed1B => vec![ + (P::NvidiaNim, "nvidia/nemotron-3-embed-1b".into()), + (P::OpenRouter, "nvidia/nemotron-3-embed-1b:free".into()), + ], + } + } +} + +pub struct EmbeddingsResponse(openai::embeddings::CreateEmbeddingResponse); + +impl IntoResponse for EmbeddingsResponse { + fn into_response(self) -> Response { + let Self(embeddings) = self; + + Json(embeddings).into_response() + } +} + +pub enum EmbeddingsError { + NoneSucceeded, +} + +impl IntoResponse for EmbeddingsError { + fn into_response(self) -> Response { + match self { + EmbeddingsError::NoneSucceeded => StatusCode::SERVICE_UNAVAILABLE.into_response(), + } + } +} + +#[tracing::instrument(skip(state, request))] +pub async fn create_embeddings( + _: V1EmbeddingsPath, + State(state): State, + Json(request): Json>, +) -> Result { + let shortcut = 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() + .url( + format!("https://api.cloudflare.com/client/v4/accounts/{account_id}/ai/v1/embeddings") + .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() + .url( + "https://generativelanguage.googleapis.com/v1beta/openai/v1/embeddings" + .parse() + .expect("TODO"), + ) + .api_key(api_key) + .build() + }) + .collect(), + Provider::Groq => state + .app_args + .groq_api_key + .iter() + .cloned() + .map(|api_key| { + Upstream::builder() + .url("https://api.groq.com/openai/v1/embeddings".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() + .url("https://router.huggingface.co/v1/embeddings".parse().expect("TODO")) + .api_key(api_key) + .build() + }) + .collect(), + Provider::LiquidAi => unreachable!(), + Provider::Mistral => state + .app_args + .mistral_api_key + .iter() + .cloned() + .map(|api_key| { + Upstream::builder() + .url("https://api.mistral.ai/v1/embeddings".parse().expect("TODO")) + .api_key(api_key) + .build() + }) + .collect(), + Provider::NvidiaNim => state + .app_args + .nvidia_nim_api_key + .iter() + .cloned() + .map(|api_key| { + Upstream::builder() + .url("https://integrate.api.nvidia.com/v1/embeddings".parse().expect("TODO")) + .api_key(api_key) + .build() + }) + .collect(), + Provider::NovitaAi => unreachable!(), + Provider::OpenRouter => state + .app_args + .openrouter_api_key + .iter() + .cloned() + .map(|api_key| { + Upstream::builder() + .url("https://openrouter.ai/api/v1/embeddings".parse().expect("TODO")) + .api_key(api_key) + .build() + }) + .collect(), + Provider::VercelAiGateway => unreachable!(), + Provider::ZAi => unreachable!(), + }; + + 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::embeddings::CreateEmbeddingRequest { + model: _shortcut, + input, + dimensions, + encoding_format, + user, + } = request.clone(); + + let request = openai::embeddings::CreateEmbeddingRequest { + model, + input, + dimensions, + encoding_format, + user, + }; + + 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(upstream.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; + }; + + let response = match upstream_response.json().await { + Ok(response) => response, + Err(json_error) => { + tracing::error!(?json_error); + continue; + } + }; + + return Ok(EmbeddingsResponse(response)); + } + Err(EmbeddingsError::NoneSucceeded) +} diff --git a/src/server/mod.rs b/src/server/mod.rs index 72c4156..f6ad60b 100644 --- a/src/server/mod.rs +++ b/src/server/mod.rs @@ -28,7 +28,10 @@ enum Provider { HuggingFace, LiquidAi, Mistral, + NovitaAi, + NvidiaNim, OpenRouter, + VercelAiGateway, ZAi, } diff --git a/src/server/system_one.rs b/src/server/system_one.rs index 780df5e..a100a92 100644 --- a/src/server/system_one.rs +++ b/src/server/system_one.rs @@ -49,6 +49,8 @@ impl Shortcut { (P::CloudflareWorkersAi, "clef-flash".into()), (P::OpenRouter, "respan/span-01-lite:free".into()), (P::OpenRouter, "respan/span-01-lite".into()), + (P::VercelAiGateway, "convaiinnovations/laya-free".into()), + (P::VercelAiGateway, "convaiinnovations/laya".into()), ] } S::M => { @@ -132,6 +134,8 @@ pub async fn create_system_one( }) .collect(), Provider::Mistral => unreachable!(), + Provider::NovitaAi => unreachable!(), + Provider::NvidiaNim => unreachable!(), Provider::OpenRouter => state .app_args .openrouter_api_key @@ -144,6 +148,18 @@ pub async fn create_system_one( .build() }) .collect(), + Provider::VercelAiGateway => state + .app_args + .vercel_ai_gateway_api_key + .iter() + .cloned() + .map(|api_key| { + Upstream::builder() + .url("https://ai-gateway.vercel.sh/typesafe/v1/systemone".parse().expect("TODO")) + .api_key(api_key) + .build() + }) + .collect(), Provider::ZAi => unreachable!(), };