Cameron CordesandClaude Opus 5.5
305eda5894
Long insight/persona chats on llama-swap failed with "exceeds the available context size" even after raising Context in the app: - llama-server fixes its window at launch (-c) and ignores num_ctx, but the history budget trusted the requested num_ctx, so raising it kept MORE history than the server could hold. LlmClient::context_window() now reports the served window (llama-swap: /upstream/<model>/props, cached 10 min) and the budget uses min(requested, served). - apply_context_budget could only drop tool-call pairs, so plain Q&A chats never shrank. It now also drops whole earlier exchanges (keeping user/assistant alternation) and, last, the current turn's tool results, never the in-progress question. - The budget only ran once, before the tool loop. ContextFitter now fits every LLM call (on a copy; the stored transcript is untouched) and, on a context-overflow error, shrinks the budget from the server's reported token counts and retries once. GET /insights/models reports each llama-swap model's context_window when it's known without loading the model (cached, or currently running), for the app's Context picker. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
1665 lines
63 KiB
Rust
1665 lines
63 KiB
Rust
// LlamaCppClient — talks to a llama-swap proxy that fronts one or more
|
||
// llama-server processes. llama-swap exposes an OpenAI-compatible HTTP
|
||
// surface (`/v1/chat/completions`, `/v1/embeddings`, `/v1/models`), so the
|
||
// wire translation mirrors `OpenRouterClient` almost exactly.
|
||
//
|
||
// Differences from OpenRouter:
|
||
// - No bearer auth or attribution headers; llama-swap is LAN-only.
|
||
// - Three model slots (`primary_model` = chat, `vision_model`, `embedding_model`)
|
||
// each map to a model id in the llama-swap config. `describe_image` and
|
||
// `generate_embeddings` issue requests with the appropriate slot id in the
|
||
// `model` field, which is how llama-swap selects which backend process to
|
||
// run.
|
||
// - `/v1/models` returns only the configured slot ids — capabilities aren't
|
||
// reported by the API, so we infer `has_vision` from a single config field
|
||
// (`vision_model`, defaulting to `"vision"`) and assume `has_tool_calling`
|
||
// is true for every slot, since llama-swap entries default to launching
|
||
// llama-server with `--jinja`.
|
||
//
|
||
// First consumer lands alongside the three-way backend dispatch in
|
||
// insight_generator / insight_chat.
|
||
#![allow(dead_code)]
|
||
|
||
use anyhow::Context as _;
|
||
use anyhow::{Result, anyhow, bail};
|
||
use async_trait::async_trait;
|
||
use opentelemetry::trace::Status;
|
||
use reqwest::Client;
|
||
use serde::Deserialize;
|
||
use serde_json::{Value, json};
|
||
use std::collections::HashMap;
|
||
use std::time::Duration;
|
||
use tracing_opentelemetry::OpenTelemetrySpanExt as _;
|
||
|
||
use crate::ai::llm_client::{
|
||
ChatMessage, LlmClient, LlmStreamEvent, ModelCapabilities, Tool, ToolCall, ToolCallFunction,
|
||
};
|
||
use futures::stream::{BoxStream, StreamExt};
|
||
|
||
const DEFAULT_BASE_URL: &str = "http://localhost:9292/v1";
|
||
const DEFAULT_PRIMARY_MODEL: &str = "chat";
|
||
const DEFAULT_VISION_MODEL: &str = "vision";
|
||
const DEFAULT_EMBEDDING_MODEL: &str = "embed";
|
||
const DEFAULT_TTS_MODEL: &str = "chatterbox";
|
||
const DEFAULT_REQUEST_TIMEOUT_SECS: u64 = 180;
|
||
/// How long a served-context-window lookup stays cached per model.
|
||
const CONTEXT_WINDOW_CACHE_TTL: Duration = Duration::from_secs(600);
|
||
|
||
/// `"<swap root>|<model>"` -> (n_ctx, fetched at). Process-wide because
|
||
/// clients are cloned per request.
|
||
fn context_window_cache() -> &'static std::sync::Mutex<HashMap<String, (usize, std::time::Instant)>>
|
||
{
|
||
static CACHE: std::sync::OnceLock<
|
||
std::sync::Mutex<HashMap<String, (usize, std::time::Instant)>>,
|
||
> = std::sync::OnceLock::new();
|
||
CACHE.get_or_init(Default::default)
|
||
}
|
||
|
||
/// OpenAI-compatible client targeting a llama-swap proxy in front of one or
|
||
/// more llama-server processes. See the module doc-comment for the slot model.
|
||
#[derive(Clone)]
|
||
pub struct LlamaCppClient {
|
||
client: Client,
|
||
pub base_url: String,
|
||
/// Chat model slot id (e.g. `"chat"`). Used for `generate` /
|
||
/// `chat_with_tools` / `chat_with_tools_stream`.
|
||
pub primary_model: String,
|
||
/// Embedding model slot id (e.g. `"embed"`). Used for
|
||
/// `generate_embeddings`.
|
||
pub embedding_model: String,
|
||
/// Vision model slot id. Used for `describe_image` routing. Defaults
|
||
/// to `primary_model` so describe_image works out of the box; override
|
||
/// via `LLAMA_SWAP_VISION_MODEL` for a dedicated vision slot.
|
||
pub vision_model: String,
|
||
/// TTS model slot id (e.g. `"chatterbox"`). Routes `text_to_speech` and
|
||
/// is the `/upstream/<id>/voices` path segment for the voice library.
|
||
/// Override via `LLAMA_SWAP_TTS_MODEL`.
|
||
pub tts_model: String,
|
||
num_ctx: Option<i32>,
|
||
temperature: Option<f32>,
|
||
top_p: Option<f32>,
|
||
top_k: Option<i32>,
|
||
min_p: Option<f32>,
|
||
/// When `Some`, forwarded to llama-server as
|
||
/// `chat_template_kwargs: {"enable_thinking": <bool>}`. The Jinja chat
|
||
/// template (e.g. Qwen3) reads this to gate its reasoning block. `None`
|
||
/// omits the key entirely, leaving the template's own default. Templates
|
||
/// that don't reference the key ignore it, so sending it is harmless.
|
||
enable_thinking: Option<bool>,
|
||
}
|
||
|
||
impl LlamaCppClient {
|
||
pub fn new(base_url: Option<String>, primary_model: Option<String>) -> Self {
|
||
let timeout_secs = std::env::var("LLAMA_SWAP_REQUEST_TIMEOUT_SECONDS")
|
||
.ok()
|
||
.and_then(|v| v.parse::<u64>().ok())
|
||
.unwrap_or(DEFAULT_REQUEST_TIMEOUT_SECS);
|
||
let pm = primary_model.unwrap_or_else(|| DEFAULT_PRIMARY_MODEL.to_string());
|
||
Self {
|
||
client: Client::builder()
|
||
.connect_timeout(Duration::from_secs(10))
|
||
.timeout(Duration::from_secs(timeout_secs))
|
||
.build()
|
||
.unwrap_or_else(|_| Client::new()),
|
||
base_url: base_url.unwrap_or_else(|| DEFAULT_BASE_URL.to_string()),
|
||
primary_model: pm.clone(),
|
||
embedding_model: DEFAULT_EMBEDDING_MODEL.to_string(),
|
||
vision_model: pm,
|
||
tts_model: DEFAULT_TTS_MODEL.to_string(),
|
||
num_ctx: None,
|
||
temperature: None,
|
||
top_p: None,
|
||
top_k: None,
|
||
min_p: None,
|
||
enable_thinking: None,
|
||
}
|
||
}
|
||
|
||
pub fn set_embedding_model(&mut self, model: String) {
|
||
self.embedding_model = model;
|
||
}
|
||
|
||
pub fn set_vision_model(&mut self, model: String) {
|
||
self.vision_model = model;
|
||
}
|
||
|
||
pub fn set_num_ctx(&mut self, num_ctx: Option<i32>) {
|
||
self.num_ctx = num_ctx;
|
||
}
|
||
|
||
/// Set the reasoning toggle forwarded as `chat_template_kwargs.enable_thinking`.
|
||
/// `None` leaves the chat template's own default in place.
|
||
pub fn set_enable_thinking(&mut self, enable_thinking: Option<bool>) {
|
||
self.enable_thinking = enable_thinking;
|
||
}
|
||
|
||
pub fn set_sampling_params(
|
||
&mut self,
|
||
temperature: Option<f32>,
|
||
top_p: Option<f32>,
|
||
top_k: Option<i32>,
|
||
min_p: Option<f32>,
|
||
) {
|
||
self.temperature = temperature;
|
||
self.top_p = top_p;
|
||
self.top_k = top_k;
|
||
self.min_p = min_p;
|
||
}
|
||
|
||
pub fn set_tts_model(&mut self, model: String) {
|
||
self.tts_model = model;
|
||
}
|
||
|
||
// --- TTS (Chatterbox behind llama-swap) ---------------------------------
|
||
//
|
||
// Speech synthesis uses the OpenAI-compatible `{base_url}/audio/speech`
|
||
// endpoint (llama-swap routes by the `model` field). The voice *library*
|
||
// (list / create cloned voices) is NOT an OpenAI endpoint — it lives on the
|
||
// upstream server directly, reached via llama-swap's passthrough at
|
||
// `{swap_root}/upstream/<tts_model>/voices`.
|
||
|
||
/// Per-request context window `model` is actually served with, read from
|
||
/// llama-server's `/props` (`default_generation_settings.n_ctx`, the
|
||
/// per-slot window) through llama-swap's passthrough. That request loads
|
||
/// the model if it isn't already, which is fine: callers ask right before
|
||
/// chatting with it. Successful lookups are cached for
|
||
/// `CONTEXT_WINDOW_CACHE_TTL` so a config.yaml edit is picked up without
|
||
/// a restart; failures aren't cached and return `None`.
|
||
pub async fn served_context_window(&self, model: &str) -> Option<usize> {
|
||
if let Some(n_ctx) = self.cached_context_window(model) {
|
||
return Some(n_ctx);
|
||
}
|
||
let key = format!("{}|{}", self.swap_root(), model);
|
||
let url = format!("{}/upstream/{}/props", self.swap_root(), model);
|
||
let props: Value = match self.client.get(&url).send().await {
|
||
Ok(resp) if resp.status().is_success() => resp.json().await.ok()?,
|
||
Ok(resp) => {
|
||
tracing::warn!(model, status = %resp.status(), "llama-swap props lookup failed");
|
||
return None;
|
||
}
|
||
Err(e) => {
|
||
tracing::warn!(model, error = %e, "llama-swap props lookup failed");
|
||
return None;
|
||
}
|
||
};
|
||
let n_ctx = props
|
||
.pointer("/default_generation_settings/n_ctx")
|
||
.or_else(|| props.get("n_ctx"))
|
||
.and_then(Value::as_u64)
|
||
.filter(|n| *n > 0)? as usize;
|
||
if let Ok(mut cache) = context_window_cache().lock() {
|
||
cache.insert(key, (n_ctx, std::time::Instant::now()));
|
||
}
|
||
Some(n_ctx)
|
||
}
|
||
|
||
/// Served context windows for `models`, for display, WITHOUT loading
|
||
/// anything: a model counts only if its window is cached from an
|
||
/// earlier lookup or llama-swap reports it as currently loaded (then
|
||
/// `/props` is answered by the running process). Others are absent.
|
||
pub async fn known_context_windows(&self, models: &[String]) -> HashMap<String, usize> {
|
||
let loaded: std::collections::HashSet<String> = async {
|
||
let url = format!("{}/models", self.base_url);
|
||
let parsed: Value = self.client.get(&url).send().await.ok()?.json().await.ok()?;
|
||
Some(
|
||
parsed
|
||
.get("data")?
|
||
.as_array()?
|
||
.iter()
|
||
.filter(|m| {
|
||
m.pointer("/status/value").and_then(Value::as_str) == Some("loaded")
|
||
})
|
||
.filter_map(|m| m.get("id")?.as_str().map(str::to_string))
|
||
.collect(),
|
||
)
|
||
}
|
||
.await
|
||
.unwrap_or_default();
|
||
|
||
let mut out = HashMap::new();
|
||
for model in models {
|
||
let cached = self.cached_context_window(model);
|
||
let n_ctx = match cached {
|
||
Some(n) => Some(n),
|
||
None if loaded.contains(model) => self.served_context_window(model).await,
|
||
None => None,
|
||
};
|
||
if let Some(n) = n_ctx {
|
||
out.insert(model.clone(), n);
|
||
}
|
||
}
|
||
out
|
||
}
|
||
|
||
fn cached_context_window(&self, model: &str) -> Option<usize> {
|
||
let key = format!("{}|{}", self.swap_root(), model);
|
||
let cache = context_window_cache().lock().ok()?;
|
||
let &(n_ctx, at) = cache.get(&key)?;
|
||
(at.elapsed() < CONTEXT_WINDOW_CACHE_TTL).then_some(n_ctx)
|
||
}
|
||
|
||
/// Root of the llama-swap proxy: `base_url` with a trailing `/v1` removed.
|
||
/// The `/upstream/...` passthrough lives here, not under `/v1`.
|
||
fn swap_root(&self) -> &str {
|
||
let b = self.base_url.trim_end_matches('/');
|
||
b.strip_suffix("/v1").unwrap_or(b)
|
||
}
|
||
|
||
/// Synthesize speech for `input` in an optional named `voice`, returning
|
||
/// the raw audio bytes (format per `response_format`, e.g. `"mp3"`/`"wav"`).
|
||
///
|
||
/// Chatterbox generation knobs are forwarded when set (caller is expected
|
||
/// to have range-clamped them): `exaggeration` (0.25–2.0, emotion),
|
||
/// `cfg_weight` (0.0–1.0, pace), `temperature` (0.05–5.0, randomness).
|
||
///
|
||
/// Callers must hold the GPU write lease (`ai::gpu::tts_lease`) across
|
||
/// this call. It is taken at the call sites in `ai::tts` rather than here
|
||
/// so the speech-job path can flip its job to `running` between acquiring
|
||
/// the GPU and sending the request.
|
||
pub async fn text_to_speech(
|
||
&self,
|
||
input: &str,
|
||
voice: Option<&str>,
|
||
response_format: &str,
|
||
exaggeration: Option<f32>,
|
||
cfg_weight: Option<f32>,
|
||
temperature: Option<f32>,
|
||
) -> Result<Vec<u8>> {
|
||
let url = format!("{}/audio/speech", self.base_url);
|
||
let mut body = json!({
|
||
"model": self.tts_model,
|
||
"input": input,
|
||
"response_format": response_format,
|
||
});
|
||
if let Some(v) = voice {
|
||
body["voice"] = Value::String(v.to_string());
|
||
}
|
||
if let Some(x) = exaggeration {
|
||
body["exaggeration"] = json!(x);
|
||
}
|
||
if let Some(x) = cfg_weight {
|
||
body["cfg_weight"] = json!(x);
|
||
}
|
||
if let Some(x) = temperature {
|
||
body["temperature"] = json!(x);
|
||
}
|
||
|
||
// TTS gets its own (longer) timeout: synthesizing a long, internally
|
||
// chunked insight can take minutes, well past the shared chat/embedding
|
||
// client timeout. Per-request `.timeout()` overrides the client default.
|
||
let tts_timeout = std::env::var("LLAMA_SWAP_TTS_REQUEST_TIMEOUT_SECONDS")
|
||
.ok()
|
||
.and_then(|v| v.parse::<u64>().ok())
|
||
.filter(|n| *n > 0)
|
||
.unwrap_or(600);
|
||
|
||
let resp = self
|
||
.client
|
||
.post(&url)
|
||
.timeout(Duration::from_secs(tts_timeout))
|
||
.json(&body)
|
||
.send()
|
||
.await
|
||
.with_context(|| format!("POST {} failed", url))?;
|
||
|
||
if !resp.status().is_success() {
|
||
let status = resp.status();
|
||
let text = resp.text().await.unwrap_or_default();
|
||
bail!("llama-swap TTS request failed: {} — {}", status, text);
|
||
}
|
||
|
||
Ok(resp
|
||
.bytes()
|
||
.await
|
||
.context("reading TTS audio bytes")?
|
||
.to_vec())
|
||
}
|
||
|
||
/// List voices in the Chatterbox voice library (raw JSON passthrough).
|
||
pub async fn list_voices(&self) -> Result<Value> {
|
||
let url = format!("{}/upstream/{}/voices", self.swap_root(), self.tts_model);
|
||
// The /upstream passthrough spins Chatterbox up (evicting the LLM),
|
||
// so it takes the exclusive GPU lease like synthesis does.
|
||
let _gpu = crate::ai::gpu::tts_lease().await;
|
||
let resp = self
|
||
.client
|
||
.get(&url)
|
||
.send()
|
||
.await
|
||
.with_context(|| format!("GET {} failed", url))?;
|
||
|
||
if !resp.status().is_success() {
|
||
let status = resp.status();
|
||
let text = resp.text().await.unwrap_or_default();
|
||
bail!("llama-swap list_voices failed: {} — {}", status, text);
|
||
}
|
||
resp.json().await.context("parsing voices response")
|
||
}
|
||
|
||
/// Register a cloned voice from raw audio bytes (multipart `voice_name` +
|
||
/// `voice_file`). Returns the upstream JSON response.
|
||
pub async fn create_voice(
|
||
&self,
|
||
voice_name: &str,
|
||
audio_bytes: Vec<u8>,
|
||
filename: &str,
|
||
mime: &str,
|
||
) -> Result<Value> {
|
||
let url = format!("{}/upstream/{}/voices", self.swap_root(), self.tts_model);
|
||
let part = reqwest::multipart::Part::bytes(audio_bytes)
|
||
.file_name(filename.to_string())
|
||
.mime_str(mime)
|
||
.context("invalid audio mime type")?;
|
||
let form = reqwest::multipart::Form::new()
|
||
.text("voice_name", voice_name.to_string())
|
||
.part("voice_file", part);
|
||
|
||
// The /upstream passthrough spins Chatterbox up (evicting the LLM),
|
||
// so it takes the exclusive GPU lease like synthesis does.
|
||
let _gpu = crate::ai::gpu::tts_lease().await;
|
||
let resp = self
|
||
.client
|
||
.post(&url)
|
||
.multipart(form)
|
||
.send()
|
||
.await
|
||
.with_context(|| format!("POST {} (multipart) failed", url))?;
|
||
|
||
if !resp.status().is_success() {
|
||
let status = resp.status();
|
||
let text = resp.text().await.unwrap_or_default();
|
||
bail!("llama-swap create_voice failed: {} — {}", status, text);
|
||
}
|
||
resp.json().await.context("parsing create_voice response")
|
||
}
|
||
|
||
/// Delete a cloned voice from the Chatterbox voice library
|
||
/// (`DELETE /voices/{name}` on the upstream, via llama-swap passthrough).
|
||
pub async fn delete_voice(&self, voice_name: &str) -> Result<Value> {
|
||
let url = format!(
|
||
"{}/upstream/{}/voices/{}",
|
||
self.swap_root(),
|
||
self.tts_model,
|
||
voice_name
|
||
);
|
||
// The /upstream passthrough spins Chatterbox up (evicting the LLM),
|
||
// so it takes the exclusive GPU lease like synthesis does.
|
||
let _gpu = crate::ai::gpu::tts_lease().await;
|
||
let resp = self
|
||
.client
|
||
.delete(&url)
|
||
.send()
|
||
.await
|
||
.with_context(|| format!("DELETE {} failed", url))?;
|
||
|
||
if !resp.status().is_success() {
|
||
let status = resp.status();
|
||
let text = resp.text().await.unwrap_or_default();
|
||
bail!("llama-swap delete_voice failed: {} — {}", status, text);
|
||
}
|
||
// Some upstreams reply with an empty body on delete.
|
||
Ok(resp
|
||
.json()
|
||
.await
|
||
.unwrap_or_else(|_| json!({ "status": "deleted" })))
|
||
}
|
||
|
||
/// Translate canonical messages to the OpenAI-compatible wire shape.
|
||
/// Behaviorally identical to `OpenRouterClient::messages_to_openai` —
|
||
/// stringify tool-call arguments, rewrite images into content-parts, attach
|
||
/// `tool_call_id` to `role=tool` messages based on the preceding assistant
|
||
/// turn's tool calls.
|
||
fn messages_to_openai(messages: &[ChatMessage]) -> Vec<Value> {
|
||
let mut out = Vec::with_capacity(messages.len());
|
||
let mut last_tool_call_ids: Vec<String> = Vec::new();
|
||
let mut next_tool_result_idx: usize = 0;
|
||
|
||
for msg in messages {
|
||
let mut obj = serde_json::Map::new();
|
||
obj.insert("role".into(), Value::String(msg.role.clone()));
|
||
|
||
// Assistant messages with tool_calls must emit `content: null`
|
||
// (not `""`) — some Jinja templates (Mistral-family) treat
|
||
// empty-string content as a regular message rather than a
|
||
// tool-calling turn, breaking role-alternation validation.
|
||
let has_tool_calls = msg.role == "assistant"
|
||
&& msg.tool_calls.as_ref().is_some_and(|tcs| !tcs.is_empty());
|
||
|
||
match &msg.images {
|
||
Some(images) if !images.is_empty() => {
|
||
let mut parts: Vec<Value> = Vec::new();
|
||
if !msg.content.is_empty() {
|
||
parts.push(json!({"type": "text", "text": msg.content}));
|
||
}
|
||
for img in images {
|
||
let url = image_to_data_url(img);
|
||
parts.push(json!({
|
||
"type": "image_url",
|
||
"image_url": { "url": url }
|
||
}));
|
||
}
|
||
obj.insert("content".into(), Value::Array(parts));
|
||
}
|
||
_ if has_tool_calls && msg.content.is_empty() => {
|
||
obj.insert("content".into(), Value::Null);
|
||
}
|
||
_ => {
|
||
obj.insert("content".into(), Value::String(msg.content.clone()));
|
||
}
|
||
}
|
||
|
||
if let Some(tcs) = &msg.tool_calls
|
||
&& msg.role == "assistant"
|
||
{
|
||
let converted: Vec<Value> = tcs
|
||
.iter()
|
||
.enumerate()
|
||
.map(|(i, call)| {
|
||
let id = call.id.clone().unwrap_or_else(|| format!("call_{}", i));
|
||
let args_str = serde_json::to_string(&call.function.arguments)
|
||
.unwrap_or_else(|_| "{}".to_string());
|
||
json!({
|
||
"id": id,
|
||
"type": "function",
|
||
"function": {
|
||
"name": call.function.name,
|
||
"arguments": args_str,
|
||
}
|
||
})
|
||
})
|
||
.collect();
|
||
last_tool_call_ids = converted
|
||
.iter()
|
||
.filter_map(|v| v.get("id").and_then(|x| x.as_str()).map(String::from))
|
||
.collect();
|
||
next_tool_result_idx = 0;
|
||
obj.insert("tool_calls".into(), Value::Array(converted));
|
||
}
|
||
|
||
if msg.role == "tool" {
|
||
let id = last_tool_call_ids
|
||
.get(next_tool_result_idx)
|
||
.cloned()
|
||
.unwrap_or_else(|| "call_0".to_string());
|
||
obj.insert("tool_call_id".into(), Value::String(id));
|
||
next_tool_result_idx += 1;
|
||
}
|
||
|
||
out.push(Value::Object(obj));
|
||
}
|
||
|
||
out
|
||
}
|
||
|
||
/// Parse an OpenAI-compatible assistant message back into canonical shape.
|
||
/// llama.cpp emits `reasoning_content` on thinking models; we drop it for
|
||
/// parity with OpenRouter (which also strips upstream reasoning fields).
|
||
fn openai_message_to_chat(msg: &Value) -> Result<ChatMessage> {
|
||
let obj = msg
|
||
.as_object()
|
||
.ok_or_else(|| anyhow!("response message is not an object"))?;
|
||
let role = obj
|
||
.get("role")
|
||
.and_then(|v| v.as_str())
|
||
.unwrap_or("assistant")
|
||
.to_string();
|
||
let content = obj
|
||
.get("content")
|
||
.and_then(|v| v.as_str())
|
||
.unwrap_or("")
|
||
.to_string();
|
||
|
||
let tool_calls = if let Some(tcs) = obj.get("tool_calls").and_then(|v| v.as_array()) {
|
||
let mut parsed = Vec::with_capacity(tcs.len());
|
||
for tc in tcs {
|
||
let id = tc.get("id").and_then(|v| v.as_str()).map(String::from);
|
||
let function = tc
|
||
.get("function")
|
||
.ok_or_else(|| anyhow!("tool_call missing function field"))?;
|
||
let name = function
|
||
.get("name")
|
||
.and_then(|v| v.as_str())
|
||
.unwrap_or_default()
|
||
.to_string();
|
||
let args_value = match function.get("arguments") {
|
||
Some(Value::String(s)) => {
|
||
serde_json::from_str::<Value>(s).unwrap_or_else(|_| json!({}))
|
||
}
|
||
Some(v @ Value::Object(_)) => v.clone(),
|
||
_ => json!({}),
|
||
};
|
||
parsed.push(ToolCall {
|
||
id,
|
||
function: ToolCallFunction {
|
||
name,
|
||
arguments: args_value,
|
||
},
|
||
});
|
||
}
|
||
Some(parsed)
|
||
} else {
|
||
None
|
||
};
|
||
|
||
Ok(ChatMessage {
|
||
role,
|
||
content,
|
||
tool_calls,
|
||
images: None,
|
||
})
|
||
}
|
||
|
||
fn build_options(&self) -> Vec<(&'static str, Value)> {
|
||
let mut v = Vec::new();
|
||
if let Some(t) = self.temperature {
|
||
v.push(("temperature", json!(t)));
|
||
}
|
||
if let Some(p) = self.top_p {
|
||
v.push(("top_p", json!(p)));
|
||
}
|
||
if let Some(k) = self.top_k {
|
||
v.push(("top_k", json!(k)));
|
||
}
|
||
if let Some(m) = self.min_p {
|
||
v.push(("min_p", json!(m)));
|
||
}
|
||
// num_ctx isn't an OpenAI param; llama-server bakes ctx in at launch
|
||
// via -c, so we silently drop the override here. The config.yaml
|
||
// entry is the source of truth for context size; chat budgets read
|
||
// it back via `context_window()`.
|
||
let _ = self.num_ctx;
|
||
// Reasoning toggle for thinking-capable templates (Qwen3 et al.).
|
||
// llama-server forwards chat_template_kwargs into the Jinja render
|
||
// (requires --jinja); templates that ignore the key are unaffected.
|
||
if let Some(think) = self.enable_thinking {
|
||
v.push(("chat_template_kwargs", json!({ "enable_thinking": think })));
|
||
}
|
||
v
|
||
}
|
||
|
||
/// Issue a chat request with an explicit model id override. Used by
|
||
/// `describe_image` to route through the vision slot without mutating
|
||
/// `self.primary_model`.
|
||
async fn chat_completion_with_model(
|
||
&self,
|
||
model: &str,
|
||
span: Option<&tracing::Span>,
|
||
messages: Vec<ChatMessage>,
|
||
tools: Vec<Tool>,
|
||
) -> Result<(ChatMessage, Option<i32>, Option<i32>)> {
|
||
let url = format!("{}/chat/completions", self.base_url);
|
||
let roles: Vec<&str> = messages.iter().map(|m| m.role.as_str()).collect();
|
||
log::debug!(
|
||
"llama-swap chat_completion: model={} roles={:?} tools={}",
|
||
model,
|
||
roles,
|
||
tools.len()
|
||
);
|
||
let mut body = serde_json::Map::new();
|
||
body.insert("model".into(), Value::String(model.to_string()));
|
||
body.insert(
|
||
"messages".into(),
|
||
Value::Array(Self::messages_to_openai(&messages)),
|
||
);
|
||
body.insert("stream".into(), Value::Bool(false));
|
||
if !tools.is_empty() {
|
||
body.insert(
|
||
"tools".into(),
|
||
serde_json::to_value(&tools).context("serializing tools")?,
|
||
);
|
||
}
|
||
for (k, v) in self.build_options() {
|
||
body.insert(k.into(), v);
|
||
}
|
||
|
||
// Wait for any TTS synthesis to release the GPU before the request
|
||
// timeout starts (see ai::gpu).
|
||
let _gpu = crate::ai::gpu::llm_lease().await;
|
||
let mut request = self.client.post(&url).json(&Value::Object(body));
|
||
if let Some(s) = span {
|
||
request = crate::otel::with_traceparent(request, s);
|
||
}
|
||
let resp = match request.send().await {
|
||
Ok(r) => r,
|
||
Err(e) => {
|
||
if let Some(s) = span {
|
||
s.set_status(Status::error("llama-swap request failed"));
|
||
}
|
||
return Err(e).with_context(|| format!("POST {} failed", url));
|
||
}
|
||
};
|
||
|
||
if !resp.status().is_success() {
|
||
let status = resp.status();
|
||
crate::otel::record_http_failure(
|
||
span,
|
||
status.as_u16(),
|
||
&format!("llama-swap chat request failed: {status}"),
|
||
);
|
||
let body = resp.text().await.unwrap_or_default();
|
||
bail!("llama-swap chat request failed: {} — {}", status, body);
|
||
}
|
||
crate::otel::set_http_status(span, resp.status().as_u16());
|
||
|
||
let parsed: Value = resp.json().await.context("parsing chat response")?;
|
||
let choice = parsed
|
||
.get("choices")
|
||
.and_then(|v| v.as_array())
|
||
.and_then(|a| a.first())
|
||
.ok_or_else(|| {
|
||
anyhow!(
|
||
"response missing choices[0]: {}",
|
||
extract_error_detail(&parsed)
|
||
)
|
||
})?;
|
||
let msg = choice.get("message").ok_or_else(|| {
|
||
anyhow!(
|
||
"choices[0] missing message: {}",
|
||
extract_error_detail(&parsed)
|
||
)
|
||
})?;
|
||
let chat_msg = Self::openai_message_to_chat(msg)?;
|
||
|
||
let usage = parsed.get("usage");
|
||
let prompt_tokens = usage
|
||
.and_then(|u| u.get("prompt_tokens"))
|
||
.and_then(|v| v.as_i64())
|
||
.map(|n| n as i32);
|
||
let completion_tokens = usage
|
||
.and_then(|u| u.get("completion_tokens"))
|
||
.and_then(|v| v.as_i64())
|
||
.map(|n| n as i32);
|
||
|
||
log_timings(&parsed, prompt_tokens, completion_tokens);
|
||
|
||
if let Some(s) = span {
|
||
if let Some(pt) = prompt_tokens {
|
||
s.set_attribute("llm.tokens.input", pt as i64);
|
||
}
|
||
if let Some(ct) = completion_tokens {
|
||
s.set_attribute("llm.tokens.output", ct as i64);
|
||
}
|
||
s.set_status(Status::Ok);
|
||
}
|
||
|
||
Ok((chat_msg, prompt_tokens, completion_tokens))
|
||
}
|
||
|
||
/// Start the `llm.request` span for an outgoing llama-swap call.
|
||
fn start_request_span(&self, operation: &str, model: &str) -> tracing::Span {
|
||
crate::otel::start_llm_request_span("llamacpp", operation, model)
|
||
}
|
||
}
|
||
|
||
#[async_trait]
|
||
impl LlmClient for LlamaCppClient {
|
||
async fn generate(
|
||
&self,
|
||
prompt: &str,
|
||
system: Option<&str>,
|
||
images: Option<Vec<String>>,
|
||
) -> Result<String> {
|
||
let mut messages: Vec<ChatMessage> = Vec::new();
|
||
if let Some(sys) = system {
|
||
messages.push(ChatMessage::system(sys));
|
||
}
|
||
let mut user = ChatMessage::user(prompt);
|
||
user.images = images;
|
||
messages.push(user);
|
||
|
||
// Delegates to chat_with_tools, which starts the llm.request span
|
||
// (operation=chat) — one span per HTTP request.
|
||
let (reply, _, _) = self.chat_with_tools(messages, Vec::new()).await?;
|
||
Ok(reply.content)
|
||
}
|
||
|
||
async fn chat_with_tools(
|
||
&self,
|
||
messages: Vec<ChatMessage>,
|
||
tools: Vec<Tool>,
|
||
) -> Result<(ChatMessage, Option<i32>, Option<i32>)> {
|
||
log::info!(
|
||
"llama-swap chat_with_tools: model={} messages={} tools={}",
|
||
self.primary_model,
|
||
messages.len(),
|
||
tools.len()
|
||
);
|
||
let span = self.start_request_span("chat", &self.primary_model);
|
||
self.chat_completion_with_model(&self.primary_model.clone(), Some(&span), messages, tools)
|
||
.await
|
||
}
|
||
|
||
async fn chat_with_tools_stream(
|
||
&self,
|
||
messages: Vec<ChatMessage>,
|
||
tools: Vec<Tool>,
|
||
) -> Result<BoxStream<'static, Result<LlmStreamEvent>>> {
|
||
// The span outlives this method — it moves into the stream
|
||
// generator and is finalized when the response drains.
|
||
let span = self.start_request_span("chat_stream", &self.primary_model);
|
||
let url = format!("{}/chat/completions", self.base_url);
|
||
let roles: Vec<&str> = messages.iter().map(|m| m.role.as_str()).collect();
|
||
log::debug!(
|
||
"llama-swap stream: model={} roles={:?} tools={}",
|
||
self.primary_model,
|
||
roles,
|
||
tools.len()
|
||
);
|
||
let mut body = serde_json::Map::new();
|
||
body.insert("model".into(), Value::String(self.primary_model.clone()));
|
||
body.insert(
|
||
"messages".into(),
|
||
Value::Array(Self::messages_to_openai(&messages)),
|
||
);
|
||
body.insert("stream".into(), Value::Bool(true));
|
||
body.insert(
|
||
"stream_options".into(),
|
||
serde_json::json!({ "include_usage": true }),
|
||
);
|
||
if !tools.is_empty() {
|
||
body.insert(
|
||
"tools".into(),
|
||
serde_json::to_value(&tools).context("serializing tools")?,
|
||
);
|
||
}
|
||
for (k, v) in self.build_options() {
|
||
body.insert(k.into(), v);
|
||
}
|
||
|
||
// Wait for any TTS synthesis to release the GPU before the request
|
||
// timeout starts (see ai::gpu). The guard is moved into the stream
|
||
// below so the lease spans the whole generation, not just the send.
|
||
let gpu = crate::ai::gpu::llm_lease().await;
|
||
let mut request = self.client.post(&url).json(&Value::Object(body));
|
||
request = crate::otel::with_traceparent(request, &span);
|
||
let resp = request
|
||
.send()
|
||
.await
|
||
.with_context(|| format!("POST {} failed", url))?;
|
||
|
||
if !resp.status().is_success() {
|
||
let status = resp.status();
|
||
crate::otel::record_http_failure(
|
||
Some(&span),
|
||
status.as_u16(),
|
||
&format!("llama-swap stream request failed: {status}"),
|
||
);
|
||
let body = resp.text().await.unwrap_or_default();
|
||
bail!("llama-swap stream request failed: {} — {}", status, body);
|
||
}
|
||
crate::otel::set_http_status(Some(&span), resp.status().as_u16());
|
||
|
||
let byte_stream = resp.bytes_stream();
|
||
let stream = async_stream::stream! {
|
||
let _gpu = gpu;
|
||
let mut byte_stream = byte_stream;
|
||
let mut buf: Vec<u8> = Vec::new();
|
||
let mut accumulated_content = String::new();
|
||
let mut tool_state = ToolCallAssembly::new();
|
||
let mut role = "assistant".to_string();
|
||
let mut prompt_tokens: Option<i32> = None;
|
||
let mut completion_tokens: Option<i32> = None;
|
||
let mut last_frame: Option<Value> = None;
|
||
let mut done_seen = false;
|
||
|
||
while let Some(chunk) = byte_stream.next().await {
|
||
let chunk = match chunk {
|
||
Ok(b) => b,
|
||
Err(e) => {
|
||
yield Err(anyhow!("stream read failed: {}", e));
|
||
return;
|
||
}
|
||
};
|
||
buf.extend_from_slice(&chunk);
|
||
|
||
while let Some(sep) = find_double_newline(&buf) {
|
||
let frame = buf.drain(..sep + 2).collect::<Vec<_>>();
|
||
let frame_str = match std::str::from_utf8(&frame) {
|
||
Ok(s) => s,
|
||
Err(_) => continue,
|
||
};
|
||
for line in frame_str.lines() {
|
||
let line = line.trim_end_matches('\r');
|
||
let payload = match line.strip_prefix("data: ") {
|
||
Some(p) => p,
|
||
None => continue,
|
||
};
|
||
if payload == "[DONE]" {
|
||
done_seen = true;
|
||
break;
|
||
}
|
||
let v: Value = match serde_json::from_str(payload) {
|
||
Ok(v) => v,
|
||
Err(e) => {
|
||
log::warn!(
|
||
"malformed llama-swap SSE frame: {} ({})",
|
||
payload,
|
||
e
|
||
);
|
||
continue;
|
||
}
|
||
};
|
||
|
||
if let Some(usage) = v.get("usage") {
|
||
prompt_tokens = usage
|
||
.get("prompt_tokens")
|
||
.and_then(|n| n.as_i64())
|
||
.map(|n| n as i32);
|
||
completion_tokens = usage
|
||
.get("completion_tokens")
|
||
.and_then(|n| n.as_i64())
|
||
.map(|n| n as i32);
|
||
last_frame = Some(v.clone());
|
||
}
|
||
|
||
let Some(choices) = v.get("choices").and_then(|c| c.as_array())
|
||
else {
|
||
continue;
|
||
};
|
||
let Some(choice) = choices.first() else { continue };
|
||
let delta = match choice.get("delta") {
|
||
Some(d) => d,
|
||
None => continue,
|
||
};
|
||
if let Some(r) = delta.get("role").and_then(|v| v.as_str()) {
|
||
role = r.to_string();
|
||
}
|
||
if let Some(content) =
|
||
delta.get("content").and_then(|v| v.as_str())
|
||
&& !content.is_empty()
|
||
{
|
||
accumulated_content.push_str(content);
|
||
yield Ok(LlmStreamEvent::TextDelta(content.to_string()));
|
||
}
|
||
if let Some(tcs) = delta.get("tool_calls").and_then(|v| v.as_array()) {
|
||
apply_tool_call_deltas(&mut tool_state, tcs);
|
||
}
|
||
}
|
||
if done_seen {
|
||
break;
|
||
}
|
||
}
|
||
if done_seen {
|
||
break;
|
||
}
|
||
}
|
||
|
||
let tool_calls = finalize_tool_calls(tool_state);
|
||
|
||
if let Some(ref frame) = last_frame {
|
||
log_timings(frame, prompt_tokens, completion_tokens);
|
||
}
|
||
|
||
let message = ChatMessage {
|
||
role,
|
||
content: accumulated_content,
|
||
tool_calls,
|
||
images: None,
|
||
};
|
||
if let Some(pt) = prompt_tokens {
|
||
span.set_attribute("llm.tokens.input", pt as i64);
|
||
}
|
||
if let Some(ct) = completion_tokens {
|
||
span.set_attribute("llm.tokens.output", ct as i64);
|
||
}
|
||
span.set_status(Status::Ok);
|
||
yield Ok(LlmStreamEvent::Done {
|
||
message,
|
||
prompt_eval_count: prompt_tokens,
|
||
eval_count: completion_tokens,
|
||
});
|
||
};
|
||
|
||
Ok(Box::pin(stream))
|
||
}
|
||
|
||
async fn generate_embeddings(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
|
||
let span = self.start_request_span("embeddings", &self.embedding_model);
|
||
// Deliberately NO GPU lease: the embed slot sits in llama-swap's
|
||
// always-resident group and never participates in a model swap, so
|
||
// leasing here would only stall searches behind a queued synthesis.
|
||
let url = format!("{}/embeddings", self.base_url);
|
||
let body = json!({
|
||
"model": self.embedding_model,
|
||
"input": texts,
|
||
});
|
||
|
||
let mut request = self.client.post(&url).json(&body);
|
||
request = crate::otel::with_traceparent(request, &span);
|
||
let resp = match request.send().await {
|
||
Ok(r) => r,
|
||
Err(e) => {
|
||
span.set_status(Status::error("llama-swap embedding request failed"));
|
||
return Err(e).with_context(|| format!("POST {} failed", url));
|
||
}
|
||
};
|
||
|
||
if !resp.status().is_success() {
|
||
let status = resp.status();
|
||
crate::otel::record_http_failure(
|
||
Some(&span),
|
||
status.as_u16(),
|
||
&format!("llama-swap embedding request failed: {status}"),
|
||
);
|
||
let body = resp.text().await.unwrap_or_default();
|
||
bail!("llama-swap embedding request failed: {} — {}", status, body);
|
||
}
|
||
crate::otel::set_http_status(Some(&span), resp.status().as_u16());
|
||
|
||
#[derive(Deserialize)]
|
||
struct EmbedResponse {
|
||
data: Vec<EmbedItem>,
|
||
}
|
||
#[derive(Deserialize)]
|
||
struct EmbedItem {
|
||
embedding: Vec<f32>,
|
||
}
|
||
|
||
let parsed: EmbedResponse = match resp.json().await {
|
||
Ok(p) => p,
|
||
Err(e) => {
|
||
span.set_status(Status::error("parsing embed response failed"));
|
||
return Err(e).context("parsing embed response");
|
||
}
|
||
};
|
||
span.set_status(Status::Ok);
|
||
Ok(parsed.data.into_iter().map(|i| i.embedding).collect())
|
||
}
|
||
|
||
async fn describe_image(&self, image_base64: &str) -> Result<String> {
|
||
// Own span with operation=describe_image — routes through the vision slot.
|
||
let span = self.start_request_span("describe_image", &self.vision_model);
|
||
let prompt = "Briefly describe what you see in this image in 1-2 sentences. \
|
||
Focus on the people, location, and activity.";
|
||
let system = "You are a scene description assistant. Be concise and factual.";
|
||
|
||
let messages = vec![
|
||
ChatMessage::system(system),
|
||
ChatMessage {
|
||
role: "user".to_string(),
|
||
content: prompt.to_string(),
|
||
tool_calls: None,
|
||
images: Some(vec![image_base64.to_string()]),
|
||
},
|
||
];
|
||
|
||
let (reply, _, _) = self
|
||
.chat_completion_with_model(
|
||
&self.vision_model.clone(),
|
||
Some(&span),
|
||
messages,
|
||
Vec::new(),
|
||
)
|
||
.await?;
|
||
Ok(reply.content)
|
||
}
|
||
|
||
async fn list_models(&self) -> Result<Vec<ModelCapabilities>> {
|
||
let url = format!("{}/models", self.base_url);
|
||
let resp = self
|
||
.client
|
||
.get(&url)
|
||
.send()
|
||
.await
|
||
.with_context(|| format!("GET {} failed", url))?;
|
||
|
||
if !resp.status().is_success() {
|
||
let status = resp.status();
|
||
let body = resp.text().await.unwrap_or_default();
|
||
bail!("llama-swap list_models failed: {} — {}", status, body);
|
||
}
|
||
|
||
let parsed: Value = resp.json().await.context("parsing models response")?;
|
||
let data = parsed
|
||
.get("data")
|
||
.and_then(|v| v.as_array())
|
||
.ok_or_else(|| anyhow!("models response missing data[]"))?;
|
||
|
||
let caps: Vec<ModelCapabilities> = data
|
||
.iter()
|
||
.map(|m| self.parse_model_capabilities(m))
|
||
.collect();
|
||
|
||
Ok(caps)
|
||
}
|
||
|
||
async fn model_capabilities(&self, model: &str) -> Result<ModelCapabilities> {
|
||
let all = self.list_models().await?;
|
||
all.into_iter()
|
||
.find(|m| m.name == model)
|
||
.ok_or_else(|| anyhow!("model '{}' not found on llama-swap", model))
|
||
}
|
||
|
||
fn primary_model(&self) -> &str {
|
||
&self.primary_model
|
||
}
|
||
|
||
/// llama-server fixes its window at launch (`-c`), so ask it via
|
||
/// llama-swap's `/upstream/<model>/props`. Cached per model; see
|
||
/// [`LlamaCppClient::served_context_window`].
|
||
async fn context_window(&self) -> Option<usize> {
|
||
self.served_context_window(&self.primary_model).await
|
||
}
|
||
}
|
||
|
||
impl LlamaCppClient {
|
||
fn parse_model_capabilities(&self, m: &Value) -> ModelCapabilities {
|
||
let name = m
|
||
.get("id")
|
||
.and_then(|v| v.as_str())
|
||
.unwrap_or_default()
|
||
.to_string();
|
||
// llama-swap doesn't expose per-model modality flags. Assume all
|
||
// configured models support vision and tool calling; a model
|
||
// without multimodal support surfaces as a chat-call error.
|
||
ModelCapabilities {
|
||
name,
|
||
has_vision: true,
|
||
has_tool_calling: true,
|
||
context_window: None,
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Extract a diagnostic fragment from a llama-swap / llama-server response
|
||
/// that doesn't match the expected `{choices: [...]}` shape. llama-server
|
||
/// returns errors as `{"error": {"message": "...", "code": N, "type": "..."}}`;
|
||
fn log_timings(parsed: &Value, prompt_tokens: Option<i32>, completion_tokens: Option<i32>) {
|
||
let timings = match parsed.get("timings") {
|
||
Some(t) => t,
|
||
None => return,
|
||
};
|
||
let prompt_tps = timings.get("prompt_per_second").and_then(|v| v.as_f64());
|
||
let gen_tps = timings.get("predicted_per_second").and_then(|v| v.as_f64());
|
||
let prompt_ms = timings.get("prompt_ms").and_then(|v| v.as_f64());
|
||
let gen_ms = timings.get("predicted_ms").and_then(|v| v.as_f64());
|
||
|
||
let mut parts: Vec<String> = Vec::new();
|
||
if let Some(c) = prompt_tokens {
|
||
let mut s = format!("prompt={} tok", c);
|
||
if let Some(ms) = prompt_ms {
|
||
s.push_str(&format!(" ({:.0} ms", ms));
|
||
if let Some(tps) = prompt_tps {
|
||
s.push_str(&format!(", {:.1} tok/s", tps));
|
||
}
|
||
s.push(')');
|
||
}
|
||
parts.push(s);
|
||
}
|
||
if let Some(c) = completion_tokens {
|
||
let mut s = format!("gen={} tok", c);
|
||
if let Some(ms) = gen_ms {
|
||
s.push_str(&format!(" ({:.0} ms", ms));
|
||
if let Some(tps) = gen_tps {
|
||
s.push_str(&format!(", {:.1} tok/s", tps));
|
||
}
|
||
s.push(')');
|
||
}
|
||
parts.push(s);
|
||
}
|
||
if !parts.is_empty() {
|
||
log::info!("llama-swap chat metrics — {}", parts.join(", "));
|
||
}
|
||
}
|
||
|
||
/// llama-swap itself sometimes wraps subprocess failures with its own
|
||
/// `{"error": "..."}` flat shape. Surface either when present, otherwise fall
|
||
/// back to a truncated raw-JSON view.
|
||
fn extract_error_detail(parsed: &Value) -> String {
|
||
if let Some(err) = parsed.get("error") {
|
||
match err {
|
||
Value::Object(_) => {
|
||
let message = err
|
||
.get("message")
|
||
.and_then(|v| v.as_str())
|
||
.unwrap_or("(no message)");
|
||
let code = err
|
||
.get("code")
|
||
.map(|v| match v {
|
||
Value::String(s) => s.clone(),
|
||
other => other.to_string(),
|
||
})
|
||
.unwrap_or_else(|| "?".to_string());
|
||
let short_message: String = message.chars().take(240).collect();
|
||
return format!("error code={} message=\"{}\"", code, short_message);
|
||
}
|
||
Value::String(s) => {
|
||
let short: String = s.chars().take(240).collect();
|
||
return format!("error=\"{}\"", short);
|
||
}
|
||
_ => {}
|
||
}
|
||
}
|
||
let raw = parsed.to_string();
|
||
raw.chars().take(300).collect()
|
||
}
|
||
|
||
/// Per-index assembly state for streamed OpenAI-style tool-call deltas:
|
||
/// `index → (id, name, concatenated argument fragments)`. BTreeMap so the
|
||
/// finalized calls come out in index order.
|
||
type ToolCallAssembly = std::collections::BTreeMap<usize, (Option<String>, Option<String>, String)>;
|
||
|
||
/// Fold one SSE frame's `delta.tool_calls` array into the assembly state.
|
||
/// Deltas carrying the same `index` merge into one call (llama.cpp streams a
|
||
/// call's argument JSON in fragments — they concatenate); distinct indexes
|
||
/// accumulate as separate calls.
|
||
fn apply_tool_call_deltas(state: &mut ToolCallAssembly, tcs: &[Value]) {
|
||
for tc_delta in tcs {
|
||
let idx = tc_delta.get("index").and_then(|n| n.as_u64()).unwrap_or(0) as usize;
|
||
let entry = state.entry(idx).or_insert((None, None, String::new()));
|
||
if let Some(id) = tc_delta.get("id").and_then(|v| v.as_str()) {
|
||
entry.0 = Some(id.to_string());
|
||
}
|
||
if let Some(func) = tc_delta.get("function") {
|
||
if let Some(name) = func.get("name").and_then(|v| v.as_str()) {
|
||
entry.1 = Some(name.to_string());
|
||
}
|
||
if let Some(args) = func.get("arguments").and_then(|v| v.as_str()) {
|
||
entry.2.push_str(args);
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Convert assembled tool-call state into canonical `ToolCall`s, parsing each
|
||
/// call's concatenated argument JSON (empty / malformed → `{}`). `None` when
|
||
/// no tool-call deltas arrived.
|
||
fn finalize_tool_calls(state: ToolCallAssembly) -> Option<Vec<ToolCall>> {
|
||
if state.is_empty() {
|
||
return None;
|
||
}
|
||
let mut v = Vec::with_capacity(state.len());
|
||
for (_idx, (id, name, args)) in state {
|
||
let arguments: Value = if args.trim().is_empty() {
|
||
Value::Object(Default::default())
|
||
} else {
|
||
serde_json::from_str(&args).unwrap_or_else(|_| Value::Object(Default::default()))
|
||
};
|
||
v.push(ToolCall {
|
||
id,
|
||
function: ToolCallFunction {
|
||
name: name.unwrap_or_default(),
|
||
arguments,
|
||
},
|
||
});
|
||
}
|
||
Some(v)
|
||
}
|
||
|
||
fn find_double_newline(buf: &[u8]) -> Option<usize> {
|
||
for i in 0..buf.len().saturating_sub(1) {
|
||
if buf[i] == b'\n' && buf[i + 1] == b'\n' {
|
||
return Some(i);
|
||
}
|
||
if i + 3 < buf.len()
|
||
&& buf[i] == b'\r'
|
||
&& buf[i + 1] == b'\n'
|
||
&& buf[i + 2] == b'\r'
|
||
&& buf[i + 3] == b'\n'
|
||
{
|
||
return Some(i + 1);
|
||
}
|
||
}
|
||
None
|
||
}
|
||
|
||
fn image_to_data_url(img: &str) -> String {
|
||
if img.starts_with("data:") {
|
||
img.to_string()
|
||
} else {
|
||
format!("data:image/jpeg;base64,{}", img)
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
/// One-shot HTTP server answering any request with `body`; returns its
|
||
/// base URL and a handle yielding the request line it received.
|
||
async fn serve_once(body: &'static str) -> (String, tokio::task::JoinHandle<String>) {
|
||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||
let addr = listener.local_addr().unwrap();
|
||
let handle = tokio::spawn(async move {
|
||
let (mut sock, _) = listener.accept().await.unwrap();
|
||
let mut buf = vec![0u8; 4096];
|
||
let n = sock.read(&mut buf).await.unwrap();
|
||
let req = String::from_utf8_lossy(&buf[..n])
|
||
.lines()
|
||
.next()
|
||
.unwrap_or("")
|
||
.to_string();
|
||
let resp = format!(
|
||
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
|
||
body.len(),
|
||
body
|
||
);
|
||
sock.write_all(resp.as_bytes()).await.unwrap();
|
||
req
|
||
});
|
||
(format!("http://{addr}/v1"), handle)
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn served_context_window_reads_props_through_upstream() {
|
||
let (base, req) =
|
||
serve_once(r#"{"default_generation_settings":{"n_ctx":32768},"total_slots":1}"#).await;
|
||
let client = LlamaCppClient::new(Some(base), Some("gemma4-31b-qat-uncensored".into()));
|
||
assert_eq!(client.context_window().await, Some(32768));
|
||
assert_eq!(
|
||
req.await.unwrap(),
|
||
"GET /upstream/gemma4-31b-qat-uncensored/props HTTP/1.1"
|
||
);
|
||
// Cached: a second lookup makes no request (the server is gone).
|
||
assert_eq!(client.context_window().await, Some(32768));
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn served_context_window_is_none_when_unreachable() {
|
||
let client = LlamaCppClient::new(Some("http://127.0.0.1:9/v1".into()), Some("nope".into()));
|
||
assert_eq!(client.context_window().await, None);
|
||
}
|
||
|
||
#[test]
|
||
fn tool_call_arguments_stringified_on_send() {
|
||
let msg = ChatMessage {
|
||
role: "assistant".into(),
|
||
content: String::new(),
|
||
tool_calls: Some(vec![ToolCall {
|
||
id: Some("call_abc".into()),
|
||
function: ToolCallFunction {
|
||
name: "search_sms".into(),
|
||
arguments: json!({"query": "hello", "limit": 5}),
|
||
},
|
||
}]),
|
||
images: None,
|
||
};
|
||
|
||
let wire = LlamaCppClient::messages_to_openai(&[msg]);
|
||
let tcs = wire[0]
|
||
.get("tool_calls")
|
||
.and_then(|v| v.as_array())
|
||
.expect("tool_calls present");
|
||
let args = tcs[0]
|
||
.get("function")
|
||
.and_then(|f| f.get("arguments"))
|
||
.and_then(|a| a.as_str())
|
||
.expect("arguments stringified");
|
||
let parsed: Value = serde_json::from_str(args).unwrap();
|
||
assert_eq!(parsed["query"], "hello");
|
||
assert_eq!(parsed["limit"], 5);
|
||
}
|
||
|
||
#[test]
|
||
fn tool_call_arguments_parsed_on_receive() {
|
||
let response_msg = json!({
|
||
"role": "assistant",
|
||
"content": "",
|
||
"tool_calls": [{
|
||
"id": "call_xyz",
|
||
"type": "function",
|
||
"function": {
|
||
"name": "get_weather",
|
||
"arguments": "{\"city\":\"Boston\",\"units\":\"celsius\"}"
|
||
}
|
||
}]
|
||
});
|
||
let parsed = LlamaCppClient::openai_message_to_chat(&response_msg).unwrap();
|
||
let tcs = parsed.tool_calls.unwrap();
|
||
assert_eq!(tcs.len(), 1);
|
||
assert_eq!(tcs[0].function.name, "get_weather");
|
||
assert_eq!(tcs[0].function.arguments["city"], "Boston");
|
||
assert_eq!(tcs[0].function.arguments["units"], "celsius");
|
||
assert_eq!(tcs[0].id.as_deref(), Some("call_xyz"));
|
||
}
|
||
|
||
#[test]
|
||
fn tool_call_arguments_accept_native_json_on_receive() {
|
||
// Some llama.cpp builds emit arguments as a JSON object directly when
|
||
// jinja's tool-output strict-string rule isn't applied — accept both.
|
||
let response_msg = json!({
|
||
"role": "assistant",
|
||
"content": "",
|
||
"tool_calls": [{
|
||
"id": "call_1",
|
||
"type": "function",
|
||
"function": {
|
||
"name": "foo",
|
||
"arguments": {"nested": {"k": 1}}
|
||
}
|
||
}]
|
||
});
|
||
let parsed = LlamaCppClient::openai_message_to_chat(&response_msg).unwrap();
|
||
let tc = &parsed.tool_calls.unwrap()[0];
|
||
assert_eq!(tc.function.arguments["nested"]["k"], 1);
|
||
}
|
||
|
||
#[test]
|
||
fn images_become_content_parts() {
|
||
let mut msg = ChatMessage::user("What is in this photo?");
|
||
msg.images = Some(vec!["BASE64DATA".into()]);
|
||
|
||
let wire = LlamaCppClient::messages_to_openai(&[msg]);
|
||
let content = wire[0].get("content").and_then(|v| v.as_array()).unwrap();
|
||
assert_eq!(content.len(), 2);
|
||
assert_eq!(content[0]["type"], "text");
|
||
assert_eq!(content[0]["text"], "What is in this photo?");
|
||
assert_eq!(content[1]["type"], "image_url");
|
||
assert_eq!(
|
||
content[1]["image_url"]["url"],
|
||
"data:image/jpeg;base64,BASE64DATA"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn data_url_images_pass_through_unchanged() {
|
||
let mut msg = ChatMessage::user("");
|
||
msg.images = Some(vec!["data:image/png;base64,ABCDEF".into()]);
|
||
let wire = LlamaCppClient::messages_to_openai(&[msg]);
|
||
let content = wire[0].get("content").and_then(|v| v.as_array()).unwrap();
|
||
assert_eq!(content.len(), 1);
|
||
assert_eq!(
|
||
content[0]["image_url"]["url"],
|
||
"data:image/png;base64,ABCDEF"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn text_only_message_stays_string() {
|
||
let msg = ChatMessage::user("hello");
|
||
let wire = LlamaCppClient::messages_to_openai(&[msg]);
|
||
assert_eq!(wire[0]["content"], "hello");
|
||
assert!(wire[0]["content"].as_str().is_some());
|
||
}
|
||
|
||
#[test]
|
||
fn tool_result_inherits_tool_call_id_from_prior_assistant() {
|
||
let assistant = ChatMessage {
|
||
role: "assistant".into(),
|
||
content: String::new(),
|
||
tool_calls: Some(vec![ToolCall {
|
||
id: Some("call_42".into()),
|
||
function: ToolCallFunction {
|
||
name: "lookup".into(),
|
||
arguments: json!({}),
|
||
},
|
||
}]),
|
||
images: None,
|
||
};
|
||
let tool_result = ChatMessage::tool_result("found it");
|
||
|
||
let wire = LlamaCppClient::messages_to_openai(&[assistant, tool_result]);
|
||
assert_eq!(wire[1]["role"], "tool");
|
||
assert_eq!(wire[1]["tool_call_id"], "call_42");
|
||
}
|
||
|
||
#[test]
|
||
fn multiple_tool_results_map_to_sequential_call_ids() {
|
||
let assistant = ChatMessage {
|
||
role: "assistant".into(),
|
||
content: String::new(),
|
||
tool_calls: Some(vec![
|
||
ToolCall {
|
||
id: Some("call_A".into()),
|
||
function: ToolCallFunction {
|
||
name: "a".into(),
|
||
arguments: json!({}),
|
||
},
|
||
},
|
||
ToolCall {
|
||
id: Some("call_B".into()),
|
||
function: ToolCallFunction {
|
||
name: "b".into(),
|
||
arguments: json!({}),
|
||
},
|
||
},
|
||
]),
|
||
images: None,
|
||
};
|
||
let r1 = ChatMessage::tool_result("a result");
|
||
let r2 = ChatMessage::tool_result("b result");
|
||
|
||
let wire = LlamaCppClient::messages_to_openai(&[assistant, r1, r2]);
|
||
assert_eq!(wire[1]["tool_call_id"], "call_A");
|
||
assert_eq!(wire[2]["tool_call_id"], "call_B");
|
||
}
|
||
|
||
#[test]
|
||
fn missing_tool_call_id_gets_synthetic_fallback() {
|
||
let assistant = ChatMessage {
|
||
role: "assistant".into(),
|
||
content: String::new(),
|
||
tool_calls: Some(vec![ToolCall {
|
||
id: None,
|
||
function: ToolCallFunction {
|
||
name: "noid".into(),
|
||
arguments: json!({}),
|
||
},
|
||
}]),
|
||
images: None,
|
||
};
|
||
let wire = LlamaCppClient::messages_to_openai(&[assistant]);
|
||
let tcs = wire[0]
|
||
.get("tool_calls")
|
||
.and_then(|v| v.as_array())
|
||
.unwrap();
|
||
assert_eq!(tcs[0]["id"], "call_0");
|
||
}
|
||
|
||
#[test]
|
||
fn all_models_report_vision_capable() {
|
||
let mut c = LlamaCppClient::new(None, Some("chat".into()));
|
||
c.set_vision_model("vision".into());
|
||
|
||
let m_chat = json!({ "id": "chat" });
|
||
let m_vision = json!({ "id": "vision" });
|
||
let m_other = json!({ "id": "embed" });
|
||
|
||
let chat = c.parse_model_capabilities(&m_chat);
|
||
let vision = c.parse_model_capabilities(&m_vision);
|
||
let other = c.parse_model_capabilities(&m_other);
|
||
|
||
assert!(chat.has_vision);
|
||
assert!(chat.has_tool_calling);
|
||
assert!(vision.has_vision);
|
||
assert!(other.has_vision);
|
||
}
|
||
|
||
#[test]
|
||
fn vision_model_defaults_to_primary() {
|
||
let c = LlamaCppClient::new(None, Some("qwen3:32b".into()));
|
||
assert_eq!(c.primary_model, "qwen3:32b");
|
||
assert_eq!(c.vision_model, "qwen3:32b");
|
||
}
|
||
|
||
#[test]
|
||
fn vision_model_explicit_override_diverges_from_primary() {
|
||
let mut c = LlamaCppClient::new(None, Some("qwen3:32b".into()));
|
||
c.set_vision_model("minicpm-v".into());
|
||
assert_eq!(c.primary_model, "qwen3:32b");
|
||
assert_eq!(c.vision_model, "minicpm-v");
|
||
}
|
||
|
||
#[test]
|
||
fn cloned_local_with_model_override_keeps_all_slots_consistent() {
|
||
// Simulates what resolve_backend does for the `local` client:
|
||
// clone the configured client, then override primary + vision
|
||
// to match the user-selected chat model. This prevents mid-turn
|
||
// model swaps in llama-swap exclusive mode.
|
||
let mut base = LlamaCppClient::new(None, Some("chat".into()));
|
||
base.set_vision_model("vision".into());
|
||
base.set_embedding_model("embed".into());
|
||
|
||
let mut local = base.clone();
|
||
let user_selected = "qwen3:32b";
|
||
local.primary_model = user_selected.to_string();
|
||
local.set_vision_model(user_selected.to_string());
|
||
|
||
// Chat generation (rerank) routes through primary_model.
|
||
assert_eq!(local.primary_model, user_selected);
|
||
// describe_image routes through vision_model.
|
||
assert_eq!(local.vision_model, user_selected);
|
||
// Embeddings stay on the dedicated slot — separate endpoint,
|
||
// no model swap conflict.
|
||
assert_eq!(local.embedding_model, "embed");
|
||
}
|
||
|
||
#[test]
|
||
fn hybrid_mode_local_client_preserves_vision_model() {
|
||
// In hybrid mode, overrides.model is an OpenRouter model id
|
||
// (e.g. "google/gemini-3-flash-preview"). The local llamacpp
|
||
// client must NOT adopt that as its vision_model — it should
|
||
// keep the configured LLAMA_SWAP_VISION_MODEL so describe_image
|
||
// hits the correct local slot instead of sending an unknown
|
||
// model name to llama-swap.
|
||
let mut base = LlamaCppClient::new(None, Some("chat".into()));
|
||
base.set_vision_model("vision".into());
|
||
base.set_embedding_model("embed".into());
|
||
|
||
// Simulate what resolve_backend SHOULD do for hybrid mode:
|
||
// clone but do NOT override primary_model / vision_model.
|
||
let local = base.clone();
|
||
|
||
// The local client keeps its configured slots.
|
||
assert_eq!(local.primary_model, "chat");
|
||
assert_eq!(local.vision_model, "vision");
|
||
assert_eq!(local.embedding_model, "embed");
|
||
}
|
||
|
||
#[test]
|
||
fn assistant_tool_calls_emit_null_content() {
|
||
let msg = ChatMessage {
|
||
role: "assistant".into(),
|
||
content: String::new(),
|
||
tool_calls: Some(vec![ToolCall {
|
||
id: Some("call_1".into()),
|
||
function: ToolCallFunction {
|
||
name: "search".into(),
|
||
arguments: json!({}),
|
||
},
|
||
}]),
|
||
images: None,
|
||
};
|
||
let wire = LlamaCppClient::messages_to_openai(&[msg]);
|
||
assert!(
|
||
wire[0]["content"].is_null(),
|
||
"empty content + tool_calls should emit null"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn assistant_with_content_and_tool_calls_preserves_content() {
|
||
let msg = ChatMessage {
|
||
role: "assistant".into(),
|
||
content: "Let me search for that.".into(),
|
||
tool_calls: Some(vec![ToolCall {
|
||
id: Some("call_1".into()),
|
||
function: ToolCallFunction {
|
||
name: "search".into(),
|
||
arguments: json!({}),
|
||
},
|
||
}]),
|
||
images: None,
|
||
};
|
||
let wire = LlamaCppClient::messages_to_openai(&[msg]);
|
||
assert_eq!(wire[0]["content"], "Let me search for that.");
|
||
}
|
||
|
||
#[test]
|
||
fn assistant_without_tool_calls_keeps_empty_string_content() {
|
||
let msg = ChatMessage {
|
||
role: "assistant".into(),
|
||
content: String::new(),
|
||
tool_calls: None,
|
||
images: None,
|
||
};
|
||
let wire = LlamaCppClient::messages_to_openai(&[msg]);
|
||
assert_eq!(wire[0]["content"], "");
|
||
}
|
||
|
||
#[test]
|
||
fn swap_root_strips_v1_suffix() {
|
||
let c = LlamaCppClient::new(Some("http://localhost:9292/v1".to_string()), None);
|
||
assert_eq!(c.swap_root(), "http://localhost:9292");
|
||
|
||
// Tolerates a trailing slash on the base URL.
|
||
let c2 = LlamaCppClient::new(Some("http://localhost:9292/v1/".to_string()), None);
|
||
assert_eq!(c2.swap_root(), "http://localhost:9292");
|
||
|
||
// No /v1 suffix → returned unchanged.
|
||
let c3 = LlamaCppClient::new(Some("http://host:1234".to_string()), None);
|
||
assert_eq!(c3.swap_root(), "http://host:1234");
|
||
}
|
||
|
||
#[test]
|
||
fn tts_model_defaults_to_chatterbox() {
|
||
let c = LlamaCppClient::new(None, None);
|
||
assert_eq!(c.tts_model, "chatterbox");
|
||
}
|
||
|
||
#[test]
|
||
fn stream_assembly_keeps_two_tool_calls_from_separate_chunks() {
|
||
// llama.cpp emits one delta per SSE frame; two calls with distinct
|
||
// `index` values arriving in separate frames must BOTH survive.
|
||
let mut state = ToolCallAssembly::new();
|
||
apply_tool_call_deltas(
|
||
&mut state,
|
||
&[json!({
|
||
"index": 0,
|
||
"id": "call_a",
|
||
"function": { "name": "get_sms_messages", "arguments": "{\"date\":\"2019-01-01\"}" }
|
||
})],
|
||
);
|
||
apply_tool_call_deltas(
|
||
&mut state,
|
||
&[json!({
|
||
"index": 1,
|
||
"id": "call_b",
|
||
"function": { "name": "reverse_geocode", "arguments": "{\"latitude\":1.0,\"longitude\":2.0}" }
|
||
})],
|
||
);
|
||
|
||
let calls = finalize_tool_calls(state).expect("two calls assembled");
|
||
assert_eq!(calls.len(), 2);
|
||
assert_eq!(calls[0].id.as_deref(), Some("call_a"));
|
||
assert_eq!(calls[0].function.name, "get_sms_messages");
|
||
assert_eq!(calls[0].function.arguments["date"], "2019-01-01");
|
||
assert_eq!(calls[1].id.as_deref(), Some("call_b"));
|
||
assert_eq!(calls[1].function.name, "reverse_geocode");
|
||
assert_eq!(calls[1].function.arguments["latitude"], 1.0);
|
||
}
|
||
|
||
#[test]
|
||
fn stream_assembly_concatenates_argument_fragments_for_same_index() {
|
||
// A single call's argument JSON streamed across frames concatenates
|
||
// into one parseable document.
|
||
let mut state = ToolCallAssembly::new();
|
||
apply_tool_call_deltas(
|
||
&mut state,
|
||
&[json!({
|
||
"index": 0,
|
||
"id": "call_x",
|
||
"function": { "name": "search_messages", "arguments": "{\"query\":" }
|
||
})],
|
||
);
|
||
apply_tool_call_deltas(
|
||
&mut state,
|
||
&[json!({
|
||
"index": 0,
|
||
"function": { "arguments": "\"dinner\"}" }
|
||
})],
|
||
);
|
||
|
||
let calls = finalize_tool_calls(state).expect("one call assembled");
|
||
assert_eq!(calls.len(), 1);
|
||
assert_eq!(calls[0].function.name, "search_messages");
|
||
assert_eq!(calls[0].function.arguments["query"], "dinner");
|
||
}
|
||
|
||
#[test]
|
||
fn stream_assembly_empty_state_finalizes_to_none() {
|
||
assert!(finalize_tool_calls(ToolCallAssembly::new()).is_none());
|
||
}
|
||
}
|