jacob/katniss-ai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: stuff

Browse Source
This commit is contained in:
Jacob
2026-10-08 01:24:23 -04:00
parent e161693587
commit 5ce36e3b8c
4 changed files with 237 additions and 199 deletions
Download Patch File Download Diff File Expand all files Collapse all files

69
src/lib.rs Normal file
View File

@@ -0,0 +1,69 @@
use clap::Parser;
use secrecy::SecretString;
use std::{convert::Infallible, net::IpAddr, str::FromStr};
use tracing_subscriber::EnvFilter;
pub mod openai;
pub mod server;
#[derive(Debug, Clone)]
pub struct CloudflareCredentials {
pub account_id: String,
pub api_key: SecretString,
}
impl FromStr for CloudflareCredentials {
type Err = Infallible; // TODO
fn from_str(s: &str) -> Result<Self, Self::Err> {
let (account_id, api_key) = s.split_once(":").expect("TODO");
let account_id = account_id.into();
let api_key = api_key.into();
Ok(Self {
account_id,
api_key,
})
}
}
#[derive(Debug, Clone, Parser)]
pub struct AppArgs {
#[arg(long, env, default_value_t = [127, 0, 0, 1].into())]
pub ip: IpAddr,
#[arg(long, env, default_value_t = 3126)]
pub port: u16,
#[arg(long)]
pub cloudflare_workers_ai_credentials: Vec<CloudflareCredentials>,
#[arg(long)]
pub google_ai_studio_api_key: Vec<SecretString>,
#[arg(long)]
pub groq_api_key: Vec<SecretString>,
#[arg(long)]
pub hugging_face_api_key: Vec<SecretString>,
#[arg(long)]
pub mistral_api_key: Vec<SecretString>,
#[arg(long)]
pub openrouter_api_key: Vec<SecretString>,
}
#[derive(Parser)]
pub struct Args {
#[arg(
long = "tracing-directives",
env = "RUST_LOG",
default_value = "warn,katniss_ai=debug"
)]
pub env_filter: EnvFilter,
#[command(flatten)]
pub app_args: AppArgs,
}

98
src/main.rs
View File

