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 futures::StreamExt as _;
use reqwest_sse::EventSource;
use secrecy::SecretString;
use tokio::{join, pin, task::LocalSet};
use tracing_subscriber::{EnvFilter, fmt::format::FmtSpan};
mod openai;
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,
}
use std::net::SocketAddr;
use tracing_subscriber::fmt::format::FmtSpan;
#[tokio::main]
async fn main() {
let args = Parser::parse();
let Args {
let katniss_ai::Args {
env_filter,
app_args,
} = args;
@@ -61,54 +17,20 @@ async fn main() {
.init();
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 (tx, rx) = async_channel::unbounded();
let app = server::create_app(
server::AppState::builder()
let app = katniss_ai::server::create_app(
katniss_ai::server::AppState::builder()
.app_args(app_args.into())
.stream_requester(tx)
.build(),
);
let listener = tokio::net::TcpListener::bind(addr).await.expect("TODO");
let stream_responses = async move {
let local_set = LocalSet::new();
local_set
.run_until(async {
pin!(rx);
while let Some((response, sse_sink)) = rx.next().await {
local_set.spawn_local(async move {
match response.events().await {
Ok(mut events) => {
while let Some(event_result) = events.next().await {
let message = event_result
.map(|event| {
axum::response::sse::Event::default().data(event.data)
})
.map_err(Into::into);
let _ = sse_sink.send(message).await;
}
}
Err(event_source_error) => {
let _ = sse_sink.send(Err(event_source_error.into())).await;
}
};
});
}
})
.await;
};
join!(stream_responses, async move {
tracing::info!(?addr, "listening");
axum::serve(listener, app).await.expect("TODO");
});
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_with::{DeserializeFromStr, SerializeDisplay};
use strum::{EnumString, VariantArray};
use std::collections::HashMap;
use strum::{EnumString, VariantArray};
// ============================================================================
// Request: Create Chat Completion
@@ -918,7 +918,9 @@ pub enum PromptCacheRetention {
}
/// 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")]
pub enum ReasoningEffort {
/// Disable reasoning.

263
src/server.rs
View File

@@ -1,4 +1,5 @@
use crate::AppArgs;
use crate::CloudflareCredentials;
use crate::openai;
use axum::Json;
use axum::Router;
@@ -7,9 +8,9 @@ use axum::http::StatusCode;
use axum::response::IntoResponse;
use axum::response::Response;
use axum::response::Sse;
use axum_extra::routing::RouterExt;
use axum_extra::routing::TypedPath;
use axum_extra::routing::{RouterExt, TypedPath};
use ecow::EcoString;
use futures::StreamExt;
use headers::Header;
use headers::authorization::Bearer;
use itertools::iproduct;
@@ -17,31 +18,26 @@ 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::Serialize;
use serde_with::DeserializeFromStr;
use serde_with::SerializeDisplay;
use strum::EnumString;
use strum::VariantArray;
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;
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)]
pub struct AppState {
app_args: Arc<AppArgs>,
stream_requester: StreamRequester,
}
#[derive(Deserialize, TypedPath)]
@@ -58,7 +54,9 @@ enum Provider {
OpenRouter,
}
#[derive(Clone, Debug, EnumString, DeserializeFromStr, strum::Display, SerializeDisplay, VariantArray)]
#[derive(
Clone, Debug, EnumString, DeserializeFromStr, strum::Display, SerializeDisplay, VariantArray,
)]
#[strum(serialize_all = "lowercase")]
enum Shortcut {
XS,
@@ -70,7 +68,6 @@ enum Shortcut {
Experimental,
}
#[derive(Debug, Clone, DeserializeFromStr, SerializeDisplay)]
struct ShortcutAndMaybeReasoning(Shortcut, Option<openai::chat::completion::ReasoningEffort>);
@@ -106,7 +103,6 @@ impl Display for ShortcutAndMaybeReasoning {
}
}
impl Shortcut {
fn suggestions(&self) -> Vec<(Provider, EcoString)> {
use Provider as P;
@@ -139,6 +135,7 @@ impl Shortcut {
(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()),
@@ -150,6 +147,10 @@ impl Shortcut {
(P::Mistral, "mistral-medium-3.5".into()),
(P::OpenRouter, "inclusionai/ling-3.0-flash-vl:free".into()),
(P::OpenRouter, "poolside/laguna-s-2.1:free".into()),
(
P::OpenRouter,
"inclusionai/ling-3.0-flash-sante:free".into(),
),
(
P::OpenRouter,
"nvidia/nemotron-3-super-120b-a12b:free".into(),
@@ -162,28 +163,29 @@ impl Shortcut {
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::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::OpenRouter, "thinkingmachines/inkling-small:free".into()),
(P::HuggingFace, "thinkingmachines/Inkling-Small".into()),
],
S::XL => vec![
(P::CloudflareWorkersAi, "@cf/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::HuggingFace, "tencent/Hy4-preview".into()),
(P::HuggingFace, "XiaomiMiMo/MiMo-V2.6-Pro-RL".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::HuggingFace, "tencent/Hy4-preview".into()),
(P::HuggingFace, "deepseek-ai/DeepSeek-V4-Pro".into()),
],
S::Experimental => vec![
(P::OpenRouter, "space-bunny-alpha".into())
]
S::Experimental => vec![(P::OpenRouter, "space-bunny-alpha".into())],
}
}
@@ -225,10 +227,7 @@ struct Upstream {
enum ChatCompletionResponse {
Immediate(openai::chat::completion::ChatCompletion),
Stream {
response: reqwest::Response,
stream_requester: StreamRequester,
},
Stream(reqwest::Response),
}
impl IntoResponse for ChatCompletionResponse {
@@ -237,17 +236,33 @@ impl IntoResponse for ChatCompletionResponse {
ChatCompletionResponse::Immediate(chat_completion) => {
Json(chat_completion).into_response()
}
ChatCompletionResponse::Stream {
response,
stream_requester,
} => {
let (sse_sink, sse_stream) = async_channel::bounded(256);
ChatCompletionResponse::Stream(response) => {
let (sink, stream) = async_channel::bounded(256);
stream_requester
.try_send((response, sse_sink))
.expect("this is an unbounded channel");
let handle = Handle::current();
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(
_: V1ChatCompletionsPath,
State(state): State<AppState>,
Json(request): Json<openai::chat::completion::CreateChatCompletionRequest<ShortcutAndMaybeReasoning>>,
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 = match provider {
Provider::CloudflareWorkersAi => vec![], // TODO
Provider::GoogleAiStudio => state
.app_args
.google_ai_studio_api_key
.iter()
.cloned()
.map(|api_key| {
Upstream::builder()
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()
.base_url(
"https://generativelanguage.googleapis.com/v1beta/openai"
format!("https://api.cloudflare.com/client/v4/accounts/{account_id}/ai")
.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(),
};
},
)
.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));
@@ -397,7 +435,11 @@ async fn create_chat_completion(
web_search_options,
} = 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 {
model,
@@ -443,7 +485,7 @@ async fn create_chat_completion(
.tls_certs_only(
webpki_root_certs::TLS_SERVER_ROOT_CERTS
.iter()
.filter_map(|cert| Certificate::from_der(&cert).ok()),
.filter_map(|cert| Certificate::from_der(cert).ok()),
)
.build()
.expect("TODO");
@@ -500,10 +542,7 @@ async fn create_chat_completion(
};
if request.stream.unwrap_or(false) {
return Ok(ChatCompletionResponse::Stream {
response: upstream_response,
stream_requester: state.stream_requester,
});
return Ok(ChatCompletionResponse::Stream(upstream_response));
} else {
let response = match upstream_response.json().await {
Ok(response) => response,
@@ -525,19 +564,25 @@ async fn create_chat_completion(
struct V1ModelsPath;
#[tracing::instrument]
async fn list_models(_: V1ModelsPath) -> Json<openai::models::ListModelsResponse<ShortcutAndMaybeReasoning>> {
let shortcuts = Shortcut::VARIANTS.into_iter().cloned();
let reasoning_efforts = openai::chat::completion::ReasoningEffort::VARIANTS.into_iter().cloned().map(Some).chain(None);
async fn list_models(
_: V1ModelsPath,
) -> 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)
.into_iter()
.map(|(shortcut, reasoning)| openai::models::Model {
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(),