chore: initialize autoresearch baseline

This commit is contained in:
Evandro Fonseca Junior committed 2026-03-26 17:28:25 -03:00
commit f83e8b83f6
42 files changed
+10550

No files matched your search

+721
View File
@@ -0,0 +1,721 @@
use std::collections::HashMap;
use anyhow::{bail, Result};
use api_contracts::{
chat::{
ChatCompletionRequest, DeltaToolCall, DeltaToolFunctionCall, FinishReason, ResponseFormat,
Role, ToolCall, ToolChoice, ToolDefinition, ToolFunctionCall, ToolType,
},
error::ApiError,
models::ModelInfo,
};
use axum::http::{HeaderMap, HeaderName, HeaderValue, StatusCode};
use serde_json::{json, Value};
use tower_http::cors::{AllowHeaders, AllowMethods, AllowOrigin, Any, CorsLayer};
const DEFAULT_SYSTEM_PROMPT: &str = concat!(
"You are a precise coding assistant running locally. ",
"Prioritize correct code, minimal speculation, explicit assumptions, and actionable fixes. ",
"When editing or generating code, prefer complete snippets, preserve requested APIs, and explain risks briefly."
);
#[derive(Debug, Clone)]
pub struct ServerConfig {
pub api_key: Option<String>,
pub auth_header: String,
pub auth_scheme: Option<String>,
pub cors_origins: Vec<String>,
pub default_system_prompt: String,
pub enable_tool_calls: bool,
pub model_catalog: ModelCatalog,
}
impl ServerConfig {
pub fn from_env(model_name: String) -> Self {
let aliases = env_csv("LOCALLLM_MODEL_ALIASES");
Self {
api_key: std::env::var("LOCALLLM_API_KEY")
.ok()
.filter(|v| !v.is_empty()),
auth_header: std::env::var("LOCALLLM_AUTH_HEADER")
.ok()
.filter(|v| !v.trim().is_empty())
.unwrap_or_else(|| "authorization".to_owned())
.to_ascii_lowercase(),
auth_scheme: std::env::var("LOCALLLM_AUTH_SCHEME")
.ok()
.map(|v| v.trim().to_owned())
.filter(|v| !v.is_empty()),
cors_origins: {
let origins = env_csv("LOCALLLM_CORS_ORIGINS");
if origins.is_empty() {
vec!["*".to_owned()]
} else {
origins
}
},
default_system_prompt: std::env::var("LOCALLLM_SYSTEM_PROMPT")
.ok()
.filter(|v| !v.trim().is_empty())
.unwrap_or_else(|| DEFAULT_SYSTEM_PROMPT.to_owned()),
enable_tool_calls: env_bool("LOCALLLM_ENABLE_TOOL_CALLS", true),
model_catalog: ModelCatalog::new(model_name, aliases),
}
}
pub fn cors_layer(&self) -> CorsLayer {
let methods = AllowMethods::any();
let headers = AllowHeaders::any();
if self.cors_origins.iter().any(|origin| origin == "*") {
return CorsLayer::new()
.allow_origin(Any)
.allow_methods(methods)
.allow_headers(headers);
}
let origins: Vec<HeaderValue> = self
.cors_origins
.iter()
.filter_map(|origin| HeaderValue::from_str(origin).ok())
.collect();
CorsLayer::new()
.allow_origin(AllowOrigin::list(origins))
.allow_methods(methods)
.allow_headers(headers)
}
pub fn resolve_model(&self, requested: &str) -> Option<&ModelDescriptor> {
self.model_catalog.resolve(requested)
}
pub fn model_list(&self, created: i64) -> Vec<ModelInfo> {
self.model_catalog.to_model_list(created)
}
}
#[derive(Debug, Clone)]
pub struct ModelCatalog {
models: Vec<ModelDescriptor>,
aliases: HashMap<String, usize>,
}
impl ModelCatalog {
pub fn new(canonical_id: String, aliases: Vec<String>) -> Self {
let mut models = vec![ModelDescriptor {
id: canonical_id.clone(),
root: canonical_id.clone(),
parent: None,
}];
let mut seen = std::collections::BTreeSet::new();
seen.insert(canonical_id.to_ascii_lowercase());
seen.insert("local-model".to_owned());
seen.insert("default".to_owned());
for alias in [
vec!["local-model".to_owned(), "default".to_owned()],
aliases,
]
.concat()
{
let trimmed = alias.trim();
if trimmed.is_empty() {
continue;
}
let lowered = trimmed.to_ascii_lowercase();
if !seen.insert(lowered) {
continue;
}
models.push(ModelDescriptor {
id: trimmed.to_owned(),
root: canonical_id.clone(),
parent: Some(canonical_id.clone()),
});
}
let aliases = models
.iter()
.enumerate()
.map(|(idx, model)| (model.id.to_ascii_lowercase(), idx))
.collect();
Self { models, aliases }
}
pub fn resolve(&self, requested: &str) -> Option<&ModelDescriptor> {
self.aliases
.get(&requested.to_ascii_lowercase())
.and_then(|idx| self.models.get(*idx))
}
pub fn to_model_list(&self, created: i64) -> Vec<ModelInfo> {
self.models
.iter()
.map(|model| ModelInfo {
id: model.id.clone(),
object: "model",
created,
owned_by: "local",
root: model.root.clone(),
parent: model.parent.clone(),
permission: Vec::new(),
})
.collect()
}
}
#[derive(Debug, Clone)]
pub struct ModelDescriptor {
pub id: String,
pub root: String,
pub parent: Option<String>,
}
#[derive(Debug, Clone)]
pub struct PreparedPrompt {
pub prompt: String,
pub streaming_mode: StreamingMode,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum StreamingMode {
LiveTokens,
Buffered,
}
#[derive(Debug, Clone)]
pub struct FinalizedAssistant {
pub content: Option<String>,
pub tool_calls: Option<Vec<ToolCall>>,
pub finish_reason: FinishReason,
}
pub fn authorize_request(
headers: &HeaderMap,
config: &ServerConfig,
) -> Result<(), (StatusCode, ApiError)> {
let Some(expected_key) = config.api_key.as_deref() else {
return Ok(());
};
let configured_header = HeaderName::from_bytes(config.auth_header.as_bytes())
.ok()
.and_then(|name| headers.get(name))
.and_then(|value| value.to_str().ok())
.map(str::trim)
.filter(|value| !value.is_empty());
let fallback_bearer = headers
.get(axum::http::header::AUTHORIZATION)
.and_then(|value| value.to_str().ok())
.map(str::trim)
.filter(|value| !value.is_empty());
let fallback_api_key = headers
.get("x-api-key")
.and_then(|value| value.to_str().ok())
.map(str::trim)
.filter(|value| !value.is_empty());
let provided = configured_header.or(fallback_bearer).or(fallback_api_key);
let Some(provided) = provided else {
return Err((
StatusCode::UNAUTHORIZED,
ApiError::new("Missing API key", "invalid_request_error").with_code("missing_api_key"),
));
};
let normalized = match config.auth_scheme.as_deref() {
Some(scheme) => provided
.strip_prefix(scheme)
.map(str::trim)
.unwrap_or(provided),
None => provided,
};
if normalized == expected_key || provided == expected_key {
Ok(())
} else {
Err((
StatusCode::UNAUTHORIZED,
ApiError::new("Invalid API key", "invalid_request_error").with_code("invalid_api_key"),
))
}
}
pub fn validate_request(
req: &ChatCompletionRequest,
config: &ServerConfig,
) -> Result<(), (StatusCode, ApiError)> {
if config.resolve_model(&req.model).is_none() {
return Err((
StatusCode::NOT_FOUND,
ApiError::new(
format!("The model '{}' does not exist", req.model),
"invalid_request_error",
)
.with_code("model_not_found")
.with_param("model"),
));
}
if req.messages.is_empty() {
return Err((
StatusCode::BAD_REQUEST,
ApiError::new("messages must not be empty", "invalid_request_error")
.with_code("empty_messages")
.with_param("messages"),
));
}
if req.n != 1 {
return Err((
StatusCode::BAD_REQUEST,
ApiError::new("Only n=1 is currently supported", "invalid_request_error")
.with_code("unsupported_value")
.with_param("n"),
));
}
if !req.tools.is_empty() && !config.enable_tool_calls {
return Err((
StatusCode::BAD_REQUEST,
ApiError::new(
"Tool calling is disabled on this server",
"invalid_request_error",
)
.with_code("tool_calls_disabled")
.with_param("tools"),
));
}
Ok(())
}
pub fn prepare_prompt(req: &ChatCompletionRequest, config: &ServerConfig) -> PreparedPrompt {
let mut prompt = String::with_capacity(1024);
let has_system = req
.messages
.iter()
.any(|message| matches!(message.role, Role::System | Role::Developer));
if !has_system {
prompt.push_str("System: ");
prompt.push_str(&config.default_system_prompt);
prompt.push('\n');
}
if let Some(format) = &req.response_format {
append_response_format_instruction(&mut prompt, format);
}
if !req.tools.is_empty() {
append_tool_instructions(&mut prompt, &req.tools, req.tool_choice.as_ref());
}
for msg in &req.messages {
let role = match msg.role {
Role::System => "System",
Role::Developer => "Developer",
Role::User => "User",
Role::Assistant => "Assistant",
Role::Tool => "Tool",
};
prompt.push_str(role);
if let Some(name) = &msg.name {
prompt.push_str("(");
prompt.push_str(name);
prompt.push(')');
}
prompt.push_str(": ");
if let Some(tool_calls) = &msg.tool_calls {
let tool_json = serde_json::to_string(tool_calls).unwrap_or_else(|_| "[]".to_owned());
prompt.push_str(&tool_json);
} else {
prompt.push_str(msg.text());
}
if let Some(tool_call_id) = &msg.tool_call_id {
prompt.push_str(" [tool_call_id=");
prompt.push_str(tool_call_id);
prompt.push(']');
}
prompt.push('\n');
}
prompt.push_str("Assistant:");
PreparedPrompt {
prompt,
streaming_mode: if req.tools.is_empty() {
StreamingMode::LiveTokens
} else {
StreamingMode::Buffered
},
}
}
pub fn finalize_assistant_message(
raw_output: String,
req: &ChatCompletionRequest,
) -> FinalizedAssistant {
if !req.tools.is_empty() {
if let Some(tool_calls) = parse_tool_calls(&raw_output) {
return FinalizedAssistant {
content: None,
tool_calls: Some(tool_calls),
finish_reason: FinishReason::ToolCalls,
};
}
}
FinalizedAssistant {
content: Some(raw_output.trim().to_owned()),
tool_calls: None,
finish_reason: FinishReason::Stop,
}
}
pub fn tool_calls_to_delta(tool_calls: &[ToolCall]) -> Vec<DeltaToolCall> {
tool_calls
.iter()
.enumerate()
.map(|(index, tool_call)| DeltaToolCall {
index: index as u32,
id: Some(tool_call.id.clone()),
tool_type: Some(tool_call.tool_type.clone()),
function: Some(DeltaToolFunctionCall {
name: Some(tool_call.function.name.clone()),
arguments: Some(tool_call.function.arguments.clone()),
}),
})
.collect()
}
fn append_response_format_instruction(prompt: &mut String, format: &ResponseFormat) {
match format.format_type.as_str() {
"json_object" => {
prompt.push_str("System: Return valid JSON only. Do not wrap it in markdown fences.\n");
}
"json_schema" => {
prompt.push_str("System: Return valid JSON only and follow this schema exactly: ");
if let Some(schema) = &format.json_schema {
prompt.push_str(&schema.to_string());
}
prompt.push('\n');
}
_ => {}
}
}
fn append_tool_instructions(
prompt: &mut String,
tools: &[ToolDefinition],
tool_choice: Option<&ToolChoice>,
) {
let tool_summaries: Vec<Value> = tools
.iter()
.map(|tool| {
json!({
"name": tool.function.name,
"description": tool.function.description,
"parameters": tool.function.parameters,
})
})
.collect();
prompt.push_str("System: Available tools: ");
prompt.push_str(&serde_json::to_string(&tool_summaries).unwrap_or_else(|_| "[]".to_owned()));
prompt.push('\n');
prompt
.push_str("System: If you need a tool, respond with JSON only using one of these shapes: ");
prompt.push_str(r#"{"tool_calls":[{"name":"tool_name","arguments":{}}]}"#);
prompt.push_str(" or ");
prompt.push_str(r#"{"name":"tool_name","arguments":{}}"#);
prompt.push_str(". If no tool is needed, answer normally.\n");
match tool_choice {
Some(ToolChoice::Mode(mode)) if mode == "required" => {
prompt.push_str("System: You must call a tool before replying.\n");
}
Some(ToolChoice::Mode(mode)) if mode == "none" => {
prompt.push_str("System: Do not call tools in this response.\n");
}
Some(ToolChoice::Specific(choice)) => {
prompt.push_str("System: You must call the tool named ");
prompt.push_str(&choice.function.name);
prompt.push_str(".\n");
}
_ => {}
}
}
fn parse_tool_calls(raw_output: &str) -> Option<Vec<ToolCall>> {
let trimmed = raw_output.trim();
if trimmed.is_empty() {
return None;
}
let value: Value = serde_json::from_str(trimmed).ok()?;
match value {
Value::Object(map) => {
if let Some(tool_calls) = map.get("tool_calls").and_then(Value::as_array) {
let parsed: Vec<ToolCall> = tool_calls
.iter()
.enumerate()
.filter_map(|(index, item)| parse_tool_call(item, index))
.collect();
if parsed.is_empty() {
None
} else {
Some(parsed)
}
} else {
parse_tool_call(&Value::Object(map), 0).map(|tool_call| vec![tool_call])
}
}
_ => None,
}
}
fn parse_tool_call(value: &Value, index: usize) -> Option<ToolCall> {
let object = value.as_object()?;
let function = object.get("function").unwrap_or(value);
let function_obj = function.as_object()?;
let name = function_obj.get("name")?.as_str()?.trim();
if name.is_empty() {
return None;
}
let arguments = function_obj
.get("arguments")
.cloned()
.unwrap_or_else(|| Value::Object(Default::default()));
let arguments = match arguments {
Value::String(s) => s,
other => serde_json::to_string(&other).ok()?,
};
Some(ToolCall {
id: format!("call_{}", index + 1),
tool_type: ToolType::Function,
function: ToolFunctionCall {
name: name.to_owned(),
arguments,
},
})
}
pub fn build_sampler(req: &ChatCompletionRequest) -> inference_core::SamplerConfig {
let mut sampler = inference_core::SamplerConfig::default();
sampler.temperature = req.temperature.clamp(0.0, 2.0);
sampler.top_p = req.top_p.clamp(0.0, 1.0);
sampler.max_tokens = req.max_tokens.max(1);
sampler.stop_sequences = req.stop.clone();
sampler.repeat_penalty = if req.tools.is_empty() { 1.05 } else { 1.0 };
sampler.seed = req.seed.unwrap_or(0);
sampler
}
fn env_csv(key: &str) -> Vec<String> {
std::env::var(key)
.ok()
.map(|value| {
value
.split(',')
.map(str::trim)
.filter(|part| !part.is_empty())
.map(ToOwned::to_owned)
.collect()
})
.unwrap_or_default()
}
fn env_bool(key: &str, default: bool) -> bool {
std::env::var(key)
.ok()
.and_then(|value| match value.trim().to_ascii_lowercase().as_str() {
"1" | "true" | "yes" | "on" => Some(true),
"0" | "false" | "no" | "off" => Some(false),
_ => None,
})
.unwrap_or(default)
}
pub fn ensure_supported_mode(req: &ChatCompletionRequest) -> Result<()> {
if matches!(req.tool_choice, Some(ToolChoice::Mode(ref mode)) if mode == "none") {
return Ok(());
}
if let Some(format) = &req.response_format {
if !matches!(
format.format_type.as_str(),
"json_object" | "json_schema" | "text"
) {
bail!("unsupported response_format type: {}", format.format_type);
}
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::{
authorize_request, finalize_assistant_message, prepare_prompt, validate_request,
FinalizedAssistant, ModelCatalog, ServerConfig, StreamingMode,
};
use api_contracts::chat::ChatCompletionRequest;
use axum::http::{HeaderMap, HeaderValue, StatusCode};
fn base_config() -> ServerConfig {
ServerConfig {
api_key: None,
auth_header: "authorization".to_owned(),
auth_scheme: Some("Bearer".to_owned()),
cors_origins: vec!["*".to_owned()],
default_system_prompt: "Default coding system prompt".to_owned(),
enable_tool_calls: true,
model_catalog: ModelCatalog::new(
"real-model".to_owned(),
vec!["alias-model".to_owned()],
),
}
}
fn parse_request(json: &str) -> ChatCompletionRequest {
serde_json::from_str(json).unwrap()
}
#[test]
fn resolves_model_aliases_case_insensitively() {
let catalog = ModelCatalog::new("real-model".to_owned(), vec!["Alias-Model".to_owned()]);
assert_eq!(catalog.resolve("alias-model").unwrap().root, "real-model");
assert_eq!(catalog.resolve("REAL-MODEL").unwrap().id, "real-model");
assert!(catalog.resolve("missing").is_none());
}
#[test]
fn validates_model_and_n() {
let config = base_config();
let request = parse_request(
r#"{
"model": "missing-model",
"messages": [{"role":"user","content":"hi"}]
}"#,
);
let err = validate_request(&request, &config).unwrap_err();
assert_eq!(err.0, StatusCode::NOT_FOUND);
let request = parse_request(
r#"{
"model": "alias-model",
"n": 2,
"messages": [{"role":"user","content":"hi"}]
}"#,
);
let err = validate_request(&request, &config).unwrap_err();
assert_eq!(err.0, StatusCode::BAD_REQUEST);
}
#[test]
fn authorizes_bearer_and_x_api_key_headers() {
let mut config = base_config();
config.api_key = Some("secret".to_owned());
let mut headers = HeaderMap::new();
headers.insert("authorization", HeaderValue::from_static("Bearer secret"));
assert!(authorize_request(&headers, &config).is_ok());
let mut headers = HeaderMap::new();
headers.insert("x-api-key", HeaderValue::from_static("secret"));
assert!(authorize_request(&headers, &config).is_ok());
}
#[test]
fn adds_default_system_prompt_and_tool_instructions() {
let config = base_config();
let request = parse_request(
r#"{
"model": "alias-model",
"messages": [{"role":"user","content":"List files"}],
"tools": [{
"type": "function",
"function": {
"name": "list_files",
"description": "List project files",
"parameters": {"type": "object"}
}
}],
"tool_choice": "required"
}"#,
);
let prepared = prepare_prompt(&request, &config);
assert_eq!(prepared.streaming_mode, StreamingMode::Buffered);
assert!(prepared.prompt.contains("Default coding system prompt"));
assert!(prepared.prompt.contains("list_files"));
assert!(prepared.prompt.contains("must call a tool"));
}
#[test]
fn extracts_tool_calls_from_model_output() {
let request = parse_request(
r#"{
"model": "alias-model",
"messages": [{"role":"user","content":"List files"}],
"tools": [{
"type": "function",
"function": {
"name": "list_files",
"parameters": {"type": "object"}
}
}]
}"#,
);
let finalized = finalize_assistant_message(
r#"{"tool_calls":[{"name":"list_files","arguments":{"path":"src"}}]}"#.to_owned(),
&request,
);
assert!(matches!(
finalized.finish_reason,
api_contracts::chat::FinishReason::ToolCalls
));
assert_eq!(finalized.content, None);
let tool_calls = finalized.tool_calls.unwrap();
assert_eq!(tool_calls[0].function.name, "list_files");
assert_eq!(tool_calls[0].function.arguments, "{\"path\":\"src\"}");
}
#[test]
fn preserves_plain_text_when_no_tool_call_is_present() {
let request = parse_request(
r#"{
"model": "alias-model",
"messages": [{"role":"user","content":"Explain borrow checker"}]
}"#,
);
let finalized =
finalize_assistant_message(" Use references and lifetimes. ".to_owned(), &request);
assert!(matches!(
finalized,
FinalizedAssistant {
content: Some(_),
tool_calls: None,
..
}
));
assert_eq!(
finalized.content.as_deref(),
Some("Use references and lifetimes.")
);
}
}
+101
View File
@@ -0,0 +1,101 @@
use std::net::SocketAddr;
use std::path::PathBuf;
use std::sync::Arc;
use anyhow::Result;
use tokio::sync::Mutex;
use tracing::info;
use tracing_subscriber::{fmt, EnvFilter};
use inference_core::{EngineConfig, InferenceEngine, RuntimeProfile};
use metrics::counters::PerfCounters;
mod compat;
mod routes;
use compat::ServerConfig;
use routes::{build_router, AppState};
#[tokio::main]
async fn main() -> Result<()> {
// ── Tracing ──────────────────────────────────────────────────────────
fmt()
.with_env_filter(
EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info")),
)
.with_target(false)
.init();
// ── Engine config (reads env vars; sane defaults for dev) ────────────
let model_path = std::env::var("LOCALLLM_MODEL")
.map(PathBuf::from)
.unwrap_or_default();
let profile = match std::env::var("LOCALLLM_PROFILE").as_deref() {
Ok("balanced_9b") => RuntimeProfile::Balanced9B,
Ok("cpu_fallback") => RuntimeProfile::CpuFallback,
_ => RuntimeProfile::Turbo4B,
};
let defaults = profile.defaults();
let config = EngineConfig {
model_path: model_path.clone(),
profile,
n_gpu_layers: env_i32("LOCALLLM_GPU_LAYERS", defaults.n_gpu_layers),
n_ctx: env_u32("LOCALLLM_CTX", defaults.n_ctx),
n_batch: env_u32("LOCALLLM_BATCH", defaults.n_batch),
n_ubatch: env_u32("LOCALLLM_UBATCH", defaults.n_ubatch),
n_threads: env_i32("LOCALLLM_THREADS", defaults.n_threads),
n_threads_batch: env_i32("LOCALLLM_THREADS_BATCH", defaults.n_threads),
flash_attn: defaults.flash_attn,
use_mmap: false,
seed: 0,
};
// ── Model name for API responses ─────────────────────────────────────
let model_name = model_path
.file_stem()
.map(|s| s.to_string_lossy().into_owned())
.unwrap_or_else(|| "local-model".into());
let server_config = Arc::new(ServerConfig::from_env(model_name.clone()));
// ── Load engine ──────────────────────────────────────────────────────
let mut engine = InferenceEngine::load(config)?;
engine.warmup()?;
let state = AppState {
engine: Arc::new(Mutex::new(engine)),
counters: PerfCounters::new(),
model_name: Arc::new(model_name),
server_config,
start_time: std::time::Instant::now(),
};
// ── Server ───────────────────────────────────────────────────────────
let addr: SocketAddr = std::env::var("LOCALLLM_ADDR")
.unwrap_or_else(|_| "0.0.0.0:8080".into())
.parse()?;
let router = build_router(state);
let listener = tokio::net::TcpListener::bind(addr).await?;
info!(addr = %addr, "LocalLLM server listening");
axum::serve(listener, router).await?;
Ok(())
}
// ── Env helpers ──────────────────────────────────────────────────────────
fn env_i32(key: &str, default: i32) -> i32 {
std::env::var(key)
.ok()
.and_then(|v| v.parse().ok())
.unwrap_or(default)
}
fn env_u32(key: &str, default: u32) -> u32 {
std::env::var(key)
.ok()
.and_then(|v| v.parse().ok())
.unwrap_or(default)
}
+420
View File
@@ -0,0 +1,420 @@
use axum::{
extract::{Json, State},
http::{HeaderMap, StatusCode},
response::{IntoResponse, Response, Sse},
routing::{get, post},
Router,
};
use std::{sync::Arc, time::SystemTime};
use tokio::sync::Mutex;
use api_contracts::{
chat::{
ChatCompletionChunk, ChatCompletionRequest, ChatCompletionResponse, ChatMessage, Choice,
ChunkChoice, Delta, FinishReason, Role, Usage,
},
error::{ApiError, ErrorResponse},
health::HealthResponse,
models::ModelList,
};
use inference_core::InferenceEngine;
use metrics::counters::PerfCounters;
use crate::compat::{
authorize_request, build_sampler, ensure_supported_mode, finalize_assistant_message,
prepare_prompt, tool_calls_to_delta, validate_request, FinalizedAssistant, ServerConfig,
StreamingMode,
};
#[derive(Clone)]
pub struct AppState {
pub engine: Arc<Mutex<InferenceEngine>>,
pub counters: Arc<PerfCounters>,
pub model_name: Arc<String>,
pub server_config: Arc<ServerConfig>,
pub start_time: std::time::Instant,
}
pub fn build_router(state: AppState) -> Router {
let cors = state.server_config.cors_layer();
Router::new()
.route("/health", get(health_handler))
.route("/v1/models", get(models_handler))
.route("/v1/chat/completions", post(chat_completions_handler))
.layer(cors)
.with_state(state)
}
fn unix_now() -> i64 {
SystemTime::now()
.duration_since(SystemTime::UNIX_EPOCH)
.unwrap_or_default()
.as_secs() as i64
}
fn error_response(status: StatusCode, error: ApiError) -> Response {
(status, Json(ErrorResponse { error })).into_response()
}
async fn health_handler(State(state): State<AppState>) -> impl IntoResponse {
Json(HealthResponse {
status: "ok",
model_loaded: true,
model_id: Some(state.model_name.as_ref().clone()),
uptime_secs: state.start_time.elapsed().as_secs(),
})
}
async fn models_handler(State(state): State<AppState>, headers: HeaderMap) -> Response {
if let Err((status, error)) = authorize_request(&headers, &state.server_config) {
return error_response(status, error);
}
Json(ModelList {
object: "list",
data: state.server_config.model_list(unix_now()),
})
.into_response()
}
async fn chat_completions_handler(
State(state): State<AppState>,
headers: HeaderMap,
Json(req): Json<ChatCompletionRequest>,
) -> Response {
if let Err((status, error)) = authorize_request(&headers, &state.server_config) {
return error_response(status, error);
}
if let Err((status, error)) = validate_request(&req, &state.server_config) {
return error_response(status, error);
}
if let Err(err) = ensure_supported_mode(&req) {
return error_response(
StatusCode::BAD_REQUEST,
ApiError::new(err.to_string(), "invalid_request_error").with_code("unsupported_mode"),
);
}
let sampler = build_sampler(&req);
let prepared = prepare_prompt(&req, &state.server_config);
if req.stream {
stream_response(
state,
req,
sampler,
prepared.prompt,
prepared.streaming_mode,
)
.await
} else {
blocking_response(state, req, sampler, prepared.prompt).await
}
}
async fn stream_response(
state: AppState,
req: ChatCompletionRequest,
sampler: inference_core::SamplerConfig,
prompt: String,
streaming_mode: StreamingMode,
) -> Response {
use axum::response::sse::{Event, KeepAlive};
let id = uuid::Uuid::new_v4().to_string();
let created = unix_now();
let resolved_model = state
.server_config
.resolve_model(&req.model)
.map(|model| model.root.clone())
.unwrap_or_else(|| state.model_name.as_ref().clone());
let engine_arc = Arc::clone(&state.engine);
let counters_arc = Arc::clone(&state.counters);
let stream = async_stream::stream! {
let first_chunk = ChatCompletionChunk {
id: id.clone(),
object: "chat.completion.chunk",
created,
model: resolved_model.clone(),
choices: vec![ChunkChoice {
index: 0,
delta: Delta { role: Some(Role::Assistant), content: None, tool_calls: None },
finish_reason: None,
}],
usage: None,
};
yield Ok::<_, std::convert::Infallible>(Event::default().data(serde_json::to_string(&first_chunk).unwrap()));
match streaming_mode {
StreamingMode::LiveTokens => {
let (tx, mut rx) = tokio::sync::mpsc::channel::<StreamItem>(256);
let prompt = prompt.clone();
tokio::task::spawn_blocking(move || {
let rt = tokio::runtime::Handle::current();
let mut engine = rt.block_on(engine_arc.lock());
let mut record = metrics::recorder::GenerationRecord::new(0);
let mut session = match engine.new_session(&prompt, sampler) {
Ok(session) => session,
Err(err) => {
tracing::error!(err = %err, "Failed to create session");
let _ = rt.block_on(tx.send(StreamItem::Error));
return;
}
};
while let Some(token) = session.next_token() {
record.tick();
counters_arc.inc_tokens(1);
let _ = rt.block_on(tx.send(StreamItem::Token(token)));
}
let usage = Usage {
prompt_tokens: session.prompt_tokens(),
completion_tokens: record.completion_tokens,
total_tokens: session.prompt_tokens() + record.completion_tokens,
tokens_per_second: record.tokens_per_second(),
};
counters_arc.inc_completed();
let _ = rt.block_on(tx.send(StreamItem::Done {
finish_reason: FinishReason::Stop,
usage,
}));
});
while let Some(item) = rx.recv().await {
match item {
StreamItem::Token(token) => {
let chunk = ChatCompletionChunk {
id: id.clone(),
object: "chat.completion.chunk",
created,
model: resolved_model.clone(),
choices: vec![ChunkChoice {
index: 0,
delta: Delta { role: None, content: Some(token), tool_calls: None },
finish_reason: None,
}],
usage: None,
};
yield Ok(Event::default().data(serde_json::to_string(&chunk).unwrap()));
}
StreamItem::Done { finish_reason, usage } => {
let final_chunk = ChatCompletionChunk {
id: id.clone(),
object: "chat.completion.chunk",
created,
model: resolved_model.clone(),
choices: vec![ChunkChoice {
index: 0,
delta: Delta { role: None, content: None, tool_calls: None },
finish_reason: Some(finish_reason),
}],
usage: req.stream_options.as_ref().filter(|options| options.include_usage).map(|_| usage),
};
yield Ok(Event::default().data(serde_json::to_string(&final_chunk).unwrap()));
yield Ok(Event::default().data("[DONE]"));
break;
}
StreamItem::Error => {
let chunk = ChatCompletionChunk {
id: id.clone(),
object: "chat.completion.chunk",
created,
model: resolved_model.clone(),
choices: vec![ChunkChoice {
index: 0,
delta: Delta { role: None, content: Some("internal error".to_owned()), tool_calls: None },
finish_reason: Some(FinishReason::Stop),
}],
usage: None,
};
yield Ok(Event::default().data(serde_json::to_string(&chunk).unwrap()));
yield Ok(Event::default().data("[DONE]"));
break;
}
}
}
}
StreamingMode::Buffered => {
match generate_full_completion(engine_arc, counters_arc, prompt, sampler, req.clone()).await {
Ok(result) => {
if let Some(tool_calls) = result.assistant.tool_calls {
let chunk = ChatCompletionChunk {
id: id.clone(),
object: "chat.completion.chunk",
created,
model: resolved_model.clone(),
choices: vec![ChunkChoice {
index: 0,
delta: Delta {
role: None,
content: None,
tool_calls: Some(tool_calls_to_delta(&tool_calls)),
},
finish_reason: None,
}],
usage: None,
};
yield Ok(Event::default().data(serde_json::to_string(&chunk).unwrap()));
} else if let Some(content) = result.assistant.content {
let chunk = ChatCompletionChunk {
id: id.clone(),
object: "chat.completion.chunk",
created,
model: resolved_model.clone(),
choices: vec![ChunkChoice {
index: 0,
delta: Delta { role: None, content: Some(content), tool_calls: None },
finish_reason: None,
}],
usage: None,
};
yield Ok(Event::default().data(serde_json::to_string(&chunk).unwrap()));
}
let final_chunk = ChatCompletionChunk {
id: id.clone(),
object: "chat.completion.chunk",
created,
model: resolved_model.clone(),
choices: vec![ChunkChoice {
index: 0,
delta: Delta { role: None, content: None, tool_calls: None },
finish_reason: Some(result.assistant.finish_reason),
}],
usage: req.stream_options.as_ref().filter(|options| options.include_usage).map(|_| result.usage),
};
yield Ok(Event::default().data(serde_json::to_string(&final_chunk).unwrap()));
yield Ok(Event::default().data("[DONE]"));
}
Err(response) => {
yield Ok(Event::default().data("[DONE]"));
let _ = response;
}
}
}
}
};
Sse::new(stream)
.keep_alive(KeepAlive::default())
.into_response()
}
async fn blocking_response(
state: AppState,
req: ChatCompletionRequest,
sampler: inference_core::SamplerConfig,
prompt: String,
) -> Response {
let id = uuid::Uuid::new_v4().to_string();
let model = state
.server_config
.resolve_model(&req.model)
.map(|descriptor| descriptor.root.clone())
.unwrap_or_else(|| state.model_name.as_ref().clone());
match generate_full_completion(
Arc::clone(&state.engine),
Arc::clone(&state.counters),
prompt,
sampler,
req,
)
.await
{
Ok(result) => Json(ChatCompletionResponse {
id,
object: "chat.completion",
created: unix_now(),
model,
choices: vec![Choice {
index: 0,
message: ChatMessage {
role: Role::Assistant,
content: result.assistant.content,
name: None,
tool_calls: result.assistant.tool_calls,
tool_call_id: None,
},
finish_reason: result.assistant.finish_reason,
}],
usage: result.usage,
})
.into_response(),
Err(response) => response,
}
}
async fn generate_full_completion(
engine_arc: Arc<Mutex<InferenceEngine>>,
counters_arc: Arc<PerfCounters>,
prompt: String,
sampler: inference_core::SamplerConfig,
req: ChatCompletionRequest,
) -> Result<GenerationResult, Response> {
let result = tokio::task::spawn_blocking(move || {
let rt = tokio::runtime::Handle::current();
let mut engine = rt.block_on(engine_arc.lock());
let mut record = metrics::recorder::GenerationRecord::new(0);
let mut content = String::new();
let mut session = engine.new_session(&prompt, sampler)?;
while let Some(tok) = session.next_token() {
record.tick();
content.push_str(&tok);
}
counters_arc.inc_tokens(record.completion_tokens as u64);
counters_arc.inc_completed();
let usage = Usage {
prompt_tokens: session.prompt_tokens(),
completion_tokens: record.completion_tokens,
total_tokens: session.prompt_tokens() + record.completion_tokens,
tokens_per_second: record.tokens_per_second(),
};
let assistant = finalize_assistant_message(content, &req);
Ok::<_, anyhow::Error>((assistant, usage))
})
.await;
match result {
Ok(Ok((assistant, usage))) => Ok(GenerationResult { assistant, usage }),
Ok(Err(err)) => {
tracing::error!(err = %err, "Generation failed");
Err(error_response(
StatusCode::INTERNAL_SERVER_ERROR,
ApiError::new("Generation failed", "server_error").with_code("generation_failed"),
))
}
Err(err) => {
tracing::error!(err = %err, "Spawn blocking panicked");
Err(error_response(
StatusCode::INTERNAL_SERVER_ERROR,
ApiError::new("Internal server error", "server_error").with_code("internal_error"),
))
}
}
}
struct GenerationResult {
assistant: FinalizedAssistant,
usage: Usage,
}
enum StreamItem {
Token(String),
Done {
finish_reason: FinishReason,
usage: Usage,
},
Error,
}