chore: initialize autoresearch baseline
This commit is contained in:
Evandro Fonseca Junior committed
commit
f83e8b83f6
42 files changed
+10550
No files matched your search
@@ -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"
|
||||
@@ -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");
|
||||
}
|
||||
@@ -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.")
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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,
|
||||
}
|
||||
Reference in new issue
Block a user