diff --git a/Cargo.lock b/Cargo.lock index f32120e..879a6fd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -540,7 +540,7 @@ dependencies = [ "encoding_rs", "enumflags2", "llama-cpp-sys-2", - "thiserror", + "thiserror 2.0.19", "tracing", "tracing-core", ] @@ -570,6 +570,7 @@ dependencies = [ "llama-cpp-2", "serde", "serde_json", + "thiserror 1.0.69", "tokio", "tokio-stream", "tower-http", @@ -958,13 +959,33 @@ version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0bf256ce5efdfa370213c1dabab5935a12e49f2c58d15e9eac2870d3b4f27263" +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl 1.0.69", +] + [[package]] name = "thiserror" version = "2.0.19" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "09a43598840e33d5b0331f38c5e30d13bb11c11210a4b58f0d9b18a5a5eefcd9" dependencies = [ - "thiserror-impl", + "thiserror-impl 2.0.19", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index c07e4a2..6a52ea7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -17,6 +17,9 @@ tower-http = { version = "0.6", features = ["cors", "trace"] } serde = { version = "1", features = ["derive"] } serde_json = "1" +# Error handling +thiserror = "1" + # Utilities tracing = "0.1" tracing-subscriber = { version = "0.3", features = ["env-filter"] } diff --git a/src/application/chat/mod.rs b/src/application/chat/mod.rs new file mode 100644 index 0000000..c5e6615 --- /dev/null +++ b/src/application/chat/mod.rs @@ -0,0 +1,3 @@ +pub mod use_cases; + +pub use use_cases::{build_prompt, build_sampler, clean_text, parse_tool_calls, SamplerParams}; diff --git a/src/application/chat/use_cases.rs b/src/application/chat/use_cases.rs new file mode 100644 index 0000000..c9fa711 --- /dev/null +++ b/src/application/chat/use_cases.rs @@ -0,0 +1,294 @@ +//! Chat completion use cases. +//! +//! Orchestrates prompt building, sampler construction, and output parsing. +//! These are pure functions with no framework dependencies. + +use llama_cpp_2::sampling::LlamaSampler; + +use crate::domain::entity::{ChatMessage, ChatRequest, ToolCall, ToolCallFunction, ToolDef}; + +/// Build a prompt string from conversation messages and optional tool definitions. +/// +/// Uses ChatML format with `<|im_start|>` / `<|im_end|>` delimiters. Tool +/// definitions are injected into the first system message. +pub fn build_prompt(messages: &[ChatMessage], tools: &Option>) -> String { + let mut prompt = String::new(); + + for (i, msg) in messages.iter().enumerate() { + match msg.role.as_str() { + "system" => { + let mut content = msg.content.clone().unwrap_or_default(); + // Inject tools into the system message (first occurrence) + if i == 0 { + if let Some(tools_list) = tools { + if !tools_list.is_empty() { + let mut tools_text = String::from( + "\n\n# Tools\n\nYou have access to the following functions:\n\n", + ); + for tool in tools_list { + tools_text.push('\n'); + tools_text.push_str( + &serde_json::to_string(tool).unwrap_or_default(), + ); + } + tools_text.push_str( + "\n\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n\n\n\nvalue_1\n\n\nThis is the value for the second parameter\nthat can span\nmultiple lines\n\n\n", + ); + content.push_str(&tools_text); + } + } + } + prompt.push_str(&format!("<|im_start|>system\n{}<|im_end|>\n", content)); + } + "user" => { + let content = msg.content.as_deref().unwrap_or(""); + if msg.tool_call_id.is_some() || msg.name.is_some() { + prompt.push_str(&format!( + "<|im_start|>user\n\n{}\n<|im_end|>\n", + content + )); + } else { + prompt.push_str(&format!("<|im_start|>user\n{}<|im_end|>\n", content)); + } + } + "assistant" => { + let content = msg.content.as_deref().unwrap_or(""); + if let Some(tcs) = &msg.tool_calls { + let mut asst = format!("<|im_start|>assistant\n{}", content); + for tc in tcs { + let args: serde_json::Value = + serde_json::from_str(&tc.function.arguments).unwrap_or_default(); + asst.push_str(&format!( + "\n\n", + tc.function.name + )); + if let Some(obj) = args.as_object() { + for (k, v) in obj { + let val = match v { + serde_json::Value::String(s) => s.clone(), + other => serde_json::to_string(other).unwrap_or_default(), + }; + asst.push_str(&format!("\n{}\n\n", k, val)); + } + } + asst.push_str("\n"); + } + asst.push_str("<|im_end|>\n"); + prompt.push_str(&asst); + } else { + prompt.push_str(&format!( + "<|im_start|>assistant\n{}<|im_end|>\n", + content + )); + } + } + _ => { + let content = msg.content.as_deref().unwrap_or(""); + prompt.push_str(&format!("<|im_start|>user\n{}<|im_end|>\n", content)); + } + } + } + + // Generation prompt: non-thinking mode + prompt.push_str("<|im_start|>assistant\n\n\n\n\n"); + prompt +} + +/// Parameters for building a [`LlamaSampler`] chain. +pub struct SamplerParams { + pub temperature: Option, + pub top_p: Option, + pub top_k: Option, + pub min_p: Option, + pub repeat_penalty: Option, + pub frequency_penalty: Option, + pub presence_penalty: Option, + pub seed: Option, +} + +impl SamplerParams { + pub fn from_request(req: &ChatRequest) -> Self { + Self { + temperature: req.temperature, + top_p: req.top_p, + top_k: req.top_k, + min_p: req.min_p, + repeat_penalty: req.repeat_penalty, + frequency_penalty: req.frequency_penalty, + presence_penalty: req.presence_penalty, + seed: req.seed, + } + } +} + +/// Build a [`LlamaSampler`] chain from [`SamplerParams`]. +pub fn build_sampler(params: &SamplerParams) -> LlamaSampler { + use llama_cpp_2::sampling::LlamaSampler as LS; + + let temperature = params.temperature; + let top_p = params.top_p; + let top_k = params.top_k; + let min_p = params.min_p; + let repeat_penalty = params.repeat_penalty; + let frequency_penalty = params.frequency_penalty; + let presence_penalty = params.presence_penalty; + let seed = params.seed; + let mut samplers: Vec = Vec::new(); + + // Repetition/frequency/presence penalties + let repeat = repeat_penalty.unwrap_or(1.0); + let freq = frequency_penalty.unwrap_or(0.0); + let present = presence_penalty.unwrap_or(0.0); + if (repeat - 1.0).abs() > 1e-6 || freq > 0.0 || present > 0.0 { + samplers.push(LS::penalties(64, repeat, freq, present)); + } + + // top_k + if let Some(k) = top_k { + samplers.push(LS::top_k(k as i32)); + } + + // top_p + if let Some(p) = top_p { + samplers.push(LS::top_p(p, 1)); + } + + // min_p + if let Some(p) = min_p { + samplers.push(LS::min_p(p, 1)); + } + + // Temperature + final selector + let temp = temperature.unwrap_or(0.0); + if temp <= 0.0 { + samplers.push(LS::greedy()); + } else { + if (temp - 1.0).abs() > 1e-6 { + samplers.push(LS::temp(temp)); + } + let s = seed.unwrap_or(0); + samplers.push(LS::dist(s)); + } + + LlamaSampler::chain_simple(samplers) +} + + +// ═══════════════════════════════════════════════════════════════ +// TEXT PROCESSING +// ═══════════════════════════════════════════════════════════════ + +/// Remove special tokens from generated text. +pub fn clean_text(text: &str) -> String { + text.replace("<|im_end|>", "") + .replace("<|im_start|>", "") + .replace("", "") + .replace("", "") + .trim() + .to_string() +} + +/// Parse tool calls from generated text in the format: +/// +/// ```xml +/// +/// +/// value +/// +/// +/// ``` +pub fn parse_tool_calls(text: &str) -> (String, Vec) { + let mut clean = text.to_string(); + let mut tool_calls: Vec = Vec::new(); + + let mut idx = 0; + loop { + let start_tag = ""; + let end_tag = ""; + + let start = match clean[idx..].find(start_tag) { + Some(s) => idx + s, + None => break, + }; + + let end = match clean[start..].find(end_tag) { + Some(e) => start + e + end_tag.len(), + None => break, + }; + + let block = &clean[start + start_tag.len()..end - end_tag.len()]; + let trimmed = block.trim(); + + // Parse function name + let func_name = trimmed + .lines() + .next() + .and_then(|l| { + let l = l.trim(); + l.strip_prefix("')) + .map(|s| s.trim().to_string()) + }) + .unwrap_or_default(); + + // Parse parameters + let mut args_map = serde_json::Map::new(); + let lines = trimmed.lines(); + let mut current_param: Option = None; + let mut current_value = String::new(); + let mut in_param = false; + + for line in lines { + let line = line.trim(); + if let Some(param) = + line.strip_prefix("')) + { + if let Some(p) = current_param.take() { + args_map.insert( + p, + serde_json::Value::String(current_value.trim().to_string()), + ); + current_value = String::new(); + } + current_param = Some(param.to_string()); + in_param = true; + } else if line == "" { + in_param = false; + } else if in_param { + if !current_value.is_empty() { + current_value.push('\n'); + } + current_value.push_str(line); + } else if line.starts_with("") { + continue; + } + } + // Save last param + if let Some(p) = current_param.take() { + args_map.insert(p, serde_json::Value::String(current_value.trim().to_string())); + } + + let args_json = serde_json::Value::Object(args_map).to_string(); + + tool_calls.push(ToolCall { + id: format!("call_{}", uuid::Uuid::new_v4().to_string().replace('-', "")), + call_type: "function".into(), + function: ToolCallFunction { + name: func_name, + arguments: args_json, + }, + }); + + idx = end; + } + + // Remove tool_call blocks from the text + clean = clean.replace("", "").replace("", ""); + // XML-like tags are already fully parsed; remaining text is the content + clean = clean.trim().to_string(); + // Strip remaining XML tags that aren't part of clean + let cleaned = clean_text(&clean); + + (cleaned, tool_calls) +} diff --git a/src/application/mod.rs b/src/application/mod.rs new file mode 100644 index 0000000..9837daa --- /dev/null +++ b/src/application/mod.rs @@ -0,0 +1,6 @@ +//! Application Layer — use cases / business orchestration. +//! +//! Contains the chat completion logic: prompt building, sampler construction, +//! and output parsing. + +pub mod chat; diff --git a/src/bootstrap/mod.rs b/src/bootstrap/mod.rs new file mode 100644 index 0000000..2710a91 --- /dev/null +++ b/src/bootstrap/mod.rs @@ -0,0 +1,70 @@ +//! Application initialization and lifecycle management. + +use std::sync::Arc; + +use axum::Router; +use tokio::net::TcpListener; +use tracing_subscriber::EnvFilter; + +use crate::config::CONFIG; +use crate::infrastructure::llama::LlamaEngine; +use crate::presentation::router::build_router; +use crate::presentation::state::AppState; + +/// The running application. +/// +/// Encapsulates the router, listener, and port so that everything can be +/// created in `build()` and then served in `run()`, enabling testability. +pub struct Application { + pub port: u16, + router: Router, + listener: TcpListener, +} + +impl Application { + /// Initialize all dependencies and build the application. + /// + /// 1. Init tracing subscriber + /// 2. Load config (triggers LazyLock — fails fast on missing vars) + /// 3. Load model and create LlamaEngine + /// 4. Build router with all routes and middleware + /// 5. Bind TCP listener + pub async fn build() -> anyhow::Result { + // Initialize tracing + let env_filter = EnvFilter::new(&CONFIG.log_level); + tracing_subscriber::fmt() + .with_env_filter(env_filter) + .init(); + tracing::info!("🚀 LLM API starting up..."); + + // Load model (fail-fast) + let engine = LlamaEngine::load().map_err(|e| { + anyhow::anyhow!("Failed to initialize LLM engine: {e}") + })?; + let engine = Arc::new(engine); + + let state = Arc::new(AppState { engine }); + + // Build router + let router = build_router(state); + + // Bind listener + let addr = format!("0.0.0.0:{}", CONFIG.server_port); + let listener = TcpListener::bind(&addr).await?; + tracing::info!( + "Server listening on {}", + listener.local_addr()? + ); + + Ok(Self { + port: CONFIG.server_port, + router, + listener, + }) + } + + /// Start serving requests. + pub async fn run(self) -> std::io::Result<()> { + axum::serve(self.listener, self.router.into_make_service()).await + } +} diff --git a/src/config/mod.rs b/src/config/mod.rs new file mode 100644 index 0000000..576eab1 --- /dev/null +++ b/src/config/mod.rs @@ -0,0 +1,59 @@ +//! Type-safe application configuration. +//! +//! Loads configuration from environment variables at startup with fail-fast behavior. + +use std::sync::LazyLock; + +const DEFAULT_MODEL_PATH: &str = "/models/MiniCPM-V-4.6-Q4_K_M.gguf"; +pub const MODEL_ID: &str = "minicpm-v-4.6"; + +/// Application configuration loaded at startup from environment variables. +#[derive(Debug, Clone)] +pub struct AppConfig { + /// Path to the GGUF model file + pub model_path: String, + + /// API key for authentication (empty = disabled) + pub api_key: String, + + /// Server port to bind to + pub server_port: u16, + + /// Log level (trace, debug, info, warn, error) + pub log_level: String, + + /// LLM context size (n_ctx) + pub n_ctx: u32, + + /// LLM batch size (n_batch) + pub n_batch: u32, + + /// Number of CPU threads for inference + pub n_threads: i32, +} + +impl AppConfig { + /// Load configuration from environment variables. + pub fn load() -> Self { + Self { + model_path: std::env::var("MODEL_PATH") + .unwrap_or_else(|_| DEFAULT_MODEL_PATH.to_string()), + api_key: std::env::var("API_KEY").unwrap_or_default(), + server_port: std::env::var("SERVER_PORT") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(8080), + log_level: std::env::var("RUST_LOG").unwrap_or_else(|_| "info".to_string()), + n_ctx: 2048, + n_batch: 512, + n_threads: 4, + } + } +} + +/// Global configuration instance, loaded once at startup. +pub static CONFIG: LazyLock = LazyLock::new(|| { + let config = AppConfig::load(); + tracing::info!("Configuration loaded: model={:?}", config.model_path); + config +}); diff --git a/src/domain/entity/mod.rs b/src/domain/entity/mod.rs new file mode 100644 index 0000000..61a70c4 --- /dev/null +++ b/src/domain/entity/mod.rs @@ -0,0 +1,193 @@ +//! Domain entities for the LLM inference API. +//! +//! Pure data structs with no framework dependencies beyond serde. +//! These represent the OpenAI-compatible API shapes used across all layers. + +use serde::{Deserialize, Serialize}; + +// ═══════════════════════════════════════════════════════════════ +// REQUEST TYPES +// ═══════════════════════════════════════════════════════════════ + +/// OpenAI-compatible chat completion request body. +#[derive(Deserialize)] +pub struct ChatRequest { + pub model: String, + pub messages: Vec, + pub max_tokens: Option, + #[serde(default)] + pub temperature: Option, + #[serde(default)] + pub top_p: Option, + #[serde(default)] + pub top_k: Option, + #[serde(default)] + pub min_p: Option, + #[serde(default)] + pub frequency_penalty: Option, + #[serde(default)] + pub presence_penalty: Option, + #[serde(default)] + pub repeat_penalty: Option, + #[serde(default)] + pub seed: Option, + #[serde(default)] + pub stream: Option, + #[serde(default)] + pub stop: Option>, + #[serde(default)] + pub tools: Option>, + #[serde(default)] + pub tool_choice: Option, +} + +/// A single message in the chat conversation. +#[derive(Deserialize)] +pub struct ChatMessage { + pub role: String, + pub content: Option, + #[serde(default)] + pub tool_calls: Option>, + #[serde(default)] + pub tool_call_id: Option, + #[serde(default)] + pub name: Option, +} + +/// Tool/function definition for function calling. +#[derive(Deserialize, Serialize, Clone)] +pub struct ToolDef { + #[serde(rename = "type")] + pub tool_type: String, + pub function: ToolFunction, +} + +#[derive(Deserialize, Serialize, Clone)] +pub struct ToolFunction { + pub name: String, + #[serde(default)] + pub description: String, + #[serde(default)] + pub parameters: serde_json::Value, +} + +// ═══════════════════════════════════════════════════════════════ +// RESPONSE TYPES +// ═══════════════════════════════════════════════════════════════ + +/// OpenAI-compatible chat completion response (non-streaming). +#[derive(Serialize)] +pub struct ChatResponse { + pub id: String, + pub object: String, + pub created: i64, + pub model: String, + pub choices: Vec, + pub usage: Usage, +} + +#[derive(Serialize)] +pub struct Choice { + pub index: u32, + pub message: ResponseMessage, + pub finish_reason: String, +} + +#[derive(Serialize)] +pub struct ResponseMessage { + pub role: String, + pub content: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub tool_calls: Option>, +} + +#[derive(Serialize)] +pub struct Usage { + pub prompt_tokens: u32, + pub completion_tokens: u32, + pub total_tokens: u32, +} + +// ═══════════════════════════════════════════════════════════════ +// TOOL CALL TYPES +// ═══════════════════════════════════════════════════════════════ + +/// A tool call in responses or streaming deltas. +#[derive(Serialize, Deserialize, Clone)] +pub struct ToolCall { + pub id: String, + #[serde(rename = "type")] + pub call_type: String, + pub function: ToolCallFunction, +} + +#[derive(Serialize, Deserialize, Clone)] +pub struct ToolCallFunction { + pub name: String, + pub arguments: String, +} + +/// For parsing tool calls from message history (has extra fields). +#[derive(Deserialize, Clone)] +pub struct ToolCallResponse { + pub id: String, + #[serde(rename = "type")] + pub call_type: String, + pub function: ToolCallFunction, +} + +// ═══════════════════════════════════════════════════════════════ +// SSE (STREAMING) TYPES +// ═══════════════════════════════════════════════════════════════ + +/// Server-Sent Event chunk for streaming responses. +#[derive(Serialize)] +pub struct SseChunk { + pub id: String, + pub object: String, + pub created: i64, + pub model: String, + pub choices: Vec, +} + +#[derive(Serialize)] +pub struct SseChoice { + pub index: u32, + pub delta: SseDelta, + #[serde(skip_serializing_if = "Option::is_none")] + pub finish_reason: Option, +} + +#[derive(Serialize)] +pub struct SseDelta { + #[serde(skip_serializing_if = "Option::is_none")] + pub role: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub content: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub tool_calls: Option>, +} + +// ═══════════════════════════════════════════════════════════════ +// MODELS & HEALTH TYPES +// ═══════════════════════════════════════════════════════════════ + +#[derive(Serialize)] +pub struct ModelsResponse { + pub object: String, + pub data: Vec, +} + +#[derive(Serialize)] +pub struct ModelInfo { + pub id: String, + pub object: String, + pub created: i64, + pub owned_by: String, +} + +#[derive(Serialize)] +pub struct HealthResponse { + pub status: String, + pub model: String, +} diff --git a/src/domain/error.rs b/src/domain/error.rs new file mode 100644 index 0000000..ce29448 --- /dev/null +++ b/src/domain/error.rs @@ -0,0 +1,38 @@ +//! Domain-level error types. +//! +//! Framework-agnostic errors that can be mapped to HTTP errors +//! at the presentation layer. + +use thiserror::Error; + +/// Errors originating from LLM inference operations. +#[derive(Error, Debug)] +pub enum LlmError { + /// Invalid request parameters + #[error("Invalid request: {0}")] + InvalidRequest(String), + + /// Model or inference error + #[error("Model error: {0}")] + Model(String), + + /// Authentication failure + #[error("Authentication failed")] + Unauthorized, + + /// Internal/unexpected error + #[error("Internal error: {0}")] + Internal(String), +} + +impl From for LlmError { + fn from(s: String) -> Self { + LlmError::Internal(s) + } +} + +impl From<&str> for LlmError { + fn from(s: &str) -> Self { + LlmError::Internal(s.to_string()) + } +} diff --git a/src/domain/mod.rs b/src/domain/mod.rs new file mode 100644 index 0000000..bb792ca --- /dev/null +++ b/src/domain/mod.rs @@ -0,0 +1,6 @@ +//! Domain Layer — pure business types, no framework dependencies. + +pub mod entity; +pub mod error; + +pub use error::LlmError; diff --git a/src/infrastructure/llama/engine.rs b/src/infrastructure/llama/engine.rs new file mode 100644 index 0000000..563453f --- /dev/null +++ b/src/infrastructure/llama/engine.rs @@ -0,0 +1,263 @@ +//! LlamaEngine — safe wrapper around llama-cpp-2 inference. +//! +//! Encapsulates model loading, context management (with the unavoidable +//! lifetime transmute), tokenization, and generation. + +use std::num::NonZeroU32; + +use llama_cpp_2::context::params::LlamaContextParams; +use llama_cpp_2::context::LlamaContext; +use llama_cpp_2::llama_backend::LlamaBackend; +use llama_cpp_2::llama_batch::LlamaBatch; +use llama_cpp_2::model::params::LlamaModelParams; +use llama_cpp_2::model::{AddBos, LlamaModel}; +use llama_cpp_2::sampling::LlamaSampler; +use llama_cpp_2::token::LlamaToken; +use llama_cpp_2::TokenToStringError; +use tokio::sync::Mutex; +use tracing::info; + +use crate::config::CONFIG; +use crate::domain::LlmError; + +// ── Thread-safe wrapper for raw llama.cpp context ── + +/// Wrapper around [`LlamaContext`] that makes it Send + Sync. +/// +/// # Safety +/// +/// The contained context has its lifetime transmuted to `'static` because it is +/// owned by [`LlamaEngine`] which lives for the entire program lifetime (held in +/// an `Arc`). The engine is only dropped at process shutdown, so no dangling +/// reference can be created. +/// Wrapper around LlamaContext with `'static` lifetime for sharing. +pub struct CtxInner { + /// Invariant: this context is dropped only when the engine is destroyed. + context: LlamaContext<'static>, +} + +unsafe impl Send for CtxInner {} +unsafe impl Sync for CtxInner {} + +/// Wrapper for [`LlamaSampler`] to make it Send + Sync. +/// +/// # Safety +/// +/// `llama-cpp-2`'s `LlamaSampler` is a C opaque pointer. The underlying +/// `llama.cpp` sampling API is reentrant for distinct contexts and thread-safe +/// when used with a single context from one thread at a time (which we enforce +/// via `Mutex`). +pub struct SendSampler(pub LlamaSampler); + +unsafe impl Send for SendSampler {} +unsafe impl Sync for SendSampler {} + +impl std::ops::Deref for SendSampler { + type Target = LlamaSampler; + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +impl std::ops::DerefMut for SendSampler { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.0 + } +} + +impl CtxInner { + pub(crate) fn clear(&mut self) { + self.context.clear_kv_cache(); + } + + pub(crate) fn prefill(&mut self, tokens: &[LlamaToken]) -> Result<(), String> { + let mut batch = LlamaBatch::new(tokens.len(), 1); + for (i, &token) in tokens.iter().enumerate() { + batch + .add(token, i as i32, &[0], i == tokens.len() - 1) + .map_err(|e| e.to_string())?; + } + self.context.decode(&mut batch).map_err(|e| e.to_string()) + } + + pub(crate) fn sample(&mut self, sampler: &mut LlamaSampler) -> LlamaToken { + sampler.sample(&self.context, -1) + } + + pub(crate) fn decode(&mut self, token: LlamaToken, pos: i32) -> Result<(), String> { + let mut batch = LlamaBatch::new(1, 1); + batch + .add(token, pos, &[0], true) + .map_err(|e| e.to_string())?; + self.context.decode(&mut batch).map_err(|e| e.to_string()) + } +} + +// ── LlamaEngine ── + +/// Safe interface to a llama.cpp model and inference context. +/// +/// All access to the underlying context is serialized through a `Mutex`, +/// so only one generation can happen at a time. This is intentional — +/// the model is designed for sequential inference. +pub struct LlamaEngine { + /// The loaded model (read-only after load, safe to share). + pub model: LlamaModel, + + /// The inference context (single-threaded access via Mutex). + pub ctx: Mutex, +} + +impl LlamaEngine { + /// Load a model from disk and create an inference context. + /// + /// # Errors + /// + /// Returns `LlmError::Model` if the model cannot be loaded or the context + /// cannot be created. + pub fn load() -> Result { + info!("Initializing llama backend..."); + let backend = LlamaBackend::init().map_err(|e| { + LlmError::Model(format!("Backend init failed: {e}")) + })?; + + // Backend must outlive model and context. We leak it to achieve 'static + // lifetime since the engine lives for the program lifetime. + let backend: &'static LlamaBackend = Box::leak(Box::new(backend)); + + info!("Loading model: {}", CONFIG.model_path); + let model = LlamaModel::load_from_file(backend, &CONFIG.model_path, &LlamaModelParams::default()) + .map_err(|e| LlmError::Model(format!("Failed to load model: {e}")))?; + info!(" Vocab: {}", model.n_vocab()); + info!(" Params: {}", model.n_params()); + info!(" Layers: {}", model.n_layer()); + + info!("Creating context..."); + let ctx_params = LlamaContextParams::default() + .with_n_ctx(NonZeroU32::new(CONFIG.n_ctx)) + .with_n_batch(CONFIG.n_batch) + .with_n_threads(CONFIG.n_threads) + .with_n_threads_batch(CONFIG.n_threads); + + let context = model + .new_context(backend, ctx_params) + .map_err(|e| LlmError::Model(format!("Failed to create context: {e}")))?; + + // SAFETY: `context` is tied to `backend`'s lifetime, which we leaked + // above to achieve `'static`. The engine owns both and lives for the + // program duration (held in a global Arc). When the engine is dropped + // at process shutdown, the leaked backend is cleaned up by the OS. + let context: LlamaContext<'static> = unsafe { std::mem::transmute(context) }; + + Ok(Self { + model, + ctx: Mutex::new(CtxInner { context }), + }) + } + + /// Tokenize a prompt string into tokens. + pub fn tokenize(&self, prompt: &str) -> Result, LlmError> { + self.model + .str_to_token(prompt, AddBos::Always) + .map_err(|e| LlmError::Model(format!("Tokenization failed: {e}"))) + } + + /// Decode a single token to its string representation. + pub fn decode_token(&self, token: LlamaToken) -> String { + let bytes = match self.model.token_to_piece_bytes(token, 32, true, None) { + Ok(b) => b, + Err(TokenToStringError::InsufficientBufferSpace(neg)) => { + let size = (-neg).max(0).try_into().unwrap_or(256); + self.model + .token_to_piece_bytes(token, size, true, None) + .unwrap_or_default() + } + _ => return String::new(), + }; + String::from_utf8(bytes).unwrap_or_default() + } + + /// Decode multiple tokens to a single string. + pub fn decode_tokens(&self, tokens: &[LlamaToken]) -> String { + let mut out = String::with_capacity(tokens.len() * 4); + for &token in tokens { + out.push_str(&self.decode_token(token)); + } + out + } + + /// Check if a token is an end-of-generation token. + pub fn is_eog(&self, token: LlamaToken) -> bool { + self.model.is_eog_token(token) + } + + /// Return a reference to the context mutex for advanced operations. + pub fn ctx(&self) -> &Mutex { + &self.ctx + } + + /// Generate tokens (non-streaming) and return output tokens, cleaned text, and tool calls. + /// + /// Locks the context mutex, prefill the prompt, then iterates sampling + decoding + /// until EOG, max_tokens, stop sequence, or tool call completion. + pub async fn generate( + &self, + input_tokens: &[LlamaToken], + sampler: &mut SendSampler, + max_tokens: u32, + stop: &[String], + ) -> Result<(Vec, String), LlmError> { + let mut inner = self.ctx.lock().await; + inner.clear(); + inner + .prefill(input_tokens) + .map_err(|e| LlmError::Model(format!("Prefill: {e}")))?; + + let mut output: Vec = Vec::new(); + let mut text_buf = String::new(); + let mut stop_now = false; + + let mut current = inner.sample(sampler); + + for _ in 0..max_tokens { + if self.model.is_eog_token(current) { + break; + } + + let pos = input_tokens.len() as i32 + output.len() as i32; + output.push(current); + + let piece = self.decode_token(current); + text_buf.push_str(&piece); + + // Check stop sequences + for s in stop { + if text_buf.contains(s) { + stop_now = true; + break; + } + } + if stop_now { + break; + } + + // Check for tool_call block completion + if text_buf.contains("") { + let close_count = text_buf.matches("").count(); + let open_count = text_buf.matches("").count(); + if open_count > 0 && close_count >= open_count { + break; + } + } + + if let Err(e) = inner.decode(current, pos) { + tracing::info!(" Decode error: {e}"); + break; + } + + current = inner.sample(sampler); + } + + Ok((output, text_buf)) + } +} diff --git a/src/infrastructure/llama/mod.rs b/src/infrastructure/llama/mod.rs new file mode 100644 index 0000000..d42d8bc --- /dev/null +++ b/src/infrastructure/llama/mod.rs @@ -0,0 +1,3 @@ +pub mod engine; + +pub use engine::{LlamaEngine, SendSampler}; diff --git a/src/infrastructure/mod.rs b/src/infrastructure/mod.rs new file mode 100644 index 0000000..35f495b --- /dev/null +++ b/src/infrastructure/mod.rs @@ -0,0 +1,5 @@ +//! Infrastructure Layer — implements external integrations. +//! +//! Contains the LLM engine wrapper around llama-cpp-2. + +pub mod llama; diff --git a/src/lib.rs b/src/lib.rs new file mode 100644 index 0000000..1ca3535 --- /dev/null +++ b/src/lib.rs @@ -0,0 +1,27 @@ +// Library root — clean architecture module structure + +// ============================================================================ +// Domain Layer — pure business logic, no framework dependencies +// ============================================================================ +pub mod domain; + +// ============================================================================ +// Application Layer — use cases / business orchestration +// ============================================================================ +pub mod application; + +// ============================================================================ +// Infrastructure Layer — implements domain ports +// ============================================================================ +pub mod infrastructure; + +// ============================================================================ +// Presentation Layer — Axum handlers, middleware, state, error +// ============================================================================ +pub mod presentation; + +// ============================================================================ +// Core Framework & Bootstrap +// ============================================================================ +pub mod bootstrap; +pub mod config; diff --git a/src/main.rs b/src/main.rs index ae48f76..90697ae 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,1011 +1,13 @@ -use axum::{ - extract::State, - http::{HeaderMap, StatusCode}, - response::{ - sse::{Event, KeepAlive, Sse}, - IntoResponse, Json, - }, - routing::{get, post}, - Router, -}; -use llama_cpp_2::{ - context::params::LlamaContextParams, - llama_backend::LlamaBackend, - llama_batch::LlamaBatch, - model::{params::LlamaModelParams, AddBos, LlamaModel}, - sampling::LlamaSampler, - token::LlamaToken, - TokenToStringError, -}; -use serde::{Deserialize, Serialize}; -use std::{num::NonZeroU32, sync::Arc}; -use tokio::sync::Mutex; -use tokio_stream::wrappers::ReceiverStream; -use tower_http::cors::CorsLayer; -use tracing::info; +//! Entry point for the LLM inference API server. +//! +//! Delegates all initialization to [`llm_api::bootstrap::Application`]. -// ── Thread-safe wrapper ── -struct CtxInner { - context: llama_cpp_2::context::LlamaContext<'static>, -} - -unsafe impl Send for CtxInner {} -unsafe impl Sync for CtxInner {} - -// LlamaSampler raw pointer is safe to Send (llama.cpp is thread-safe) -struct SendSampler(LlamaSampler); -unsafe impl Send for SendSampler {} -unsafe impl Sync for SendSampler {} - -impl std::ops::Deref for SendSampler { - type Target = LlamaSampler; - fn deref(&self) -> &Self::Target { &self.0 } -} -impl std::ops::DerefMut for SendSampler { - fn deref_mut(&mut self) -> &mut Self::Target { &mut self.0 } -} - -impl CtxInner { - fn clear(&mut self) { - self.context.clear_kv_cache(); - } - - fn prefill(&mut self, tokens: &[LlamaToken]) -> Result<(), String> { - let mut batch = LlamaBatch::new(tokens.len(), 1); - for (i, &token) in tokens.iter().enumerate() { - batch - .add(token, i as i32, &[0], i == tokens.len() - 1) - .map_err(|e| e.to_string())?; - } - self.context.decode(&mut batch).map_err(|e| e.to_string()) - } - - fn sample(&mut self, sampler: &mut LlamaSampler) -> LlamaToken { - sampler.sample(&self.context, -1) - } - - fn decode(&mut self, token: LlamaToken, pos: i32) -> Result<(), String> { - let mut batch = LlamaBatch::new(1, 1); - batch - .add(token, pos, &[0], true) - .map_err(|e| e.to_string())?; - self.context.decode(&mut batch).map_err(|e| e.to_string()) - } -} - -struct AppState { - model: LlamaModel, - ctx: Mutex, -} - -// ── OpenAI Chat Request ── -#[derive(Deserialize)] -struct ChatRequest { - model: String, - messages: Vec, - max_tokens: Option, - #[serde(default)] - temperature: Option, - #[serde(default)] - top_p: Option, - #[serde(default)] - top_k: Option, - #[serde(default)] - min_p: Option, - #[serde(default)] - frequency_penalty: Option, - #[serde(default)] - presence_penalty: Option, - #[serde(default)] - repeat_penalty: Option, - #[serde(default)] - seed: Option, - #[serde(default)] - stream: Option, - #[serde(default)] - stop: Option>, - #[serde(default)] - tools: Option>, - #[serde(default)] - tool_choice: Option, -} - -#[derive(Deserialize)] -struct ChatMessage { - role: String, - content: Option, // null for tool calls - #[serde(default)] - tool_calls: Option>, - #[serde(default)] - tool_call_id: Option, - #[serde(default)] - name: Option, -} - -#[derive(Deserialize, Serialize, Clone)] -struct ToolDef { - #[serde(rename = "type")] - tool_type: String, - function: ToolFunction, -} - -#[derive(Deserialize, Serialize, Clone)] -struct ToolFunction { - name: String, - #[serde(default)] - description: String, - #[serde(default)] - parameters: serde_json::Value, -} - -// ── OpenAI Chat Response ── -#[derive(Serialize)] -struct ChatResponse { - id: String, - object: String, - created: i64, - model: String, - choices: Vec, - usage: Usage, -} - -#[derive(Serialize)] -struct Choice { - index: u32, - message: ResponseMessage, - finish_reason: String, -} - -#[derive(Serialize)] -struct ResponseMessage { - role: String, - content: Option, - #[serde(skip_serializing_if = "Option::is_none")] - tool_calls: Option>, -} - -#[derive(Serialize)] -struct Usage { - prompt_tokens: u32, - completion_tokens: u32, - total_tokens: u32, -} - -// ── Tool Call Types ── -#[derive(Serialize, Deserialize, Clone)] -struct ToolCall { - id: String, - #[serde(rename = "type")] - call_type: String, - function: ToolCallFunction, -} - -#[derive(Serialize, Deserialize, Clone)] -struct ToolCallFunction { - name: String, - arguments: String, // JSON string -} - -// For parsing tool calls from history -#[derive(Deserialize, Clone)] -struct ToolCallResponse { - id: String, - #[serde(rename = "type")] - call_type: String, - function: ToolCallFunction, -} - -// ── SSE Chunk Types ── -#[derive(Serialize)] -struct SseChunk { - id: String, - object: String, - created: i64, - model: String, - choices: Vec, -} - -#[derive(Serialize)] -struct SseChoice { - index: u32, - delta: SseDelta, - #[serde(skip_serializing_if = "Option::is_none")] - finish_reason: Option, -} - -#[derive(Serialize)] -struct SseDelta { - #[serde(skip_serializing_if = "Option::is_none")] - role: Option, - #[serde(skip_serializing_if = "Option::is_none")] - content: Option, - #[serde(skip_serializing_if = "Option::is_none")] - tool_calls: Option>, -} - -#[derive(Serialize)] -struct ModelsResponse { - object: String, - data: Vec, -} - -#[derive(Serialize)] -struct ModelInfo { - id: String, - object: String, - created: i64, - owned_by: String, -} - -#[derive(Serialize)] -struct HealthResponse { - status: String, - model: String, -} - -const DEFAULT_MODEL_PATH: &str = "/models/MiniCPM-V-4.6-Q4_K_M.gguf"; -const MODEL_ID: &str = "minicpm-v-4.6"; - -// ═══════════════════════════════════════════ -// SAMPLER BUILDER -// ═══════════════════════════════════════════ - -fn build_sampler(req: &ChatRequest) -> LlamaSampler { - let mut samplers: Vec = Vec::new(); - - // Repetition/frequency/presence penalties - let repeat = req.repeat_penalty.unwrap_or(1.0); - let freq = req.frequency_penalty.unwrap_or(0.0); - let present = req.presence_penalty.unwrap_or(0.0); - if (repeat - 1.0).abs() > 1e-6 || freq > 0.0 || present > 0.0 { - samplers.push(LlamaSampler::penalties(64, repeat, freq, present)); - } - - // top_k - if let Some(k) = req.top_k { - samplers.push(LlamaSampler::top_k(k as i32)); - } - - // top_p - if let Some(p) = req.top_p { - samplers.push(LlamaSampler::top_p(p, 1)); - } - - // min_p - if let Some(p) = req.min_p { - samplers.push(LlamaSampler::min_p(p, 1)); - } - - // Temperature + final selector - let temp = req.temperature.unwrap_or(0.0); - if temp <= 0.0 { - samplers.push(LlamaSampler::greedy()); - } else { - if (temp - 1.0).abs() > 1e-6 { - samplers.push(LlamaSampler::temp(temp)); - } - let seed = req.seed.unwrap_or(0); - samplers.push(LlamaSampler::dist(seed)); - } - - LlamaSampler::chain_simple(samplers) -} - -fn build_sampler_params( - temperature: Option, - top_p: Option, - top_k: Option, - min_p: Option, - repeat_penalty: Option, - frequency_penalty: Option, - presence_penalty: Option, - seed: Option, -) -> LlamaSampler { - let mut samplers: Vec = Vec::new(); - - let repeat = repeat_penalty.unwrap_or(1.0); - let freq = frequency_penalty.unwrap_or(0.0); - let present = presence_penalty.unwrap_or(0.0); - if (repeat - 1.0).abs() > 1e-6 || freq > 0.0 || present > 0.0 { - samplers.push(LlamaSampler::penalties(64, repeat, freq, present)); - } - - if let Some(k) = top_k { - samplers.push(LlamaSampler::top_k(k as i32)); - } - - if let Some(p) = top_p { - samplers.push(LlamaSampler::top_p(p, 1)); - } - - if let Some(p) = min_p { - samplers.push(LlamaSampler::min_p(p, 1)); - } - - let temp = temperature.unwrap_or(0.0); - if temp <= 0.0 { - samplers.push(LlamaSampler::greedy()); - } else { - if (temp - 1.0).abs() > 1e-6 { - samplers.push(LlamaSampler::temp(temp)); - } - let s = seed.unwrap_or(0); - samplers.push(LlamaSampler::dist(s)); - } - - LlamaSampler::chain_simple(samplers) -} - -// ═══════════════════════════════════════════ -// PROMPT / TOOL BUILDERS -// ═══════════════════════════════════════════ - -fn build_prompt(messages: &[ChatMessage], tools: &Option>) -> String { - let mut prompt = String::new(); - - for (i, msg) in messages.iter().enumerate() { - match msg.role.as_str() { - "system" => { - let mut content = msg.content.clone().unwrap_or_default(); - // Inject tools into the system message (first occurrence) - if i == 0 { - if let Some(tools_list) = tools { - if !tools_list.is_empty() { - let mut tools_text = String::from( - "\n\n# Tools\n\nYou have access to the following functions:\n\n", - ); - for tool in tools_list { - tools_text.push('\n'); - tools_text.push_str( - &serde_json::to_string(tool).unwrap_or_default(), - ); - } - tools_text.push_str( - "\n\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n\n\n\nvalue_1\n\n\nThis is the value for the second parameter\nthat can span\nmultiple lines\n\n\n", - ); - content.push_str(&tools_text); - } - } - } - prompt.push_str(&format!("<|im_start|>system\n{}<|im_end|>\n", content)); - } - "user" => { - let content = msg.content.as_deref().unwrap_or(""); - // Handle tool results - if msg.tool_call_id.is_some() || msg.name.is_some() { - prompt.push_str(&format!( - "<|im_start|>user\n\n{}\n<|im_end|>\n", - content - )); - } else { - prompt.push_str(&format!("<|im_start|>user\n{}<|im_end|>\n", content)); - } - } - "assistant" => { - let content = msg.content.as_deref().unwrap_or(""); - if let Some(tcs) = &msg.tool_calls { - // Assistant message with tool calls - let mut asst = format!("<|im_start|>assistant\n{}", content); - for tc in tcs { - let args: serde_json::Value = - serde_json::from_str(&tc.function.arguments).unwrap_or_default(); - asst.push_str(&format!( - "\n\n", - tc.function.name - )); - if let Some(obj) = args.as_object() { - for (k, v) in obj { - let val = match v { - serde_json::Value::String(s) => s.clone(), - other => serde_json::to_string(other).unwrap_or_default(), - }; - asst.push_str(&format!("\n{}\n\n", k, val)); - } - } - asst.push_str("\n"); - } - asst.push_str("<|im_end|>\n"); - prompt.push_str(&asst); - } else { - prompt.push_str(&format!( - "<|im_start|>assistant\n{}<|im_end|>\n", - content - )); - } - } - _ => { - let content = msg.content.as_deref().unwrap_or(""); - prompt.push_str(&format!("<|im_start|>user\n{}<|im_end|>\n", content)); - } - } - } - - // Generation prompt: non-thinking mode - prompt.push_str("<|im_start|>assistant\n\n\n\n\n"); - prompt -} - -// ═══════════════════════════════════════════ -// TOOL CALL PARSER -// ═══════════════════════════════════════════ - -fn parse_tool_calls(text: &str) -> (String, Vec) { - let mut clean = text.to_string(); - let mut tool_calls: Vec = Vec::new(); - - // Find all ... blocks - let mut idx = 0; - loop { - let start_tag = ""; - let end_tag = ""; - - let start = match clean[idx..].find(start_tag) { - Some(s) => idx + s, - None => break, - }; - - let end = match clean[start..].find(end_tag) { - Some(e) => start + e + end_tag.len(), - None => break, - }; - - let block = &clean[start + start_tag.len()..end - end_tag.len()]; - let trimmed = block.trim(); - - // Parse function name - let func_name = trimmed - .lines() - .next() - .and_then(|l| { - let l = l.trim(); - l.strip_prefix("')) - .map(|s| s.trim().to_string()) - }) - .unwrap_or_default(); - - // Parse parameters - let mut args_map = serde_json::Map::new(); - let lines = trimmed.lines(); - let mut current_param: Option = None; - let mut current_value = String::new(); - let mut in_param = false; - - for line in lines { - let line = line.trim(); - if let Some(param) = line.strip_prefix("')) - { - // Save previous param - if let Some(p) = current_param.take() { - args_map.insert(p, serde_json::Value::String(current_value.trim().to_string())); - current_value = String::new(); - } - current_param = Some(param.to_string()); - in_param = true; - } else if line == "" { - in_param = false; - } else if in_param { - if !current_value.is_empty() { - current_value.push('\n'); - } - current_value.push_str(line); - } else if line.starts_with("") { - continue; - } - } - // Save last param - if let Some(p) = current_param.take() { - args_map.insert(p, serde_json::Value::String(current_value.trim().to_string())); - } - - let args_json = serde_json::Value::Object(args_map).to_string(); - - tool_calls.push(ToolCall { - id: format!("call_{}", uuid::Uuid::new_v4().to_string().replace('-', "")), - call_type: "function".into(), - function: ToolCallFunction { - name: func_name, - arguments: args_json, - }, - }); - - idx = end; - } - - // Remove tool_call blocks from the text - clean = clean.replace("", "").replace("", ""); - // Also remove leftover function/parameter XML - clean = clean.replace(r#""#, ""); - clean = clean.replace("", ""); - clean = clean.replace(r#""#, ""); - clean = clean.replace("", ""); - - (clean_text(&clean), tool_calls) -} - -// ═══════════════════════════════════════════ -// TOKEN DECODING -// ═══════════════════════════════════════════ - -fn decode_token_piece(model: &LlamaModel, token: LlamaToken) -> String { - let bytes = match model.token_to_piece_bytes(token, 32, true, None) { - Ok(b) => b, - Err(TokenToStringError::InsufficientBufferSpace(neg)) => { - let size = (-neg).max(0).try_into().unwrap_or(256); - model - .token_to_piece_bytes(token, size, true, None) - .unwrap_or_default() - } - _ => return String::new(), - }; - String::from_utf8(bytes).unwrap_or_default() -} - -fn clean_text(text: &str) -> String { - text.replace("<|im_end|>", "") - .replace("<|im_start|>", "") - .replace("", "") - .replace("", "") - .trim() - .to_string() -} - -fn decode_tokens(model: &LlamaModel, tokens: &[LlamaToken]) -> String { - let mut out = String::with_capacity(tokens.len() * 4); - for &token in tokens { - out.push_str(&decode_token_piece(model, token)); - } - clean_text(&out) -} - -// ═══════════════════════════════════════════ -// AUTH -// ═══════════════════════════════════════════ - -fn check_auth(headers: &HeaderMap) -> Result<(), (StatusCode, String)> { - let api_key = std::env::var("API_KEY").unwrap_or_default(); - if api_key.is_empty() { - return Ok(()); - } - let header = headers - .get("authorization") - .and_then(|v| v.to_str().ok()) - .unwrap_or(""); - let expected = format!("Bearer {api_key}"); - if header == expected || header == api_key { - return Ok(()); - } - Err(( - StatusCode::UNAUTHORIZED, - "{\"error\":\"unauthorized\",\"message\":\"Invalid API key\"}".into(), - )) -} - -// ═══════════════════════════════════════════ -// MAIN -// ═══════════════════════════════════════════ +use llm_api::bootstrap::Application; #[tokio::main] -async fn main() { - tracing_subscriber::fmt() - .with_env_filter("info") - .init(); +async fn main() -> anyhow::Result<()> { + let app = Application::build().await?; + app.run().await?; - info!("Initializing backend..."); - let backend = LlamaBackend::init().expect("Backend init failed"); - - info!("Loading model..."); - let model_path = std::env::var("MODEL_PATH").unwrap_or_else(|_| DEFAULT_MODEL_PATH.to_string()); - info!(" Model: {model_path}"); - let model = LlamaModel::load_from_file(&backend, &model_path, &LlamaModelParams::default()) - .expect("Failed to load model"); - info!(" Vocab: {}", model.n_vocab()); - info!(" Params: {}", model.n_params()); - info!(" Layers: {}", model.n_layer()); - - info!("Creating context..."); - let ctx_params = LlamaContextParams::default() - .with_n_ctx(NonZeroU32::new(2048)) - .with_n_batch(512) - .with_n_threads(4) - .with_n_threads_batch(4); - - let context = model - .new_context(&backend, ctx_params) - .expect("Failed to create context"); - let context: llama_cpp_2::context::LlamaContext<'static> = - unsafe { std::mem::transmute(context) }; - - let state = Arc::new(AppState { - model, - ctx: Mutex::new(CtxInner { context }), - }); - - info!("Server ready on :8080"); - - let app = Router::new() - .route("/health", get(health)) - .route("/v1/models", get(list_models)) - .route("/v1/chat/completions", post(chat_completions)) - .layer(CorsLayer::permissive()) - .with_state(state); - - let listener = tokio::net::TcpListener::bind("0.0.0.0:8080") - .await - .expect("Failed to bind"); - - axum::serve(listener, app).await.expect("Server failed"); -} - -async fn health() -> Json { - Json(HealthResponse { - status: "ok".into(), - model: format!("{MODEL_ID}-q4_k_m"), - }) -} - -async fn list_models() -> Json { - Json(ModelsResponse { - object: "list".into(), - data: vec![ModelInfo { - id: MODEL_ID.into(), - object: "model".into(), - created: chrono::Utc::now().timestamp(), - owned_by: "asepharyana".into(), - }], - }) -} - -// ═══════════════════════════════════════════ -// GENERATE TOKENS (shared by stream & non-stream) -// ═══════════════════════════════════════════ - -fn generate_tokens( - state: &AppState, - inner: &mut CtxInner, - input_tokens: &[LlamaToken], - sampler: &mut LlamaSampler, - max_tokens: u32, - stop: &[String], -) -> (Vec, String, Vec) { - let mut output: Vec = Vec::new(); - let mut text_buf = String::new(); - let mut stop_now = false; - - let mut current = inner.sample(sampler); - - for _ in 0..max_tokens { - if state.model.is_eog_token(current) { - break; - } - - let pos = input_tokens.len() as i32 + output.len() as i32; - output.push(current); - - if let Err(e) = inner.decode(current, pos) { - info!(" Decode error: {e}"); - break; - } - - // Decode this token for stop checking - let piece = decode_token_piece(&state.model, current); - text_buf.push_str(&piece); - - // Check stop sequences - for s in stop { - if text_buf.contains(s) { - stop_now = true; - break; - } - } - if stop_now { - break; - } - - // Check for tool_call start - if text_buf.contains("") { - // Keep generating until tool_call block is closed - let close_tag = ""; - let close_count = text_buf.matches(close_tag).count(); - let open_count = text_buf.matches("").count(); - if open_count > 0 && close_count >= open_count { - // All tool call blocks are closed - break; - } - } - - current = inner.sample(sampler); - } - - let (clean, tool_calls) = parse_tool_calls(&clean_text(&text_buf)); - (output, clean, tool_calls) -} - -// ═══════════════════════════════════════════ -// CHAT COMPLETIONS HANDLER -// ═══════════════════════════════════════════ - -async fn chat_completions( - State(state): State>, - headers: HeaderMap, - Json(req): Json, -) -> Result { - check_auth(&headers)?; - - let chat_id = format!("chatcmpl-{}", uuid::Uuid::new_v4()); - let created = chrono::Utc::now().timestamp(); - let max_tokens = req.max_tokens.unwrap_or(256).min(1024); - let stop = req.stop.clone().unwrap_or_default(); - let prompt = build_prompt(&req.messages, &req.tools); - let has_tools = req - .tools - .as_ref() - .is_some_and(|t| !t.is_empty()); - - // Tokenize - let input_tokens = state - .model - .str_to_token(&prompt, AddBos::Always) - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; - - let prompt_tokens = input_tokens.len() as u32; - info!( - " Chat: {} prompt tokens, max_tokens={}, tools={}", - prompt_tokens, - max_tokens, - has_tools - ); - - if req.stream.unwrap_or(false) { - // ── STREAMING ── - let state = state.clone(); - let model_name = req.model.clone(); - let stop_clone = stop.clone(); - let (tx, rx) = tokio::sync::mpsc::channel::>(64); - - // Extract sampling params for the spawned task - let temp = req.temperature; - let top_p = req.top_p; - let top_k = req.top_k; - let min_p = req.min_p; - let freq_penalty = req.frequency_penalty; - let pres_penalty = req.presence_penalty; - let rep_penalty = req.repeat_penalty; - let seed = req.seed; - let _max_tokens_s = max_tokens; - - tokio::spawn(async move { - // Role chunk - let role_chunk = serde_json::to_string(&SseChunk { - id: chat_id.clone(), - object: "chat.completion.chunk".into(), - created, - model: model_name.clone(), - choices: vec![SseChoice { - index: 0, - delta: SseDelta { - role: Some("assistant".into()), - content: None, - tool_calls: None, - }, - finish_reason: None, - }], - }) - .unwrap(); - if tx.send(Ok(Event::default().data(role_chunk))).await.is_err() { - return; - } - - // Lock context - let mut inner = state.ctx.lock().await; - inner.clear(); - if let Err(e) = inner.prefill(&input_tokens) { - info!(" Prefill error: {e}"); - return; - } - - let mut sampler = SendSampler(build_sampler_params( - temp, top_p, top_k, min_p, - rep_penalty, freq_penalty, pres_penalty, seed, - )); - let mut count = 0u32; - let mut text_buf = String::new(); - let mut current = inner.sample(&mut sampler); - - loop { - if count >= max_tokens { - let chunk = serde_json::to_string(&SseChunk { - id: chat_id.clone(), - object: "chat.completion.chunk".into(), - created, - model: model_name.clone(), - choices: vec![SseChoice { - index: 0, - delta: SseDelta { - role: None, - content: None, - tool_calls: None, - }, - finish_reason: Some("length".into()), - }], - }) - .unwrap(); - let _ = tx.send(Ok(Event::default().data(chunk))).await; - break; - } - - if state.model.is_eog_token(current) { - let reason = if has_tools && text_buf.contains("") { - "tool_calls" - } else { - "stop" - }; - let chunk = serde_json::to_string(&SseChunk { - id: chat_id.clone(), - object: "chat.completion.chunk".into(), - created, - model: model_name.clone(), - choices: vec![SseChoice { - index: 0, - delta: SseDelta { - role: None, - content: None, - tool_calls: None, - }, - finish_reason: Some(reason.into()), - }], - }) - .unwrap(); - let _ = tx.send(Ok(Event::default().data(chunk))).await; - break; - } - - let piece = decode_token_piece(&state.model, current); - let content = clean_text(&piece); - - if !content.is_empty() { - let chunk = serde_json::to_string(&SseChunk { - id: chat_id.clone(), - object: "chat.completion.chunk".into(), - created, - model: model_name.clone(), - choices: vec![SseChoice { - index: 0, - delta: SseDelta { - role: None, - content: Some(content.clone()), - tool_calls: None, - }, - finish_reason: None, - }], - }) - .unwrap(); - if tx.send(Ok(Event::default().data(chunk))).await.is_err() { - break; - } - } - - text_buf.push_str(&piece); - - // Check stop - let mut stop_now = false; - for s in &stop_clone { - if text_buf.contains(s) { - stop_now = true; - break; - } - } - if stop_now { - let chunk = serde_json::to_string(&SseChunk { - id: chat_id.clone(), - object: "chat.completion.chunk".into(), - created, - model: model_name.clone(), - choices: vec![SseChoice { - index: 0, - delta: SseDelta { - role: None, - content: None, - tool_calls: None, - }, - finish_reason: Some("stop".into()), - }], - }) - .unwrap(); - let _ = tx.send(Ok(Event::default().data(chunk))).await; - break; - } - - // Check tool call completeness - if has_tools && text_buf.contains("") { - let open = text_buf.matches("").count(); - let close = text_buf.matches("").count(); - if close >= open { - let chunk = serde_json::to_string(&SseChunk { - id: chat_id.clone(), - object: "chat.completion.chunk".into(), - created, - model: model_name.clone(), - choices: vec![SseChoice { - index: 0, - delta: SseDelta { - role: None, - content: None, - tool_calls: None, - }, - finish_reason: Some("tool_calls".into()), - }], - }) - .unwrap(); - let _ = tx.send(Ok(Event::default().data(chunk))).await; - break; - } - } - - let pos = input_tokens.len() as i32 + count as i32; - if let Err(e) = inner.decode(current, pos) { - info!(" Decode error: {e}"); - break; - } - count += 1; - current = inner.sample(&mut sampler); - } - }); - - let stream = ReceiverStream::new(rx); - let sse = Sse::new(stream).keep_alive(KeepAlive::default()); - return Ok(sse.into_response()); - } - - // ── NON-STREAMING ── - let mut inner = state.ctx.lock().await; - inner.clear(); - inner.prefill(&input_tokens).map_err(|e| { - (StatusCode::INTERNAL_SERVER_ERROR, format!("Prefill: {e}")) - })?; - - let mut sampler = build_sampler(&req); - let (output_tokens, output_text, tool_calls) = - generate_tokens(&state, &mut inner, &input_tokens, &mut sampler, max_tokens, &stop); - - let completion_tokens = output_tokens.len() as u32; - info!(" {} generated tokens", completion_tokens); - - let finish_reason = if has_tools && !tool_calls.is_empty() { - "tool_calls" - } else if completion_tokens < max_tokens { - "stop" - } else { - "length" - }; - - Ok(Json(ChatResponse { - id: chat_id, - object: "chat.completion".into(), - created, - model: req.model, - choices: vec![Choice { - index: 0, - message: ResponseMessage { - role: "assistant".into(), - content: if tool_calls.is_empty() { - Some(output_text) - } else { - Some(output_text) - }, - tool_calls: if tool_calls.is_empty() { - None - } else { - Some(tool_calls) - }, - }, - finish_reason: finish_reason.into(), - }], - usage: Usage { - prompt_tokens, - completion_tokens, - total_tokens: prompt_tokens + completion_tokens, - }, - }) - .into_response()) + Ok(()) } diff --git a/src/presentation/dto/common.rs b/src/presentation/dto/common.rs new file mode 100644 index 0000000..01186ba --- /dev/null +++ b/src/presentation/dto/common.rs @@ -0,0 +1,30 @@ +//! Common API response types. + +use serde::Serialize; + +#[derive(Serialize)] +pub struct ApiResponse { + pub success: bool, + #[serde(skip_serializing_if = "Option::is_none")] + pub message: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub data: Option, +} + +impl ApiResponse { + pub fn success(data: T) -> Self { + Self { + success: true, + message: None, + data: Some(data), + } + } + + pub fn error(message: String) -> Self { + Self { + success: false, + message: Some(message), + data: None, + } + } +} diff --git a/src/presentation/dto/mod.rs b/src/presentation/dto/mod.rs new file mode 100644 index 0000000..34994bf --- /dev/null +++ b/src/presentation/dto/mod.rs @@ -0,0 +1 @@ +pub mod common; diff --git a/src/presentation/error.rs b/src/presentation/error.rs new file mode 100644 index 0000000..4438845 --- /dev/null +++ b/src/presentation/error.rs @@ -0,0 +1,98 @@ +//! Application-level HTTP error handling. +//! +//! Maps domain errors into HTTP responses. + +use axum::http::StatusCode; +use axum::response::IntoResponse; +use axum::Json; +use serde::Serialize; +use thiserror::Error; + +use crate::domain::LlmError; + +/// Top-level HTTP error returned by all API handlers. +#[derive(Error, Debug)] +pub enum AppError { + #[error("Bad request: {0}")] + BadRequest(String), + + #[error("Not found: {0}")] + NotFound(String), + + #[error("Authentication failed")] + Unauthorized, + + #[error("Model error: {0}")] + LlmError(String), + + #[error("Internal error: {0}")] + Internal(String), +} + +// ── From impls — convert domain errors to AppError ── + +impl From for AppError { + fn from(err: LlmError) -> Self { + match err { + LlmError::InvalidRequest(msg) => AppError::BadRequest(msg), + LlmError::Model(msg) => AppError::LlmError(msg), + LlmError::Unauthorized => AppError::Unauthorized, + LlmError::Internal(msg) => AppError::Internal(msg), + } + } +} + +impl From for AppError { + fn from(s: String) -> Self { + AppError::Internal(s) + } +} + +impl From<&str> for AppError { + fn from(s: &str) -> Self { + AppError::Internal(s.to_string()) + } +} + +// ── Error Response DTO ── + +#[derive(Serialize)] +struct ErrorBody { + error: String, + message: String, +} + +// ── IntoResponse — render AppError as HTTP response ── + +impl IntoResponse for AppError { + fn into_response(self) -> axum::response::Response { + let (status, error_msg, detail_msg) = match &self { + AppError::BadRequest(msg) => (StatusCode::BAD_REQUEST, "bad_request", msg.as_str()), + AppError::NotFound(msg) => (StatusCode::NOT_FOUND, "not_found", msg.as_str()), + AppError::Unauthorized => (StatusCode::UNAUTHORIZED, "unauthorized", "Invalid API key"), + AppError::LlmError(_msg) => { + tracing::error!(%self, "LLM error"); + ( + StatusCode::INTERNAL_SERVER_ERROR, + "internal_error", + "Internal server error", + ) + } + AppError::Internal(_) => { + tracing::error!(%self, "Internal error"); + ( + StatusCode::INTERNAL_SERVER_ERROR, + "internal_error", + "Internal server error", + ) + } + }; + + let body = Json(ErrorBody { + error: error_msg.to_string(), + message: detail_msg.to_string(), + }); + + (status, body).into_response() + } +} diff --git a/src/presentation/handler/chat-ui/index.html b/src/presentation/handler/chat-ui/index.html new file mode 100644 index 0000000..a04d87c --- /dev/null +++ b/src/presentation/handler/chat-ui/index.html @@ -0,0 +1,246 @@ + + + + + +AI Chat + + + +
+

AI Chat

+ llm-api +
+
+
+
+ + +
+ + + + diff --git a/src/presentation/handler/chat.rs b/src/presentation/handler/chat.rs new file mode 100644 index 0000000..e36f5e1 --- /dev/null +++ b/src/presentation/handler/chat.rs @@ -0,0 +1,310 @@ +//! Chat completions endpoint — streaming and non-streaming. + +use std::convert::Infallible; +use std::sync::Arc; + +use axum::extract::State; +use axum::response::sse::{Event, KeepAlive, Sse}; +use axum::{Json, response::{IntoResponse, Response}}; +use chrono::Utc; +use tokio::sync::mpsc; +use tokio_stream::wrappers::ReceiverStream; +use tracing::info; + +use crate::application::chat; +use crate::domain::entity::{ + ChatRequest, ChatResponse, Choice, ResponseMessage, SseChunk, SseChoice, SseDelta, Usage, +}; +use crate::infrastructure::llama::SendSampler; +use crate::presentation::error::AppError; +use crate::presentation::state::AppState; + +/// POST /v1/chat/completions +pub async fn chat_completions( + State(state): State>, + Json(req): Json, +) -> Result { + let max_tokens = req.max_tokens.unwrap_or(256).min(1024); + let stop = req.stop.clone().unwrap_or_default(); + let prompt = chat::build_prompt(&req.messages, &req.tools); + + // Tokenize + let input_tokens = state + .engine + .tokenize(&prompt) + .map_err(|e| AppError::LlmError(e.to_string()))?; + + let prompt_tokens = input_tokens.len() as u32; + info!( + " Chat: {} prompt tokens, max_tokens={}, tools={}", + prompt_tokens, + max_tokens, + req.tools.as_ref().is_some_and(|t| !t.is_empty()) + ); + + let response = if req.stream.unwrap_or(false) { + handle_streaming(state.clone(), req, max_tokens, stop, input_tokens).await? + } else { + handle_non_streaming(state.clone(), req, max_tokens, stop, input_tokens).await? + }; + Ok(response.into_response()) +} + +// ── Non-streaming path ── + +async fn handle_non_streaming( + state: Arc, + req: ChatRequest, + max_tokens: u32, + stop: Vec, + input_tokens: Vec, +) -> Result { + let chat_id = format!("chatcmpl-{}", uuid::Uuid::new_v4()); + let created = Utc::now().timestamp(); + let prompt_tokens = input_tokens.len() as u32; + let has_tools = req.tools.as_ref().is_some_and(|t| !t.is_empty()); + let params = chat::SamplerParams::from_request(&req); + let mut sampler = SendSampler(chat::build_sampler(¶ms)); + + let (output_tokens, raw_text) = state + .engine + .generate(&input_tokens, &mut sampler, max_tokens, &stop) + .await?; + + let (output_text, tool_calls) = chat::parse_tool_calls(&chat::clean_text(&raw_text)); + + let completion_tokens = output_tokens.len() as u32; + info!(" {} generated tokens", completion_tokens); + + let finish_reason = if has_tools && !tool_calls.is_empty() { + "tool_calls" + } else if completion_tokens < max_tokens { + "stop" + } else { + "length" + }; + + Ok(Json(ChatResponse { + id: chat_id, + object: "chat.completion".into(), + created, + model: req.model, + choices: vec![Choice { + index: 0, + message: ResponseMessage { + role: "assistant".into(), + content: Some(output_text), + tool_calls: if tool_calls.is_empty() { + None + } else { + Some(tool_calls) + }, + }, + finish_reason: finish_reason.into(), + }], + usage: Usage { + prompt_tokens, + completion_tokens, + total_tokens: prompt_tokens + completion_tokens, + }, + }) + .into_response()) +} + +// ── Streaming path ── + +async fn handle_streaming( + state: Arc, + req: ChatRequest, + max_tokens: u32, + stop: Vec, + input_tokens: Vec, +) -> Result { + let chat_id = format!("chatcmpl-{}", uuid::Uuid::new_v4()); + let created = Utc::now().timestamp(); + let has_tools = req.tools.as_ref().is_some_and(|t| !t.is_empty()); + let model_name = req.model.clone(); + let params = chat::SamplerParams::from_request(&req); + let (tx, rx) = mpsc::channel::>(64); + + tokio::spawn(async move { + // Role chunk + let role_chunk = serde_json::to_string(&SseChunk { + id: chat_id.clone(), + object: "chat.completion.chunk".into(), + created, + model: model_name.clone(), + choices: vec![SseChoice { + index: 0, + delta: SseDelta { + role: Some("assistant".into()), + content: None, + tool_calls: None, + }, + finish_reason: None, + }], + }) + .unwrap(); + if tx.send(Ok(Event::default().data(role_chunk))).await.is_err() { + return; + } + + // Build sampler + let mut sampler = SendSampler(chat::build_sampler(¶ms)); + + // Lock context + let mut inner = state.engine.ctx().lock().await; + inner.clear(); + if let Err(e) = inner.prefill(&input_tokens) { + info!(" Prefill error: {e}"); + return; + } + + let mut count = 0u32; + let mut text_buf = String::new(); + let mut current = inner.sample(&mut sampler); + + loop { + if count >= max_tokens { + let chunk = serde_json::to_string(&SseChunk { + id: chat_id.clone(), + object: "chat.completion.chunk".into(), + created, + model: model_name.clone(), + choices: vec![SseChoice { + index: 0, + delta: SseDelta { + role: None, + content: None, + tool_calls: None, + }, + finish_reason: Some("length".into()), + }], + }) + .unwrap(); + let _ = tx.send(Ok(Event::default().data(chunk))).await; + break; + } + + if state.engine.is_eog(current) { + let reason = if has_tools && text_buf.contains("") { + "tool_calls" + } else { + "stop" + }; + let chunk = serde_json::to_string(&SseChunk { + id: chat_id.clone(), + object: "chat.completion.chunk".into(), + created, + model: model_name.clone(), + choices: vec![SseChoice { + index: 0, + delta: SseDelta { + role: None, + content: None, + tool_calls: None, + }, + finish_reason: Some(reason.into()), + }], + }) + .unwrap(); + let _ = tx.send(Ok(Event::default().data(chunk))).await; + break; + } + + let piece = state.engine.decode_token(current); + let content = chat::clean_text(&piece); + + if !content.is_empty() { + let chunk = serde_json::to_string(&SseChunk { + id: chat_id.clone(), + object: "chat.completion.chunk".into(), + created, + model: model_name.clone(), + choices: vec![SseChoice { + index: 0, + delta: SseDelta { + role: None, + content: Some(content.clone()), + tool_calls: None, + }, + finish_reason: None, + }], + }) + .unwrap(); + if tx.send(Ok(Event::default().data(chunk))).await.is_err() { + break; + } + } + + text_buf.push_str(&piece); + + // Check stop sequences + let mut stop_now = false; + for s in &stop { + if text_buf.contains(s) { + stop_now = true; + break; + } + } + if stop_now { + let chunk = serde_json::to_string(&SseChunk { + id: chat_id.clone(), + object: "chat.completion.chunk".into(), + created, + model: model_name.clone(), + choices: vec![SseChoice { + index: 0, + delta: SseDelta { + role: None, + content: None, + tool_calls: None, + }, + finish_reason: Some("stop".into()), + }], + }) + .unwrap(); + let _ = tx.send(Ok(Event::default().data(chunk))).await; + break; + } + + // Check tool call completeness + if has_tools && text_buf.contains("") { + let open = text_buf.matches("").count(); + let close = text_buf.matches("").count(); + if close >= open { + let chunk = serde_json::to_string(&SseChunk { + id: chat_id.clone(), + object: "chat.completion.chunk".into(), + created, + model: model_name.clone(), + choices: vec![SseChoice { + index: 0, + delta: SseDelta { + role: None, + content: None, + tool_calls: None, + }, + finish_reason: Some("tool_calls".into()), + }], + }) + .unwrap(); + let _ = tx.send(Ok(Event::default().data(chunk))).await; + break; + } + } + + let pos = input_tokens.len() as i32 + count as i32; + if let Err(e) = inner.decode(current, pos) { + info!(" Decode error: {e}"); + break; + } + count += 1; + current = inner.sample(&mut sampler); + } + }); + + let stream = ReceiverStream::new(rx); + let sse = Sse::new(stream).keep_alive(KeepAlive::default()); + Ok(sse.into_response()) +} diff --git a/src/presentation/handler/chat_ui.rs b/src/presentation/handler/chat_ui.rs new file mode 100644 index 0000000..451a59f --- /dev/null +++ b/src/presentation/handler/chat_ui.rs @@ -0,0 +1,16 @@ +//! Chat UI — simple web interface for interacting with the LLM. + +use axum::http::{header, HeaderValue, StatusCode}; +use axum::response::{IntoResponse, Response}; + +const HTML: &str = include_str!("chat-ui/index.html"); + +/// GET / — serve the chat UI page. +pub async fn chat_ui() -> Response { + ( + StatusCode::OK, + [(header::CONTENT_TYPE, HeaderValue::from_static("text/html; charset=utf-8"))], + HTML, + ) + .into_response() +} diff --git a/src/presentation/handler/health.rs b/src/presentation/handler/health.rs new file mode 100644 index 0000000..663627b --- /dev/null +++ b/src/presentation/handler/health.rs @@ -0,0 +1,13 @@ +//! Health check endpoint. + +use axum::Json; + +use crate::config::MODEL_ID; +use crate::domain::entity::HealthResponse; + +pub async fn health_check() -> Json { + Json(HealthResponse { + status: "ok".into(), + model: format!("{MODEL_ID}-q4_k_m"), + }) +} diff --git a/src/presentation/handler/mod.rs b/src/presentation/handler/mod.rs new file mode 100644 index 0000000..409ac03 --- /dev/null +++ b/src/presentation/handler/mod.rs @@ -0,0 +1,4 @@ +pub mod chat; +pub mod chat_ui; +pub mod health; +pub mod models; diff --git a/src/presentation/handler/models.rs b/src/presentation/handler/models.rs new file mode 100644 index 0000000..ca769b6 --- /dev/null +++ b/src/presentation/handler/models.rs @@ -0,0 +1,19 @@ +//! Models list endpoint. + +use axum::Json; +use chrono::Utc; + +use crate::config::MODEL_ID; +use crate::domain::entity::{ModelInfo, ModelsResponse}; + +pub async fn list_models() -> Json { + Json(ModelsResponse { + object: "list".into(), + data: vec![ModelInfo { + id: MODEL_ID.into(), + object: "model".into(), + created: Utc::now().timestamp(), + owned_by: "asepharyana".into(), + }], + }) +} diff --git a/src/presentation/middleware/auth.rs b/src/presentation/middleware/auth.rs new file mode 100644 index 0000000..d2cd085 --- /dev/null +++ b/src/presentation/middleware/auth.rs @@ -0,0 +1,39 @@ +//! Authentication middleware. +//! +//! Checks for a valid Bearer token in the Authorization header. +//! Only applied to routes that require authentication. + +use axum::extract::Request; +use axum::http::StatusCode; +use axum::middleware::Next; +use axum::response::Response; + +use crate::config::CONFIG; + +/// Middleware that validates the Bearer token in the Authorization header. +/// +/// If `API_KEY` is not set (empty), authentication is disabled and +/// all requests pass through. If set, the middleware rejects requests +/// without a matching token. +pub async fn auth_middleware(request: Request, next: Next) -> Result { + let api_key = &CONFIG.api_key; + if api_key.is_empty() { + return Ok(next.run(request).await); + } + + let header = request + .headers() + .get("authorization") + .and_then(|v| v.to_str().ok()) + .unwrap_or(""); + + let expected = format!("Bearer {api_key}"); + if header == expected || header == api_key { + return Ok(next.run(request).await); + } + + Err(( + StatusCode::UNAUTHORIZED, + "{\"error\":\"unauthorized\",\"message\":\"Invalid API key\"}".into(), + )) +} diff --git a/src/presentation/middleware/mod.rs b/src/presentation/middleware/mod.rs new file mode 100644 index 0000000..0e4a05d --- /dev/null +++ b/src/presentation/middleware/mod.rs @@ -0,0 +1 @@ +pub mod auth; diff --git a/src/presentation/mod.rs b/src/presentation/mod.rs new file mode 100644 index 0000000..6825463 --- /dev/null +++ b/src/presentation/mod.rs @@ -0,0 +1,8 @@ +//! Presentation Layer — Axum handlers, middleware, state, and error handling. + +pub mod dto; +pub mod error; +pub mod handler; +pub mod middleware; +pub mod router; +pub mod state; diff --git a/src/presentation/router.rs b/src/presentation/router.rs new file mode 100644 index 0000000..6d4ef02 --- /dev/null +++ b/src/presentation/router.rs @@ -0,0 +1,27 @@ +//! Axum router assembly. + +use std::sync::Arc; + +use axum::middleware; +use axum::routing::{get, post}; +use axum::Router; +use tower_http::cors::CorsLayer; + +use super::handler::{chat, chat_ui, health, models}; +use crate::presentation::middleware::auth::auth_middleware; +use crate::presentation::state::AppState; + +/// Build the main application router with all routes and middleware. +pub fn build_router(state: Arc) -> Router { + Router::new() + // Public routes (no auth) + .route("/", get(chat_ui::chat_ui)) + .route("/health", get(health::health_check)) + .route("/v1/models", get(models::list_models)) + // Chat completions (auth-protected) + .route("/v1/chat/completions", post(chat::chat_completions)) + .route_layer(middleware::from_fn(auth_middleware)) + // Global middleware + .layer(CorsLayer::permissive()) + .with_state(state) +} diff --git a/src/presentation/state.rs b/src/presentation/state.rs new file mode 100644 index 0000000..1c414f8 --- /dev/null +++ b/src/presentation/state.rs @@ -0,0 +1,13 @@ +//! Application state shared across all handlers. + +use std::sync::Arc; + +use crate::infrastructure::llama::LlamaEngine; + +/// Shared application state injected into every handler via Axum State. +/// +/// Contains the infrastructure dependencies that handlers need. +#[derive(Clone)] +pub struct AppState { + pub engine: Arc, +}