fix(chat): budget chat history against the model's served context window
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>
This commit is contained in:
1 parent
1d2848db05
commit
305eda5894
8 files changed
+654
-104
No files matched your search
@@ -798,10 +798,11 @@ TTS_PRONUNCIATIONS_PATH=tts_pronunciations.json # JSON map of pronunciation over
|
|||||||
# Insight Chat Continuation
|
# Insight Chat Continuation
|
||||||
AGENTIC_CHAT_MAX_ITERATIONS=6 # Cap on tool-calling iterations per chat turn (default 6)
|
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
|
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
|
# when a chat request omits num_ctx AND the backend can't
|
||||||
# the smallest context among the chat models actually served;
|
# report its served window (llama-swap reports it via /props,
|
||||||
# too small silently guts replayed history every turn (and
|
# so this mostly matters for Ollama/OpenRouter). Too small
|
||||||
# destroys llama.cpp KV-cache prefix reuse).
|
# silently guts replayed history every turn (and destroys
|
||||||
|
# llama.cpp KV-cache prefix reuse).
|
||||||
```
|
```
|
||||||
|
|
||||||
**AI Insights Fallback Behavior:**
|
**AI Insights Fallback Behavior:**
|
||||||
@@ -924,18 +925,29 @@ Persistence:
|
|||||||
Per-`(library_id, file_path)` async mutex (`AppState.insight_chat.chat_locks`)
|
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.
|
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
|
Context management (`ContextFitter` in `insight_chat.rs`): the history
|
||||||
`num_ctx - 2048` tokens (cheap 4-byte/token heuristic; `num_ctx` defaults
|
budget is `min(num_ctx, served window) - 2048` tokens (cheap 4-byte/token
|
||||||
to `AGENTIC_CHAT_DEFAULT_NUM_CTX`, 32768, when the request omits it), the
|
heuristic). The served window comes from `LlmClient::context_window()`:
|
||||||
oldest assistant-tool_call + tool_result pairs are dropped until under budget. The
|
llama-swap models fix it at launch (`-c`) and ignore the request's
|
||||||
initial user message (with any images) and system prompt are always preserved.
|
`num_ctx`, so `LlamaCppClient` reads it from `/upstream/<model>/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.
|
The `truncated` event / flag is surfaced to the client when a drop occurred.
|
||||||
|
|
||||||
Configurable env:
|
Configurable env:
|
||||||
- `AGENTIC_CHAT_MAX_ITERATIONS` — cap on tool-calling iterations per turn
|
- `AGENTIC_CHAT_MAX_ITERATIONS` — cap on tool-calling iterations per turn
|
||||||
(default 6). Per-request `max_iterations` is clamped to this cap.
|
(default 6). Per-request `max_iterations` is clamped to this cap.
|
||||||
- `AGENTIC_CHAT_DEFAULT_NUM_CTX` — assumed context window for the truncation
|
- `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):**
|
**Apollo Places integration (optional):**
|
||||||
|
|
||||||
|
|||||||
@@ -995,6 +995,8 @@ pub async fn generate_agentic_insight_handler(
|
|||||||
/// configured `LLAMA_SWAP_VISION_MODEL` slot id, and `has_tool_calling` is
|
/// configured `LLAMA_SWAP_VISION_MODEL` slot id, and `has_tool_calling` is
|
||||||
/// reported as true for every slot (llama-server is launched with `--jinja`
|
/// reported as true for every slot (llama-server is launched with `--jinja`
|
||||||
/// by convention — a misconfigured slot surfaces as a chat-call error).
|
/// 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")]
|
#[get("/insights/models")]
|
||||||
pub async fn get_available_models_handler(
|
pub async fn get_available_models_handler(
|
||||||
_claims: Claims,
|
_claims: Claims,
|
||||||
@@ -1005,6 +1007,10 @@ pub async fn get_available_models_handler(
|
|||||||
if crate::ai::local_backend_is_llamacpp()
|
if crate::ai::local_backend_is_llamacpp()
|
||||||
&& let Some(lc) = app_state.llamacpp.as_ref()
|
&& 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<ModelCapabilities> = app_state
|
let models: Vec<ModelCapabilities> = app_state
|
||||||
.llamacpp_allowed_models
|
.llamacpp_allowed_models
|
||||||
.iter()
|
.iter()
|
||||||
@@ -1012,6 +1018,7 @@ pub async fn get_available_models_handler(
|
|||||||
name: name.clone(),
|
name: name.clone(),
|
||||||
has_vision: true,
|
has_vision: true,
|
||||||
has_tool_calling: true,
|
has_tool_calling: true,
|
||||||
|
context_window: windows.get(name).copied(),
|
||||||
})
|
})
|
||||||
.collect();
|
.collect();
|
||||||
let primary = ServerModels {
|
let primary = ServerModels {
|
||||||
|
|||||||
+455
-87
@@ -21,12 +21,12 @@ use tracing_opentelemetry::OpenTelemetrySpanExt;
|
|||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
|
|
||||||
pub const DEFAULT_MAX_ITERATIONS: usize = 6;
|
pub const DEFAULT_MAX_ITERATIONS: usize = 6;
|
||||||
/// Assumed context window when the request doesn't specify `num_ctx`.
|
/// Assumed context window when the request doesn't specify `num_ctx` and
|
||||||
/// The llama-swap chat slots serve 20k-131k contexts and real conversations
|
/// the backend can't report the window it serves (llama-swap can; see
|
||||||
/// rarely pass ~16k tokens, so 32k keeps the truncation pass from gutting
|
/// `effective_num_ctx`). Real conversations rarely pass ~16k tokens, so 32k
|
||||||
/// history that the server could comfortably hold (which also destroys the
|
/// keeps the truncation pass from gutting history the server could hold
|
||||||
/// server's KV-cache prefix reuse). Override per-deploy with
|
/// (which also destroys the server's KV-cache prefix reuse). Override
|
||||||
/// AGENTIC_CHAT_DEFAULT_NUM_CTX if the serving models change shape.
|
/// per-deploy with AGENTIC_CHAT_DEFAULT_NUM_CTX.
|
||||||
const DEFAULT_NUM_CTX: i32 = 32768;
|
const DEFAULT_NUM_CTX: i32 = 32768;
|
||||||
/// Headroom reserved for the model's response, deducted from the context
|
/// Headroom reserved for the model's response, deducted from the context
|
||||||
/// budget when deciding whether to truncate the replayed history.
|
/// budget when deciding whether to truncate the replayed history.
|
||||||
@@ -323,12 +323,10 @@ impl InsightChatService {
|
|||||||
None
|
None
|
||||||
};
|
};
|
||||||
|
|
||||||
// 6. Apply truncation budget. Drops oldest tool_call+tool pairs
|
// 6. Apply truncation budget (see `apply_context_budget`), sized to
|
||||||
// (preserves system + first user including any images).
|
// the window the model is actually served with.
|
||||||
let budget_tokens = (req.num_ctx.unwrap_or_else(env_default_num_ctx) as usize)
|
let mut fitter = ContextFitter::for_client(backend.chat(), req.num_ctx).await;
|
||||||
.saturating_sub(RESPONSE_HEADROOM_TOKENS);
|
let truncated = fitter.trim(&mut messages);
|
||||||
let budget_bytes = budget_tokens.saturating_mul(BYTES_PER_TOKEN);
|
|
||||||
let truncated = apply_context_budget(&mut messages, budget_bytes);
|
|
||||||
if truncated {
|
if truncated {
|
||||||
span.set_attribute("history_truncated", true);
|
span.set_attribute("history_truncated", true);
|
||||||
}
|
}
|
||||||
@@ -371,9 +369,8 @@ impl InsightChatService {
|
|||||||
iterations_used = iteration + 1;
|
iterations_used = iteration + 1;
|
||||||
log::info!("Chat iteration {}/{}", iterations_used, max_iterations);
|
log::info!("Chat iteration {}/{}", iterations_used, max_iterations);
|
||||||
|
|
||||||
let (response, prompt_tokens, eval_tokens) = backend
|
let (response, prompt_tokens, eval_tokens) = fitter
|
||||||
.chat()
|
.chat(backend.chat(), &messages, tools.clone())
|
||||||
.chat_with_tools(messages.clone(), tools.clone())
|
|
||||||
.instrument(loop_span.clone())
|
.instrument(loop_span.clone())
|
||||||
.await?;
|
.await?;
|
||||||
last_prompt_eval_count = prompt_tokens;
|
last_prompt_eval_count = prompt_tokens;
|
||||||
@@ -429,9 +426,8 @@ impl InsightChatService {
|
|||||||
iterations_used
|
iterations_used
|
||||||
);
|
);
|
||||||
let synthetic_idx = push_synthetic_final_prompt(&mut messages);
|
let synthetic_idx = push_synthetic_final_prompt(&mut messages);
|
||||||
let (final_response, prompt_tokens, eval_tokens) = backend
|
let (final_response, prompt_tokens, eval_tokens) = fitter
|
||||||
.chat()
|
.chat(backend.chat(), &messages, vec![])
|
||||||
.chat_with_tools(messages.clone(), vec![])
|
|
||||||
.instrument(loop_span.clone())
|
.instrument(loop_span.clone())
|
||||||
.await?;
|
.await?;
|
||||||
last_prompt_eval_count = prompt_tokens;
|
last_prompt_eval_count = prompt_tokens;
|
||||||
@@ -1021,10 +1017,8 @@ impl InsightChatService {
|
|||||||
None
|
None
|
||||||
};
|
};
|
||||||
|
|
||||||
let budget_tokens = (req.num_ctx.unwrap_or_else(env_default_num_ctx) as usize)
|
let mut fitter = ContextFitter::for_client(backend.chat(), req.num_ctx).await;
|
||||||
.saturating_sub(RESPONSE_HEADROOM_TOKENS);
|
let truncated = fitter.trim(&mut messages);
|
||||||
let budget_bytes = budget_tokens.saturating_mul(BYTES_PER_TOKEN);
|
|
||||||
let truncated = apply_context_budget(&mut messages, budget_bytes);
|
|
||||||
if truncated {
|
if truncated {
|
||||||
let _ = entry.push_event(ChatStreamEvent::Truncated).await;
|
let _ = entry.push_event(ChatStreamEvent::Truncated).await;
|
||||||
}
|
}
|
||||||
@@ -1050,6 +1044,7 @@ impl InsightChatService {
|
|||||||
&active_persona,
|
&active_persona,
|
||||||
max_iterations,
|
max_iterations,
|
||||||
&entry,
|
&entry,
|
||||||
|
&mut fitter,
|
||||||
)
|
)
|
||||||
.await?;
|
.await?;
|
||||||
let AgenticLoopOutcome {
|
let AgenticLoopOutcome {
|
||||||
@@ -1255,6 +1250,7 @@ impl InsightChatService {
|
|||||||
user_msg.images = Some(vec![img.clone()]);
|
user_msg.images = Some(vec![img.clone()]);
|
||||||
}
|
}
|
||||||
let mut messages = vec![system_msg, user_msg];
|
let mut messages = vec![system_msg, user_msg];
|
||||||
|
let mut fitter = ContextFitter::for_client(backend.chat(), req.num_ctx).await;
|
||||||
|
|
||||||
let outcome = self
|
let outcome = self
|
||||||
.run_streaming_agentic_loop_with_entry(
|
.run_streaming_agentic_loop_with_entry(
|
||||||
@@ -1267,6 +1263,7 @@ impl InsightChatService {
|
|||||||
&active_persona,
|
&active_persona,
|
||||||
max_iterations,
|
max_iterations,
|
||||||
&entry,
|
&entry,
|
||||||
|
&mut fitter,
|
||||||
)
|
)
|
||||||
.await?;
|
.await?;
|
||||||
let AgenticLoopOutcome {
|
let AgenticLoopOutcome {
|
||||||
@@ -1354,6 +1351,7 @@ impl InsightChatService {
|
|||||||
active_persona: &str,
|
active_persona: &str,
|
||||||
max_iterations: usize,
|
max_iterations: usize,
|
||||||
entry: &Arc<TurnEntry>,
|
entry: &Arc<TurnEntry>,
|
||||||
|
fitter: &mut ContextFitter,
|
||||||
) -> Result<AgenticLoopOutcome> {
|
) -> Result<AgenticLoopOutcome> {
|
||||||
crate::ai::insight_chat::run_streaming_agentic_loop_with_entry(
|
crate::ai::insight_chat::run_streaming_agentic_loop_with_entry(
|
||||||
&self.generator,
|
&self.generator,
|
||||||
@@ -1366,6 +1364,7 @@ impl InsightChatService {
|
|||||||
active_persona,
|
active_persona,
|
||||||
max_iterations,
|
max_iterations,
|
||||||
entry,
|
entry,
|
||||||
|
fitter,
|
||||||
)
|
)
|
||||||
.await
|
.await
|
||||||
}
|
}
|
||||||
@@ -1494,10 +1493,8 @@ impl InsightChatService {
|
|||||||
};
|
};
|
||||||
|
|
||||||
// Truncate before appending the new user turn.
|
// Truncate before appending the new user turn.
|
||||||
let budget_tokens = (req.num_ctx.unwrap_or_else(env_default_num_ctx) as usize)
|
let mut fitter = ContextFitter::for_client(backend.chat(), req.num_ctx).await;
|
||||||
.saturating_sub(RESPONSE_HEADROOM_TOKENS);
|
let truncated = fitter.trim(&mut messages);
|
||||||
let budget_bytes = budget_tokens.saturating_mul(BYTES_PER_TOKEN);
|
|
||||||
let truncated = apply_context_budget(&mut messages, budget_bytes);
|
|
||||||
if truncated {
|
if truncated {
|
||||||
let _ = tx.send(ChatStreamEvent::Truncated).await;
|
let _ = tx.send(ChatStreamEvent::Truncated).await;
|
||||||
}
|
}
|
||||||
@@ -1520,6 +1517,7 @@ impl InsightChatService {
|
|||||||
&active_persona,
|
&active_persona,
|
||||||
max_iterations,
|
max_iterations,
|
||||||
&tx,
|
&tx,
|
||||||
|
&mut fitter,
|
||||||
)
|
)
|
||||||
.await?;
|
.await?;
|
||||||
let AgenticLoopOutcome {
|
let AgenticLoopOutcome {
|
||||||
@@ -1740,6 +1738,7 @@ impl InsightChatService {
|
|||||||
user_msg.images = Some(vec![img.clone()]);
|
user_msg.images = Some(vec![img.clone()]);
|
||||||
}
|
}
|
||||||
let mut messages = vec![system_msg, user_msg];
|
let mut messages = vec![system_msg, user_msg];
|
||||||
|
let mut fitter = ContextFitter::for_client(backend.chat(), req.num_ctx).await;
|
||||||
|
|
||||||
let outcome = self
|
let outcome = self
|
||||||
.run_streaming_agentic_loop(
|
.run_streaming_agentic_loop(
|
||||||
@@ -1752,6 +1751,7 @@ impl InsightChatService {
|
|||||||
&active_persona,
|
&active_persona,
|
||||||
max_iterations,
|
max_iterations,
|
||||||
&tx,
|
&tx,
|
||||||
|
&mut fitter,
|
||||||
)
|
)
|
||||||
.await?;
|
.await?;
|
||||||
let AgenticLoopOutcome {
|
let AgenticLoopOutcome {
|
||||||
@@ -1833,6 +1833,7 @@ impl InsightChatService {
|
|||||||
active_persona: &str,
|
active_persona: &str,
|
||||||
max_iterations: usize,
|
max_iterations: usize,
|
||||||
tx: &tokio::sync::mpsc::Sender<ChatStreamEvent>,
|
tx: &tokio::sync::mpsc::Sender<ChatStreamEvent>,
|
||||||
|
fitter: &mut ContextFitter,
|
||||||
) -> Result<AgenticLoopOutcome> {
|
) -> Result<AgenticLoopOutcome> {
|
||||||
let mut tool_calls_made = 0usize;
|
let mut tool_calls_made = 0usize;
|
||||||
let mut iterations_used = 0usize;
|
let mut iterations_used = 0usize;
|
||||||
@@ -1849,9 +1850,8 @@ impl InsightChatService {
|
|||||||
})
|
})
|
||||||
.await;
|
.await;
|
||||||
|
|
||||||
let mut stream = backend
|
let mut stream = fitter
|
||||||
.chat()
|
.chat_stream(backend.chat(), messages, tools.clone())
|
||||||
.chat_with_tools_stream(messages.clone(), tools.clone())
|
|
||||||
.await?;
|
.await?;
|
||||||
|
|
||||||
let mut final_message: Option<ChatMessage> = None;
|
let mut final_message: Option<ChatMessage> = None;
|
||||||
@@ -1941,10 +1941,7 @@ impl InsightChatService {
|
|||||||
// pending_tools at this position (wiping the prior tool
|
// pending_tools at this position (wiping the prior tool
|
||||||
// calls from the final assistant render).
|
// calls from the final assistant render).
|
||||||
let synthetic_idx = push_synthetic_final_prompt(messages);
|
let synthetic_idx = push_synthetic_final_prompt(messages);
|
||||||
let mut stream = backend
|
let mut stream = fitter.chat_stream(backend.chat(), messages, vec![]).await?;
|
||||||
.chat()
|
|
||||||
.chat_with_tools_stream(messages.clone(), vec![])
|
|
||||||
.await?;
|
|
||||||
let mut final_message: Option<ChatMessage> = None;
|
let mut final_message: Option<ChatMessage> = None;
|
||||||
while let Some(ev) = stream.next().await {
|
while let Some(ev) = stream.next().await {
|
||||||
let ev = ev?;
|
let ev = ev?;
|
||||||
@@ -2270,6 +2267,145 @@ pub(crate) fn env_default_num_ctx() -> i32 {
|
|||||||
.max(RESPONSE_HEADROOM_TOKENS as i32 + 1024)
|
.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<i32>,
|
||||||
|
) -> 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<i32>,
|
||||||
|
) -> 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<ChatMessage>) -> bool {
|
||||||
|
apply_context_budget(messages, self.budget_bytes)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn fit(&self, messages: &[ChatMessage]) -> Vec<ChatMessage> {
|
||||||
|
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<Tool>,
|
||||||
|
) -> Result<(ChatMessage, Option<i32>, Option<i32>)> {
|
||||||
|
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<Tool>,
|
||||||
|
) -> Result<futures::stream::BoxStream<'static, Result<LlmStreamEvent>>> {
|
||||||
|
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<Option<(usize, usize)>> {
|
||||||
|
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<usize> {
|
||||||
|
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
|
/// Append a per-turn iteration-budget reminder to the replayed system
|
||||||
/// message so the model knows how many tool-calling rounds this turn gets.
|
/// message so the model knows how many tool-calling rounds this turn gets.
|
||||||
/// Returns the original `content` so the caller can restore it before
|
/// 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,
|
active_persona: &str,
|
||||||
max_iterations: usize,
|
max_iterations: usize,
|
||||||
entry: &Arc<TurnEntry>,
|
entry: &Arc<TurnEntry>,
|
||||||
|
fitter: &mut ContextFitter,
|
||||||
) -> Result<AgenticLoopOutcome> {
|
) -> Result<AgenticLoopOutcome> {
|
||||||
let mut tool_calls_made = 0usize;
|
let mut tool_calls_made = 0usize;
|
||||||
let mut iterations_used = 0usize;
|
let mut iterations_used = 0usize;
|
||||||
@@ -2350,9 +2487,8 @@ pub async fn run_streaming_agentic_loop_with_entry(
|
|||||||
})
|
})
|
||||||
.await;
|
.await;
|
||||||
|
|
||||||
let mut stream = backend
|
let mut stream = fitter
|
||||||
.chat()
|
.chat_stream(backend.chat(), messages, tools.clone())
|
||||||
.chat_with_tools_stream(messages.clone(), tools.clone())
|
|
||||||
.await?;
|
.await?;
|
||||||
|
|
||||||
let mut final_message: Option<ChatMessage> = None;
|
let mut final_message: Option<ChatMessage> = None;
|
||||||
@@ -2431,10 +2567,7 @@ pub async fn run_streaming_agentic_loop_with_entry(
|
|||||||
|
|
||||||
if final_content.is_empty() {
|
if final_content.is_empty() {
|
||||||
let synthetic_idx = push_synthetic_final_prompt(messages);
|
let synthetic_idx = push_synthetic_final_prompt(messages);
|
||||||
let mut stream = backend
|
let mut stream = fitter.chat_stream(backend.chat(), messages, vec![]).await?;
|
||||||
.chat()
|
|
||||||
.chat_with_tools_stream(messages.clone(), vec![])
|
|
||||||
.await?;
|
|
||||||
let mut final_message: Option<ChatMessage> = None;
|
let mut final_message: Option<ChatMessage> = None;
|
||||||
while let Some(ev) = stream.next().await {
|
while let Some(ev) = stream.next().await {
|
||||||
let ev = ev?;
|
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
|
/// Trim history to fit within `budget_bytes` of serialized JSON. Always
|
||||||
/// the system message and the first user message (with its base64 images
|
/// preserves the system message and the first user message (with its base64
|
||||||
/// intact, since dropping those would invalidate the model's prior visual
|
/// images intact, since dropping those would invalidate the model's prior
|
||||||
/// reasoning). Drops the oldest assistant-tool_call + corresponding
|
/// visual reasoning) plus that first exchange's final reply, i.e. the
|
||||||
/// tool-result pair on each pass until the budget is met or only the
|
/// insight itself.
|
||||||
/// preserved prefix remains.
|
///
|
||||||
|
/// 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.
|
/// Returns true when at least one message was dropped.
|
||||||
pub(crate) fn apply_context_budget(messages: &mut Vec<ChatMessage>, budget_bytes: usize) -> bool {
|
pub(crate) fn apply_context_budget(messages: &mut Vec<ChatMessage>, budget_bytes: usize) -> bool {
|
||||||
if budget_bytes == 0 {
|
if budget_bytes == 0 || estimate_bytes(messages) <= budget_bytes {
|
||||||
return false;
|
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;
|
let mut dropped_any = false;
|
||||||
loop {
|
while estimate_bytes(messages) > budget_bytes {
|
||||||
if estimate_bytes(messages) <= budget_bytes {
|
let Some(range) = next_budget_drop(messages) else {
|
||||||
break;
|
break;
|
||||||
}
|
};
|
||||||
// Find the oldest assistant-with-tool_calls strictly after the
|
messages.drain(range);
|
||||||
// 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);
|
|
||||||
dropped_any = true;
|
dropped_any = true;
|
||||||
}
|
}
|
||||||
|
|
||||||
dropped_any
|
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<std::ops::Range<usize>> {
|
||||||
|
let user_idxs: Vec<usize> = 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<std::ops::Range<usize>> = 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<usize>| {
|
||||||
|
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<usize>| {
|
||||||
|
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,
|
/// Estimate the serialized byte size of `messages` for the truncation budget,
|
||||||
/// EXCLUDING inlined base64 image payloads. Images are charged a flat
|
/// EXCLUDING inlined base64 image payloads. Images are charged a flat
|
||||||
/// `IMAGE_TOKENS_EACH` instead: their base64 is hundreds of KB of characters
|
/// `IMAGE_TOKENS_EACH` instead: their base64 is hundreds of KB of characters
|
||||||
@@ -2913,6 +3080,207 @@ mod tests {
|
|||||||
assert!(!dropped);
|
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<usize>,
|
||||||
|
fail_first: usize,
|
||||||
|
error: String,
|
||||||
|
sent: Mutex<Vec<usize>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait::async_trait]
|
||||||
|
impl crate::ai::llm_client::LlmClient for ScriptedClient {
|
||||||
|
async fn generate(
|
||||||
|
&self,
|
||||||
|
_: &str,
|
||||||
|
_: Option<&str>,
|
||||||
|
_: Option<Vec<String>>,
|
||||||
|
) -> Result<String> {
|
||||||
|
unimplemented!()
|
||||||
|
}
|
||||||
|
async fn chat_with_tools(
|
||||||
|
&self,
|
||||||
|
messages: Vec<ChatMessage>,
|
||||||
|
_: Vec<Tool>,
|
||||||
|
) -> Result<(ChatMessage, Option<i32>, Option<i32>)> {
|
||||||
|
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<ChatMessage>,
|
||||||
|
_: Vec<Tool>,
|
||||||
|
) -> Result<futures::stream::BoxStream<'static, Result<LlmStreamEvent>>> {
|
||||||
|
unimplemented!()
|
||||||
|
}
|
||||||
|
async fn generate_embeddings(&self, _: &[&str]) -> Result<Vec<Vec<f32>>> {
|
||||||
|
unimplemented!()
|
||||||
|
}
|
||||||
|
async fn describe_image(&self, _: &str) -> Result<String> {
|
||||||
|
unimplemented!()
|
||||||
|
}
|
||||||
|
async fn list_models(&self) -> Result<Vec<crate::ai::llm_client::ModelCapabilities>> {
|
||||||
|
unimplemented!()
|
||||||
|
}
|
||||||
|
async fn model_capabilities(
|
||||||
|
&self,
|
||||||
|
_: &str,
|
||||||
|
) -> Result<crate::ai::llm_client::ModelCapabilities> {
|
||||||
|
unimplemented!()
|
||||||
|
}
|
||||||
|
fn primary_model(&self) -> &str {
|
||||||
|
"scripted"
|
||||||
|
}
|
||||||
|
async fn context_window(&self) -> Option<usize> {
|
||||||
|
self.window
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn scripted(window: Option<usize>, 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]
|
#[test]
|
||||||
fn render_tree_path_detects_fork_at_tool_dispatch_child() {
|
fn render_tree_path_detects_fork_at_tool_dispatch_child() {
|
||||||
// Regression: a regeneration where the model replies with a tool call
|
// Regression: a regeneration where the model replies with a tool call
|
||||||
|
|||||||
+149
-1
@@ -27,6 +27,7 @@ use opentelemetry::trace::Status;
|
|||||||
use reqwest::Client;
|
use reqwest::Client;
|
||||||
use serde::Deserialize;
|
use serde::Deserialize;
|
||||||
use serde_json::{Value, json};
|
use serde_json::{Value, json};
|
||||||
|
use std::collections::HashMap;
|
||||||
use std::time::Duration;
|
use std::time::Duration;
|
||||||
use tracing_opentelemetry::OpenTelemetrySpanExt as _;
|
use tracing_opentelemetry::OpenTelemetrySpanExt as _;
|
||||||
|
|
||||||
@@ -41,6 +42,18 @@ const DEFAULT_VISION_MODEL: &str = "vision";
|
|||||||
const DEFAULT_EMBEDDING_MODEL: &str = "embed";
|
const DEFAULT_EMBEDDING_MODEL: &str = "embed";
|
||||||
const DEFAULT_TTS_MODEL: &str = "chatterbox";
|
const DEFAULT_TTS_MODEL: &str = "chatterbox";
|
||||||
const DEFAULT_REQUEST_TIMEOUT_SECS: u64 = 180;
|
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
|
/// 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.
|
/// 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
|
// upstream server directly, reached via llama-swap's passthrough at
|
||||||
// `{swap_root}/upstream/<tts_model>/voices`.
|
// `{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.
|
/// Root of the llama-swap proxy: `base_url` with a trailing `/v1` removed.
|
||||||
/// The `/upstream/...` passthrough lives here, not under `/v1`.
|
/// The `/upstream/...` passthrough lives here, not under `/v1`.
|
||||||
fn swap_root(&self) -> &str {
|
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
|
// 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
|
// 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;
|
let _ = self.num_ctx;
|
||||||
// Reasoning toggle for thinking-capable templates (Qwen3 et al.).
|
// Reasoning toggle for thinking-capable templates (Qwen3 et al.).
|
||||||
// llama-server forwards chat_template_kwargs into the Jinja render
|
// llama-server forwards chat_template_kwargs into the Jinja render
|
||||||
@@ -939,6 +1033,13 @@ impl LlmClient for LlamaCppClient {
|
|||||||
fn primary_model(&self) -> &str {
|
fn primary_model(&self) -> &str {
|
||||||
&self.primary_model
|
&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 {
|
impl LlamaCppClient {
|
||||||
@@ -955,6 +1056,7 @@ impl LlamaCppClient {
|
|||||||
name,
|
name,
|
||||||
has_vision: true,
|
has_vision: true,
|
||||||
has_tool_calling: true,
|
has_tool_calling: true,
|
||||||
|
context_window: None,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -1113,6 +1215,52 @@ fn image_to_data_url(img: &str) -> String {
|
|||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
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]
|
#[test]
|
||||||
fn tool_call_arguments_stringified_on_send() {
|
fn tool_call_arguments_stringified_on_send() {
|
||||||
let msg = ChatMessage {
|
let msg = ChatMessage {
|
||||||
|
|||||||
@@ -63,6 +63,14 @@ pub trait LlmClient: Send + Sync {
|
|||||||
|
|
||||||
/// Primary model identifier this client was constructed with.
|
/// Primary model identifier this client was constructed with.
|
||||||
fn primary_model(&self) -> &str;
|
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<usize> {
|
||||||
|
None
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Events emitted by streaming `chat_with_tools_stream`. A stream is a
|
/// Events emitted by streaming `chat_with_tools_stream`. A stream is a
|
||||||
@@ -174,6 +182,11 @@ pub struct ModelCapabilities {
|
|||||||
pub name: String,
|
pub name: String,
|
||||||
pub has_vision: bool,
|
pub has_vision: bool,
|
||||||
pub has_tool_calling: 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<usize>,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// A single node in the conversation tree. Wraps ChatMessage with structural
|
/// A single node in the conversation tree. Wraps ChatMessage with structural
|
||||||
|
|||||||
@@ -315,6 +315,7 @@ impl OllamaClient {
|
|||||||
name: model_name.to_string(),
|
name: model_name.to_string(),
|
||||||
has_vision,
|
has_vision,
|
||||||
has_tool_calling,
|
has_tool_calling,
|
||||||
|
context_window: None,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -346,6 +347,7 @@ impl OllamaClient {
|
|||||||
name: model_name,
|
name: model_name,
|
||||||
has_vision: false,
|
has_vision: false,
|
||||||
has_tool_calling: false,
|
has_tool_calling: false,
|
||||||
|
context_window: None,
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -901,6 +901,7 @@ fn parse_model_capabilities(m: &Value) -> ModelCapabilities {
|
|||||||
name,
|
name,
|
||||||
has_vision,
|
has_vision,
|
||||||
has_tool_calling,
|
has_tool_calling,
|
||||||
|
context_window: None,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -37,8 +37,8 @@ use uuid::Uuid;
|
|||||||
use crate::ai::backend::{BackendKind, SamplingOverrides};
|
use crate::ai::backend::{BackendKind, SamplingOverrides};
|
||||||
use crate::ai::handlers::ReplayQuery;
|
use crate::ai::handlers::ReplayQuery;
|
||||||
use crate::ai::insight_chat::{
|
use crate::ai::insight_chat::{
|
||||||
BYTES_PER_TOKEN, ChatStreamEvent, DEFAULT_MAX_ITERATIONS, ForkInfo, RESPONSE_HEADROOM_TOKENS,
|
ChatStreamEvent, ContextFitter, DEFAULT_MAX_ITERATIONS, ForkInfo, env_max_iterations,
|
||||||
apply_context_budget, env_default_num_ctx, env_max_iterations, render_tree_path,
|
render_tree_path,
|
||||||
};
|
};
|
||||||
use crate::ai::insight_generator::InsightGenerator;
|
use crate::ai::insight_generator::InsightGenerator;
|
||||||
use crate::ai::llm_client::{ChatHistoryStore, ChatMessage, StoredChatNode};
|
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
|
// Trim to the model's context window before the turn. Without this a
|
||||||
// rolling persona transcript grows unbounded until the backend
|
// rolling persona transcript grows unbounded until the backend
|
||||||
// silently drops the front of it mid-conversation.
|
// silently drops the front of it mid-conversation.
|
||||||
let budget_tokens = (req.num_ctx.unwrap_or_else(env_default_num_ctx) as usize)
|
let mut fitter = ContextFitter::for_client(backend.chat(), req.num_ctx).await;
|
||||||
.saturating_sub(RESPONSE_HEADROOM_TOKENS);
|
let truncated = fitter.trim(&mut messages);
|
||||||
let budget_bytes = budget_tokens.saturating_mul(BYTES_PER_TOKEN);
|
|
||||||
let truncated = apply_context_budget(&mut messages, budget_bytes);
|
|
||||||
if truncated {
|
if truncated {
|
||||||
let _ = entry.push_event(ChatStreamEvent::Truncated).await;
|
let _ = entry.push_event(ChatStreamEvent::Truncated).await;
|
||||||
}
|
}
|
||||||
@@ -880,6 +878,7 @@ impl PersonaChatSession {
|
|||||||
&persona.persona_id,
|
&persona.persona_id,
|
||||||
max_iterations,
|
max_iterations,
|
||||||
&entry,
|
&entry,
|
||||||
|
&mut fitter,
|
||||||
)
|
)
|
||||||
.await
|
.await
|
||||||
.map_err(|e| anyhow!("persona chat loop failed: {e:?}"))?;
|
.map_err(|e| anyhow!("persona chat loop failed: {e:?}"))?;
|
||||||
|
|||||||
Reference in new issue
Block a user