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

+10
View File
@@ -0,0 +1,10 @@
[package]
name = "api-contracts"
version.workspace = true
edition.workspace = true
[dependencies]
serde.workspace = true
serde_json.workspace = true
uuid.workspace = true
chrono.workspace = true
+361
View File
@@ -0,0 +1,361 @@
use serde::{Deserialize, Deserializer, Serialize};
use serde_json::Value;
#[derive(Debug, Clone, Deserialize)]
pub struct ChatCompletionRequest {
pub model: String,
pub messages: Vec<ChatMessage>,
#[serde(default = "default_max_tokens", alias = "max_completion_tokens")]
pub max_tokens: u32,
#[serde(default = "default_temperature")]
pub temperature: f32,
#[serde(default = "default_top_p")]
pub top_p: f32,
#[serde(default)]
pub stream: bool,
#[serde(default, deserialize_with = "deserialize_stop_sequences")]
pub stop: Vec<String>,
#[serde(default = "default_n")]
pub n: u32,
#[serde(default)]
pub seed: Option<u32>,
#[serde(default)]
pub user: Option<String>,
#[serde(default)]
pub response_format: Option<ResponseFormat>,
#[serde(default)]
pub tools: Vec<ToolDefinition>,
#[serde(default)]
pub tool_choice: Option<ToolChoice>,
#[serde(default)]
pub parallel_tool_calls: Option<bool>,
#[serde(default)]
pub stream_options: Option<StreamOptions>,
}
fn default_max_tokens() -> u32 {
1024
}
fn default_temperature() -> f32 {
0.15
}
fn default_top_p() -> f32 {
0.95
}
fn default_n() -> u32 {
1
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatMessage {
pub role: Role,
#[serde(
default,
deserialize_with = "deserialize_message_content",
skip_serializing_if = "Option::is_none"
)]
pub content: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub name: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub tool_calls: Option<Vec<ToolCall>>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub tool_call_id: Option<String>,
}
impl ChatMessage {
pub fn text(&self) -> &str {
self.content.as_deref().unwrap_or("")
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum Role {
System,
Developer,
User,
Assistant,
Tool,
}
#[derive(Debug, Clone, Serialize)]
pub struct ChatCompletionResponse {
pub id: String,
pub object: &'static str,
pub created: i64,
pub model: String,
pub choices: Vec<Choice>,
pub usage: Usage,
}
#[derive(Debug, Clone, Serialize)]
pub struct Choice {
pub index: u32,
pub message: ChatMessage,
pub finish_reason: FinishReason,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum FinishReason {
Stop,
Length,
ToolCalls,
}
#[derive(Debug, Clone, Serialize)]
pub struct ChatCompletionChunk {
pub id: String,
pub object: &'static str,
pub created: i64,
pub model: String,
pub choices: Vec<ChunkChoice>,
#[serde(skip_serializing_if = "Option::is_none")]
pub usage: Option<Usage>,
}
#[derive(Debug, Clone, Serialize)]
pub struct ChunkChoice {
pub index: u32,
pub delta: Delta,
pub finish_reason: Option<FinishReason>,
}
#[derive(Debug, Clone, Serialize)]
pub struct Delta {
#[serde(skip_serializing_if = "Option::is_none")]
pub role: Option<Role>,
#[serde(skip_serializing_if = "Option::is_none")]
pub content: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tool_calls: Option<Vec<DeltaToolCall>>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolDefinition {
#[serde(rename = "type")]
pub tool_type: ToolType,
pub function: FunctionDefinition,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum ToolType {
Function,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FunctionDefinition {
pub name: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub description: Option<String>,
#[serde(default)]
pub parameters: Value,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolCall {
pub id: String,
#[serde(rename = "type")]
pub tool_type: ToolType,
pub function: ToolFunctionCall,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolFunctionCall {
pub name: String,
pub arguments: String,
}
#[derive(Debug, Clone, Serialize)]
pub struct DeltaToolCall {
pub index: u32,
#[serde(skip_serializing_if = "Option::is_none")]
pub id: Option<String>,
#[serde(rename = "type", skip_serializing_if = "Option::is_none")]
pub tool_type: Option<ToolType>,
#[serde(skip_serializing_if = "Option::is_none")]
pub function: Option<DeltaToolFunctionCall>,
}
#[derive(Debug, Clone, Serialize)]
pub struct DeltaToolFunctionCall {
#[serde(skip_serializing_if = "Option::is_none")]
pub name: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub arguments: Option<String>,
}
#[derive(Debug, Clone, Deserialize)]
#[serde(untagged)]
pub enum ToolChoice {
Mode(String),
Specific(ToolChoiceObject),
}
#[derive(Debug, Clone, Deserialize)]
pub struct ToolChoiceObject {
#[serde(rename = "type")]
pub choice_type: ToolType,
pub function: ToolChoiceFunction,
}
#[derive(Debug, Clone, Deserialize)]
pub struct ToolChoiceFunction {
pub name: String,
}
#[derive(Debug, Clone, Deserialize)]
pub struct ResponseFormat {
#[serde(rename = "type")]
pub format_type: String,
#[serde(default)]
pub json_schema: Option<Value>,
}
#[derive(Debug, Clone, Deserialize)]
pub struct StreamOptions {
#[serde(default)]
pub include_usage: bool,
}
#[derive(Debug, Clone, Serialize, Default)]
pub struct Usage {
pub prompt_tokens: u32,
pub completion_tokens: u32,
pub total_tokens: u32,
pub tokens_per_second: f64,
}
fn deserialize_stop_sequences<'de, D>(deserializer: D) -> Result<Vec<String>, D::Error>
where
D: Deserializer<'de>,
{
let value = Option::<Value>::deserialize(deserializer)?;
let Some(value) = value else {
return Ok(Vec::new());
};
match value {
Value::Null => Ok(Vec::new()),
Value::String(s) => Ok(vec![s]),
Value::Array(items) => items
.into_iter()
.map(|item| match item {
Value::String(s) => Ok(s),
_ => Err(serde::de::Error::custom("stop items must be strings")),
})
.collect(),
_ => Err(serde::de::Error::custom(
"stop must be a string, array of strings, or null",
)),
}
}
fn deserialize_message_content<'de, D>(deserializer: D) -> Result<Option<String>, D::Error>
where
D: Deserializer<'de>,
{
let value = Option::<Value>::deserialize(deserializer)?;
let Some(value) = value else {
return Ok(None);
};
match value {
Value::Null => Ok(None),
Value::String(s) => Ok(Some(s)),
Value::Array(items) => {
let mut text_parts = Vec::new();
for item in items {
match item {
Value::String(s) => text_parts.push(s),
Value::Object(map) => {
let item_type = map.get("type").and_then(Value::as_str).unwrap_or("text");
if matches!(item_type, "text" | "input_text" | "output_text") {
if let Some(text) = map.get("text").and_then(Value::as_str) {
text_parts.push(text.to_owned());
} else if let Some(text) = map.get("content").and_then(Value::as_str) {
text_parts.push(text.to_owned());
}
}
}
_ => {
return Err(serde::de::Error::custom(
"message content array items must be strings or text parts",
));
}
}
}
if text_parts.is_empty() {
Ok(None)
} else {
Ok(Some(text_parts.join("\n")))
}
}
_ => Err(serde::de::Error::custom(
"message content must be a string, array, or null",
)),
}
}
#[cfg(test)]
mod tests {
use super::{ChatCompletionRequest, Role, ToolChoice};
#[test]
fn deserializes_openai_style_message_parts_and_aliases() {
let request: ChatCompletionRequest = serde_json::from_str(
r#"{
"model": "alias-model",
"messages": [
{
"role": "developer",
"content": [
{ "type": "text", "text": "Write code" },
{ "type": "input_text", "text": "Use Rust" }
]
},
{
"role": "user",
"content": "Build a parser"
}
],
"max_completion_tokens": 77,
"stop": "</tool_call>",
"tool_choice": "auto"
}"#,
)
.unwrap();
assert_eq!(request.max_tokens, 77);
assert_eq!(request.stop, vec!["</tool_call>"]);
assert!(matches!(request.messages[0].role, Role::Developer));
assert_eq!(request.messages[0].text(), "Write code\nUse Rust");
assert!(matches!(request.tool_choice, Some(ToolChoice::Mode(ref mode)) if mode == "auto"));
}
#[test]
fn deserializes_stop_arrays_and_null_content() {
let request: ChatCompletionRequest = serde_json::from_str(
r#"{
"model": "alias-model",
"messages": [
{ "role": "assistant", "content": null },
{ "role": "tool", "tool_call_id": "call_1", "content": "done" }
],
"stop": ["END", "STOP"]
}"#,
)
.unwrap();
assert_eq!(request.stop, vec!["END", "STOP"]);
assert_eq!(request.messages[0].content, None);
assert_eq!(request.messages[1].tool_call_id.as_deref(), Some("call_1"));
}
}
+38
View File
@@ -0,0 +1,38 @@
use serde::Serialize;
#[derive(Debug, Clone, Serialize)]
pub struct ErrorResponse {
pub error: ApiError,
}
#[derive(Debug, Clone, Serialize)]
pub struct ApiError {
pub message: String,
#[serde(rename = "type")]
pub error_type: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub param: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub code: Option<String>,
}
impl ApiError {
pub fn new(message: impl Into<String>, error_type: impl Into<String>) -> Self {
Self {
message: message.into(),
error_type: error_type.into(),
param: None,
code: None,
}
}
pub fn with_code(mut self, code: impl Into<String>) -> Self {
self.code = Some(code.into());
self
}
pub fn with_param(mut self, param: impl Into<String>) -> Self {
self.param = Some(param.into());
self
}
}
+9
View File
@@ -0,0 +1,9 @@
use serde::Serialize;
#[derive(Debug, Clone, Serialize)]
pub struct HealthResponse {
pub status: &'static str,
pub model_loaded: bool,
pub model_id: Option<String>,
pub uptime_secs: u64,
}
+6
View File
@@ -0,0 +1,6 @@
// api-contracts: shared request/response types (OpenAI-compatible)
pub mod chat;
pub mod error;
pub mod health;
pub mod models;
+21
View File
@@ -0,0 +1,21 @@
use serde::Serialize;
use serde_json::Value;
#[derive(Debug, Clone, Serialize)]
pub struct ModelList {
pub object: &'static str,
pub data: Vec<ModelInfo>,
}
#[derive(Debug, Clone, Serialize)]
pub struct ModelInfo {
pub id: String,
pub object: &'static str,
pub created: i64,
pub owned_by: &'static str,
pub root: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub parent: Option<String>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub permission: Vec<Value>,
}
+17
View File
@@ -0,0 +1,17 @@
[package]
name = "inference-core"
version.workspace = true
edition.workspace = true
[dependencies]
llama-sys = { path = "../llama-sys" }
anyhow.workspace = true
thiserror.workspace = true
tracing.workspace = true
serde.workspace = true
serde_json.workspace = true
tokio = { workspace = true }
futures.workspace = true
tokio-stream.workspace = true
uuid.workspace = true
chrono.workspace = true
+305
View File
@@ -0,0 +1,305 @@
//! InferenceEngine: loads a GGUF model via llama-sys FFI and manages the
//! lifetime of `llama_model` + `llama_context`.
//!
//! Design decisions:
//! - Single model, single context per engine instance (Phase 2 design).
//! - `LlamaLib` is loaded once and stored in an `Arc` so future session spawning
//! can share it without re-loading the dynamic library.
//! - All raw pointers are wrapped in RAII guards that call the matching
//! `llama_free` / `llama_model_free` on drop.
//! - The engine is `Send` but NOT `Sync` – callers must coordinate access
//! (the scheduler in a later phase will own this exclusively).
use std::ffi::CString;
use std::path::PathBuf;
use std::sync::Arc;
use anyhow::{bail, Context, Result};
use tracing::{info, warn};
use llama_sys::{
llama_context, llama_model, LlamaLib, LLAMA_FLASH_ATTN_TYPE_DISABLED,
LLAMA_FLASH_ATTN_TYPE_ENABLED,
};
use crate::profile::RuntimeProfile;
use crate::sampler::SamplerConfig;
use crate::session::GenerationSession;
// ─── Configuration ─────────────────────────────────────────────────────────
/// Configuration for the inference engine.
#[derive(Debug, Clone)]
pub struct EngineConfig {
/// Path to the GGUF model file.
pub model_path: PathBuf,
/// Runtime performance profile (used to fill in defaults).
pub profile: RuntimeProfile,
/// How many GPU layers to offload (999 = all).
pub n_gpu_layers: i32,
/// Context window size in tokens.
pub n_ctx: u32,
/// Physical batch size (prompt processing).
pub n_batch: u32,
/// Micro-batch size (generation step).
pub n_ubatch: u32,
/// Number of CPU threads for generation.
pub n_threads: i32,
/// Number of CPU threads for batch/prefill.
pub n_threads_batch: i32,
/// Enable flash attention.
pub flash_attn: bool,
/// Use mmap for model loading (false = load fully into RAM/VRAM).
pub use_mmap: bool,
/// RNG seed for sampling (0 = random).
pub seed: u32,
}
impl Default for EngineConfig {
fn default() -> Self {
let prof = RuntimeProfile::Turbo4B;
let d = prof.defaults();
Self {
model_path: PathBuf::new(),
profile: prof,
n_gpu_layers: d.n_gpu_layers,
n_ctx: d.n_ctx,
n_batch: d.n_batch,
n_ubatch: d.n_ubatch,
n_threads: d.n_threads,
n_threads_batch: d.n_threads,
flash_attn: d.flash_attn,
use_mmap: false, // --no-mmap recommended for CUDA
seed: 0,
}
}
}
impl EngineConfig {
/// Build from a profile, applying its defaults.
pub fn from_profile(profile: RuntimeProfile, model_path: PathBuf) -> Self {
let d = profile.defaults();
Self {
model_path,
profile,
n_gpu_layers: d.n_gpu_layers,
n_ctx: d.n_ctx,
n_batch: d.n_batch,
n_ubatch: d.n_ubatch,
n_threads: d.n_threads,
n_threads_batch: d.n_threads,
flash_attn: d.flash_attn,
use_mmap: false,
seed: 0,
}
}
}
// ─── RAII wrappers ─────────────────────────────────────────────────────────
/// Owns a `*mut llama_model` and frees it on drop.
struct OwnedModel {
ptr: *mut llama_model,
lib: Arc<LlamaLib>,
}
impl Drop for OwnedModel {
fn drop(&mut self) {
if !self.ptr.is_null() {
unsafe { (self.lib.model_free)(self.ptr) };
}
}
}
// SAFETY: llama_model is reference-counted internally; we are the only owner.
unsafe impl Send for OwnedModel {}
/// Owns a `*mut llama_context` and frees it on drop.
pub(crate) struct OwnedContext {
pub ptr: *mut llama_context,
pub lib: Arc<LlamaLib>,
}
impl Drop for OwnedContext {
fn drop(&mut self) {
if !self.ptr.is_null() {
unsafe { (self.lib.free)(self.ptr) };
}
}
}
// SAFETY: we enforce single-threaded access through InferenceEngine.
unsafe impl Send for OwnedContext {}
// ─── Engine ────────────────────────────────────────────────────────────────
/// Core inference engine.
///
/// - Owns a loaded model and its associated context.
/// - Provides `new_session()` to start a streaming generation.
/// - NOT `Sync` – the scheduler must ensure single-threaded use.
pub struct InferenceEngine {
pub config: EngineConfig,
pub(crate) lib: Arc<LlamaLib>,
_model: OwnedModel,
pub(crate) ctx: OwnedContext,
}
impl InferenceEngine {
/// Load a GGUF model and initialise a llama_context.
///
/// Blocks until the model is fully loaded.
pub fn load(config: EngineConfig) -> Result<Self> {
if config.model_path.as_os_str().is_empty() {
bail!("EngineConfig.model_path is empty");
}
if !config.model_path.exists() {
bail!("Model file not found: {}", config.model_path.display());
}
info!(
model = %config.model_path.display(),
profile = ?config.profile,
n_gpu_layers = config.n_gpu_layers,
n_ctx = config.n_ctx,
n_batch = config.n_batch,
n_ubatch = config.n_ubatch,
flash_attn = config.flash_attn,
use_mmap = config.use_mmap,
"Loading model via llama-sys"
);
// 1. Load dynamic library
let lib = Arc::new(LlamaLib::load_auto().context("Failed to load llama.dll")?);
// 2. Backend init
unsafe {
(lib.backend_load_all)();
(lib.backend_init)();
}
info!("llama backend initialised");
// 3. Model params
let mut mparams = unsafe { (lib.model_default_params)() };
mparams.n_gpu_layers = config.n_gpu_layers;
mparams.use_mmap = config.use_mmap;
// 4. Load model from file
let model_path_c = CString::new(config.model_path.to_str().context("Non-UTF8 model path")?)
.context("CString::new failed")?;
let model_ptr = unsafe { (lib.model_load_from_file)(model_path_c.as_ptr(), mparams) };
if model_ptr.is_null() {
bail!("llama_model_load_from_file returned NULL – check model path and VRAM");
}
// Log model info
{
let mut desc = vec![0u8; 256];
unsafe {
(lib.model_desc)(model_ptr, desc.as_mut_ptr() as *mut _, desc.len());
}
let desc_str = std::ffi::CStr::from_bytes_until_nul(&desc)
.map(|s| s.to_string_lossy().into_owned())
.unwrap_or_else(|_| "<unknown>".into());
let n_params = unsafe { (lib.model_n_params)(model_ptr) };
let is_hybrid = unsafe { (lib.model_is_hybrid)(model_ptr) };
let is_recurrent = unsafe { (lib.model_is_recurrent)(model_ptr) };
info!(
desc = %desc_str,
params_b = n_params as f64 / 1e9,
is_hybrid,
is_recurrent,
"Model loaded"
);
}
let model = OwnedModel {
ptr: model_ptr,
lib: Arc::clone(&lib),
};
// 5. Context params
let mut cparams = unsafe { (lib.context_default_params)() };
cparams.n_ctx = config.n_ctx;
cparams.n_batch = config.n_batch;
cparams.n_ubatch = config.n_ubatch;
cparams.n_threads = config.n_threads;
cparams.n_threads_batch = config.n_threads_batch;
cparams.flash_attn_type = if config.flash_attn {
LLAMA_FLASH_ATTN_TYPE_ENABLED
} else {
LLAMA_FLASH_ATTN_TYPE_DISABLED
};
cparams.offload_kqv = true;
cparams.no_perf = false;
// 6. Init context
let ctx_ptr = unsafe { (lib.init_from_model)(model_ptr, cparams) };
if ctx_ptr.is_null() {
bail!("llama_init_from_model returned NULL – insufficient VRAM or bad params");
}
let actual_ctx = unsafe { (lib.n_ctx)(ctx_ptr) };
info!(n_ctx = actual_ctx, "Context initialised");
let ctx = OwnedContext {
ptr: ctx_ptr,
lib: Arc::clone(&lib),
};
Ok(Self {
config,
lib,
_model: model,
ctx,
})
}
/// Warm up: run a single dummy decode to pre-allocate CUDA buffers.
///
/// Call once after `load()` before accepting real requests.
pub fn warmup(&mut self) -> Result<()> {
info!("Running engine warmup (dummy decode)…");
unsafe {
// Mark context as warmup so llama.cpp skips perf counters
if let Ok(f) = std::panic::catch_unwind(|| {}) {
let _ = f;
}
// set_warmup may not exist in all builds; call decode with 1 token
let dummy_token: llama_sys::llama_token = 1; // BOS-ish token
let mut tokens = vec![dummy_token];
let batch = (self.lib.batch_get_one)(tokens.as_mut_ptr(), 1);
let ret = (self.lib.decode)(self.ctx.ptr, batch);
if ret != 0 {
warn!(ret, "Warmup decode returned non-zero (non-fatal)");
}
}
info!("Warmup complete");
Ok(())
}
/// Create a new generation session for a single request.
///
/// `prompt` is the fully-formatted string (system + user turns already joined).
/// The engine must not be accessed concurrently while the session is active.
pub fn new_session(
&mut self,
prompt: &str,
sampler_cfg: SamplerConfig,
) -> Result<GenerationSession> {
let vocab = unsafe { (self.lib.model_get_vocab)(self._model.ptr) };
GenerationSession::with_vocab(
Arc::clone(&self.lib),
self.ctx.ptr,
vocab,
prompt,
sampler_cfg,
)
}
}
// SAFETY: we are the sole owner of the pointers; no concurrent access.
unsafe impl Send for InferenceEngine {}
+11
View File
@@ -0,0 +1,11 @@
// inference-core: engine abstraction over llama-sys
pub mod engine;
pub mod profile;
pub mod sampler;
pub mod session;
pub use engine::{EngineConfig, InferenceEngine};
pub use profile::RuntimeProfile;
pub use sampler::SamplerConfig;
pub use session::GenerationSession;
+59
View File
@@ -0,0 +1,59 @@
/// Predefined runtime profiles matching the plan targets.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RuntimeProfile {
/// Qwen3.5-4B – maximum TPS, text-only, flash attention, moderate context.
Turbo4B,
/// Qwen3.5-9B – balanced quality/speed, tuned KV cache.
Balanced9B,
/// CPU-only fallback for machines without a GPU.
CpuFallback,
}
impl RuntimeProfile {
/// Returns sensible defaults for each profile.
pub fn defaults(&self) -> ProfileDefaults {
match self {
Self::Turbo4B => ProfileDefaults {
n_gpu_layers: 999,
n_ctx: 4096,
n_batch: 512,
n_ubatch: 128,
flash_attn: true,
n_threads: 8,
},
Self::Balanced9B => ProfileDefaults {
n_gpu_layers: 999,
n_ctx: 2048,
n_batch: 256,
n_ubatch: 64,
flash_attn: true,
n_threads: 8,
},
Self::CpuFallback => ProfileDefaults {
n_gpu_layers: 0,
n_ctx: 2048,
n_batch: 128,
n_ubatch: 32,
flash_attn: false,
n_threads: num_cpus(),
},
}
}
}
fn num_cpus() -> i32 {
// Simple fallback; replace with `num_cpus` crate if needed.
std::thread::available_parallelism()
.map(|n| n.get() as i32)
.unwrap_or(4)
}
#[derive(Debug, Clone)]
pub struct ProfileDefaults {
pub n_gpu_layers: i32,
pub n_ctx: u32,
pub n_batch: u32,
pub n_ubatch: u32,
pub flash_attn: bool,
pub n_threads: i32,
}
+37
View File
@@ -0,0 +1,37 @@
/// Sampling hyper-parameters for a single generation request.
#[derive(Debug, Clone)]
pub struct SamplerConfig {
pub temperature: f32,
pub top_p: f32,
pub top_k: i32,
pub repeat_penalty: f32,
pub max_tokens: u32,
pub stop_sequences: Vec<String>,
pub seed: u32,
}
impl Default for SamplerConfig {
fn default() -> Self {
Self {
temperature: 0.7,
top_p: 0.9,
top_k: 40,
repeat_penalty: 1.1,
max_tokens: 512,
stop_sequences: vec![],
seed: 0,
}
}
}
impl SamplerConfig {
pub fn from_request(temp: f32, top_p: f32, max_tokens: u32, stop: Vec<String>) -> Self {
Self {
temperature: temp,
top_p,
max_tokens,
stop_sequences: stop,
..Default::default()
}
}
}
+376
View File
@@ -0,0 +1,376 @@
//! GenerationSession: single-request token-by-token generation loop.
//!
//! Flow:
//! 1. Tokenize the prompt (including system + user messages).
//! 2. Prefill: feed all prompt tokens through llama_decode() in one batch.
//! 3. Generation loop: sample → decode → yield token text → check stop.
//!
//! The session borrows raw pointers from InferenceEngine and must not outlive it.
//! In Phase 7+ this will be refactored to use a shared context pool.
use std::ffi::{CStr, CString};
use std::sync::Arc;
use std::time::Instant;
use anyhow::{bail, Context, Result};
use tracing::{debug, info};
use llama_sys::{llama_context, llama_token, llama_vocab, LlamaLib};
use crate::sampler::SamplerConfig;
// ─── GenerationSession ─────────────────────────────────────────────────────
pub struct GenerationSession {
lib: Arc<LlamaLib>,
ctx: *mut llama_context,
vocab: *const llama_vocab,
sampler: *mut llama_sys::llama_sampler,
stop_sequences: Vec<String>,
// generation state
n_decoded: i32, // number of tokens decoded so far (prefill + gen)
prompt_tokens: u32,
pub max_tokens: u32,
pub tokens_generated: u32,
pub finished: bool,
// timing
pub prefill_ms: Option<f64>,
gen_start: Option<Instant>,
pub tokens_per_sec: f64,
}
// SAFETY: pointers come from InferenceEngine which is Send.
unsafe impl Send for GenerationSession {}
impl GenerationSession {
/// Tokenize prompt and run prefill decode. Returns a session ready to generate.
pub fn new(
_lib: Arc<LlamaLib>,
_ctx: *mut llama_context,
_sampler_cfg: SamplerConfig,
) -> Result<Self> {
// Get vocab from context's model (via a workaround: we pass vocab separately)
// For now we'll get it indirectly — see Engine::new_session which passes vocab.
// This constructor signature will be refined; use the vocab variant below.
bail!("Use GenerationSession::with_vocab() instead")
}
/// Full constructor called by InferenceEngine::new_session().
pub fn with_vocab(
lib: Arc<LlamaLib>,
ctx: *mut llama_context,
vocab: *const llama_vocab,
prompt: &str,
sampler_cfg: SamplerConfig,
) -> Result<Self> {
// ── 1. Tokenize prompt ─────────────────────────────────────────────
let prompt_c = CString::new(prompt).context("Invalid UTF-8 in prompt")?;
// First call: get token count
let ret = unsafe {
(lib.tokenize)(
vocab,
prompt_c.as_ptr(),
prompt.len() as i32,
std::ptr::null_mut(),
0,
true, // add_special (BOS)
true, // parse_special
)
};
let n_tokens_needed = if ret < 0 { -ret } else { ret };
if n_tokens_needed == 0 {
bail!("llama_tokenize failed to return token count");
}
let mut tokens: Vec<llama_token> = vec![0; n_tokens_needed as usize];
let n_tokens = unsafe {
(lib.tokenize)(
vocab,
prompt_c.as_ptr(),
prompt.len() as i32,
tokens.as_mut_ptr(),
tokens.len() as i32,
true,
true,
)
};
if n_tokens < 0 {
bail!(
"llama_tokenize (second call) failed (returned {})",
n_tokens
);
}
tokens.truncate(n_tokens as usize);
let n_prompt_tokens = tokens.len() as i32;
debug!(n_tokens = n_prompt_tokens, "Prompt tokenised");
// ── 2. Build sampler chain ─────────────────────────────────────────
let sc_params = unsafe { (lib.sampler_chain_default_params)() };
let sampler_chain = unsafe { (lib.sampler_chain_init)(sc_params) };
if sampler_chain.is_null() {
bail!("sampler_chain_init returned NULL");
}
unsafe {
if sampler_cfg.temperature <= 0.0 || sampler_cfg.temperature < 1e-6 {
// Greedy
let s = (lib.sampler_init_greedy)();
(lib.sampler_chain_add)(sampler_chain, s);
} else {
// top-k → temp → top-p → dist (for randomness)
let s_topk = (lib.sampler_init_top_k)(sampler_cfg.top_k);
(lib.sampler_chain_add)(sampler_chain, s_topk);
let s_temp = (lib.sampler_init_temp)(sampler_cfg.temperature);
(lib.sampler_chain_add)(sampler_chain, s_temp);
let s_topp = (lib.sampler_init_top_p)(sampler_cfg.top_p, 1);
(lib.sampler_chain_add)(sampler_chain, s_topp);
let s_dist = (lib.sampler_init_dist)(sampler_cfg.seed);
(lib.sampler_chain_add)(sampler_chain, s_dist);
}
}
// ── 3. Prefill decode ──────────────────────────────────────────────
let prefill_start = Instant::now();
let ret = unsafe {
let batch = (lib.batch_get_one)(tokens.as_mut_ptr(), n_tokens);
(lib.decode)(ctx, batch)
};
if ret != 0 {
unsafe { (lib.sampler_free)(sampler_chain) };
bail!(
"llama_decode (prefill) returned {} — context may be full",
ret
);
}
let prefill_ms = prefill_start.elapsed().as_secs_f64() * 1000.0;
let prefill_tps = n_prompt_tokens as f64 / (prefill_ms / 1000.0);
info!(
"Prefill finished in {:.2} ms ({:.2} t/s)",
prefill_ms, prefill_tps
);
// create session state
Ok(Self {
lib,
ctx,
vocab,
sampler: sampler_chain, // use sampler_chain here
stop_sequences: sampler_cfg.stop_sequences.clone(),
n_decoded: n_prompt_tokens,
prompt_tokens: n_prompt_tokens as u32,
max_tokens: sampler_cfg.max_tokens,
tokens_generated: 0,
finished: false,
prefill_ms: Some(prefill_ms),
gen_start: None,
tokens_per_sec: 0.0,
})
}
/// Sample and decode the next token. Returns `None` when generation ends.
///
/// Each call advances the context by one token.
pub fn next_token(&mut self) -> Option<String> {
if self.finished {
return None;
}
if self.tokens_generated >= self.max_tokens {
self.finished = true;
self.finalize_timing();
return None;
}
// Start timing on first generated token
if self.gen_start.is_none() {
self.gen_start = Some(Instant::now());
}
unsafe {
// Sample the next token
let new_token = (self.lib.sampler_sample)(self.sampler, self.ctx, -1);
// Check end-of-generation
if (self.lib.vocab_is_eog)(self.vocab, new_token) {
self.finished = true;
self.finalize_timing();
return None;
}
// Decode (advance context with the new token)
let mut token_arr = [new_token];
let batch = (self.lib.batch_get_one)(token_arr.as_mut_ptr(), 1);
let ret = (self.lib.decode)(self.ctx, batch);
if ret != 0 {
debug!(ret, "llama_decode (gen step) error — stopping");
self.finished = true;
self.finalize_timing();
return None;
}
self.n_decoded += 1;
self.tokens_generated += 1;
// Detokenize
let token_text = self.detokenize_token(new_token);
// Check stop sequences
// (simple check: if generated text so far ends with a stop seq)
// For streaming we check per-token — good enough for Phase 2.
for stop in &self.stop_sequences {
if token_text.trim_end() == stop.as_str() {
self.finished = true;
self.finalize_timing();
return None;
}
}
Some(token_text)
}
}
fn finalize_timing(&mut self) {
if let Some(start) = self.gen_start.take() {
let elapsed = start.elapsed().as_secs_f64();
if elapsed > 0.0 {
self.tokens_per_sec = self.tokens_generated as f64 / elapsed;
}
}
info!(
tokens_generated = self.tokens_generated,
tps = format!("{:.1}", self.tokens_per_sec),
prefill_ms = self
.prefill_ms
.map(|v| format!("{:.1}", v))
.unwrap_or_default(),
"Generation complete"
);
}
pub fn tokens_generated(&self) -> u32 {
self.tokens_generated
}
pub fn prompt_tokens(&self) -> u32 {
self.prompt_tokens
}
pub fn tps(&self) -> f64 {
self.tokens_per_sec
}
pub fn prefill_ms(&self) -> Option<f64> {
self.prefill_ms
}
fn detokenize_token(&self, token: llama_token) -> String {
if let Some(token_to_piece) = self.lib.token_to_piece {
return decode_token_piece_with(|buf, len| unsafe {
token_to_piece(self.vocab, token, buf, len, 0, true)
});
}
let text_ptr = unsafe { (self.lib.vocab_get_text)(self.vocab, token) };
if text_ptr.is_null() {
String::new()
} else {
unsafe { CStr::from_ptr(text_ptr) }
.to_string_lossy()
.into_owned()
}
}
}
fn decode_token_piece_with<F>(mut call: F) -> String
where
F: FnMut(*mut std::ffi::c_char, i32) -> i32,
{
let mut capacity = 64usize;
loop {
let mut buffer = vec![0u8; capacity];
let written = call(
buffer.as_mut_ptr() as *mut std::ffi::c_char,
capacity as i32,
);
if written < 0 {
capacity = (-written) as usize;
continue;
}
let written = written as usize;
return String::from_utf8_lossy(&buffer[..written]).into_owned();
}
}
#[cfg(test)]
mod tests {
use super::decode_token_piece_with;
#[test]
fn retries_when_piece_buffer_is_too_small() {
let mut calls = 0;
let decoded = decode_token_piece_with(|buf, len| {
calls += 1;
if len < 17 {
return -17;
}
let bytes = b" Quantum computer";
unsafe {
std::ptr::copy_nonoverlapping(bytes.as_ptr(), buf as *mut u8, bytes.len());
}
bytes.len() as i32
});
assert_eq!(decoded, " Quantum computer");
assert_eq!(calls, 1);
}
#[test]
fn decodes_piece_after_resizing_buffer() {
let mut calls = 0;
let decoded = decode_token_piece_with(|buf, len| {
calls += 1;
if calls == 1 {
return -18;
}
assert!(len >= 18);
let bytes = b" Quantum computer";
unsafe {
std::ptr::copy_nonoverlapping(bytes.as_ptr(), buf as *mut u8, bytes.len());
}
bytes.len() as i32
});
assert_eq!(decoded, " Quantum computer");
assert_eq!(calls, 2);
}
}
impl Drop for GenerationSession {
fn drop(&mut self) {
if !self.sampler.is_null() {
unsafe { (self.lib.sampler_free)(self.sampler) };
self.sampler = std::ptr::null_mut();
}
}
}
+30
View File
@@ -0,0 +1,30 @@
[package]
name = "inference-server"
version.workspace = true
edition.workspace = true
build = "build.rs"
[[bin]]
name = "localllm"
path = "src/main.rs"
[dependencies]
inference-core = { path = "../inference-core" }
api-contracts = { path = "../api-contracts" }
metrics = { path = "../metrics" }
tokio.workspace = true
axum.workspace = true
hyper.workspace = true
tower.workspace = true
tower-http.workspace = true
serde.workspace = true
serde_json.workspace = true
anyhow.workspace = true
thiserror.workspace = true
tracing.workspace = true
tracing-subscriber.workspace = true
futures.workspace = true
tokio-stream.workspace = true
bytes.workspace = true
uuid.workspace = true
async-stream = "0.3"
+38
View File
@@ -0,0 +1,38 @@
fn main() {
let out_dir = std::env::var("OUT_DIR").unwrap();
let target_dir = std::path::Path::new(&out_dir)
.ancestors()
.nth(3) // Volta 3 níveis de `target/debug/build/inference-server-XYZ/out` para `target/debug/` ou `target/release/`
.unwrap();
let vendor_dir = std::path::PathBuf::from(std::env::var("CARGO_MANIFEST_DIR").unwrap())
.parent()
.unwrap()
.parent()
.unwrap()
.join("vendor")
.join("llama-bin");
if vendor_dir.exists() {
println!(
"cargo:warning=Copying DLLs from {:?} to {:?}",
vendor_dir, target_dir
);
// Vamos copiar todos os arquivos .dll
if let Ok(entries) = std::fs::read_dir(&vendor_dir) {
for entry in entries.flatten() {
let path = entry.path();
if path.is_file() && path.extension().and_then(|s| s.to_str()) == Some("dll") {
let file_name = path.file_name().unwrap();
let dest = target_dir.join(file_name);
let _ = std::fs::copy(&path, &dest);
}
}
}
}
// Rebuildar apenas se houver mudanças nessa build file ou nos bins da vendor
println!("cargo:rerun-if-changed=build.rs");
println!("cargo:rerun-if-changed=../../vendor/llama-bin");
}
+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,
}
+21
View File
@@ -0,0 +1,21 @@
[package]
name = "llama-sys"
version.workspace = true
edition.workspace = true
[dependencies]
anyhow.workspace = true
libloading.workspace = true
once_cell.workspace = true
tracing.workspace = true
[build-dependencies]
# bindgen e cmake ficam para quando compilarmos do source (Phase 3+)
# bindgen = "0.72"
# cmake = "0.1"
[features]
default = ["cuda"]
cuda = []
metal = []
cpu-only = []
+14
View File
@@ -0,0 +1,14 @@
/// Build script for llama-sys.
///
/// Phase 2: Dynamic loading via libloading — no .lib needed.
/// The real llama.dll path is communicated to the crate via the
/// LLAMA_DLL_PATH env var or a sensible default.
///
/// Phase 3+ will optionally compile llama.cpp from source.
fn main() {
// Nothing to link here — we use libloading at runtime.
// Rerun if the env var changes.
println!("cargo:rerun-if-env-changed=LLAMA_DLL_PATH");
println!("cargo:rerun-if-changed=build.rs");
}
+552
View File
@@ -0,0 +1,552 @@
//! llama-sys: Dynamic FFI bindings to llama.cpp via libloading.
//!
//! Loads `llama.dll` (Windows) / `libllama.so` (Linux) / `libllama.dylib` (macOS)
//! at runtime from the path specified by `LLAMA_DLL_PATH` env var, or from a
//! set of well-known locations.
//!
//! This approach requires no `.lib` import library and no source compilation.
#![allow(non_upper_case_globals)]
#![allow(non_camel_case_types)]
#![allow(non_snake_case)]
#![allow(dead_code)]
#![allow(clippy::all)]
use std::ffi::{c_char, c_void};
use std::path::{Path, PathBuf};
pub use libloading::Library;
// ─── Primitive type aliases ────────────────────────────────────────────────
pub type llama_pos = i32;
pub type llama_token = i32;
pub type llama_seq_id = i32;
// ─── Opaque struct handles ─────────────────────────────────────────────────
#[repr(C)]
pub struct llama_model {
_private: [u8; 0],
}
#[repr(C)]
pub struct llama_context {
_private: [u8; 0],
}
#[repr(C)]
pub struct llama_vocab {
_private: [u8; 0],
}
#[repr(C)]
pub struct llama_sampler {
_private: [u8; 0],
}
// ─── Model params ──────────────────────────────────────────────────────────
#[repr(C)]
pub struct llama_model_params {
pub devices: *mut c_void,
pub tensor_buft_overrides: *const c_void,
pub n_gpu_layers: i32,
pub split_mode: i32,
pub main_gpu: i32,
pub tensor_split: *const f32,
pub progress_callback: Option<extern "C" fn(f32, *mut c_void) -> bool>,
pub progress_callback_user_data: *mut c_void,
pub kv_overrides: *const c_void,
pub vocab_only: bool,
pub use_mmap: bool,
pub use_direct_io: bool,
pub use_mlock: bool,
pub check_tensors: bool,
pub use_extra_bufts: bool,
pub no_host: bool,
pub no_alloc: bool,
}
// ─── Context params ────────────────────────────────────────────────────────
#[repr(C)]
pub struct llama_context_params {
pub n_ctx: u32,
pub n_batch: u32,
pub n_ubatch: u32,
pub n_seq_max: u32,
pub n_threads: i32,
pub n_threads_batch: i32,
pub rope_scaling_type: i32,
pub pooling_type: i32,
pub attention_type: i32,
pub flash_attn_type: i32,
pub rope_freq_base: f32,
pub rope_freq_scale: f32,
pub yarn_ext_factor: f32,
pub yarn_attn_factor: f32,
pub yarn_beta_fast: f32,
pub yarn_beta_slow: f32,
pub yarn_orig_ctx: u32,
pub defrag_thold: f32,
pub cb_eval: Option<extern "C" fn(*mut c_void, bool) -> bool>,
pub cb_eval_user_data: *mut c_void,
pub type_k: i32,
pub type_v: i32,
pub abort_callback: Option<extern "C" fn(*mut c_void) -> bool>,
pub abort_callback_data: *mut c_void,
pub embeddings: bool,
pub offload_kqv: bool,
pub no_perf: bool,
pub op_offload: bool,
pub swa_full: bool,
pub kv_unified: bool,
pub samplers: *mut c_void,
pub n_samplers: usize,
}
// ─── Batch ─────────────────────────────────────────────────────────────────
#[repr(C)]
pub struct llama_batch {
pub n_tokens: i32,
pub token: *mut llama_token,
pub embd: *mut f32,
pub pos: *mut llama_pos,
pub n_seq_id: *mut i32,
pub seq_id: *mut *mut llama_seq_id,
pub logits: *mut i8,
}
// ─── Sampler chain params ──────────────────────────────────────────────────
#[repr(C)]
pub struct llama_sampler_chain_params {
pub no_perf: bool,
}
// ─── Flash attn type ───────────────────────────────────────────────────────
pub const LLAMA_FLASH_ATTN_TYPE_AUTO: i32 = -1;
pub const LLAMA_FLASH_ATTN_TYPE_DISABLED: i32 = 0;
pub const LLAMA_FLASH_ATTN_TYPE_ENABLED: i32 = 1;
// ─── Dynamic loader ────────────────────────────────────────────────────────
/// Container for all resolved function pointers from llama.dll.
pub struct LlamaLib {
// Keep the library alive as long as this struct is alive.
_lib: Library,
// ── Lifecycle ──────────────────────────────────────────────────────────
pub backend_init: unsafe extern "C" fn(),
pub backend_free: unsafe extern "C" fn(),
pub backend_load_all: unsafe extern "C" fn(),
// ── Defaults ──────────────────────────────────────────────────────────
pub model_default_params: unsafe extern "C" fn() -> llama_model_params,
pub context_default_params: unsafe extern "C" fn() -> llama_context_params,
pub sampler_chain_default_params: unsafe extern "C" fn() -> llama_sampler_chain_params,
// ── Model ─────────────────────────────────────────────────────────────
pub model_load_from_file:
unsafe extern "C" fn(*const c_char, llama_model_params) -> *mut llama_model,
pub model_free: unsafe extern "C" fn(*mut llama_model),
pub model_get_vocab: unsafe extern "C" fn(*const llama_model) -> *const llama_vocab,
pub model_n_ctx_train: unsafe extern "C" fn(*const llama_model) -> i32,
pub model_is_recurrent: unsafe extern "C" fn(*const llama_model) -> bool,
pub model_is_hybrid: unsafe extern "C" fn(*const llama_model) -> bool,
pub model_desc: unsafe extern "C" fn(*const llama_model, *mut c_char, usize) -> i32,
pub model_n_params: unsafe extern "C" fn(*const llama_model) -> u64,
pub model_size: unsafe extern "C" fn(*const llama_model) -> u64,
// ── Context ───────────────────────────────────────────────────────────
pub init_from_model:
unsafe extern "C" fn(*mut llama_model, llama_context_params) -> *mut llama_context,
pub free: unsafe extern "C" fn(*mut llama_context),
pub n_ctx: unsafe extern "C" fn(*const llama_context) -> u32,
// ── Batch ─────────────────────────────────────────────────────────────
pub batch_get_one: unsafe extern "C" fn(*mut llama_token, i32) -> llama_batch,
pub batch_init: unsafe extern "C" fn(i32, i32, i32) -> llama_batch,
pub batch_free: unsafe extern "C" fn(llama_batch),
// ── Decode ────────────────────────────────────────────────────────────
pub decode: unsafe extern "C" fn(*mut llama_context, llama_batch) -> i32,
pub set_n_threads: unsafe extern "C" fn(*mut llama_context, i32, i32),
pub set_warmup: unsafe extern "C" fn(*mut llama_context, bool),
pub synchronize: unsafe extern "C" fn(*mut llama_context),
// ── Logits ────────────────────────────────────────────────────────────
pub get_logits_ith: unsafe extern "C" fn(*mut llama_context, i32) -> *mut f32,
// ── Vocab / Tokenizer ─────────────────────────────────────────────────
pub vocab_n_tokens: unsafe extern "C" fn(*const llama_vocab) -> i32,
pub vocab_is_eog: unsafe extern "C" fn(*const llama_vocab, llama_token) -> bool,
pub vocab_get_text: unsafe extern "C" fn(*const llama_vocab, llama_token) -> *const c_char,
pub token_to_piece: Option<
unsafe extern "C" fn(*const llama_vocab, llama_token, *mut c_char, i32, i32, bool) -> i32,
>,
pub tokenize: unsafe extern "C" fn(
*const llama_vocab,
*const c_char,
i32,
*mut llama_token,
i32,
bool,
bool,
) -> i32,
pub token_bos: unsafe extern "C" fn(*const llama_vocab) -> llama_token,
pub token_eos: unsafe extern "C" fn(*const llama_vocab) -> llama_token,
// ── Sampler ───────────────────────────────────────────────────────────
pub sampler_chain_init: unsafe extern "C" fn(llama_sampler_chain_params) -> *mut llama_sampler,
pub sampler_chain_add: unsafe extern "C" fn(*mut llama_sampler, *mut llama_sampler),
pub sampler_free: unsafe extern "C" fn(*mut llama_sampler),
pub sampler_sample:
unsafe extern "C" fn(*mut llama_sampler, *mut llama_context, i32) -> llama_token,
pub sampler_init_greedy: unsafe extern "C" fn() -> *mut llama_sampler,
pub sampler_init_temp: unsafe extern "C" fn(f32) -> *mut llama_sampler,
pub sampler_init_top_p: unsafe extern "C" fn(f32, usize) -> *mut llama_sampler,
pub sampler_init_top_k: unsafe extern "C" fn(i32) -> *mut llama_sampler,
pub sampler_init_min_p: unsafe extern "C" fn(f32, usize) -> *mut llama_sampler,
pub sampler_init_dist: unsafe extern "C" fn(u32) -> *mut llama_sampler,
// ── Misc ──────────────────────────────────────────────────────────────
pub supports_gpu_offload: unsafe extern "C" fn() -> bool,
pub time_us: unsafe extern "C" fn() -> i64,
}
impl LlamaLib {
/// Load the llama shared library from the given path.
pub fn load(path: &Path) -> anyhow::Result<Self> {
use anyhow::Context;
let lib = unsafe { Library::new(path) }
.with_context(|| format!("Failed to load llama library from {:?}", path))?;
macro_rules! sym {
($name:literal) => {{
*unsafe { lib.get::<unsafe extern "C" fn()>(concat!($name, "\0").as_bytes()) }
.with_context(|| format!("Missing symbol: {}", $name))?
}};
}
// We need transmutes to specific fn types — use a helper macro:
macro_rules! load {
($field:ident, $ty:ty, $name:literal) => {{
let raw: libloading::Symbol<unsafe extern "C" fn()> = unsafe {
lib.get(concat!($name, "\0").as_bytes())
.with_context(|| format!("Missing symbol: {}", $name))?
};
unsafe { std::mem::transmute::<_, $ty>(*raw) }
}};
}
Ok(Self {
backend_init: load!(backend_init, unsafe extern "C" fn(), "llama_backend_init"),
backend_free: load!(backend_free, unsafe extern "C" fn(), "llama_backend_free"),
backend_load_all: {
let sym_name = b"ggml_backend_load_all\0";
if let Ok(raw) = unsafe { lib.get::<unsafe extern "C" fn()>(sym_name) } {
unsafe { std::mem::transmute(*raw) }
} else {
let ggml_path = path.with_file_name("ggml.dll");
if let Ok(ggml_lib) = unsafe { libloading::Library::new(&ggml_path) } {
// Faz a extração do raw *antes* de dar leak no Box, desfazendo ownership borrow issue do Rust
let raw_ptr = if let Ok(raw) =
unsafe { ggml_lib.get::<unsafe extern "C" fn()>(sym_name) }
{
Some(unsafe { std::mem::transmute::<_, unsafe extern "C" fn()>(*raw) })
} else {
None
};
// Faz leak no handle do módulo pra manter vivo
std::mem::forget(ggml_lib);
raw_ptr.unwrap_or_else(|| {
panic!("Missing symbol ggml_backend_load_all in both llama.dll and ggml.dll")
})
} else {
anyhow::bail!("Missing symbol ggml_backend_load_all in llama.dll and could not load ggml.dll");
}
}
},
model_default_params: load!(
model_default_params,
unsafe extern "C" fn() -> llama_model_params,
"llama_model_default_params"
),
context_default_params: load!(
context_default_params,
unsafe extern "C" fn() -> llama_context_params,
"llama_context_default_params"
),
sampler_chain_default_params: load!(
sampler_chain_default_params,
unsafe extern "C" fn() -> llama_sampler_chain_params,
"llama_sampler_chain_default_params"
),
model_load_from_file: load!(
model_load_from_file,
unsafe extern "C" fn(*const c_char, llama_model_params) -> *mut llama_model,
"llama_model_load_from_file"
),
model_free: load!(
model_free,
unsafe extern "C" fn(*mut llama_model),
"llama_model_free"
),
model_get_vocab: load!(
model_get_vocab,
unsafe extern "C" fn(*const llama_model) -> *const llama_vocab,
"llama_model_get_vocab"
),
model_n_ctx_train: load!(
model_n_ctx_train,
unsafe extern "C" fn(*const llama_model) -> i32,
"llama_model_n_ctx_train"
),
model_is_recurrent: load!(
model_is_recurrent,
unsafe extern "C" fn(*const llama_model) -> bool,
"llama_model_is_recurrent"
),
model_is_hybrid: load!(
model_is_hybrid,
unsafe extern "C" fn(*const llama_model) -> bool,
"llama_model_is_hybrid"
),
model_desc: load!(
model_desc,
unsafe extern "C" fn(*const llama_model, *mut c_char, usize) -> i32,
"llama_model_desc"
),
model_n_params: load!(
model_n_params,
unsafe extern "C" fn(*const llama_model) -> u64,
"llama_model_n_params"
),
model_size: load!(
model_size,
unsafe extern "C" fn(*const llama_model) -> u64,
"llama_model_size"
),
init_from_model: load!(
init_from_model,
unsafe extern "C" fn(*mut llama_model, llama_context_params) -> *mut llama_context,
"llama_init_from_model"
),
free: load!(free, unsafe extern "C" fn(*mut llama_context), "llama_free"),
n_ctx: load!(
n_ctx,
unsafe extern "C" fn(*const llama_context) -> u32,
"llama_n_ctx"
),
batch_get_one: load!(
batch_get_one,
unsafe extern "C" fn(*mut llama_token, i32) -> llama_batch,
"llama_batch_get_one"
),
batch_init: load!(
batch_init,
unsafe extern "C" fn(i32, i32, i32) -> llama_batch,
"llama_batch_init"
),
batch_free: load!(
batch_free,
unsafe extern "C" fn(llama_batch),
"llama_batch_free"
),
decode: load!(
decode,
unsafe extern "C" fn(*mut llama_context, llama_batch) -> i32,
"llama_decode"
),
set_n_threads: load!(
set_n_threads,
unsafe extern "C" fn(*mut llama_context, i32, i32),
"llama_set_n_threads"
),
set_warmup: load!(
set_warmup,
unsafe extern "C" fn(*mut llama_context, bool),
"llama_set_warmup"
),
synchronize: load!(
synchronize,
unsafe extern "C" fn(*mut llama_context),
"llama_synchronize"
),
get_logits_ith: load!(
get_logits_ith,
unsafe extern "C" fn(*mut llama_context, i32) -> *mut f32,
"llama_get_logits_ith"
),
vocab_n_tokens: load!(
vocab_n_tokens,
unsafe extern "C" fn(*const llama_vocab) -> i32,
"llama_vocab_n_tokens"
),
vocab_is_eog: load!(
vocab_is_eog,
unsafe extern "C" fn(*const llama_vocab, llama_token) -> bool,
"llama_vocab_is_eog"
),
vocab_get_text: load!(
vocab_get_text,
unsafe extern "C" fn(*const llama_vocab, llama_token) -> *const c_char,
"llama_vocab_get_text"
),
token_to_piece: unsafe {
lib.get::<unsafe extern "C" fn(
*const llama_vocab,
llama_token,
*mut c_char,
i32,
i32,
bool,
) -> i32>(b"llama_token_to_piece\0")
.ok()
.map(|raw| *raw)
},
tokenize: load!(
tokenize,
unsafe extern "C" fn(
*const llama_vocab,
*const c_char,
i32,
*mut llama_token,
i32,
bool,
bool,
) -> i32,
"llama_tokenize"
),
token_bos: load!(
token_bos,
unsafe extern "C" fn(*const llama_vocab) -> llama_token,
"llama_vocab_bos"
),
token_eos: load!(
token_eos,
unsafe extern "C" fn(*const llama_vocab) -> llama_token,
"llama_vocab_eos"
),
sampler_chain_init: load!(
sampler_chain_init,
unsafe extern "C" fn(llama_sampler_chain_params) -> *mut llama_sampler,
"llama_sampler_chain_init"
),
sampler_chain_add: load!(
sampler_chain_add,
unsafe extern "C" fn(*mut llama_sampler, *mut llama_sampler),
"llama_sampler_chain_add"
),
sampler_free: load!(
sampler_free,
unsafe extern "C" fn(*mut llama_sampler),
"llama_sampler_free"
),
sampler_sample: load!(
sampler_sample,
unsafe extern "C" fn(*mut llama_sampler, *mut llama_context, i32) -> llama_token,
"llama_sampler_sample"
),
sampler_init_greedy: load!(
sampler_init_greedy,
unsafe extern "C" fn() -> *mut llama_sampler,
"llama_sampler_init_greedy"
),
sampler_init_temp: load!(
sampler_init_temp,
unsafe extern "C" fn(f32) -> *mut llama_sampler,
"llama_sampler_init_temp"
),
sampler_init_top_p: load!(
sampler_init_top_p,
unsafe extern "C" fn(f32, usize) -> *mut llama_sampler,
"llama_sampler_init_top_p"
),
sampler_init_top_k: load!(
sampler_init_top_k,
unsafe extern "C" fn(i32) -> *mut llama_sampler,
"llama_sampler_init_top_k"
),
sampler_init_min_p: load!(
sampler_init_min_p,
unsafe extern "C" fn(f32, usize) -> *mut llama_sampler,
"llama_sampler_init_min_p"
),
sampler_init_dist: load!(
sampler_init_dist,
unsafe extern "C" fn(u32) -> *mut llama_sampler,
"llama_sampler_init_dist"
),
supports_gpu_offload: load!(
supports_gpu_offload,
unsafe extern "C" fn() -> bool,
"llama_supports_gpu_offload"
),
time_us: load!(time_us, unsafe extern "C" fn() -> i64, "llama_time_us"),
_lib: lib,
})
}
/// Resolve the DLL path from env var or known locations.
pub fn find_dll() -> anyhow::Result<PathBuf> {
// 1. Explicit env override
if let Ok(p) = std::env::var("LLAMA_DLL_PATH") {
let path = PathBuf::from(p);
if path.exists() {
return Ok(path);
}
}
// 2. Load directly relying on the Windows Library loader searching the .exe dir
// since our build.rs now copies all DLLs to target/{debug,release}/
#[cfg(target_os = "windows")]
{
return Ok(PathBuf::from("llama.dll"));
}
// 3. Unix-like: rely on LD_LIBRARY_PATH / rpath
#[cfg(not(target_os = "windows"))]
{
for name in &["libllama.so", "libllama.dylib"] {
if let Ok(lib) = Library::new(name) {
drop(lib);
return Ok(PathBuf::from(name));
}
}
}
anyhow::bail!(
"Cannot find llama shared library. Set LLAMA_DLL_PATH to the full path of llama.dll"
)
}
/// Convenience: find + load in one call.
pub fn load_auto() -> anyhow::Result<Self> {
let path = Self::find_dll()?;
tracing::info!(path = %path.display(), "Loading llama library");
Self::load(&path)
}
}
// SAFETY: LlamaLib only holds a Library and function pointers.
// All calls must be externally synchronized (enforced by InferenceEngine).
unsafe impl Send for LlamaLib {}
unsafe impl Sync for LlamaLib {}
+11
View File
@@ -0,0 +1,11 @@
[package]
name = "metrics"
version.workspace = true
edition.workspace = true
[dependencies]
serde.workspace = true
serde_json.workspace = true
tracing.workspace = true
chrono.workspace = true
tokio.workspace = true
+46
View File
@@ -0,0 +1,46 @@
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
/// Global atomic performance counters.
#[derive(Debug, Default)]
pub struct PerfCounters {
pub tokens_generated: AtomicU64,
pub requests_completed: AtomicU64,
pub requests_failed: AtomicU64,
pub total_prompt_tokens: AtomicU64,
}
impl PerfCounters {
pub fn new() -> Arc<Self> {
Arc::new(Self::default())
}
pub fn inc_tokens(&self, n: u64) {
self.tokens_generated.fetch_add(n, Ordering::Relaxed);
}
pub fn inc_completed(&self) {
self.requests_completed.fetch_add(1, Ordering::Relaxed);
}
pub fn inc_failed(&self) {
self.requests_failed.fetch_add(1, Ordering::Relaxed);
}
pub fn snapshot(&self) -> CountersSnapshot {
CountersSnapshot {
tokens_generated: self.tokens_generated.load(Ordering::Relaxed),
requests_completed: self.requests_completed.load(Ordering::Relaxed),
requests_failed: self.requests_failed.load(Ordering::Relaxed),
total_prompt_tokens: self.total_prompt_tokens.load(Ordering::Relaxed),
}
}
}
#[derive(Debug, Clone, serde::Serialize)]
pub struct CountersSnapshot {
pub tokens_generated: u64,
pub requests_completed: u64,
pub requests_failed: u64,
pub total_prompt_tokens: u64,
}
+4
View File
@@ -0,0 +1,4 @@
// metrics: performance counters and structured telemetry
pub mod counters;
pub mod recorder;
+34
View File
@@ -0,0 +1,34 @@
use std::time::Instant;
/// Records a single generation run and computes TPS.
#[derive(Debug)]
pub struct GenerationRecord {
pub start: Instant,
pub prompt_tokens: u32,
pub completion_tokens: u32,
}
impl GenerationRecord {
pub fn new(prompt_tokens: u32) -> Self {
Self {
start: Instant::now(),
prompt_tokens,
completion_tokens: 0,
}
}
pub fn tick(&mut self) {
self.completion_tokens += 1;
}
pub fn elapsed_secs(&self) -> f64 {
self.start.elapsed().as_secs_f64()
}
pub fn tokens_per_second(&self) -> f64 {
if self.completion_tokens == 0 {
return 0.0;
}
self.completion_tokens as f64 / self.elapsed_secs()
}
}