@@ -1,55 +1,11 @@
use std::net::{IpAddr, SocketAddr};
use clap::Parser; use clap::Parser;
use futures::StreamExt as _; use std::net::SocketAddr;
use reqwest_sse::EventSource; use tracing_subscriber::fmt::format::FmtSpan;
use secrecy::SecretString;
use tokio::{join, pin, task::LocalSet};
use tracing_subscriber::{EnvFilter, fmt::format::FmtSpan};
mod openai;
mod server;
#[derive(Parser, Debug)]
struct AppArgs {
#[arg(long, env, default_value_t = [127, 0, 0, 1].into())]
ip: IpAddr,
#[arg(long, env, default_value_t = 3126)]
port: u16,
#[arg(long)]
cloudflare_workers_ai_api_key: Vec<SecretString>,
#[arg(long)]
google_ai_studio_api_key: Vec<SecretString>,
#[arg(long)]
groq_api_key: Vec<SecretString>,
#[arg(long)]
hugging_face_api_key: Vec<SecretString>,
#[arg(long)]
mistral_api_key: Vec<SecretString>,
#[arg(long)]
openrouter_api_key: Vec<SecretString>,
}
#[derive(Parser)]
struct Args {
#[arg(long = "tracing-directives", env = "RUST_LOG", default_value = "warn,katniss_ai=debug")]
env_filter: EnvFilter,
#[command(flatten)]
app_args: AppArgs,
}
#[tokio::main] #[tokio::main]
async fn main() { async fn main() {
let args = Parser::parse(); let args = Parser::parse();
let Args { let katniss_ai::Args {
env_filter, env_filter,
app_args, app_args,
} = args; } = args;
@@ -61,54 +17,20 @@ async fn main() {
.init(); .init();
tracing::debug!(?app_args, "using"); tracing::debug!(?app_args, "using");
rustls::crypto::aws_lc_rs::default_provider().install_default().unwrap(); rustls::crypto::aws_lc_rs::default_provider()
.install_default()
.unwrap();
let addr = SocketAddr::new(app_args.ip, app_args.port); let addr = SocketAddr::new(app_args.ip, app_args.port);
let (tx, rx) = async_channel::unbounded(); let app = katniss_ai::server::create_app(
katniss_ai::server::AppState::builder()
let app = server::create_app(
server::AppState::builder()
.app_args(app_args.into()) .app_args(app_args.into())
.stream_requester(tx)
.build(), .build(),
); );
let listener = tokio::net::TcpListener::bind(addr).await.expect("TODO"); let listener = tokio::net::TcpListener::bind(addr).await.expect("TODO");
let stream_responses = async move { tracing::info!(?addr, "listening");
let local_set = LocalSet::new(); axum::serve(listener, app).await.expect("TODO");
local_set
.run_until(async {
pin!(rx);
while let Some((response, sse_sink)) = rx.next().await {
local_set.spawn_local(async move {
match response.events().await {
Ok(mut events) => {
while let Some(event_result) = events.next().await {
let message = event_result
.map(|event| {
axum::response::sse::Event::default().data(event.data)
})
.map_err(Into::into);
let _ = sse_sink.send(message).await;
}
}
Err(event_source_error) => {
let _ = sse_sink.send(Err(event_source_error.into())).await;
}
};
});
}
})
.await;
};
join!(stream_responses, async move {
tracing::info!(?addr, "listening");
axum::serve(listener, app).await.expect("TODO");
});
} }

6
src/openai/chat/completion.rs
View File

@@ -1,7 +1,7 @@
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use serde_with::{DeserializeFromStr, SerializeDisplay}; use serde_with::{DeserializeFromStr, SerializeDisplay};
use strum::{EnumString, VariantArray};
use std::collections::HashMap; use std::collections::HashMap;
use strum::{EnumString, VariantArray};
// ============================================================================ // ============================================================================
// Request: Create Chat Completion // Request: Create Chat Completion
@@ -918,7 +918,9 @@ pub enum PromptCacheRetention {
} }
/// Constrains reasoning effort for reasoning models. /// Constrains reasoning effort for reasoning models.
#[derive(Clone, Debug, EnumString, DeserializeFromStr, strum::Display, SerializeDisplay, VariantArray)] #[derive(
Clone, Debug, EnumString, DeserializeFromStr, strum::Display, SerializeDisplay, VariantArray,
)]
#[strum(serialize_all = "lowercase")] #[strum(serialize_all = "lowercase")]
pub enum ReasoningEffort { pub enum ReasoningEffort {
/// Disable reasoning. /// Disable reasoning.

263
src/server.rs
View File

@@ -1,4 +1,5 @@
use crate::AppArgs; use crate::AppArgs;
use crate::CloudflareCredentials;
use crate::openai; use crate::openai;
use axum::Json; use axum::Json;
use axum::Router; use axum::Router;
@@ -7,9 +8,9 @@ use axum::http::StatusCode;
use axum::response::IntoResponse; use axum::response::IntoResponse;
use axum::response::Response; use axum::response::Response;
use axum::response::Sse; use axum::response::Sse;
use axum_extra::routing::RouterExt; use axum_extra::routing::{RouterExt, TypedPath};
use axum_extra::routing::TypedPath;
use ecow::EcoString; use ecow::EcoString;
use futures::StreamExt;
use headers::Header; use headers::Header;
use headers::authorization::Bearer; use headers::authorization::Bearer;
use itertools::iproduct; use itertools::iproduct;
@@ -17,31 +18,26 @@ use mitsein::vec1::Vec1;
use rand::Rng; use rand::Rng;
use rand::RngExt; use rand::RngExt;
use reqwest::Certificate; use reqwest::Certificate;
use reqwest_sse::EventSource as _;
use secrecy::ExposeSecret; use secrecy::ExposeSecret;
use secrecy::SecretString; use secrecy::SecretString;
use serde::Deserialize; use serde::Deserialize;
use serde::Serialize;
use serde_with::DeserializeFromStr; use serde_with::DeserializeFromStr;
use serde_with::SerializeDisplay; use serde_with::SerializeDisplay;
use strum::EnumString;
use strum::VariantArray;
use std::collections::VecDeque; use std::collections::VecDeque;
use std::convert::Infallible; use std::convert::Infallible;
use std::fmt::Display; use std::fmt::Display;
use std::str::FromStr; use std::str::FromStr;
use std::sync::Arc; use std::sync::Arc;
use strum::EnumString;
use strum::VariantArray;
use tokio::runtime::Handle;
use typed_builder::TypedBuilder; use typed_builder::TypedBuilder;
use url::Url; use url::Url;
type SSEEventResult = Result<axum::response::sse::Event, Box<dyn std::error::Error + Sync + Send>>;
type StreamRequest = (reqwest::Response, async_channel::Sender<SSEEventResult>);
type StreamRequester = async_channel::Sender<StreamRequest>;
#[derive(Debug, Clone, TypedBuilder)] #[derive(Debug, Clone, TypedBuilder)]
pub struct AppState { pub struct AppState {
app_args: Arc<AppArgs>, app_args: Arc<AppArgs>,
stream_requester: StreamRequester,
} }
#[derive(Deserialize, TypedPath)] #[derive(Deserialize, TypedPath)]
@@ -58,7 +54,9 @@ enum Provider {
OpenRouter, OpenRouter,
} }
#[derive(Clone, Debug, EnumString, DeserializeFromStr, strum::Display, SerializeDisplay, VariantArray)] #[derive(
Clone, Debug, EnumString, DeserializeFromStr, strum::Display, SerializeDisplay, VariantArray,
)]
#[strum(serialize_all = "lowercase")] #[strum(serialize_all = "lowercase")]
enum Shortcut { enum Shortcut {
XS, XS,
@@ -70,7 +68,6 @@ enum Shortcut {
Experimental, Experimental,
} }
#[derive(Debug, Clone, DeserializeFromStr, SerializeDisplay)] #[derive(Debug, Clone, DeserializeFromStr, SerializeDisplay)]
struct ShortcutAndMaybeReasoning(Shortcut, Option<openai::chat::completion::ReasoningEffort>); struct ShortcutAndMaybeReasoning(Shortcut, Option<openai::chat::completion::ReasoningEffort>);
@@ -106,7 +103,6 @@ impl Display for ShortcutAndMaybeReasoning {
} }
} }
impl Shortcut { impl Shortcut {
fn suggestions(&self) -> Vec<(Provider, EcoString)> { fn suggestions(&self) -> Vec<(Provider, EcoString)> {
use Provider as P; use Provider as P;
@@ -139,6 +135,7 @@ impl Shortcut {
(P::HuggingFace, "Qwen/Qwen3.8-27B".into()), (P::HuggingFace, "Qwen/Qwen3.8-27B".into()),
(P::HuggingFace, "meta-models/Muse-Glimmer-30B".into()), (P::HuggingFace, "meta-models/Muse-Glimmer-30B".into()),
(P::OpenRouter, "google/gemma-4-31b-it:free".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::Mistral, "mistral-small-2603".into()),
(P::Groq, "groq/compound".into()), (P::Groq, "groq/compound".into()),
(P::OpenRouter, "openrouter/fusion".into()), (P::OpenRouter, "openrouter/fusion".into()),
@@ -150,6 +147,10 @@ impl Shortcut {
(P::Mistral, "mistral-medium-3.5".into()), (P::Mistral, "mistral-medium-3.5".into()),
(P::OpenRouter, "inclusionai/ling-3.0-flash-vl:free".into()), (P::OpenRouter, "inclusionai/ling-3.0-flash-vl:free".into()),
(P::OpenRouter, "poolside/laguna-s-2.1:free".into()), (P::OpenRouter, "poolside/laguna-s-2.1:free".into()),
(
P::OpenRouter,
"inclusionai/ling-3.0-flash-sante:free".into(),
),
( (
P::OpenRouter, P::OpenRouter,
"nvidia/nemotron-3-super-120b-a12b:free".into(), "nvidia/nemotron-3-super-120b-a12b:free".into(),
@@ -162,28 +163,29 @@ impl Shortcut {
S::L => vec![ S::L => vec![
(P::CloudflareWorkersAi, "@cf/zai-org/glm-5.3-flash".into()), (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".into()),
(P::HuggingFace, "zai-org/GLM-5.3-Flash-BF16".into()),
(P::HuggingFace, "deepseek-ai/DeepSeek-V4.1-Flash".into()), (P::HuggingFace, "deepseek-ai/DeepSeek-V4.1-Flash".into()),
(P::OpenRouter, "nex-agi/nex-n2.5-pro:free".into()), (P::HuggingFace, "XiaomiMiMo/MiMo-V2.6-Flash-RL".into()),
(P::OpenRouter, "inclusionai/ling-3.1-flash".into()),
(P::OpenRouter, "dots-studio/dots-3-note-preview:free".into()),
(P::HuggingFace, "MiniMaxAI/MiniMax-M3".into()), (P::HuggingFace, "MiniMaxAI/MiniMax-M3".into()),
(P::OpenRouter, "thinkingmachines/inkling-small:free".into()), (P::OpenRouter, "thinkingmachines/inkling-small:free".into()),
(P::HuggingFace, "thinkingmachines/Inkling-Small".into()), (P::HuggingFace, "thinkingmachines/Inkling-Small".into()),
], ],
S::XL => vec![ S::XL => vec![
(P::CloudflareWorkersAi, "@cf/zai-org/glm-5.3".into()),
(P::HuggingFace, "zai-org/GLM-5.3".into()), (P::HuggingFace, "zai-org/GLM-5.3".into()),
(P::CloudflareWorkersAi, "@cf/zai-org/glm-5.3".into()),
(P::Mistral, "zai-glm-5-3".into()), (P::Mistral, "zai-glm-5-3".into()),
(P::HuggingFace, "tencent/Hy4-preview".into()),
(P::HuggingFace, "XiaomiMiMo/MiMo-V2.6-Pro-RL".into()),
(P::HuggingFace, "moonshotai/Kimi-K3".into()), (P::HuggingFace, "moonshotai/Kimi-K3".into()),
(P::HuggingFace, "Qwen/Qwen3.8-2.4T-A95B".into()), (P::HuggingFace, "Qwen/Qwen3.8-2.4T-A95B".into()),
(P::OpenRouter, "thinkingmachines/inkling:free".into()), (P::OpenRouter, "thinkingmachines/inkling:free".into()),
(P::HuggingFace, "thinkingmachines/Inkling".into()), (P::HuggingFace, "thinkingmachines/Inkling".into()),
(P::HuggingFace, "tencent/Hy4-preview".into()),
(P::HuggingFace, "deepseek-ai/DeepSeek-V4-Pro".into()),
], ],
S::Experimental => vec![ S::Experimental => vec![(P::OpenRouter, "space-bunny-alpha".into())],
(P::OpenRouter, "space-bunny-alpha".into())
]
} }
} }
@@ -225,10 +227,7 @@ struct Upstream {
enum ChatCompletionResponse { enum ChatCompletionResponse {
Immediate(openai::chat::completion::ChatCompletion), Immediate(openai::chat::completion::ChatCompletion),
Stream { Stream(reqwest::Response),
response: reqwest::Response,
stream_requester: StreamRequester,
},
} }
impl IntoResponse for ChatCompletionResponse { impl IntoResponse for ChatCompletionResponse {
@@ -237,17 +236,33 @@ impl IntoResponse for ChatCompletionResponse {
ChatCompletionResponse::Immediate(chat_completion) => { ChatCompletionResponse::Immediate(chat_completion) => {
Json(chat_completion).into_response() Json(chat_completion).into_response()
} }
ChatCompletionResponse::Stream { ChatCompletionResponse::Stream(response) => {
response, let (sink, stream) = async_channel::bounded(256);
stream_requester,
} => {
let (sse_sink, sse_stream) = async_channel::bounded(256);
stream_requester let handle = Handle::current();
.try_send((response, sse_sink))
.expect("this is an unbounded channel");
Sse::new(sse_stream).into_response() std::thread::spawn(move || {
handle.block_on(async move {
match response.events().await {
Ok(mut events) => {
while let Some(event_result) = events.next().await {
let message = event_result
.map(|event| {
axum::response::sse::Event::default().data(event.data)
})
.map_err(Box::<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()
} }
} }
} }
@@ -270,79 +285,102 @@ impl IntoResponse for ChatCompletionError {
async fn create_chat_completion( async fn create_chat_completion(
_: V1ChatCompletionsPath, _: V1ChatCompletionsPath,
State(state): State<AppState>, State(state): State<AppState>,
Json(request): Json<openai::chat::completion::CreateChatCompletionRequest<ShortcutAndMaybeReasoning>>, Json(request): Json<
openai::chat::completion::CreateChatCompletionRequest<ShortcutAndMaybeReasoning>,
>,
) -> Result<ChatCompletionResponse, ChatCompletionError> { ) -> Result<ChatCompletionResponse, ChatCompletionError> {
let ShortcutAndMaybeReasoning(shortcut, reasoning_in_id) = request.model.clone(); let ShortcutAndMaybeReasoning(shortcut, reasoning_in_id) = request.model.clone();
let mut model_upstreams = VecDeque::new(); let mut model_upstreams = VecDeque::new();
for (provider, model) in shortcut.suggestions() { for (provider, model) in shortcut.suggestions() {
let upstreams = match provider { let upstreams: Vec<Upstream> =
Provider::CloudflareWorkersAi => vec![], // TODO match provider {
Provider::GoogleAiStudio => state Provider::CloudflareWorkersAi => state
.app_args .app_args
.google_ai_studio_api_key .cloudflare_workers_ai_credentials
.iter() .iter()
.cloned() .cloned()
.map(|api_key| { .map(
Upstream::builder() |CloudflareCredentials {
account_id,
api_key,
}| {
Upstream::builder()
.base_url( .base_url(
"https://generativelanguage.googleapis.com/v1beta/openai" format!("https://api.cloudflare.com/client/v4/accounts/{account_id}/ai")
.parse() .parse()
.expect("TODO"), .expect("TODO"),
) )
.api_key(api_key) .api_key(api_key)
.build() .build()
}) },
.collect(), )
Provider::Groq => state .collect(),
.app_args Provider::GoogleAiStudio => state
.groq_api_key .app_args
.iter() .google_ai_studio_api_key
.cloned() .iter()
.map(|api_key| { .cloned()
Upstream::builder() .map(|api_key| {
.base_url("https://api.groq.com/openai".parse().expect("TODO")) Upstream::builder()
.api_key(api_key) .base_url(
.build() "https://generativelanguage.googleapis.com/v1beta/openai"
}) .parse()
.collect(), .expect("TODO"),
Provider::HuggingFace => state )
.app_args .api_key(api_key)
.hugging_face_api_key .build()
.iter() })
.cloned() .collect(),
.map(|api_key| { Provider::Groq => state
Upstream::builder() .app_args
.base_url("https://router.huggingface.co".parse().expect("TODO")) .groq_api_key
.api_key(api_key) .iter()
.build() .cloned()
}) .map(|api_key| {
.collect(), Upstream::builder()
Provider::Mistral => state .base_url("https://api.groq.com/openai".parse().expect("TODO"))
.app_args .api_key(api_key)
.mistral_api_key .build()
.iter() })
.cloned() .collect(),
.map(|api_key| { Provider::HuggingFace => state
Upstream::builder() .app_args
.base_url("https://api.mistral.ai".parse().expect("TODO")) .hugging_face_api_key
.api_key(api_key) .iter()
.build() .cloned()
}) .map(|api_key| {
.collect(), Upstream::builder()
Provider::OpenRouter => state .base_url("https://router.huggingface.co".parse().expect("TODO"))
.app_args .api_key(api_key)
.openrouter_api_key .build()
.iter() })
.cloned() .collect(),
.map(|api_key| { Provider::Mistral => state
Upstream::builder() .app_args
.base_url("https://openrouter.ai/api".parse().expect("TODO")) .mistral_api_key
.api_key(api_key) .iter()
.build() .cloned()
}) .map(|api_key| {
.collect(), 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) { if let Ok(upstreams) = Vec1::try_from(upstreams) {
model_upstreams.push_back((model, upstreams)); model_upstreams.push_back((model, upstreams));
@@ -397,7 +435,11 @@ async fn create_chat_completion(
web_search_options, web_search_options,
} = request.clone(); } = request.clone();
let reasoning_effort = Some(reasoning_effort.or(reasoning_in_id.clone()).unwrap_or(shortcut.default_reasoning())); let reasoning_effort = Some(
reasoning_effort
.or(reasoning_in_id.clone())
.unwrap_or(shortcut.default_reasoning()),
);
let request = openai::chat::completion::CreateChatCompletionRequest { let request = openai::chat::completion::CreateChatCompletionRequest {
model, model,
@@ -443,7 +485,7 @@ async fn create_chat_completion(
.tls_certs_only( .tls_certs_only(
webpki_root_certs::TLS_SERVER_ROOT_CERTS webpki_root_certs::TLS_SERVER_ROOT_CERTS
.iter() .iter()
.filter_map(|cert| Certificate::from_der(&cert).ok()), .filter_map(|cert| Certificate::from_der(cert).ok()),
) )
.build() .build()
.expect("TODO"); .expect("TODO");
@@ -500,10 +542,7 @@ async fn create_chat_completion(
}; };
if request.stream.unwrap_or(false) { if request.stream.unwrap_or(false) {
return Ok(ChatCompletionResponse::Stream { return Ok(ChatCompletionResponse::Stream(upstream_response));
response: upstream_response,
stream_requester: state.stream_requester,
});
} else { } else {
let response = match upstream_response.json().await { let response = match upstream_response.json().await {
Ok(response) => response, Ok(response) => response,
@@ -525,19 +564,25 @@ async fn create_chat_completion(
struct V1ModelsPath; struct V1ModelsPath;
#[tracing::instrument] #[tracing::instrument]
async fn list_models(_: V1ModelsPath) -> Json<openai::models::ListModelsResponse<ShortcutAndMaybeReasoning>> { async fn list_models(
let shortcuts = Shortcut::VARIANTS.into_iter().cloned(); _: V1ModelsPath,
let reasoning_efforts = openai::chat::completion::ReasoningEffort::VARIANTS.into_iter().cloned().map(Some).chain(None); ) -> Json<openai::models::ListModelsResponse<ShortcutAndMaybeReasoning>> {
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) let models = iproduct!(shortcuts, reasoning_efforts).map(|(shortcut, reasoning)| {
.into_iter() openai::models::Model {
.map(|(shortcut, reasoning)| openai::models::Model {
id: ShortcutAndMaybeReasoning(shortcut, reasoning), id: ShortcutAndMaybeReasoning(shortcut, reasoning),
created: 0, created: 0,
object: openai::models::ModelObject::Model, object: openai::models::ModelObject::Model,
owned_by: "nobody".into(), owned_by: "nobody".into(),
shutdown_date: None, shutdown_date: None,
}); }
});
Json(openai::models::ListModelsResponse { Json(openai::models::ListModelsResponse {
data: models.collect(), data: models.collect(),