diff --git a/CLAUDE.md b/CLAUDE.md index e3edd51..9e6f3b1 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -798,10 +798,11 @@ TTS_PRONUNCIATIONS_PATH=tts_pronunciations.json # JSON map of pronunciation over # Insight Chat Continuation AGENTIC_CHAT_MAX_ITERATIONS=6 # Cap on tool-calling iterations per chat turn (default 6) AGENTIC_CHAT_DEFAULT_NUM_CTX=32768 # Assumed context window for the history-truncation budget - # when a chat request omits num_ctx (default 32768). Size to - # the smallest context among the chat models actually served; - # too small silently guts replayed history every turn (and - # destroys llama.cpp KV-cache prefix reuse). + # when a chat request omits num_ctx AND the backend can't + # report its served window (llama-swap reports it via /props, + # so this mostly matters for Ollama/OpenRouter). Too small + # silently guts replayed history every turn (and destroys + # llama.cpp KV-cache prefix reuse). ``` **AI Insights Fallback Behavior:** @@ -924,18 +925,29 @@ Persistence: Per-`(library_id, file_path)` async mutex (`AppState.insight_chat.chat_locks`) serialises concurrent turns on the same insight so the JSON blob doesn't race. -Context management is a soft bound: if the serialized history exceeds -`num_ctx - 2048` tokens (cheap 4-byte/token heuristic; `num_ctx` defaults -to `AGENTIC_CHAT_DEFAULT_NUM_CTX`, 32768, when the request omits it), the -oldest assistant-tool_call + tool_result pairs are dropped until under budget. The -initial user message (with any images) and system prompt are always preserved. +Context management (`ContextFitter` in `insight_chat.rs`): the history +budget is `min(num_ctx, served window) - 2048` tokens (cheap 4-byte/token +heuristic). The served window comes from `LlmClient::context_window()`: +llama-swap models fix it at launch (`-c`) and ignore the request's +`num_ctx`, so `LlamaCppClient` reads it from `/upstream//props` +(cached 10 min). When neither is known the budget uses +`AGENTIC_CHAT_DEFAULT_NUM_CTX` (32768). Over budget, `apply_context_budget` +drops, oldest first: tool-call scaffolding in completed exchanges, then +whole earlier exchanges (keeps user/assistant alternation), then the +current turn's own tool results. The system prompt, the first user message +(with any images) and the insight reply are always preserved, as is the +in-progress question. Every LLM call in the loop is re-fitted (tool results +grow mid-turn) on a copy, so the stored transcript is never trimmed. If the +server still rejects the prompt as too long, the budget is shrunk from its +reported token counts and the call retried once. The `truncated` event / flag is surfaced to the client when a drop occurred. Configurable env: - `AGENTIC_CHAT_MAX_ITERATIONS` — cap on tool-calling iterations per turn (default 6). Per-request `max_iterations` is clamped to this cap. - `AGENTIC_CHAT_DEFAULT_NUM_CTX` — assumed context window for the truncation - budget when the request omits `num_ctx` (default 32768). + budget when the request omits `num_ctx` and the backend can't report its + served window (default 32768). **Apollo Places integration (optional):** diff --git a/src/ai/handlers.rs b/src/ai/handlers.rs index 7482058..2be4e85 100644 --- a/src/ai/handlers.rs +++ b/src/ai/handlers.rs @@ -995,6 +995,8 @@ pub async fn generate_agentic_insight_handler( /// configured `LLAMA_SWAP_VISION_MODEL` slot id, and `has_tool_calling` is /// reported as true for every slot (llama-server is launched with `--jinja` /// by convention — a misconfigured slot surfaces as a chat-call error). +/// `context_window` is each slot's served `-c` window when known without +/// loading the model (cached from a recent chat, or the model is running). #[get("/insights/models")] pub async fn get_available_models_handler( _claims: Claims, @@ -1005,6 +1007,10 @@ pub async fn get_available_models_handler( if crate::ai::local_backend_is_llamacpp() && let Some(lc) = app_state.llamacpp.as_ref() { + // Served windows for the context picker; never loads a model. + let windows = lc + .known_context_windows(&app_state.llamacpp_allowed_models) + .await; let models: Vec = app_state .llamacpp_allowed_models .iter() @@ -1012,6 +1018,7 @@ pub async fn get_available_models_handler( name: name.clone(), has_vision: true, has_tool_calling: true, + context_window: windows.get(name).copied(), }) .collect(); let primary = ServerModels { diff --git a/src/ai/insight_chat.rs b/src/ai/insight_chat.rs index b7e119b..41c4599 100644 --- a/src/ai/insight_chat.rs +++ b/src/ai/insight_chat.rs @@ -21,12 +21,12 @@ use tracing_opentelemetry::OpenTelemetrySpanExt; use uuid::Uuid; pub const DEFAULT_MAX_ITERATIONS: usize = 6; -/// Assumed context window when the request doesn't specify `num_ctx`. -/// The llama-swap chat slots serve 20k-131k contexts and real conversations -/// rarely pass ~16k tokens, so 32k keeps the truncation pass from gutting -/// history that the server could comfortably hold (which also destroys the -/// server's KV-cache prefix reuse). Override per-deploy with -/// AGENTIC_CHAT_DEFAULT_NUM_CTX if the serving models change shape. +/// Assumed context window when the request doesn't specify `num_ctx` and +/// the backend can't report the window it serves (llama-swap can; see +/// `effective_num_ctx`). Real conversations rarely pass ~16k tokens, so 32k +/// keeps the truncation pass from gutting history the server could hold +/// (which also destroys the server's KV-cache prefix reuse). Override +/// per-deploy with AGENTIC_CHAT_DEFAULT_NUM_CTX. const DEFAULT_NUM_CTX: i32 = 32768; /// Headroom reserved for the model's response, deducted from the context /// budget when deciding whether to truncate the replayed history. @@ -323,12 +323,10 @@ impl InsightChatService { None }; - // 6. Apply truncation budget. Drops oldest tool_call+tool pairs - // (preserves system + first user including any images). - let budget_tokens = (req.num_ctx.unwrap_or_else(env_default_num_ctx) as usize) - .saturating_sub(RESPONSE_HEADROOM_TOKENS); - let budget_bytes = budget_tokens.saturating_mul(BYTES_PER_TOKEN); - let truncated = apply_context_budget(&mut messages, budget_bytes); + // 6. Apply truncation budget (see `apply_context_budget`), sized to + // the window the model is actually served with. + let mut fitter = ContextFitter::for_client(backend.chat(), req.num_ctx).await; + let truncated = fitter.trim(&mut messages); if truncated { span.set_attribute("history_truncated", true); } @@ -371,9 +369,8 @@ impl InsightChatService { iterations_used = iteration + 1; log::info!("Chat iteration {}/{}", iterations_used, max_iterations); - let (response, prompt_tokens, eval_tokens) = backend - .chat() - .chat_with_tools(messages.clone(), tools.clone()) + let (response, prompt_tokens, eval_tokens) = fitter + .chat(backend.chat(), &messages, tools.clone()) .instrument(loop_span.clone()) .await?; last_prompt_eval_count = prompt_tokens; @@ -429,9 +426,8 @@ impl InsightChatService { iterations_used ); let synthetic_idx = push_synthetic_final_prompt(&mut messages); - let (final_response, prompt_tokens, eval_tokens) = backend - .chat() - .chat_with_tools(messages.clone(), vec![]) + let (final_response, prompt_tokens, eval_tokens) = fitter + .chat(backend.chat(), &messages, vec![]) .instrument(loop_span.clone()) .await?; last_prompt_eval_count = prompt_tokens; @@ -1021,10 +1017,8 @@ impl InsightChatService { None }; - let budget_tokens = (req.num_ctx.unwrap_or_else(env_default_num_ctx) as usize) - .saturating_sub(RESPONSE_HEADROOM_TOKENS); - let budget_bytes = budget_tokens.saturating_mul(BYTES_PER_TOKEN); - let truncated = apply_context_budget(&mut messages, budget_bytes); + let mut fitter = ContextFitter::for_client(backend.chat(), req.num_ctx).await; + let truncated = fitter.trim(&mut messages); if truncated { let _ = entry.push_event(ChatStreamEvent::Truncated).await; } @@ -1050,6 +1044,7 @@ impl InsightChatService { &active_persona, max_iterations, &entry, + &mut fitter, ) .await?; let AgenticLoopOutcome { @@ -1255,6 +1250,7 @@ impl InsightChatService { user_msg.images = Some(vec![img.clone()]); } let mut messages = vec![system_msg, user_msg]; + let mut fitter = ContextFitter::for_client(backend.chat(), req.num_ctx).await; let outcome = self .run_streaming_agentic_loop_with_entry( @@ -1267,6 +1263,7 @@ impl InsightChatService { &active_persona, max_iterations, &entry, + &mut fitter, ) .await?; let AgenticLoopOutcome { @@ -1354,6 +1351,7 @@ impl InsightChatService { active_persona: &str, max_iterations: usize, entry: &Arc, + fitter: &mut ContextFitter, ) -> Result { crate::ai::insight_chat::run_streaming_agentic_loop_with_entry( &self.generator, @@ -1366,6 +1364,7 @@ impl InsightChatService { active_persona, max_iterations, entry, + fitter, ) .await } @@ -1494,10 +1493,8 @@ impl InsightChatService { }; // Truncate before appending the new user turn. - let budget_tokens = (req.num_ctx.unwrap_or_else(env_default_num_ctx) as usize) - .saturating_sub(RESPONSE_HEADROOM_TOKENS); - let budget_bytes = budget_tokens.saturating_mul(BYTES_PER_TOKEN); - let truncated = apply_context_budget(&mut messages, budget_bytes); + let mut fitter = ContextFitter::for_client(backend.chat(), req.num_ctx).await; + let truncated = fitter.trim(&mut messages); if truncated { let _ = tx.send(ChatStreamEvent::Truncated).await; } @@ -1520,6 +1517,7 @@ impl InsightChatService { &active_persona, max_iterations, &tx, + &mut fitter, ) .await?; let AgenticLoopOutcome { @@ -1740,6 +1738,7 @@ impl InsightChatService { user_msg.images = Some(vec![img.clone()]); } let mut messages = vec![system_msg, user_msg]; + let mut fitter = ContextFitter::for_client(backend.chat(), req.num_ctx).await; let outcome = self .run_streaming_agentic_loop( @@ -1752,6 +1751,7 @@ impl InsightChatService { &active_persona, max_iterations, &tx, + &mut fitter, ) .await?; let AgenticLoopOutcome { @@ -1833,6 +1833,7 @@ impl InsightChatService { active_persona: &str, max_iterations: usize, tx: &tokio::sync::mpsc::Sender, + fitter: &mut ContextFitter, ) -> Result { let mut tool_calls_made = 0usize; let mut iterations_used = 0usize; @@ -1849,9 +1850,8 @@ impl InsightChatService { }) .await; - let mut stream = backend - .chat() - .chat_with_tools_stream(messages.clone(), tools.clone()) + let mut stream = fitter + .chat_stream(backend.chat(), messages, tools.clone()) .await?; let mut final_message: Option = None; @@ -1941,10 +1941,7 @@ impl InsightChatService { // pending_tools at this position (wiping the prior tool // calls from the final assistant render). let synthetic_idx = push_synthetic_final_prompt(messages); - let mut stream = backend - .chat() - .chat_with_tools_stream(messages.clone(), vec![]) - .await?; + let mut stream = fitter.chat_stream(backend.chat(), messages, vec![]).await?; let mut final_message: Option = None; while let Some(ev) = stream.next().await { let ev = ev?; @@ -2270,6 +2267,145 @@ pub(crate) fn env_default_num_ctx() -> i32 { .max(RESPONSE_HEADROOM_TOKENS as i32 + 1024) } +/// Context window (tokens) to budget a chat turn against: the smaller of what +/// the request asked for and what the server actually serves. llama-swap +/// models fix their window at launch (`-c` in config.yaml) and ignore the +/// per-request `num_ctx`, so budgeting against a bigger requested value just +/// keeps more history than the server can hold. +pub(crate) async fn effective_num_ctx( + client: &dyn crate::ai::llm_client::LlmClient, + requested: Option, +) -> usize { + let requested = requested.filter(|n| *n > 0).map(|n| n as usize); + match (requested, client.context_window().await) { + (Some(r), Some(served)) => r.min(served), + (None, Some(served)) => served, + (Some(r), None) => r, + (None, None) => env_default_num_ctx() as usize, + } +} + +/// Keeps each LLM call of a chat turn inside the context budget. Every call +/// goes through [`ContextFitter::chat`] / [`ContextFitter::chat_stream`], +/// which send a trimmed copy of `messages` (the caller's vector — and so what +/// gets persisted — is untouched) and, if the server still rejects the prompt +/// as too long (the byte-per-token estimate is only a heuristic), shrink the +/// budget from the server's own token counts and retry once. The shrunk +/// budget sticks for the rest of the turn. +pub struct ContextFitter { + budget_bytes: usize, +} + +impl ContextFitter { + pub(crate) fn from_num_ctx(num_ctx: usize) -> Self { + let budget_tokens = num_ctx.saturating_sub(RESPONSE_HEADROOM_TOKENS); + Self { + budget_bytes: budget_tokens.saturating_mul(BYTES_PER_TOKEN), + } + } + + /// Budget for `client`'s model; see [`effective_num_ctx`]. + pub(crate) async fn for_client( + client: &dyn crate::ai::llm_client::LlmClient, + requested: Option, + ) -> Self { + Self::from_num_ctx(effective_num_ctx(client, requested).await) + } + + /// Trim `messages` in place to the budget. Used once on the replayed + /// history before the turn starts; returns true when anything was dropped. + pub(crate) fn trim(&self, messages: &mut Vec) -> bool { + apply_context_budget(messages, self.budget_bytes) + } + + fn fit(&self, messages: &[ChatMessage]) -> Vec { + let mut fitted = messages.to_vec(); + apply_context_budget(&mut fitted, self.budget_bytes); + fitted + } + + /// Shrink the budget after a context-overflow error. Returns false (and + /// leaves the budget alone) when `err` isn't one. + fn shrink_after(&mut self, err: &anyhow::Error, sent: &[ChatMessage]) -> bool { + let Some(counts) = parse_context_overflow(&format!("{err:#}")) else { + return false; + }; + let sent_bytes = estimate_bytes(sent).min(self.budget_bytes); + let shrunk = match counts { + // Scale what we sent by the server's measured overshoot, with 10% + // slack since trimming lands on message boundaries. + Some((prompt_tokens, n_ctx)) if prompt_tokens > 0 => { + let target = n_ctx.saturating_sub(RESPONSE_HEADROOM_TOKENS); + sent_bytes.saturating_mul(target) / prompt_tokens * 9 / 10 + } + _ => sent_bytes / 2, + }; + tracing::warn!( + old_budget_bytes = self.budget_bytes, + new_budget_bytes = shrunk, + "chat prompt exceeded the model's context window; trimming and retrying" + ); + self.budget_bytes = shrunk; + true + } + + pub(crate) async fn chat( + &mut self, + client: &dyn crate::ai::llm_client::LlmClient, + messages: &[ChatMessage], + tools: Vec, + ) -> Result<(ChatMessage, Option, Option)> { + let sent = self.fit(messages); + match client.chat_with_tools(sent.clone(), tools.clone()).await { + Err(e) if self.shrink_after(&e, &sent) => { + client.chat_with_tools(self.fit(messages), tools).await + } + result => result, + } + } + + pub(crate) async fn chat_stream( + &mut self, + client: &dyn crate::ai::llm_client::LlmClient, + messages: &[ChatMessage], + tools: Vec, + ) -> Result>> { + let sent = self.fit(messages); + match client + .chat_with_tools_stream(sent.clone(), tools.clone()) + .await + { + Err(e) if self.shrink_after(&e, &sent) => { + client + .chat_with_tools_stream(self.fit(messages), tools) + .await + } + result => result, + } + } +} + +/// Recognise a "prompt longer than the context window" error from any +/// backend. `None` = some other error. `Some(Some((prompt_tokens, n_ctx)))` +/// when the body carries llama-server's counts +/// (`exceed_context_size_error` with `n_prompt_tokens` / `n_ctx`), +/// `Some(None)` when it's an overflow without usable numbers. +fn parse_context_overflow(err: &str) -> Option> { + let lower = err.to_lowercase(); + let is_overflow = lower.contains("exceed_context_size") + || lower.contains("context_length_exceeded") + || lower.contains("maximum context length") + || (lower.contains("exceed") && lower.contains("context")); + if !is_overflow { + return None; + } + let num = |key: &str| -> Option { + let re = regex::Regex::new(&format!(r#""{key}"\s*:\s*(\d+)"#)).ok()?; + re.captures(err)?.get(1)?.as_str().parse().ok() + }; + Some(num("n_prompt_tokens").zip(num("n_ctx"))) +} + /// Append a per-turn iteration-budget reminder to the replayed system /// message so the model knows how many tool-calling rounds this turn gets. /// Returns the original `content` so the caller can restore it before @@ -2323,6 +2459,7 @@ pub async fn run_streaming_agentic_loop_with_entry( active_persona: &str, max_iterations: usize, entry: &Arc, + fitter: &mut ContextFitter, ) -> Result { let mut tool_calls_made = 0usize; let mut iterations_used = 0usize; @@ -2350,9 +2487,8 @@ pub async fn run_streaming_agentic_loop_with_entry( }) .await; - let mut stream = backend - .chat() - .chat_with_tools_stream(messages.clone(), tools.clone()) + let mut stream = fitter + .chat_stream(backend.chat(), messages, tools.clone()) .await?; let mut final_message: Option = None; @@ -2431,10 +2567,7 @@ pub async fn run_streaming_agentic_loop_with_entry( if final_content.is_empty() { let synthetic_idx = push_synthetic_final_prompt(messages); - let mut stream = backend - .chat() - .chat_with_tools_stream(messages.clone(), vec![]) - .await?; + let mut stream = fitter.chat_stream(backend.chat(), messages, vec![]).await?; let mut final_message: Option = None; while let Some(ev) = stream.next().await { let ev = ev?; @@ -2729,65 +2862,99 @@ pub(crate) fn truncate_tool_result(s: &str) -> (String, bool) { } } -/// Trim history to fit within `budget_bytes` of serialized JSON. Preserves -/// the system message and the first user message (with its base64 images -/// intact, since dropping those would invalidate the model's prior visual -/// reasoning). Drops the oldest assistant-tool_call + corresponding -/// tool-result pair on each pass until the budget is met or only the -/// preserved prefix remains. +/// Trim history to fit within `budget_bytes` of serialized JSON. Always +/// preserves the system message and the first user message (with its base64 +/// images intact, since dropping those would invalidate the model's prior +/// visual reasoning) plus that first exchange's final reply, i.e. the +/// insight itself. +/// +/// Drops, oldest first, in three phases until the budget is met: +/// 1. tool-call scaffolding (an assistant `tool_calls` turn + its tool +/// results) inside *completed* exchanges — cheapest to lose, the final +/// reply that used them stays; +/// 2. whole completed exchanges (user turn through its final reply) after +/// the first. Dropping a whole exchange keeps user/assistant alternation +/// intact, which strict chat templates (Gemma) enforce; +/// 3. tool-call scaffolding in the in-progress exchange (the current turn's +/// own tool results, which grow during the agentic loop). +/// +/// An exchange is "completed" when it contains an assistant reply without +/// tool calls. The current turn never is while the loop runs, so the user's +/// question is never dropped — even after `push_synthetic_final_prompt` +/// appends a later user message. /// /// Returns true when at least one message was dropped. pub(crate) fn apply_context_budget(messages: &mut Vec, budget_bytes: usize) -> bool { - if budget_bytes == 0 { + if budget_bytes == 0 || estimate_bytes(messages) <= budget_bytes { return false; } - if estimate_bytes(messages) <= budget_bytes { - return false; - } - - // Find the index past the protected prefix: system messages + the first - // user message. Everything after is droppable in pairs. - let first_user_idx = messages.iter().position(|m| m.role == "user"); - let preserve_through = match first_user_idx { - Some(i) => i, // keep [0..=i] - None => return false, - }; - let mut dropped_any = false; - loop { - if estimate_bytes(messages) <= budget_bytes { + while estimate_bytes(messages) > budget_bytes { + let Some(range) = next_budget_drop(messages) else { break; - } - // Find the oldest assistant-with-tool_calls strictly after the - // preserved prefix. Drop it together with the following tool turn(s) - // until we hit the next assistant or user turn. - let drop_start = (preserve_through + 1..messages.len()).find(|&i| { - let m = &messages[i]; - m.role == "assistant" - && m.tool_calls - .as_ref() - .map(|c| !c.is_empty()) - .unwrap_or(false) - }); - let Some(start) = drop_start else { break }; - // Determine end: drop the assistant turn plus any contiguous tool - // result turns that follow. - let mut end = start + 1; - while end < messages.len() && messages[end].role == "tool" { - end += 1; - } - // Stop if dropping these would leave the just-appended user turn at - // the end alone with no preceding context — we still want it kept. - if end > messages.len() { - break; - } - messages.drain(start..end); + }; + messages.drain(range); dropped_any = true; } - dropped_any } +/// Pick the next span of `messages` for [`apply_context_budget`] to drop, +/// following its phase order. `None` when nothing droppable remains. +fn next_budget_drop(messages: &[ChatMessage]) -> Option> { + let user_idxs: Vec = messages + .iter() + .enumerate() + .filter(|(_, m)| m.role == "user") + .map(|(i, _)| i) + .collect(); + let first_user = *user_idxs.first()?; + + // Exchanges: each user turn through the message before the next one. + let exchanges: Vec> = user_idxs + .iter() + .enumerate() + .map(|(n, &u)| u..user_idxs.get(n + 1).copied().unwrap_or(messages.len())) + .collect(); + let completed = |r: &std::ops::Range| { + messages[r.clone()] + .iter() + .any(|m| m.role == "assistant" && !has_tool_calls(m)) + }; + // Oldest tool-call turn (+ trailing tool results) inside `r`. + let tool_pair_in = |r: &std::ops::Range| { + let start = (r.start + 1..r.end) + .find(|&i| messages[i].role == "assistant" && has_tool_calls(&messages[i]))?; + let mut end = start + 1; + while end < r.end && messages[end].role == "tool" { + end += 1; + } + Some(start..end) + }; + + // Phase 1: tool scaffolding in completed exchanges. + if let Some(r) = exchanges + .iter() + .filter(|r| completed(r)) + .find_map(tool_pair_in) + { + return Some(r); + } + // Phase 2: whole completed exchanges after the first. + if let Some(r) = exchanges + .iter() + .find(|r| r.start != first_user && completed(r)) + { + return Some(r.clone()); + } + // Phase 3: tool scaffolding in the in-progress exchange(s). + exchanges.iter().find_map(tool_pair_in) +} + +fn has_tool_calls(m: &ChatMessage) -> bool { + m.tool_calls.as_ref().is_some_and(|c| !c.is_empty()) +} + /// Estimate the serialized byte size of `messages` for the truncation budget, /// EXCLUDING inlined base64 image payloads. Images are charged a flat /// `IMAGE_TOKENS_EACH` instead: their base64 is hundreds of KB of characters @@ -2913,6 +3080,207 @@ mod tests { assert!(!dropped); } + fn roles(msgs: &[ChatMessage]) -> Vec<&str> { + msgs.iter().map(|m| m.role.as_str()).collect() + } + + #[test] + fn truncation_drops_old_plain_exchanges_when_no_tool_pairs_left() { + // Long chat of plain Q&A: nothing but whole exchanges to drop. The + // old pass could only drop tool pairs, so this never shrank. + let filler = "x".repeat(2_000); + let mut msgs = vec![ + ChatMessage::system("sys"), + ChatMessage::user("first"), + assistant_text("the insight"), + ChatMessage::user(format!("q1 {filler}")), + assistant_text(&format!("a1 {filler}")), + ChatMessage::user(format!("q2 {filler}")), + assistant_text(&format!("a2 {filler}")), + ChatMessage::user("current question"), + ]; + let budget = estimate_bytes(&msgs) - 1_000; + assert!(apply_context_budget(&mut msgs, budget)); + assert!(estimate_bytes(&msgs) <= budget); + // Oldest exchange (q1/a1) went; the insight exchange, q2/a2 and the + // current question stay, and roles still alternate. + assert_eq!( + roles(&msgs), + ["system", "user", "assistant", "user", "assistant", "user"] + ); + assert_eq!(msgs[2].content, "the insight"); + assert!(msgs[3].content.starts_with("q2")); + assert_eq!(msgs[5].content, "current question"); + } + + #[test] + fn truncation_drops_tool_scaffolding_before_whole_exchanges() { + let mut msgs = vec![ + ChatMessage::system("sys"), + ChatMessage::user("first"), + assistant_text("the insight"), + ChatMessage::user("q1"), + assistant_with_tool_call("lookup"), + ChatMessage::tool_result("r".repeat(4_000)), + assistant_text("a1"), + ChatMessage::user("current question"), + ]; + let budget = estimate_bytes(&msgs) - 1_000; + assert!(apply_context_budget(&mut msgs, budget)); + // Only the tool pair was needed; q1/a1 survive. + assert_eq!( + roles(&msgs), + ["system", "user", "assistant", "user", "assistant", "user"] + ); + assert_eq!(msgs[4].content, "a1"); + } + + #[test] + fn truncation_never_drops_the_in_progress_question() { + // Mid-loop with the synthetic final prompt appended: the real + // question's exchange has no final reply yet, so phase 2 must skip + // it; only its tool results (phase 3) can go. + let mut msgs = vec![ + ChatMessage::system("sys"), + ChatMessage::user("first"), + assistant_text("the insight"), + ChatMessage::user("real question"), + assistant_with_tool_call("lookup"), + ChatMessage::tool_result("r".repeat(8_000)), + ChatMessage::user("synthetic: answer now"), + ]; + assert!(apply_context_budget(&mut msgs, 600)); + assert!(msgs.iter().any(|m| m.content == "real question")); + assert!(msgs.iter().all(|m| m.role != "tool")); + } + + #[test] + fn parse_context_overflow_reads_llama_server_counts() { + let err = r#"llama-swap chat request failed: 400 Bad Request — {"error":{"code":400,"message":"the request exceeds the available context size, try increasing it","type":"exceed_context_size_error","n_prompt_tokens":40312,"n_ctx":32768}}"#; + assert_eq!(parse_context_overflow(err), Some(Some((40312, 32768)))); + assert_eq!( + parse_context_overflow("400 — This model's maximum context length is 128000 tokens"), + Some(None) + ); + assert_eq!(parse_context_overflow("500 Internal Server Error"), None); + } + + /// Scripted client for `ContextFitter`: fails the first `fail_first` + /// calls with `error`, records the size of every prompt it's sent. + struct ScriptedClient { + window: Option, + fail_first: usize, + error: String, + sent: Mutex>, + } + + #[async_trait::async_trait] + impl crate::ai::llm_client::LlmClient for ScriptedClient { + async fn generate( + &self, + _: &str, + _: Option<&str>, + _: Option>, + ) -> Result { + unimplemented!() + } + async fn chat_with_tools( + &self, + messages: Vec, + _: Vec, + ) -> Result<(ChatMessage, Option, Option)> { + let mut sent = self.sent.lock().unwrap(); + sent.push(messages.len()); + if sent.len() <= self.fail_first { + bail!("{}", self.error); + } + Ok((assistant_text("ok"), None, None)) + } + async fn chat_with_tools_stream( + &self, + _: Vec, + _: Vec, + ) -> Result>> { + unimplemented!() + } + async fn generate_embeddings(&self, _: &[&str]) -> Result>> { + unimplemented!() + } + async fn describe_image(&self, _: &str) -> Result { + unimplemented!() + } + async fn list_models(&self) -> Result> { + unimplemented!() + } + async fn model_capabilities( + &self, + _: &str, + ) -> Result { + unimplemented!() + } + fn primary_model(&self) -> &str { + "scripted" + } + async fn context_window(&self) -> Option { + self.window + } + } + + fn scripted(window: Option, fail_first: usize, error: &str) -> ScriptedClient { + ScriptedClient { + window, + fail_first, + error: error.to_string(), + sent: Mutex::new(Vec::new()), + } + } + + #[tokio::test] + async fn effective_num_ctx_caps_request_at_served_window() { + let served_32k = scripted(Some(32_768), 0, ""); + assert_eq!(effective_num_ctx(&served_32k, Some(65_536)).await, 32_768); + assert_eq!(effective_num_ctx(&served_32k, Some(8_192)).await, 8_192); + assert_eq!(effective_num_ctx(&served_32k, None).await, 32_768); + let unknown = scripted(None, 0, ""); + assert_eq!(effective_num_ctx(&unknown, Some(65_536)).await, 65_536); + } + + #[tokio::test] + async fn fitter_retries_overflow_with_a_smaller_prompt() { + let filler = "x".repeat(2_000); + let mut msgs = vec![ + ChatMessage::system("sys"), + ChatMessage::user("first"), + assistant_text("insight"), + ]; + for i in 0..10 { + msgs.push(ChatMessage::user(format!("q{i} {filler}"))); + msgs.push(assistant_text(&format!("a{i} {filler}"))); + } + msgs.push(ChatMessage::user("current")); + let client = scripted( + None, + 1, + r#"400 — {"error":{"type":"exceed_context_size_error","n_prompt_tokens":20000,"n_ctx":10000}}"#, + ); + let mut fitter = ContextFitter::from_num_ctx(1_000_000); + let (reply, _, _) = fitter.chat(&client, &msgs, vec![]).await.unwrap(); + assert_eq!(reply.content, "ok"); + let sent = client.sent.lock().unwrap().clone(); + assert_eq!(sent.len(), 2, "one retry"); + assert_eq!(sent[0], msgs.len(), "first attempt fit the generous budget"); + assert!(sent[1] < sent[0], "retry sent a trimmed prompt"); + } + + #[tokio::test] + async fn fitter_does_not_retry_other_errors() { + let client = scripted(None, 5, "502 Bad Gateway"); + let mut fitter = ContextFitter::from_num_ctx(32_768); + let msgs = vec![ChatMessage::system("s"), ChatMessage::user("u")]; + assert!(fitter.chat(&client, &msgs, vec![]).await.is_err()); + assert_eq!(client.sent.lock().unwrap().len(), 1); + } + #[test] fn render_tree_path_detects_fork_at_tool_dispatch_child() { // Regression: a regeneration where the model replies with a tool call diff --git a/src/ai/llamacpp.rs b/src/ai/llamacpp.rs index b69cf0f..9bcd24d 100644 --- a/src/ai/llamacpp.rs +++ b/src/ai/llamacpp.rs @@ -27,6 +27,7 @@ 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 _; @@ -41,6 +42,18 @@ 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); + +/// `"|"` -> (n_ctx, fetched at). Process-wide because +/// clients are cloned per request. +fn context_window_cache() -> &'static std::sync::Mutex> +{ + static CACHE: std::sync::OnceLock< + std::sync::Mutex>, + > = 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. @@ -145,6 +158,86 @@ impl LlamaCppClient { // upstream server directly, reached via llama-swap's passthrough at // `{swap_root}/upstream//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 { + 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 { + let loaded: std::collections::HashSet = 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 { + 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 { @@ -472,7 +565,8 @@ impl LlamaCppClient { } // 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. + // 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 @@ -939,6 +1033,13 @@ impl LlmClient for LlamaCppClient { fn primary_model(&self) -> &str { &self.primary_model } + + /// llama-server fixes its window at launch (`-c`), so ask it via + /// llama-swap's `/upstream//props`. Cached per model; see + /// [`LlamaCppClient::served_context_window`]. + async fn context_window(&self) -> Option { + self.served_context_window(&self.primary_model).await + } } impl LlamaCppClient { @@ -955,6 +1056,7 @@ impl LlamaCppClient { name, has_vision: true, has_tool_calling: true, + context_window: None, } } } @@ -1113,6 +1215,52 @@ fn image_to_data_url(img: &str) -> String { 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) { + 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 { diff --git a/src/ai/llm_client.rs b/src/ai/llm_client.rs index c8b0966..657c2fa 100644 --- a/src/ai/llm_client.rs +++ b/src/ai/llm_client.rs @@ -63,6 +63,14 @@ pub trait LlmClient: Send + Sync { /// Primary model identifier this client was constructed with. fn primary_model(&self) -> &str; + + /// Context window (tokens) the server actually serves for the primary + /// model, when the backend can report one that overrides the request's + /// `num_ctx`. `None` = unknown or request-controlled (Ollama honours + /// `num_ctx` per request); callers fall back to the requested value. + async fn context_window(&self) -> Option { + None + } } /// Events emitted by streaming `chat_with_tools_stream`. A stream is a @@ -174,6 +182,11 @@ pub struct ModelCapabilities { pub name: String, pub has_vision: bool, pub has_tool_calling: bool, + /// Context window (tokens) the model is served with, when the server + /// fixes it regardless of the request's `num_ctx` (llama-swap) and it's + /// known without loading the model. Omitted otherwise. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub context_window: Option, } /// A single node in the conversation tree. Wraps ChatMessage with structural diff --git a/src/ai/ollama.rs b/src/ai/ollama.rs index 335b047..bf7c454 100644 --- a/src/ai/ollama.rs +++ b/src/ai/ollama.rs @@ -315,6 +315,7 @@ impl OllamaClient { name: model_name.to_string(), has_vision, has_tool_calling, + context_window: None, }) } @@ -346,6 +347,7 @@ impl OllamaClient { name: model_name, has_vision: false, has_tool_calling: false, + context_window: None, }); } } diff --git a/src/ai/openrouter.rs b/src/ai/openrouter.rs index 01e226b..3c450ee 100644 --- a/src/ai/openrouter.rs +++ b/src/ai/openrouter.rs @@ -901,6 +901,7 @@ fn parse_model_capabilities(m: &Value) -> ModelCapabilities { name, has_vision, has_tool_calling, + context_window: None, } } diff --git a/src/ai/persona_chat.rs b/src/ai/persona_chat.rs index 83eda37..73be4dd 100644 --- a/src/ai/persona_chat.rs +++ b/src/ai/persona_chat.rs @@ -37,8 +37,8 @@ use uuid::Uuid; use crate::ai::backend::{BackendKind, SamplingOverrides}; use crate::ai::handlers::ReplayQuery; use crate::ai::insight_chat::{ - BYTES_PER_TOKEN, ChatStreamEvent, DEFAULT_MAX_ITERATIONS, ForkInfo, RESPONSE_HEADROOM_TOKENS, - apply_context_budget, env_default_num_ctx, env_max_iterations, render_tree_path, + ChatStreamEvent, ContextFitter, DEFAULT_MAX_ITERATIONS, ForkInfo, env_max_iterations, + render_tree_path, }; use crate::ai::insight_generator::InsightGenerator; use crate::ai::llm_client::{ChatHistoryStore, ChatMessage, StoredChatNode}; @@ -842,10 +842,8 @@ impl PersonaChatSession { // Trim to the model's context window before the turn. Without this a // rolling persona transcript grows unbounded until the backend // silently drops the front of it mid-conversation. - let budget_tokens = (req.num_ctx.unwrap_or_else(env_default_num_ctx) as usize) - .saturating_sub(RESPONSE_HEADROOM_TOKENS); - let budget_bytes = budget_tokens.saturating_mul(BYTES_PER_TOKEN); - let truncated = apply_context_budget(&mut messages, budget_bytes); + let mut fitter = ContextFitter::for_client(backend.chat(), req.num_ctx).await; + let truncated = fitter.trim(&mut messages); if truncated { let _ = entry.push_event(ChatStreamEvent::Truncated).await; } @@ -880,6 +878,7 @@ impl PersonaChatSession { &persona.persona_id, max_iterations, &entry, + &mut fitter, ) .await .map_err(|e| anyhow!("persona chat loop failed: {e:?}"))?;