717 lines
27 KiB
Rust
717 lines
27 KiB
Rust
use crate::CloudflareCredentials;
|
|
use crate::openai;
|
|
use crate::server::AppState;
|
|
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<openai::chat::completion::ReasoningEffort>,
|
|
);
|
|
|
|
impl FromStr for ShortcutAndMaybeReasoning {
|
|
type Err = Infallible; // TODO
|
|
|
|
fn from_str(s: &str) -> Result<Self, Self::Err> {
|
|
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(())
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
enum Provider {
|
|
CloudflareWorkersAi,
|
|
Cohere,
|
|
GoogleAiStudio,
|
|
Groq,
|
|
HuggingFace,
|
|
Kilo,
|
|
Mistral,
|
|
NovitaAi,
|
|
NvidiaNim,
|
|
Ollama,
|
|
OpenCode,
|
|
OpenRouter,
|
|
Poolside,
|
|
VercelAiGateway,
|
|
ZAi,
|
|
}
|
|
|
|
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::VercelAiGateway, "openai/gpt-5.4-nano".into()),
|
|
(P::HuggingFace, "Qwen/Qwen3.6-35B-A3B".into()),
|
|
(P::OpenRouter, "xingchen-agi/xing4.0-29b-a4b:free".into()),
|
|
(P::NovitaAi, "apodex/apodex-1.1-mini".into()),
|
|
(P::OpenRouter, "apodex/apodex-1.1-mini:free".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::Kilo, "google-gemma-4-26b-a4b-it-free".into()),
|
|
(P::HuggingFace, "google/gemma-4-26B-A4B-it".into()),
|
|
(P::Poolside, "poolside/laguna-xs-2.1".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::Cohere, "north-mini-code-1-0".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::OpenCode, "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::Ollama, "nemotron-3-nano:30b".into()),
|
|
(P::Groq, "openai/gpt-oss-20b".into()),
|
|
(P::Ollama, "gpt-oss:20b".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::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::Kilo, "kilo-auto/free".into()),
|
|
(P::OpenRouter, "openrouter/free".into()),
|
|
],
|
|
|
|
S::S => vec![
|
|
(P::Groq, "qwen/qwen3.8-27b".into()),
|
|
(P::CloudflareWorkersAi, "@cf/qwen/qwen3.8-27b".into()),
|
|
(P::OpenRouter, "qwen/qwen3.8-27b:free".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::Ollama, "gemma4:31b".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::NvidiaNim,
|
|
"nvidia/nemotron-3.5-super-vl-120b-a12b".into(),
|
|
),
|
|
(P::OpenRouter, "nvidia/nemotron-3.5-super-vl:free".into()),
|
|
(P::OpenCode, "nemotron-3.5-super-vl-free".into()),
|
|
(P::Mistral, "mistral-medium-3.5".into()),
|
|
(P::Kilo, "inclusionai-ling-3-0-flash-free".into()),
|
|
(P::Kilo, "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::Poolside, "poolside/laguna-s-2.1".into()),
|
|
(P::Kilo, "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::VercelAiGateway, "inclusionai/ling-3.0-flash-fin".into()),
|
|
|
|
(
|
|
P::VercelAiGateway,
|
|
"inclusionai/ling-3.0-flash-sante".into(),
|
|
),
|
|
(P::VercelAiGateway, "mistral/mistral-large-3".into()),
|
|
(P::Ollama, "nemotron-3-super".into()),
|
|
(P::Cohere, "command-a-plus-05-2026".into()),
|
|
(P::NvidiaNim, "nvidia/nemotron-3-super-120b-a12b".into()),
|
|
(P::Kilo, "nvidia-nemotron-3-super-120b-a12b-free".into()),
|
|
(
|
|
P::OpenRouter,
|
|
"nvidia/nemotron-3-super-120b-a12b:free".into(),
|
|
),
|
|
(P::Groq, "openai/gpt-oss-120b".into()),
|
|
(P::Ollama, "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::OpenCode, "glm-5.3-flash-free".into()),
|
|
(P::HuggingFace, "zai-org/GLM-5.3-Flash".into()),
|
|
(P::HuggingFace, "zai-org/GLM-5.3-Flash-BF16".into()),
|
|
(P::Kilo, "step-5-preview-free".into()),
|
|
(P::NvidiaNim, "stepfun-ai/step-5-preview".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::OpenCode, "deepseek-v4.1-flash-free".into()),
|
|
(P::VercelAiGateway, "minimax/minimax-m3.1-flash".into()),
|
|
(P::OpenRouter, "minimax/minimax-m3.1-flash:free".into()),
|
|
(P::VercelAiGateway, "inclusionai/ling-3.1-flash".into()),
|
|
(P::HuggingFace, "inclusionAI/Ling-3.1-flash".into()),
|
|
(P::VercelAiGateway, "reflection/beam-free".into()),
|
|
(P::Ollama, "beam".into()),
|
|
(P::OpenCode, "beam-free".into()),
|
|
(P::VercelAiGateway, "xiaomi/mimo-v2.6-flash".into()),
|
|
(P::OpenCode, "mimo-v2.6-flash-free".into()),
|
|
(P::HuggingFace, "XiaomiMiMo/MiMo-V2.6-Flash-RL".into()),
|
|
(P::Kilo, "nex-agi-nex-n2-5-pro-free".into()),
|
|
(P::VercelAiGateway, "xiaomi/mimo-v2.5".into()),
|
|
(P::Kilo, "dots-studio-dots-3-note-preview-free".into()),
|
|
(P::OpenRouter, "dots-studio/dots-3-note-preview:free".into()),
|
|
(P::Kilo, "minimax-minimax-m3-free".into()),
|
|
(P::HuggingFace, "MiniMaxAI/MiniMax-M3".into()),
|
|
(P::VercelAiGateway, "tencent/hy3".into()),
|
|
(P::Kilo, "tencent-hy3-free".into()),
|
|
(P::Kilo, "tencent-hy3-preview-free".into()),
|
|
(P::OpenRouter, "thinkingmachines/inkling-small:free".into()),
|
|
(P::HuggingFace, "thinkingmachines/Inkling-Small".into()),
|
|
(P::Ollama, "nemotron-3-ultra".into()),
|
|
(P::NvidiaNim, "nvidia/nemotron-3-ultra-550b-a55b".into()),
|
|
(P::Kilo, "nvidia-nemotron-3-ultra-550b-a55b-free".into()),
|
|
(P::OpenCode, "nemotron-3-ultra-free".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::VercelAiGateway, "deepseek/deepseek-v4.1-flash-1017".into()),
|
|
(P::OpenCode, "muse-spark-1.3-contributor-free".into()),
|
|
(P::NvidiaNim, "moonshotai/kimi-k3".into()),
|
|
(P::HuggingFace, "moonshotai/Kimi-K3".into()),
|
|
(P::HuggingFace, "tencent/Hy4-preview".into()),
|
|
(P::HuggingFace, "XiaomiMiMo/MiMo-V2.6-Pro-RL".into()),
|
|
(P::Mistral, "mistral-large-4".into()),
|
|
(P::HuggingFace, "Qwen/Qwen3.8-2.4T-A95B".into()),
|
|
(P::VercelAiGateway, "xiaomi/mimo-v2.5-pro".into()),
|
|
(P::OpenCode, "longcat-2.5-preview-free".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::OpenCode, "space-bunny-free".into()),
|
|
(P::OpenCode, "big-pickle".into()),
|
|
(P::VercelAiGateway, "stealth/glyph-cluster".into()),
|
|
],
|
|
}
|
|
}
|
|
}
|
|
|
|
pub enum ChatCompletionResponse {
|
|
Immediate(openai::chat::completion::ChatCompletion<String>),
|
|
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::<dyn std::error::Error + Send + Sync>::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<AppState>,
|
|
Json(request): Json<
|
|
openai::chat::completion::CreateChatCompletionRequest<ShortcutAndMaybeReasoning>,
|
|
>,
|
|
) -> Result<ChatCompletionResponse, ChatCompletionError> {
|
|
let ShortcutAndMaybeReasoning(shortcut, reasoning_in_id) = request.model.clone();
|
|
let mut model_upstreams = VecDeque::new();
|
|
|
|
for (provider, model) in shortcut.suggestions() {
|
|
let upstreams: Vec<Upstream> =
|
|
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::Cohere => state
|
|
.app_args
|
|
.cohere_api_key
|
|
.iter()
|
|
.cloned()
|
|
.map(|api_key| {
|
|
Upstream::builder()
|
|
.url(
|
|
"https://api.cohere.com/compatibility/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::Kilo => state
|
|
.app_args
|
|
.kilo_api_key
|
|
.iter()
|
|
.cloned()
|
|
.map(|api_key| {
|
|
Upstream::builder()
|
|
.url("https://api.kilo.ai/api/gateway/v1/chat/completions".parse().expect("TODO"))
|
|
.api_key(api_key)
|
|
.build()
|
|
})
|
|
.collect(),
|
|
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::Ollama => state
|
|
.app_args
|
|
.ollama_cloud_api_key
|
|
.iter()
|
|
.cloned()
|
|
.map(|api_key| {
|
|
Upstream::builder()
|
|
.url("https://ollama.com/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::OpenCode => state
|
|
.app_args
|
|
.opencode_console_api_key
|
|
.iter()
|
|
.cloned()
|
|
.map(|api_key| {
|
|
Upstream::builder()
|
|
.url("https://opencode.ai/inference/openai/v1/chat/completions".parse().expect("TODO"))
|
|
.api_key(api_key)
|
|
.build()
|
|
})
|
|
.collect(),
|
|
Provider::Poolside => state
|
|
.app_args
|
|
.poolside_api_key
|
|
.iter()
|
|
.cloned()
|
|
.map(|api_key| {
|
|
Upstream::builder()
|
|
.url("https://inference.poolside.ai/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::<Bearer>::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)
|
|
}
|