mirror of
https://github.com/unslothai/unsloth.git
synced 2026-08-23 07:44:06 +00:00
Studio: keep and search the turns rolling context evicts (#9074)
* Studio: add rolling context windows for local GGUF chat
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: return rolling context metadata for non-stream chats
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: keep original context when rolling fit fails
* Studio: keep instruction groups independently protected
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: preserve rolling context metadata across retries
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: count sanitized rolling context prompts
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: refit rolling context after respawn
* Studio: scope middle truncation to passthrough
* Studio: refit tool prompts after respawn
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: retain later choice truncation metadata
* Studio: report clipping-only context truncation
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: expose the turns a rolling fit evicts
Rolling context eviction currently drops turns with no way for a caller to
learn which ones went. Extract the turn grouping so the eviction unit is
reusable, add an identity-based diff of what a fit removed, and let a caller
reserve room for content it intends to add back after fitting.
The reserve deliberately does not affect whether trimming happens at all, only
how far it goes once it is already required, so a conversation that fits today
is still returned untouched.
* Studio: archive evicted turns into a per-thread rag scope
An evicted turn is currently gone for the rest of the session, so the model
will state the conversation began wherever its visible context begins. Keep the
turns the rolling window drops in a searchable scope built on the existing
store, chunker, embedder and hybrid retrieval.
The archive is cumulative: every compaction adds to it and nothing is cleared,
so a later compaction can still find what an earlier one evicted. It lives in
its own scope rather than the thread's document scope, because thread documents
are injected in full on every request and would re-inject the whole history.
Idempotent by content hash, since the same turns are evicted again on every
later request. Every entry point degrades to a no-op rather than raising.
* Studio: recall archived turns on the turn that evicted them
Given only a search tool, a model decides for itself whether to look, and
mostly does not. Measured on MRCR v2, a 35B declined on 56% of rows, scoring
0.099 when it skipped against 0.461 when it searched. Forcing one retrieval on
the compaction turn took tool-only 0.258 to 0.604, and the model then called
the tool on 0% of rows, so the common path costs nothing extra.
Recall fires at most once per request, since the tool loop refits on every
iteration. The tool loop renders it as an ordinary tool exchange through the
builder shared with document auto-inject; the plain path prefixes the latest
user message instead, because it sends no tools array and a tool role without
one breaks strict chat templates.
Only scalar counts join the context_truncated event, so no message content
reaches the wire.
* Studio: let the model search a compacted conversation
Forced recall answers the turn that evicted, but a later turn can refer to
something the forced pass had no reason to fetch. Add search_conversation so
the model can go looking, scoped to the thread's own archive and sharing the
admission slot with document search so the two cannot race for the embedder.
The tool is offered only once a thread has actually had turns archived, so an
ordinary short chat never pays for the schema, and it is classified read-only
so auto mode does not prompt on every call. A matching system-prompt note tells
the model the session was compacted, since otherwise it assumes the
conversation began where its visible context begins.
Deleting a thread now drops its archive rather than leaking a scope per chat.
* Studio: read thread_id defensively when selecting tools
_select_request_tools also serves the token-count request model, which has no
thread_id field, so the archive gate raised there.
* Studio: retrieve archived turns lexically first
Recalling your own conversation is mostly an exact-match problem: a name, a
number, an identifier someone pasted twenty turns ago. Those live or die on
rare-token matching, and hybrid fusion was losing them.
Measured on a 30-turn walkthrough of a 230k-character document at a 16k window,
where every turn shared the same wrapper text. The chunk holding the needle
ranked 3rd lexically at any k, was never returned by dense retrieval at all,
and RRF pushed it to 16th because it had 30 useless dense hits to fuse with.
End to end the model answered with the exact code once lexical leads, and could
not answer at all before.
Dense still fills whatever the lexical pass leaves, for paraphrased recall.
* Studio: keep recall on the branch the user is actually on
Editing an earlier message rewinds a thread and continues down a new branch,
but the archive is append-only and still holds everything the abandoned
continuation produced. Verified against a live build: after rewinding past a
turn, querying its distinctive text still returned it, so the model could be
handed a turn that on this branch never happened.
Recall now drops archived turns that are absent from the thread's saved
transcript. Threads with no saved transcript are left unfiltered, since an API
caller may pass a thread_id without persisting messages and an empty transcript
is absence of evidence rather than evidence the turns are gone.
Containment on a normalised prefix rather than a digest, because the archived
copy is rendered from the inference projection and the saved copy comes back
through the message store.
* Studio: show a persistent compaction notice on the turn that compacted
The only signal that a long chat had been compacted was a toast, which vanishes after
a few seconds and does not survive a reload. A user who scrolls back later has no way
to find out why the model seemed to forget the start of the conversation.
The notice renders from metadata.custom.contextTruncation inside the assistant
message's own container, so it is not part of the conversation sent to the model, is
not editable, and is not exported as content, but it stays attached to the turn it
describes. It reports how many messages were dropped and, when the conversation
archive is on, that they are still searchable and how many passages were recalled.
* Studio: show the compaction notice once, on the turn it started
A thread that has outgrown its window compacts on every turn from then on, not
just the first, so a notice per compacted turn was a notice on every reply for the
rest of the conversation. The user needs telling once.
The notice is now gated on being the first compacted assistant turn in the thread,
found by walking the thread rather than assuming the compacted turns are contiguous
or that this is the last one, so a rollback that removes the turn it was on moves it
to whichever turn now compacts first. The wording follows: it describes the state
the conversation is in from here on, and carries the counts from the turn it began
on in parentheses.
* Studio: make compaction an occasional event instead of every turn
The fit was stateless. The client re-sends the whole saved transcript on every
request, so "keep the newest N tokens" recomputes from scratch each time and slides
forward a turn or two at a time. Measured on a 40-turn thread against an 8k window,
the eviction boundary moved on 12 of 40 turns, which means every few replies quietly
lost a little more of the conversation, llama-server's prefix cache was thrown away
each time the head of the prompt moved, and there was no such thing as a compaction
event to tell the user about.
Two changes make it discrete. The fit now reads back the boundary the thread last
compacted to, from the newest assistant turn's own persisted truncation, and reapplies
it before deciding anything, so nothing new is stored and it survives a restart. And
when the boundary does have to move, the trim takes a further ROLLING_COMPACTION_
HEADROOM_RATIO of the budget out (default 0.25) rather than skimming to the brim, so
the new boundary has room to stay put. Same 40-turn thread, same window: 4 compactions
over 60 turns instead of 14, keeping about 82 percent of the usable budget.
Both are gated on the prompt not already fitting, exactly as the recall reserve is. A
conversation inside its window is never evicted to satisfy either, and a stale boundary
from a branch that was rolled back cannot evict a chat that now fits.
The notice follows: it is shown when dropped_messages rises above the last turn that
reported it, so it appears once per compaction and stays quiet in between.
* Studio: pin that the compaction notice can never become conversation
The notice is a sidecar rendered from metadata, not a message, and the ways that
could quietly stop being true are all one careless edit away: moving it inside
MessagePrimitive.Parts would make it a content part, and everything that walks parts
would then replay it to the model, copy it and export it.
Asserts it renders as a sibling of the content parts, that neither the outbound
message builder nor the assistant replay serialiser reads the key it renders from
(bounded to those function bodies, since the streaming handler reads the same key
legitimately on the way in), that the markdown export does not mention it, and that
it is suppressed while editing so it cannot be typed into the textarea and saved
back as text.
* Studio: say which part is too long when nothing can make a turn fit
If the message just sent is itself bigger than the window, no amount of eviction
helps. The fit already handled that correctly, returning the conversation untouched
so the request reaches llama-server's normal context-length error, but what the user
was then told was actively misleading: the error reports the size of the WHOLE
conversation and advises shortening it, when the history has already been evicted and
the single message is the part that does not fit. Measured at a 4096-token window: a
5000-token message produces 'Message too long: 10290 tokens ... shorten the
conversation', and shortening it cannot possibly work.
The fit now returns a fits:false diagnosis instead of a bare None, carrying what the
conversation could not be reduced below and how much of that is the latest turn.
Every consumer already gated on fits, so this is inert wherever a truncation is
treated as a compaction; the streaming paths now forward it so the client can use it.
The toast reads the diagnosis and, when the latest turn alone exceeds the window,
says so with the numbers instead of offering advice that leads nowhere. The merge
drops the diagnosis once a later refit succeeds, by delete rather than by assigning
undefined, so an ordinary response keeps exactly the shape it had before.
* Studio: fix the review findings on the conversation archive
Fifteen items, each reproduced against the code before changing anything.
Correctness in the request path. The fit reported a successful recall-capable fit
whenever the result was under the prompt budget, but protected messages can stop the
trim reaching the reserve target, and the recall then went in anyway: reproduced at
ctx 8000, the fit accepted at 6900 and recall took the request to 8948, past the
window it had just been made to fit. Recall is now sized from the room the fit
actually obtained, and skipped when there is none. Separately, the sticky boundary
describes the original transcript, so re-applying it on a later tool-loop fit evicted
another boundary-sized block of live history: measured, a second fit dropped 28 of the
30 surviving messages instead of 14, and the summed count persisted an inflated
boundary for the next request. It is now spent after the first fit of a request.
Availability and privacy. enabled() trusted RAG_AVAILABLE, which only records that
import sqlite_vec worked; rag_available() exists because the native vec0 library it
loads is a separate file a venv can lack. On such a machine the fit held a recall
reserve back, evicting extra history, and then both the archive write and the recall
failed, so the user paid for content they never got. And a temporary chat is never
written to studio.db, yet the frontend still sends its thread_id and the request
carries no incognito flag, so its turns were archived to a scope no deletion flow
could reach. Archival now requires the thread to be persisted, which is the same rule
that keeps every archive reachable by a delete. Clear-history and project deletion
drop archives too; only DELETE /threads did.
Archive integrity. The document was committed before its chunks, so a failed chunk
write left an empty row marked completed that document_by_hash then skipped forever;
both now go in one transaction. The live-branch filter accepted a turn on its first
matching line, so editing only the assistant half kept serving the old answer; the
whole turn must be present. Retrieval fetched exactly k before that filter, so stale
turns could starve live ones and recall returned nothing; it over-fetches first. Tool
turns archived only the tool name, which cannot answer what was actually run, so a
bounded rendering of the arguments and the assistant text goes in as well.
The tool surface. Studio always sends an explicit enabled_tools array and has no
reason to name an internal tool, so the allowlist filter removed search_conversation
before the archive gate ran and the tool, plus the compaction nudge gated on it, never
appeared in a Studio chat. It now follows the archive rather than the allowlist. The
forced recall rendered a tool exchange even when the tool was absent from the
catalogue, which is the strict-template hazard the plain path avoids; it picks inline
in that case, and always inline for the final-answer request, which sends no tools at
all. A model-supplied top_k reached a slice as out[:-1] and returned nearly the whole
candidate pool, so it is clamped. Both retrieval tools now share the per-turn search
cap; only the knowledge-base one was counted.
The UI. A fits:false diagnosis is the fitter reporting it could NOT fit, so toasting
that older turns were removed was untrue and burned the once-per-thread flag a later
real compaction needed. The too-long advice compared the latest turn against the raw
context length rather than the prompt budget, so a 3,500-token message in a
4,096-token window was still told to start a new chat, which fails identically.
* Studio: fix the follow-up findings on the archive fixes
Four items, three of them about last round's own fixes.
Archived tool turns had become permanently unrecallable. render_turn now writes
'assistant called X: args' and 'tool result: ...' lines, while assistant-ui persists a
tool call as a structured tool-call content part that the transcript flattener dropped;
with every archived line required to appear in that transcript, no tool turn could ever
match. The transcript now flattens toolName, args and result, and the probe strips the
'assistant called <name>:' label, which is ours rather than the stored message's.
The branch probe compared only the first 160 normalized characters of each line, so an
edit to the tail of a long answer left the stale copy eligible. Ordinary lines are now
compared whole; tool results keep a prefix, since render_turn deliberately truncates
those and the archived copy is not meant to equal the stored one.
The forced recall sized itself by dividing the remaining budget by CHUNK_TOKENS, which
is an embedding-token limit rather than the chat template's cost, and prices none of
the wrappers around the injection. It is now recounted with the same tokenizer the fit
used and dropped if it overshoots, so the estimate can no longer eat the reply reserve.
An omitted top_k on search_conversation defaulted to the clamp ceiling of eight rather
than the configured recall default, so an ordinary search could return eight archived
turns into the protected current exchange that rolling truncation cannot evict.
* Studio: filter conversation recall to the active branch
A thread's stored rows are the whole message DAG. Retry and regenerate keep the
replaced response as a sibling on purpose, so filtering recall against the whole
thread cannot tell a live turn from one the user replaced, and an archived copy of
the abandoned response could be recalled into a branch where it never happened.
Filter against the messages the request was actually sent with instead, which is
one branch by construction, and hand the same branch to search_conversation so the
model cannot ask for what the forced recall refused. Falls back to the thread-wide
blob for a caller with no branch to offer.
Also retunes two respawn-refit fixtures whose windows no longer produced two
compactions after compaction started trimming a headroom margin below the budget.
* Studio: survive a delete mid-archive, and read the boundary off the active branch
Deleting a chat cancels its generation, but cancellation is cooperative and the
chunk-and-embed pass between the archive's liveness check and its commit does not
observe it. A delete landing in that window drops the thread's rows and sweeps its
scope before the commit puts rows back, leaving content the user deleted in a scope
no later delete can reach. Re-check after the commit and drop the scope: the delete
route removes rows first and sweeps archives last, so either order converges.
The sticky compaction boundary had the same thread-wide read as recall did. The
stored rows are the whole DAG ordered by creation time, so after a Retry the newest
assistant turn can be the sibling the user switched away from, and its boundary is
sized for history the active branch does not have. Resolve it against the request's
own messages instead.
* Studio: give the safetensors loop the same conversation-search guards
search_conversation is advertised by thread, not by backend: the tool selector is
shared, so a chat compacted under a GGUF model still offers it after the user
switches to a safetensors one. The safetensors loop had neither guard the GGUF loop
applies to it. It passed no active branch, so a search there fell back to the
thread-wide rows and could answer from a branch Retry left behind, and it capped
only search_knowledge_base, so paraphrased conversation searches could append
archived passages into the protected current exchange on every iteration until the
window failed. The shared set of capped retrieval tools now lives beside the cap.
The forced recall also derived its query from the loop conversation, which on a
later iteration can end with an internal user-role re-prompt rather than anything
the user wrote. It reads the request branch's own latest user turn instead.
* Studio: widen recall past an abandoned branch, and name whose turn overflowed
One over-fetch is not enough for the live-branch filter. Rewinding or retrying a
continuation that had already been compacted leaves enough stale turns to fill any
fixed candidate window, and the whole page is then rejected while the live match
sitting just below it is never examined, so recall reports nothing although the
answer is in the archive. Widen and re-ask instead, stopping as soon as there are
enough live hits, when the archive stops yielding candidates, or at a bound.
The irreducible-fit diagnosis also carried the size of the last message without
saying whose it was. A tool loop refits with the tool result appended, so that turn
is often output the user never wrote and cannot edit, and the client told them to
shorten it. It now reports the role, and the advice splits on it.
* Studio: shrink an over-budget recall, and keep truncated tool turns on their branch
The exact recount is the right gate, but dropping the whole recall when it fails is
the wrong response: with the shipped defaults a full top-K of long turns lands just
over the reserve once the wrappers are priced, which would disable the forced
retrieval on exactly the long conversations it exists for. Halve the number of turns
and re-ask instead, down to one, before giving up.
The live-branch probe keyed its prefix rule off the tool-result label, but a long
tool result is one appended string containing many newlines, so only its first line
carries that label. Continuation lines were compared in full, including the last one,
which is the only line the truncation marker is on and can therefore never appear in
a transcript: every archived tool turn over the cap was rejected as rolled back.
Key off the marker instead, which also covers truncated tool arguments and compares
the full text everywhere else.
* Studio: retire an edited turn's whole archived copy, not just the edited chunk
A turn longer than CHUNK_TOKENS is stored as several chunks of one document, and the
live-branch filter ran per chunk. Editing the second half of a long answer therefore
retired only the chunks carrying the edit, and an untouched earlier chunk of the same
retired turn stayed eligible on its own. The unit that was archived is the turn, which
is what the filter already claimed to enforce.
Validate every chunk of the candidate's document before admitting any of it, cached
per call since candidates from one turn share a document. A query failure falls back
to the per-chunk answer rather than failing the recall.
* Studio: do not hold a recall reserve back for a chat that is never archived
archive_turns refuses a thread with no saved messages, because a temporary chat must
not be persisted into a scope no deletion flow can reach. The reserve did not follow
that rule: it was granted whenever the RAG stack was available, so an incognito chat,
or an API client sending a thread_id without saving anything, paid a full 2,048-token
reserve for content that can never arrive. The fit subtracts the reserve from its trim
target, so that room is bought with evicted history. Measured on a 4K window: 15
tokens of conversation survived a compaction instead of 1,615.
Both now ask one predicate, can_archive, backed by an existence probe rather than a
row load.
* Studio: match archived turns in order, and delete them without sqlite-vec
Independent line membership accepts a turn whose lines were merely rearranged: every
probe still occurs somewhere, so the pre-edit ordering stayed eligible and would be
served back as what happened. Match the probes in order instead.
That only works if both sides agree on the order, and they did not: render_turn writes
a tool call before any assistant text on the same message and the result after it,
while both transcript builders wrote the text first, so a tool turn carrying both was
rejected outright. One flattener now lays out a message the way render_turn does, for
the request shape and the stored shape alike, and it offers both JSON spacings for a
stored call's arguments, whose object form is not the string the model emitted.
Deletion no longer depends on the optional native extension. An archive is only written
while vec0 loads, but the library can stop loading afterwards, and a delete that quietly
did nothing left a deleted conversation's turns on disk to answer again once it loaded.
The fallback removes the text-bearing rows over a metadata connection; the embedding
rows it cannot reach carry no text and resolve through tables that are gone.
* Studio: give the provider loops the branch, and keep the thread across a respawn
The provider tool loops take their catalogue from the same selector as the local ones,
so search_conversation is advertised there too once a thread has an archive, but the
loop passed no active branch and its searches fell back to the whole stored DAG, where
Retry keeps the response it replaced. This is the third loop to need the same wiring.
The respawn retry also dropped the thread. It re-enables context_overflow on purpose,
to refit an already-compacted prompt for a replacement window, so it is the one path
that deliberately compacts a second time: without the thread those extra evictions were
archived nowhere, nothing was recalled in their place, and the fit held back no reserve
and re-applied no boundary.
* Studio: keep the reply that follows a recall, and archive respawn refits
group_turns keeps an assistant tool call, its result and the reply that follows in one
group, and the archiver rejected any group containing one of our own injections. So on
a turn that forced a recall, the model's actual answer was thrown away with it: the
question was archived from its own group and the answer was not, and a later search
could find what was asked and never what was said. The injections come out now and the
rest of the turn stays; retrieved passages are still kept out of the index they came
from.
The two respawn refits inside the tool loop also evicted without archiving. They run
against a smaller replacement window, so those turns are simply gone otherwise. They
archive only, deliberately: no reserve is held back on that path, so injecting a recall
there is what would push the retry back over the window, and the next request can still
recall what this one archived.
* Studio: bound the branch check to one turn, and anchor only what recall injected
The branch check searched one flattened transcript, so a line an edit removed could be
supplied by any later message that happened to repeat the words. Short answers repeat
constantly: an archived "Should I deploy? / No" survived its answer being edited to
"Yes" because a later turn said "No", and the stale pair stayed recallable. The
transcript is now one normalised string per message, and a turn matches only if its
lines appear in order within a run of adjacent messages no longer than the turn itself.
Recall anchoring took the last two messages of the conversation. That is right for the
tool style, which appends a synthetic pair, and wrong for the inline style, which
appends nothing and rewrites the latest user message in place: it also pinned the
assistant turn before it, and with it an eviction unit the fit was entitled to drop,
which can fail a later iteration that would otherwise have fit. The injection now
reports exactly what it added or rewrote, and only that is anchored.
* Studio: validate a turn's chunks against one run of the branch, not each on its own
The chunks of an archived turn are consecutive slices of a single rendering, and each was
checked for itself. That let a turn be reassembled out of parts that never sat together:
the head matching the question and the answer it has now, the tail matching some later
message that happens to repeat the passage an edit removed, so the stale association
stayed recallable. Reproduced before the change, with each chunk passing on its own.
All chunks of the document must now be found within one run of adjacent messages, bounded
by the document's own line count, which is at least the number of messages the turn was
rendered from, so a turn that really is still there always fits.
* Studio: budget conversation searches and disambiguate identical replies
search_conversation clamped the model's top_k only against a fixed ceiling of 8.
Eight chunks is roughly 4,000 tokens once wrapped, and the result lands in the
current tool exchange, which rolling truncation protects and cannot evict, so on
a small context the search itself made the turn unsendable. The GGUF loop now
passes the room the window actually has left, the tool clamps top_k by it, and a
search with no room says so instead of returning a result that cannot be sent.
The sticky compaction boundary took the newest on-branch assistant row, but the
branch check is textual, so two siblings whose replies read the same ("Done.")
are indistinguishable from there, and Retry is exactly what produces them.
Taking the first match applied a boundary measured on a different, deeper branch
and evicted live history. Where the text cannot separate them, the smallest
boundary is now used: too small costs one more compaction, too large loses turns
the branch still has.
* Studio: price a conversation search against the whole prompt
Three gaps in the budget the previous commit introduced.
The GGUF loop subtracted an estimate of the messages alone, leaving the tool
catalogue out of the prompt entirely. A large catalogue is thousands of tokens,
so the request could already be near its budget while the search was still told
there was room for several 500-token chunks. The preflight already prices the
request exactly, catalogue and template included, so the difference between that
count and the estimate of the same messages is carried forward and the estimate
only covers what the loop appends after it.
The safetensors loop advertises and executes search_conversation for a thread
compacted under a GGUF model, but forwarded no budget, so the clamp in the tool
was skipped there. It now passes the room this model has left. With no known
context length the argument is omitted rather than sent as zero, which would
refuse every search.
The sticky boundary filtered candidates with a substring test, which is right
for archived chunks and wrong for whole messages: an abandoned "Done" rode in on
a live "Not done yet", and having no live twin it then decided the boundary
alone. Where any candidate matches a live message exactly, only the exact ones
are considered; where none does, the old behaviour stands, since a stored row is
not always byte-identical to what the client re-sends.
* Studio: keep the overflow diagnosis on the tool path
The irreducible-overflow branch counted the newest message on its own, and a
tool loop reaches that branch with a tool result last. A tool result by itself
is not a conversation: templates that require it to follow its assistant tool
call refuse to render one, and the exception escaped the fit entirely, so the
caller fell back to the untrimmed request and the client was told nothing at
all. That is the one path this diagnosis was added for.
The count is now attempted and falls back to the estimator when the template
refuses. An approximate number is worth more here than a diagnosis that never
arrives.
* Studio: budget the conversation search in the provider loop
The third tool loop forwarded the active branch but no budget, so the clamp in
the tool was skipped there and a model-chosen top_k of 8 could append roughly 4K
tokens to a prompt this loop replays on its next call.
Studio knows no window for an external model: the request carries no context
length and there is no registry to look one up in, and a custom
OpenAI-compatible endpoint can be a small local server. So rather than a
measured budget, this path spends no more than one ordinary recall's worth,
which is the same amount the compaction turn itself is sized for.
* Studio: bound an archived turn to its own messages, and archive it once
Two problems in the archive.
The whole-document branch check bounded the run by the number of LINES the turn
produced, and gave each chunk its own start. A turn is two or three messages
however long it is, so a hundred-line answer got a hundred-message window, and
the tail of an edited answer could be satisfied by a message well outside the
turn. The run is now bounded by the messages the turn was rendered from, counted
from the labels render_turn writes, and the chunks are scanned as one pass: each
continues where the previous one stopped rather than restarting. The cursor
inside a message is deliberately not carried over, because chunks overlap and a
continuation chunk repeats the tail of the one before it.
The hash check that skips an already-archived turn ran long before the insert,
with the embedding pass in between, and the index on (scope, sha256) is not
unique. Two generations compacting the same thread both cleared it and both
wrote, so the turn was stored twice and its copies took two of the few recall
slots. The check is now repeated under a write lock immediately before the
insert. Reproduced with two concurrent archive passes: two documents before, one
after.
* Studio: record how many messages an archived turn came from
The run an archived turn is allowed to occupy on the branch was bounded by
counting the role labels in its rendered text. That counts lines the user wrote
as well as the ones the renderer did: a pasted chat log carries lines that look
exactly the same, and each one widens the run by a message, which is enough for
the message after an edited turn to supply the passage the edit removed.
The group's size is now recorded on the document when it is archived, in a
nullable column added the way the other lazy upgrades are. Archives written
before this have NULL and fall back to the label count, so nothing needs
backfilling to keep working.
* Studio: persist a compaction boundary the next request can use
The boundary was read back from dropped_messages, which counts what each fit
removed from the conversation in front of it. The tool loop refits on every
iteration and the client sums those counts, so a long agent run added the tool
exchanges the turn itself created, and the next request applied the total to its
saved transcript. Reproduced with six tool calls on a 4K window: three fits of 4
summed to 12, on a branch that only ever had 4 evictable messages.
The boundary is now carried separately, measured against the messages the
request was sent with, so it is absolute and re-sending it cannot advance it.
The client keeps the latest value rather than summing, and turns saved before
this fall back to the dropped count, which is the same number for a turn that
fit once.
The recall reserve is also dropped once archiving has failed. sqlite-vec can be
present and the thread saved while the embedder cannot start: archive_turns
swallows that and recall injects nothing, so the room the fit held back was pure
loss on every compaction, and the failure mode forgot more history than having
the feature off. A failed write marks the archive degraded and the next
successful one clears it.
* Studio: count the boundary past the system prompt, and notice it
Two faults in the boundary added in the previous commit.
It stopped at the first message still present, and a Studio request always
starts with a system prompt that a fit never evicts, so every compaction
recorded a boundary of zero. That is the same as having no boundary: the next
request would move the eviction point again and invalidate the prefix cache on
every turn. Instruction messages are now skipped rather than treated as the
front of the branch, and the newest turn is excluded because it is never evicted
and an inline recall rewrites it in place.
The compaction notice still keyed off dropped_messages, which is the
accumulated count. A tool-heavy turn reporting 12 while the boundary moved to 4
set a high-water mark that silenced the next two real advances. It now reads the
boundary, falling back to the dropped count only for turns saved before the
boundary was recorded.
* Studio: tighten comments in the conversation archive and rolling context window
* Studio: give the pasted-text import an extension on this branch too
The node test runner cannot resolve an extensionless relative import, so
delete-chat-files-preference fails to load the preferences store. main fixed
this in d43892ea7; this branch predates it, and the same line is on the branch
below it. The change is identical to main's, so it disappears on merge.
* Studio: size a conversation search by what it renders to
CHUNK_TOKENS is what the chunker aims at, not what a chunk weighs: chunks
overlap, the chunker's tokenizer is not the model's, and the rendered block adds
markup, source metadata and the tool framing around it. Dividing the budget by
CHUNK_TOKENS therefore permitted a result far larger than the room measured, and
it lands in the current exchange, which the rolling window cannot evict.
Measured on a 500-token budget: one chunk came back at 1,256 estimated tokens.
The count is now halved until the rendered result fits, the same backoff the
forced recall uses, and a single chunk that still does not fit is refused.
The late archive cleanup also spared a recreated thread. DELETE removes the rows,
awaits the sandbox pass, and only then sweeps the archive; another tab can POST
the same id in that window and its generation can archive turns under it. The
sandbox pass re-checks for exactly that, and this now does too, so the recreated
chat keeps its memory.
* Studio: keep the configured default when top_k is omitted
Budgeting an omitted top_k by dividing the whole budget treated the room as a
target rather than a cap: on a 128K chat with most of its window free it asked
the archive for 200 passages, past the configured default of 4 and past the
ceiling of 8 that the model's own value is held to. The default now applies as
before, with the room and the ceiling capping it.
* Studio: cut the boundary at the newest user turn
Excluding only the branch's last message assumed the newest user turn is last,
and a continued assistant message puts a prefill after it. The user turn then
stayed in the identity scan, and since inline recall rewrites it into a new dict
it read as evicted, inflating the boundary by one and costing the next request a
live message. The scan now stops at the newest user turn, which is what is
protected in any case.
* Studio: require an archived turn to end where the live message does
Editing a reply by keeping it and adding to it ("No" becoming "No, correction:
yes") left every probe matching, so the pre-edit copy stayed eligible and a
search could return "No" as the answer with the correction nowhere in it. The
scan now reports where it finished inside the message, and the turn is accepted
only if it reaches the end of it.
Tool results and tool arguments are exempt: render_turn cuts them, so their
probes are prefixes by design and demanding the end would retire every turn that
carried one.
* Studio: check both ends of an archived turn, and charge dense text properly
An edit that keeps the old text and adds to it leaves every probe matching
whichever side it adds on. The end anchor caught "No" becoming "No, correction:
yes"; it did not catch "Correction: no", which ends exactly where the archived
copy does. The scan now also reports where the run into the final message began,
and the turn is accepted only if it covers that message end to end. render_turn
only ever cuts the tail, so the start is required unconditionally.
A conversation search sized itself with the shared estimator, which charges four
characters per token. That is an English rule: CJK and emoji run closer to one
token per character, so a result could be accepted at a quarter of its real size
and then land in the current tool exchange, which the window cannot evict. No
exact counter is reachable there, the provider loop having no tokenizer at all,
so non-ASCII characters are charged one token each and the rest at the usual
rate. Measured on 320 CJK characters: 87 tokens claimed, 320 charged.
* Studio: earn the compaction headroom, and keep this turn's tool results
The headroom cut every overflowing prompt to about 75% of its budget, including
on paths that can never put the boundary back: an incognito chat, an API request
with no persisted thread, or one whose turns are not saved. There the deeper cut
buys nothing and cannot be recalled, so it is simply less history than plain
eviction would have kept, and turning the archive off did not restore the old
behaviour. The fit now takes the headroom only when the caller can restore the
boundary next request.
The branch the archive is filtered against was the messages the client sent. A
long agent run evicts, and archives, tool exchanges it created itself, and those
were then refused as an abandoned branch: the model could not search back a tool
result it still needed to answer. The branch is now accumulated across the
request, so this turn's own exchanges count as live, while a sibling response
from a Retry still does not.
* Studio: read the branch as replies when restoring a boundary
The rows the rolling window checks are assistant replies, but the branch they
were checked against was every message of it, flattened without roles. An
abandoned "Done" left by a Retry therefore rode in on a live user message that
merely contains it ("not done yet I think"), and with no exact match to prefer,
its much larger boundary was applied to a branch that never had that reply.
The branch representation now takes a role filter and the boundary asks for
assistant messages only. A branch with no reply of its own has no boundary to
restore, so it reports none rather than falling back to the unfiltered check.
* Studio: account for every message an archived turn claims
The anchors covered the final message only, so an edit to the question
underneath it left the archived copy eligible, whichever side the text was added
on. The scan now checks each message the run touches: it must be matched from
the first character, and nothing may be left over when the run moves on.
Tool calls are exempt, and the exemption covers the whole message. The store
keeps a call as a structured part, so the live text carries the tool name and
both spellings of the arguments, spaced and compact, while the archived copy has
one line of one of them. Nothing there lines up character for character, and
demanding it would retire every tool turn in the archive.
* Studio: re-embed archived turns after a model change, and batch the first pass
Deduplication was by hash alone, but dense search only reads documents whose
recorded embedder matches the query's. A turn archived under the previous model
was therefore skipped on every later compaction while being invisible to every
paraphrased search, permanently. The check now compares the identity too and
replaces the copy, which is what ingestion does with a re-uploaded file.
Embedding also ran once per evicted turn. A first compaction of a long chat sent
dozens of one-item jobs back to back, and both backends serialise them, so the
reply waited for all of them. The turns are chunked first and embedded in one
pass, with the per-turn write and its duplicate check unchanged. Measured on 40
turns: 40 embedding calls before, 1 after.
* Studio: re-pin the permission read the deep-research contract lost to #8686
The contract still asserted the literal `const permissionMode = loadPermissionMode();`,
which #8686 removed when it put a chat-scoped override in front of that read. This branch
already carries the refactored store, so Repo tests (CPU) cannot pass here or on anything
stacked above it.
Ported main's replacement rather than dropping the check. The read is still the contract,
and it is scoped to the deep-research setter so it cannot be satisfied by the initial-state
constant, which is a different property and would keep passing if the read were deleted.
* Studio: re-derive the compaction boundary by position, not by count
`boundary_messages` is an absolute count against the transcript it was counted
on. Delete an already-evicted prompt from a compacted thread and the front of
that transcript gets shorter, but the count is replayed unchanged, so the cut
lands that many messages too deep and evicts turns that are still live.
Reproduced on `fit_rolling_context` with a 30-turn chat at ctx 2000: a fresh fit
keeps from turn 16; deleting one already-evicted pair and replaying the stored
count keeps from turn 17, one live turn lost. The loss is then baked in, since
the next turn records the boundary against the shortened transcript.
Stamp `boundary_anchor`, the text of the first message the fit KEPT, next to the
count at all five sites that write it, and look it up on the next request's own
branch. The anchor is only ever allowed to make the boundary SHALLOWER, so a
stale, ambiguous or repeated anchor costs one extra compaction and can never
evict a live turn. Rows written before the key exists behave exactly as today.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: say what the archive kill switch does and does not cover
RAG_CONVERSATION_ARCHIVE=0 stops indexing, recall and the recall reserve. It
does not turn the rolling window back into what it was before it learned to
compact in discrete events: the headroom and the sticky boundary belong to the
window, and the headroom has its own knob. "Off restores plain eviction" read
as though it covered all of it.
Comment only. Gating the window's own behaviour on this flag would make a host
with no sqlite-vec silently compact differently from one that has it.
* Studio: charge dense text honestly, and reject an edit inside an archived turn
Two fixes.
The safetensors search budget priced what is already in the prompt at four
characters per token. That is about right for English and roughly half the truth
for CJK and emoji: measured on an 81-message Japanese chat, 1295 estimated
against 2737 real, reporting 1777 tokens of room where 335 remained. The tool
result is already charged a token per non-ASCII character; this is the same rule
applied to the spend, which was the missing half. That path runs no rolling fit,
so nothing downstream recovers once the exchange lands. The GGUF site takes it
only on its fallback leg, since the fit's exact tokenizer count needs no
correction, and eviction keeps the flat estimate: making eviction pessimistic
would drop history a request could have kept.
The probe scan anchored the start and the end of a run but not the middle, so a
correction inserted between two archived lines matched both probes with the new
line sitting unexamined in the gap, and the pre-edit turn stayed recallable.
Verified on "A\nB" becoming "A\ncorrection\nB". A gap is now allowed only where
it is a label render_turn wrote and the probe had therefore stripped, since a
pasted chat log legitimately carries its own "user:" lines, and the tool-call
exemption still covers the rest of its message.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: end the tool-call anchor exemption where the call ends
A stored tool call cannot line up character for character with the live text,
because the store keeps arguments as an object and offers both JSON spellings,
so the cursor after one is not exact and the character anchors have to relax.
That relaxation was then held for the rest of the message.
An assistant turn carrying both a call and text therefore stayed matched after a
correction was appended to the text: "old answer" becoming "old answer,
correction: new answer" left the archived pre-edit turn eligible for recall.
Cleared once an ordinary text probe matches, since the cursor is exact again from
there. Verified on that exact pair: intact still eligible, edited no longer.
* Studio: count a video part as media in the rolling preflight
`_inject_video_part` writes llama.cpp's own `input_video` part into the same
message list the fit is then handed, and the media predicate did not list it.
Video prompts therefore ran the rolling preflight, which is skipped for media
precisely because `/apply-template` token counting does not include the sampled
video tokens. The prompt could be certified as fitting when it does not, or lose
history it never needed to lose, and still fail with context_length_exceeded.
* Studio: keep which tool was called in the branch probe
`render_turn` writes "assistant called <name>: <args>", and the whole label was
stripped before probing, name included, so only the arguments decided. A retry
that kept the arguments and the result and changed only the tool left the
archived pre-edit turn matching: verified on the same arguments moving from
terminal to python, which stayed eligible and could be recalled as though the
old call had happened on this branch.
The name is now its own probe, ahead of the arguments, which is the order
`_probe_text` renders a live call in. The label itself still has to go, since it
exists only in the archived copy, but the name does not: the live text carries
it. `render_turn`'s "tool" fallback for a nameless call is excluded, because the
live text has no name there either.
Not an exact anchor: the name is matched as a substring of the message, so a
tool whose name appears inside the arguments of another can still match. This
closes the ordinary retry case rather than every construction of it.
* Studio: do not archive a search the model asked for
`_is_injected` recognises the ids this feature and the RAG auto-inject generate,
and a search the MODEL issues carries neither: the parser gives it an ordinary
`call_N` id. Both the call and the passages it retrieved were therefore indexed
as fresh conversation, and a second search archived the first one's output inside
its own, one nesting level per distinct search, each copy competing for the four
recall slots.
Dropped by tool NAME instead, reusing `RAG_SEARCH_TOOLS`, which both loops
already share. Only the retrieval parts go: a reply that follows a search is real
conversation and is still archived, and an assistant message carrying a retrieval
call beside an ordinary one keeps the ordinary one and its result.
* Delete the archived conversation even when its thread id comes back
The late archive sweep skipped the whole scope when another tab had recreated
the id, which spared the recreated chat's memory but kept the deleted
conversation's too. The scope is keyed by thread id alone, so those turns stayed
recallable in the new chat with nothing left to sweep them, and the endpoint
reported success.
Take a cutoff before the rows go and bound the delete by created_at when the id
has come back: everything archived before the delete was accepted belongs to the
deleted conversation, everything after to the new one. Applied at DELETE
/threads, project delete and clear all, and on the no-vec0 path so it cannot
delete more than the vec0 path does.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: tighten rolling context window comments
* Stop folding case when validating an archived turn against its branch
A turn corrected only in capitalisation still matched its archived copy, so the
branch filter kept the pre-edit document live and a recall could answer with the
spelling the user had just fixed. Nothing else retires it: the edit changes the
digest, so the corrected turn is written as a new document and the branch filter
is the only thing that could have dropped the old one.
Trim both ends rather than only the right. Keeping leading whitespace looks
tighter and is worse, since render_turn strips the whole message and a live turn
opening on a space or a newline then starts its run at a non-zero offset and is
retired outright. Measured: ' hello there' and a pasted block opening on a
newline both went from live to retired. Indentation-only edits stay tolerated in
both directions as a result, because probes are matched per line and as
substrings; closing that needs the splitter to carry offsets.
Also fixes an unrelated retirement found while measuring this. When
render_turn's 4000 character cut lands exactly on a newline the marker becomes a
line of its own, strips to empty, and was dropped along with the truncation flag,
so the last real probe read as complete and an unedited over-cap tool result was
retired. Measured on a 900 line result: no query could return it.
Note for anyone testing an upgrade in place: boundary anchors written by an
earlier build of this branch were stored folded, so they will not match until the
next compaction rewrites them. The anchor can only ever move the boundary
shallower, so the cost is one extra compaction.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Let search_conversation through the Anthropic gate, and probe a tool call before its answer
Adding the schema to ALL_TOOLS made the Anthropic selector pick it while
_ANTHROPIC_UNPROMPTED_SAFE_TOOLS still listed only web_search and
search_knowledge_base, so the pre-switch guard classified a read-only tool as
confirmation-gated. Measured: enable_tools with enabled_tools
['search_conversation'] returned 400 on auto and on the omitted default, with
the terminal/python message, even though is_potentially_unsafe_tool_call marks
it always safe and that set is documented as mirroring it. Pre-PR the same
request was served, so this is a regression the PR introduced.
Separately, _probe_text bucketed a whole message as call, text, result, which
only holds while the text came before the call. A persisted row whose tool call
is followed by the model's final answer, the ordinary agent turn, renders as
call, result, answer, so _scan_probes advanced past the answer to find the
result and could not find it again. Measured end to end on the persisted-row
transcript: recall returned the user's question alone and the document holding
the answer was filtered out. The buckets now flush when text arrives after a
call, the same rule the replay serializer uses; text written before a call still
rides ahead of it.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Record the transcript span, and price the recall budget exactly when nothing was trimmed
archive_messages bounds the branch check's run, and it was recording the number
of messages ARCHIVED rather than the span of the turn. The two differ whenever
_archivable drops something: an assistant batch that called search_conversation
alongside an ordinary tool archives three messages while the live transcript
holds four. Measured, validation failed at three and passed at four, so a
perfectly valid ordinary-tool exchange and the answer that followed were
rejected as off-branch and could never be recalled.
fit_rolling_context returns None when it drops nothing, so a prompt that simply
fits, after a context-length increase or on a shorter branch, left the recall
budget to a character estimate that cannot see the template's own framing.
Measured on a request whose real prompt is about 2800 tokens of a 3584 budget:
the budget came back 3512, nearly the whole window, so the recall could append a
passage that does not fit and the next iteration cannot evict it again, since the
current tool exchange is protected. Priced exactly instead, once per request and
only when the model actually reaches for a retrieval tool on a request that did
not truncate. A failure falls back to the old estimate.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Price the recall result with the model tokenizer, recount every search, and anchor on tool calls
Count the search result with the caller's own token counter instead of the ASCII
heuristic, recompute the exact prompt count on every search rather than once per
turn, and take the boundary anchor from a message's tool calls when it has no text,
so a tool-call message no longer records an empty anchor and disables the rebase.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Widen a reused turn's window to its longest span, and price every search exactly
A three-message tool exchange and a four-message batch holding the same exchange render
identically, so the second is skipped as a duplicate and inherited the shorter window,
which the branch check then used to reject it. Grow the stored span instead, upwards
only. Count the prompt exactly on every conversation search rather than only when the
fit dropped nothing.
* Let an overlapping chunk restart far enough back to cover its carried messages
A chunk's leading overlap is the previous chunk's tail, and it can carry a whole short
message with it, so resuming the scan at the message the previous chunk finished in
could never match it and retired unedited turns as off-branch. The forward position is
still tried first and the walk back stops at the previous chunk's own opening.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Widen the span on the locked duplicate path as well
Both turns can arrive in one compaction: the pre-check clears both before either is
written, the shorter is written first, and the longer then met the re-check under the
write lock and left the window at the shorter figure.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Bound the persisted anchor, and spend any room at all on one recall attempt
The anchor rides in every truncation event and every assistant turn's metadata while the
boundary stays sticky, so a large pasted message was copied across the thread; cap it at
a head on both sides, where the read side already clamps only shallower. And a budget
below one chunk is not no budget: CHUNK_TOKENS is a ceiling, not the size of a turn, so
try one and let the exact recount reject it.
---------
Co-authored-by: alkinun <alkinunl@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: danielhanchen <unslothshared@gmail.com>
This commit is contained in:
Daniel Han
GitHub
parent
645ca5c7bd
commit
18b97f8b72
30 changed files with 7877 additions and 198 deletions
|
|
@ -6,11 +6,21 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from collections.abc import Callable
|
||||
from typing import Any, Optional
|
||||
|
||||
_OMITTED_TOOL_EXCHANGE = "[Earlier tool exchange omitted from the rolling context window.]"
|
||||
|
||||
# How far BELOW the prompt budget a compaction trims, as a fraction of that budget.
|
||||
# Trimming to exactly the budget puts the next turn over it again, so the boundary creeps
|
||||
# forward every turn: llama-server's prefix cache dies each time and there is no discrete
|
||||
# compaction event to report. Taking a chunk out in one go buys a stretch of turns with a
|
||||
# fixed head, at the cost of the headroom itself, hence a minority of the budget.
|
||||
_COMPACTION_HEADROOM_RATIO = max(
|
||||
0.0, min(0.9, float(os.environ.get("ROLLING_COMPACTION_HEADROOM_RATIO", "0.25")))
|
||||
)
|
||||
|
||||
|
||||
def estimate_message_tokens(message: dict) -> int:
|
||||
try:
|
||||
|
|
@ -23,21 +33,39 @@ def estimate_messages_tokens(messages: list[dict]) -> int:
|
|||
return sum(estimate_message_tokens(message) for message in messages)
|
||||
|
||||
|
||||
def truncate_oldest_messages(
|
||||
messages: list[dict],
|
||||
keep_ratio: float,
|
||||
*,
|
||||
protected_message_ids: Optional[set[int]] = None,
|
||||
) -> tuple[list[dict], int]:
|
||||
"""Drop complete oldest turns while preserving system messages and the latest turn.
|
||||
def estimate_messages_tokens_dense(messages: list[dict]) -> int:
|
||||
"""The same estimate, but honest about text that tokenises densely.
|
||||
|
||||
Normal user/assistant turns stay together. Each assistant tool call starts a
|
||||
separate group containing its tool results, so long agent runs can evict old
|
||||
exchanges without orphaning results or losing the task that initiated them.
|
||||
Four characters per token is about right for English and badly wrong for CJK and
|
||||
emoji, which run closer to one token per character. Measured on an 81-message CJK
|
||||
chat: 1295 estimated against 2737 real, a 2.1x undercount, which a caller sizing a
|
||||
search budget spends as room it does not have. `_conversation_search_tokens` already
|
||||
charges non-ASCII a token each when pricing the result; this is the same rule applied
|
||||
to the spend, which is the half that was missing.
|
||||
|
||||
Deliberately NOT the default. `truncate_oldest_messages` evicts on the flat estimate,
|
||||
and making eviction pessimistic would drop history a request could have kept. Only a
|
||||
caller deciding how much room is LEFT wants this, where erring high is the safe side.
|
||||
"""
|
||||
if not messages or keep_ratio >= 1.0:
|
||||
return messages, 0
|
||||
total = 0
|
||||
for message in messages:
|
||||
try:
|
||||
text = json.dumps(message, ensure_ascii = False)
|
||||
except Exception:
|
||||
total += 1
|
||||
continue
|
||||
dense = sum(1 for char in text if ord(char) > 127)
|
||||
total += max(1, dense + (len(text) - dense) // 4)
|
||||
return total
|
||||
|
||||
|
||||
def group_turns(messages: list[dict]) -> list[list[dict]]:
|
||||
"""Split messages into the turn groups the rolling window evicts as single units.
|
||||
|
||||
Each assistant tool call starts its own group holding its tool results, so long agent
|
||||
runs evict old exchanges without orphaning results. Exposed so callers that act on
|
||||
evicted turns use the same unit the evictor does.
|
||||
"""
|
||||
groups: list[list[dict]] = []
|
||||
for message in messages:
|
||||
starts_tool_exchange = message.get("role") == "assistant" and bool(
|
||||
|
|
@ -54,6 +82,35 @@ def truncate_oldest_messages(
|
|||
groups.append([message])
|
||||
else:
|
||||
groups[-1].append(message)
|
||||
return groups
|
||||
|
||||
|
||||
def evicted_messages(before: list[dict], after: list[dict]) -> list[dict]:
|
||||
"""Messages present in ``before`` and absent from ``after``, in their original order.
|
||||
|
||||
Identity, not equality: the truncation helpers reuse the same dict objects, and a
|
||||
chat can contain two byte-identical turns ("continue" twice) that equality collapses.
|
||||
"""
|
||||
kept = {id(message) for message in after}
|
||||
return [message for message in before if id(message) not in kept]
|
||||
|
||||
|
||||
def truncate_oldest_messages(
|
||||
messages: list[dict],
|
||||
keep_ratio: float,
|
||||
*,
|
||||
protected_message_ids: Optional[set[int]] = None,
|
||||
min_dropped: int = 0,
|
||||
) -> tuple[list[dict], int]:
|
||||
"""Drop complete oldest turns while preserving system messages and the latest turn.
|
||||
|
||||
``min_dropped`` keeps evicting past the point where the prompt fits, so a thread can
|
||||
re-apply the boundary it already compacted to instead of one that slides every turn.
|
||||
"""
|
||||
if not messages or (keep_ratio >= 1.0 and min_dropped <= 0):
|
||||
return messages, 0
|
||||
|
||||
groups = group_turns(messages)
|
||||
|
||||
if len(groups) <= 1:
|
||||
return messages, 0
|
||||
|
|
@ -102,7 +159,7 @@ def truncate_oldest_messages(
|
|||
|
||||
dropped_groups: set[int] = set()
|
||||
for unit in eviction_units:
|
||||
if current_estimate <= target_estimate:
|
||||
if current_estimate <= target_estimate and dropped >= min_dropped:
|
||||
break
|
||||
dropped_groups.update(unit)
|
||||
for group_index in unit:
|
||||
|
|
@ -117,8 +174,8 @@ def truncate_oldest_messages(
|
|||
for index, group in enumerate(groups):
|
||||
if index not in dropped_groups:
|
||||
if kept and kept[-1].get("role") == "user" and group and group[0].get("role") == "user":
|
||||
# Strict chat templates reject adjacent user turns. This occurs when
|
||||
# an internal tool re-prompt follows an evicted exchange.
|
||||
# Strict chat templates reject adjacent user turns, which an internal
|
||||
# tool re-prompt after an evicted exchange would produce.
|
||||
kept.append({"role": "assistant", "content": _OMITTED_TOOL_EXCHANGE})
|
||||
kept.extend(group)
|
||||
return kept, dropped
|
||||
|
|
@ -132,11 +189,9 @@ def messages_have_media(messages: list[dict]) -> bool:
|
|||
for part in content:
|
||||
if not isinstance(part, dict):
|
||||
continue
|
||||
# `input_video` is llama.cpp's own part type, written by `_inject_video_part`
|
||||
# into the same message list this fit then sees. Missing from the list, a video
|
||||
# prompt ran the rolling preflight even though `/apply-template` does not count
|
||||
# the sampled video tokens, so it could be certified as fitting when it does
|
||||
# not, or lose history it did not need to lose.
|
||||
# `input_video` is llama.cpp's own part type (written by `_inject_video_part`);
|
||||
# missing it here would let a video prompt take the rolling preflight, whose
|
||||
# `/apply-template` count omits the sampled video tokens.
|
||||
if part.get("type") in (
|
||||
"image_url",
|
||||
"input_audio",
|
||||
|
|
@ -148,6 +203,33 @@ def messages_have_media(messages: list[dict]) -> bool:
|
|||
return False
|
||||
|
||||
|
||||
def prompt_budget(context_length: int, max_tokens: Optional[int]) -> int:
|
||||
"""Tokens available to the PROMPT, once room for the reply is set aside.
|
||||
|
||||
Exported so the recall sizing and the client's over-long-request explanation share
|
||||
this formula rather than each keeping a copy that drifts from the actual fit.
|
||||
"""
|
||||
if context_length <= 1:
|
||||
return context_length
|
||||
requested = max_tokens if max_tokens is not None and max_tokens > 0 else context_length
|
||||
return context_length - min(requested, max(1, context_length // 4))
|
||||
|
||||
|
||||
def _latest_turn_tokens(messages: list[dict], count_tokens: Callable[[list[dict]], int]) -> int:
|
||||
"""Tokens in the newest message, estimated if the template refuses to render it.
|
||||
|
||||
A tool loop can reach the does-not-fit diagnosis with a tool result last, which strict
|
||||
templates reject on its own; letting that raise would abort the fit and tell the user
|
||||
nothing. An approximate number beats a diagnosis that never arrives.
|
||||
"""
|
||||
if not messages:
|
||||
return 0
|
||||
try:
|
||||
return count_tokens(messages[-1:])
|
||||
except Exception:
|
||||
return estimate_messages_tokens(messages[-1:])
|
||||
|
||||
|
||||
def fit_rolling_context(
|
||||
messages: list[dict],
|
||||
*,
|
||||
|
|
@ -155,27 +237,72 @@ def fit_rolling_context(
|
|||
max_tokens: Optional[int],
|
||||
count_tokens: Callable[[list[dict]], int],
|
||||
protected_message_ids: Optional[set[int]] = None,
|
||||
reserve_tokens: int = 0,
|
||||
sticky_dropped: int = 0,
|
||||
keeps_boundary: bool = False,
|
||||
) -> tuple[list[dict], Optional[dict[str, Any]]]:
|
||||
"""Fit a chat into its real context by dropping oldest complete turns.
|
||||
|
||||
The exact tokenizer/template count decides whether trimming is needed. The
|
||||
inexpensive estimator only chooses candidate turns; exact recounts verify the
|
||||
result. The current turn is never clipped, so an irreducibly large request still
|
||||
reaches llama-server's normal context-length error.
|
||||
The exact tokenizer/template count decides whether trimming is needed; the cheap
|
||||
estimator only picks candidate turns. The current turn is never clipped, so an
|
||||
irreducibly large request still reaches llama-server's context-length error.
|
||||
|
||||
``reserve_tokens`` leaves room for what the caller adds back after fitting (recalled
|
||||
turns). It deliberately does not affect whether to trim at all, so a chat that fits
|
||||
today is never evicted just because the reserve would not fit alongside it.
|
||||
|
||||
``sticky_dropped`` is the boundary this thread last compacted to, in messages,
|
||||
re-applied before anything else and moved only if what is left still does not fit.
|
||||
Without it the fit is stateless (the client re-sends the whole transcript each turn)
|
||||
so the boundary slides every reply; with it plus ``_COMPACTION_HEADROOM_RATIO`` of
|
||||
slack, compaction is an occasional event the prefix cache can survive.
|
||||
|
||||
Acceptance is still checked against the untightened ``prompt_target``: falling short
|
||||
of the headroom is not a failure to fit.
|
||||
"""
|
||||
if context_length <= 1:
|
||||
return messages, None
|
||||
|
||||
requested_headroom = max_tokens if max_tokens is not None and max_tokens > 0 else context_length
|
||||
output_headroom = min(requested_headroom, max(1, context_length // 4))
|
||||
prompt_target = context_length - output_headroom
|
||||
prompt_target = prompt_budget(context_length, max_tokens)
|
||||
fitted = list(messages)
|
||||
initial_tokens = count_tokens(fitted)
|
||||
current_tokens = initial_tokens
|
||||
dropped_total = 0
|
||||
|
||||
while current_tokens > prompt_target:
|
||||
keep_ratio = min(0.95, prompt_target / max(1, current_tokens))
|
||||
# Phase one: put the boundary back where this thread already had it, so a compacted
|
||||
# thread stops compacting further every turn. Gated on the prompt not already
|
||||
# fitting: a saved boundary describes the branch it was measured on, and after a
|
||||
# rollback it would evict most of a chat that comfortably fits.
|
||||
if sticky_dropped > 0 and initial_tokens > prompt_target:
|
||||
candidate, dropped = truncate_oldest_messages(
|
||||
fitted,
|
||||
1.0,
|
||||
protected_message_ids = protected_message_ids,
|
||||
min_dropped = sticky_dropped,
|
||||
)
|
||||
if dropped:
|
||||
fitted = candidate
|
||||
dropped_total = dropped
|
||||
current_tokens = count_tokens(fitted)
|
||||
|
||||
# Phase two, only if what is left still does not fit: move the boundary, taking a
|
||||
# chunk out rather than skimming to the brim so it can stay put for a while.
|
||||
trim_target = prompt_target
|
||||
if current_tokens > prompt_target:
|
||||
# Summed, not max()'d: the reserve is spent immediately on recalled passages, so
|
||||
# counting it as headroom would hand back room that is already taken.
|
||||
#
|
||||
# And only for a caller that can put the boundary back next request. The headroom
|
||||
# buys quiet turns between compactions by cutting deeper than needed, which is a
|
||||
# bargain only if the deeper cut is remembered. An incognito chat, an API request
|
||||
# with no persisted thread, or a request whose turns are not saved gets neither
|
||||
# the boundary nor a recall of what went, so there it is simply 25% less history
|
||||
# than plain eviction would have kept.
|
||||
headroom = int(prompt_target * _COMPACTION_HEADROOM_RATIO) if keeps_boundary else 0
|
||||
trim_target = max(1, prompt_target - reserve_tokens - headroom)
|
||||
|
||||
while current_tokens > trim_target:
|
||||
keep_ratio = min(0.95, trim_target / max(1, current_tokens))
|
||||
candidate, dropped = truncate_oldest_messages(
|
||||
fitted,
|
||||
keep_ratio,
|
||||
|
|
@ -187,7 +314,29 @@ def fit_rolling_context(
|
|||
dropped_total += dropped
|
||||
current_tokens = count_tokens(fitted)
|
||||
|
||||
if dropped_total == 0 or current_tokens > prompt_target:
|
||||
if current_tokens > prompt_target:
|
||||
# Evicted everything evictable and it still does not fit. Return the ORIGINAL
|
||||
# messages, not the partial eviction: the request is refused either way, so
|
||||
# dropping turns off a doomed request loses them for nothing. The diagnosis is
|
||||
# still worth returning; llama-server's own error reports the size of the WHOLE
|
||||
# conversation and advises shortening it, which cannot work when what is left is
|
||||
# the system prompt plus the latest turn. Consumers all gate on `fits`.
|
||||
return messages, {
|
||||
"fits": False,
|
||||
"dropped_messages": 0,
|
||||
"prompt_tokens_before": initial_tokens,
|
||||
"prompt_tokens_after": initial_tokens,
|
||||
# Floor for the conversation, and how much of it is the message just sent:
|
||||
# together they say whether the chat or the single message is the problem.
|
||||
"irreducible_tokens": current_tokens,
|
||||
"latest_turn_tokens": _latest_turn_tokens(messages, count_tokens),
|
||||
# ...and whose message that is. In a tool loop the last message is often a
|
||||
# tool result, which the user did not write and cannot shorten.
|
||||
"latest_turn_role": str(messages[-1].get("role") or "") if messages else "",
|
||||
"context_length": context_length,
|
||||
"prompt_target": prompt_target,
|
||||
}
|
||||
if dropped_total == 0:
|
||||
return messages, None
|
||||
return fitted, {
|
||||
"dropped_messages": dropped_total,
|
||||
|
|
|
|||
|
|
@ -1035,6 +1035,9 @@ class InferenceBackend:
|
|||
rag_scope = rag_scope,
|
||||
reasoning_prefilled = reasoning_prefilled,
|
||||
continue_final_message = continue_final_message,
|
||||
# So a conversation search can be sized against what this model can hold.
|
||||
context_length = _model_info.get("context_length"),
|
||||
max_tokens = max_new_tokens,
|
||||
)
|
||||
|
||||
def generate_chat_response(
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -1878,6 +1878,9 @@ class InferenceOrchestrator:
|
|||
permission_mode = permission_mode,
|
||||
reasoning_prefilled = reasoning_prefilled,
|
||||
continue_final_message = continue_final_message,
|
||||
# So a conversation search can be sized against what this model can hold.
|
||||
context_length = _model_info.get("context_length"),
|
||||
max_tokens = max_new_tokens,
|
||||
)
|
||||
|
||||
def generate_with_adapter_control(
|
||||
|
|
|
|||
|
|
@ -33,6 +33,7 @@ from core.inference.tool_call_parser import (
|
|||
NUDGE_TOOL_CALLS_STATUS,
|
||||
RAG_MAX_SEARCHES_PER_TURN,
|
||||
RAG_SEARCH_CAP_NUDGE,
|
||||
RAG_SEARCH_TOOLS,
|
||||
StreamingMarkupStripper,
|
||||
TOOL_XML_SIGNALS,
|
||||
is_reprompt_repeat,
|
||||
|
|
@ -432,22 +433,26 @@ def _tool_event_provenance(**flags: object) -> dict[str, object]:
|
|||
return tool_event_provenance(**flags)
|
||||
|
||||
|
||||
def _accepts_output_callback(func: Callable[..., str]) -> bool:
|
||||
"""Whether an injectable ``execute_tool`` supports ``output_callback``.
|
||||
def _accepts_kwarg(func: Callable[..., str], name: str) -> bool:
|
||||
"""Whether an injectable ``execute_tool`` supports the keyword ``name``.
|
||||
|
||||
The loop's ``execute_tool`` is a parameter (tests inject fakes), so forward
|
||||
the live-output kwarg only when the callable declares it or takes ``**kwargs``.
|
||||
an optional kwarg only when the callable declares it or takes ``**kwargs``.
|
||||
"""
|
||||
try:
|
||||
sig = inspect.signature(func)
|
||||
except (TypeError, ValueError):
|
||||
return False
|
||||
params = sig.parameters
|
||||
if "output_callback" in params:
|
||||
if name in params:
|
||||
return True
|
||||
return any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values())
|
||||
|
||||
|
||||
def _accepts_output_callback(func: Callable[..., str]) -> bool:
|
||||
return _accepts_kwarg(func, "output_callback")
|
||||
|
||||
|
||||
def _call_single_turn(single_turn, conversation: list, active_tools: list[dict]):
|
||||
"""Call a single-turn generator with active tool schemas when supported."""
|
||||
try:
|
||||
|
|
@ -479,6 +484,8 @@ def run_safetensors_tool_loop(
|
|||
continue_final_message: bool = False,
|
||||
markup = None,
|
||||
renderable_tools = None,
|
||||
context_length: Optional[int] = None,
|
||||
max_tokens: Optional[int] = None,
|
||||
) -> Generator[dict, None, None]:
|
||||
"""Drive an agentic tool loop on top of a cumulative-text generator.
|
||||
|
||||
|
|
@ -503,6 +510,11 @@ def run_safetensors_tool_loop(
|
|||
* ``{"type": "tool_end", "tool_name", "tool_call_id", "result"}``
|
||||
"""
|
||||
conversation = list(messages)
|
||||
# The branch this request is on, before the loop appends anything. A GGUF-compacted
|
||||
# thread keeps its archive across a switch to safetensors, so search_conversation is
|
||||
# advertised here too and needs the same filtering: the stored rows are the whole
|
||||
# DAG, and Retry leaves the replaced response in them.
|
||||
request_branch = list(messages)
|
||||
|
||||
# Mirrors the GGUF loop: "full" and bypass_permissions are the same switch;
|
||||
# unset defaults to "auto", unknown falls back to the stricter "ask"; "off"
|
||||
|
|
@ -1277,7 +1289,7 @@ def run_safetensors_tool_loop(
|
|||
eff_timeout = None if tool_call_timeout >= 9999 else tool_call_timeout
|
||||
# RAG: cap paraphrased KB re-searches that slip past the dup guard.
|
||||
if (
|
||||
decision.tool_name == "search_knowledge_base"
|
||||
decision.tool_name in RAG_SEARCH_TOOLS
|
||||
and kb_search_count >= RAG_MAX_SEARCHES_PER_TURN
|
||||
):
|
||||
result = RAG_SEARCH_CAP_NUDGE
|
||||
|
|
@ -1295,6 +1307,31 @@ def run_safetensors_tool_loop(
|
|||
rag_scope = rag_scope,
|
||||
disable_sandbox = bypass_permissions,
|
||||
)
|
||||
if _accepts_kwarg(execute_tool, "conversation_branch"):
|
||||
kwargs["conversation_branch"] = request_branch
|
||||
# And the room the model has left, as the GGUF loop does: without a
|
||||
# budget the tool's clamp is skipped and a model-chosen top_k of 8
|
||||
# appends roughly 4K tokens to an already full prompt.
|
||||
if context_length and _accepts_kwarg(
|
||||
execute_tool, "conversation_budget_tokens"
|
||||
):
|
||||
from core.inference.context_window import (
|
||||
estimate_messages_tokens_dense,
|
||||
prompt_budget,
|
||||
)
|
||||
|
||||
# Dense, unlike the eviction estimator: four characters per token
|
||||
# undercounts CJK and emoji by about half, and this path has no
|
||||
# rolling fit to recover if the tool exchange it sizes then puts
|
||||
# the next prompt over the window. Measured on an 81-message CJK
|
||||
# chat: 1295 estimated against 2737 real, reporting 1777 tokens of
|
||||
# room where 335 remained.
|
||||
kwargs["conversation_budget_tokens"] = max(
|
||||
0,
|
||||
prompt_budget(int(context_length), max_tokens)
|
||||
- estimate_messages_tokens_dense(conversation)
|
||||
- estimate_messages_tokens_dense(tools or []),
|
||||
)
|
||||
if _accepts_output_callback(execute_tool):
|
||||
kwargs["output_callback"] = _output_callback
|
||||
return execute_tool(_decision.tool_name, _decision.arguments, **kwargs)
|
||||
|
|
@ -1309,7 +1346,7 @@ def run_safetensors_tool_loop(
|
|||
except Exception as exc:
|
||||
logger.exception("Tool %s raised: %s", decision.tool_name, exc)
|
||||
result = f"Error: tool raised an exception: {exc}"
|
||||
if decision.tool_name == "search_knowledge_base":
|
||||
if decision.tool_name in RAG_SEARCH_TOOLS:
|
||||
kb_search_count += 1
|
||||
|
||||
completion = tool_controller.record_result(decision, result)
|
||||
|
|
|
|||
|
|
@ -72,6 +72,7 @@ from core.inference.tool_loop_controller import (
|
|||
)
|
||||
from core.inference.tool_stream_exec import (
|
||||
TOOL_HEARTBEAT_INTERVAL_S,
|
||||
accepts_kwarg,
|
||||
accepts_output_callback,
|
||||
stream_tool_execution,
|
||||
)
|
||||
|
|
@ -729,6 +730,8 @@ async def stream_with_studio_tools(
|
|||
) -> AsyncIterator[str]:
|
||||
"""Stream a provider, execute requested Studio tools, continue to a final answer."""
|
||||
conversation = [dict(message) for message in run.messages]
|
||||
# Kept before the loop appends anything: this is the branch the request is on.
|
||||
request_branch = list(run.messages)
|
||||
remaining = policy.max_calls
|
||||
unlimited = remaining >= 9999
|
||||
session_id = run.session_id
|
||||
|
|
@ -1220,6 +1223,24 @@ async def stream_with_studio_tools(
|
|||
"rag_scope": rag_scope,
|
||||
"disable_sandbox": bypass_permissions,
|
||||
}
|
||||
# Provider loops share the local catalogue selector, so
|
||||
# search_conversation is advertised here too once a thread has an archive
|
||||
# and needs the same branch: the stored rows are the whole DAG, and Retry
|
||||
# leaves the replaced response in them.
|
||||
if accepts_kwarg(execute_tool, "conversation_branch"):
|
||||
kwargs["conversation_branch"] = request_branch
|
||||
# And a budget, so the tool's clamp is not skipped. Studio cannot measure
|
||||
# an external model's window, and a custom OpenAI-compatible endpoint can
|
||||
# be a small local server, so a model-chosen 8 chunks is roughly 4K tokens
|
||||
# replayed on every later call. Unmeasurable means one recall's worth.
|
||||
if accepts_kwarg(execute_tool, "conversation_budget_tokens"):
|
||||
try:
|
||||
from core.rag import config as rag_config
|
||||
kwargs["conversation_budget_tokens"] = max(
|
||||
1, int(rag_config.CHUNK_TOKENS)
|
||||
) * max(1, int(rag_config.CONVERSATION_ARCHIVE_TOP_K))
|
||||
except Exception:
|
||||
pass
|
||||
if accepts_output_callback(execute_tool):
|
||||
kwargs["output_callback"] = output_callback
|
||||
return execute_tool(call.tool_name, call.arguments, **kwargs)
|
||||
|
|
|
|||
|
|
@ -158,6 +158,10 @@ BUDGET_EXHAUSTED_NUDGE = (
|
|||
|
||||
# The exact-args dup guard misses paraphrased re-searches, so also cap KB searches per turn.
|
||||
RAG_MAX_SEARCHES_PER_TURN = 3
|
||||
# Both retrieval tools share that cap. Their top-K passages land in the current
|
||||
# exchange, which the rolling window protects and cannot evict, so an uncapped search
|
||||
# only ends the turn in a context-length error. Here so both tool loops agree on it.
|
||||
RAG_SEARCH_TOOLS = frozenset({"search_knowledge_base", "search_conversation"})
|
||||
RAG_SEARCH_CAP_NUDGE = (
|
||||
"You have already searched the knowledge base several times this turn. "
|
||||
"Do not search again. Answer the question using the passages already "
|
||||
|
|
|
|||
|
|
@ -32,22 +32,26 @@ from loggers import get_logger
|
|||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
def accepts_output_callback(func: Callable[..., str]) -> bool:
|
||||
"""Whether an injectable ``execute_tool`` supports ``output_callback``.
|
||||
def accepts_kwarg(func: Callable[..., str], name: str) -> bool:
|
||||
"""Whether an injectable ``execute_tool`` supports the keyword ``name``.
|
||||
|
||||
``execute_tool`` is replaceable (tests inject fakes / the pre-PR signature),
|
||||
so forward the kwarg only when the callable declares it or takes ``**kwargs``
|
||||
so forward a kwarg only when the callable declares it or takes ``**kwargs``
|
||||
(passing it unconditionally would ``TypeError`` on an old signature).
|
||||
"""
|
||||
try:
|
||||
params = inspect.signature(func).parameters
|
||||
except (TypeError, ValueError):
|
||||
return False
|
||||
if "output_callback" in params:
|
||||
if name in params:
|
||||
return True
|
||||
return any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values())
|
||||
|
||||
|
||||
def accepts_output_callback(func: Callable[..., str]) -> bool:
|
||||
return accepts_kwarg(func, "output_callback")
|
||||
|
||||
|
||||
# Cadence of heartbeat events while a tool blocks with no output. Well under
|
||||
# common proxy idle caps (Cloudflare ~100 s, nginx default 60 s).
|
||||
TOOL_HEARTBEAT_INTERVAL_S = 10.0
|
||||
|
|
|
|||
|
|
@ -4530,7 +4530,9 @@ def _render_html_reaches_network(arguments: dict) -> bool:
|
|||
# Tools that are read-only regardless of their arguments, so auto mode never has
|
||||
# to pause them and their safety needs no argument scan. render_html is handled
|
||||
# separately above because a networked canvas does need approval.
|
||||
_ALWAYS_SAFE_TOOLS = frozenset({"web_search", "search_knowledge_base"})
|
||||
# search_conversation only reads this chat's own past turns, so auto mode would otherwise
|
||||
# prompt for approval on every call.
|
||||
_ALWAYS_SAFE_TOOLS = frozenset({"web_search", "search_knowledge_base", "search_conversation"})
|
||||
|
||||
|
||||
def is_always_safe_tool(name: str) -> bool:
|
||||
|
|
@ -9772,6 +9774,33 @@ SEARCH_KNOWLEDGE_BASE_TOOL = {
|
|||
},
|
||||
}
|
||||
|
||||
SEARCH_CONVERSATION_TOOL = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "search_conversation",
|
||||
"description": (
|
||||
"Search earlier turns of THIS conversation that were removed from your "
|
||||
"context when it grew too long. Use it whenever the user refers to something "
|
||||
"discussed earlier that you cannot see, instead of saying you have no record "
|
||||
"of it."
|
||||
),
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"query": {
|
||||
"type": "string",
|
||||
"description": "Natural-language search query.",
|
||||
},
|
||||
"top_k": {
|
||||
"type": "integer",
|
||||
"description": "Max earlier turns to return.",
|
||||
},
|
||||
},
|
||||
"required": ["query"],
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
ALL_TOOLS = [
|
||||
WEB_SEARCH_TOOL,
|
||||
PYTHON_TOOL,
|
||||
|
|
@ -9779,6 +9808,7 @@ ALL_TOOLS = [
|
|||
EDIT_FILE_TOOL,
|
||||
RENDER_HTML_TOOL,
|
||||
SEARCH_KNOWLEDGE_BASE_TOOL,
|
||||
SEARCH_CONVERSATION_TOOL,
|
||||
]
|
||||
|
||||
|
||||
|
|
@ -9974,6 +10004,9 @@ def execute_tool(
|
|||
disable_sandbox: bool = False,
|
||||
output_callback = None,
|
||||
website_policy: dict | None = None,
|
||||
conversation_branch: list[dict] | None = None,
|
||||
conversation_budget_tokens: int | None = None,
|
||||
conversation_token_counter = None,
|
||||
) -> str:
|
||||
"""Execute a tool by name with the given arguments; returns a string.
|
||||
|
||||
|
|
@ -10001,6 +10034,21 @@ def execute_tool(
|
|||
effective_timeout,
|
||||
cancel_event,
|
||||
)
|
||||
if name == "search_conversation":
|
||||
# Scoped by thread id alone: the archive is this chat's own evicted turns, so it
|
||||
# works with or without a document rag_scope.
|
||||
return _search_knowledge_base_with_budget(
|
||||
arguments,
|
||||
{
|
||||
"thread_id": thread_id,
|
||||
"branch_messages": conversation_branch,
|
||||
"budget_tokens": conversation_budget_tokens,
|
||||
"token_counter": conversation_token_counter,
|
||||
},
|
||||
effective_timeout,
|
||||
cancel_event,
|
||||
search_fn = _search_conversation,
|
||||
)
|
||||
if name == "render_html":
|
||||
return _render_html_result(arguments)
|
||||
if name.startswith(MCP_TOOL_PREFIX):
|
||||
|
|
@ -10140,12 +10188,159 @@ def _search_knowledge_base(arguments: dict, rag_scope: dict | None) -> str:
|
|||
return text
|
||||
|
||||
|
||||
# Ceiling for a model-supplied top_k. Small on purpose: this returns whole archived turns
|
||||
# into a protected exchange the rolling window cannot trim.
|
||||
_MAX_CONVERSATION_SEARCH_TOP_K = 8
|
||||
|
||||
|
||||
def _search_conversation(arguments: dict, rag_scope: dict | None) -> str:
|
||||
"""Search this thread's archived turns. ``rag_scope`` carries only the thread id here;
|
||||
the model supplies ``query``/``top_k``."""
|
||||
scope = rag_scope or {}
|
||||
thread_id = scope.get("thread_id")
|
||||
query = (arguments or {}).get("query", "")
|
||||
if not query or not str(query).strip():
|
||||
return "Error: query is empty."
|
||||
if not thread_id:
|
||||
return "There is no earlier conversation to search."
|
||||
try:
|
||||
from core.rag import conversation_archive
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.warning("Conversation archive unavailable: %s", exc)
|
||||
return "Searching earlier conversation is unavailable on this server."
|
||||
if not conversation_archive.enabled():
|
||||
return "Searching earlier conversation is unavailable on this server."
|
||||
|
||||
# Clamped, not trusted: top_k comes from the model, and a negative value reaches a
|
||||
# Python slice as out[:-1], returning nearly the whole candidate pool as a ~30k-token
|
||||
# tool result that the protected current exchange cannot evict.
|
||||
requested = _opt_int((arguments or {}).get("top_k"))
|
||||
# None, not the ceiling: an omitted top_k must fall through to the configured recall
|
||||
# default. Defaulting to the maximum returned eight chunks into that same protected
|
||||
# exchange, enough to fail the next pass on a 4K chat.
|
||||
top_k = (
|
||||
None if requested is None else max(1, min(_MAX_CONVERSATION_SEARCH_TOP_K, int(requested)))
|
||||
)
|
||||
# Then against the room actually left: the fixed cap bounds what the model may ask
|
||||
# for, not what fits. Eight chunks is roughly 4,000 tokens once wrapped, so on a 4K
|
||||
# chat overshooting here is a context-length error no later preflight can recover.
|
||||
budget = scope.get("budget_tokens")
|
||||
if budget is not None:
|
||||
default_k = 1
|
||||
try:
|
||||
from core.rag import config as rag_config
|
||||
affordable = max(0, int(budget)) // max(1, int(rag_config.CHUNK_TOKENS))
|
||||
default_k = max(1, int(rag_config.CONVERSATION_ARCHIVE_TOP_K))
|
||||
except Exception:
|
||||
affordable = 0
|
||||
if affordable <= 0:
|
||||
return "There is no room left in this context to search earlier conversation."
|
||||
# An omitted top_k still means the configured default; room is a cap on it, not a
|
||||
# target. Taking the room itself asked a 128K chat for 200 passages, past both the
|
||||
# default and the ceiling the model's own value is held to.
|
||||
top_k = (
|
||||
min(default_k, _MAX_CONVERSATION_SEARCH_TOP_K, affordable)
|
||||
if top_k is None
|
||||
else max(1, min(top_k, affordable))
|
||||
)
|
||||
|
||||
# The branch this request is on, so a response replaced by Retry cannot be searched
|
||||
# back out of the archive. Absent callers fall back to the whole stored thread.
|
||||
def _recall(k):
|
||||
return conversation_archive.recall(
|
||||
str(thread_id),
|
||||
str(query),
|
||||
top_k = k,
|
||||
branch_messages = scope.get("branch_messages"),
|
||||
)
|
||||
|
||||
found = _recall(top_k)
|
||||
if not found:
|
||||
return "No earlier turns of this conversation matched that query."
|
||||
|
||||
# Then against what the result actually costs. CHUNK_TOKENS is what the chunker AIMS
|
||||
# at, not what a chunk weighs: chunks overlap, the chunker's tokenizer is not the
|
||||
# model's, and the rendered block adds markup, source metadata and the tool framing
|
||||
# around it. Measured on a 500-token budget: one chunk came back as 1,256 estimated
|
||||
# tokens. So the count is halved until the rendered result fits, the same backoff the
|
||||
# forced recall uses, and a single chunk that still does not fit is refused rather
|
||||
# than appended to an exchange the window is not allowed to evict.
|
||||
if budget is not None:
|
||||
counter = scope.get("token_counter")
|
||||
attempt = max(1, int(top_k or 1))
|
||||
while True:
|
||||
rendered = _rendered_conversation_search(found)
|
||||
if _conversation_search_cost(rendered, counter) <= int(budget):
|
||||
return rendered
|
||||
if attempt <= 1:
|
||||
return "There is no room left in this context to search earlier conversation."
|
||||
attempt = max(1, attempt // 2)
|
||||
found = _recall(attempt)
|
||||
if not found:
|
||||
return "No earlier turns of this conversation matched that query."
|
||||
return _rendered_conversation_search(found)
|
||||
|
||||
|
||||
# What a `tool` message costs beyond its own text: the role, the call id and whatever the
|
||||
# template wraps them in. Small, fixed, and left out entirely before, which is the wrong
|
||||
# direction on a check whose whole job is to refuse a result that will not fit.
|
||||
_TOOL_MESSAGE_FRAMING_TOKENS = 8
|
||||
|
||||
|
||||
def _conversation_search_cost(text: str, counter = None) -> int:
|
||||
"""What admitting this result really costs, exactly when the caller has a tokenizer.
|
||||
|
||||
The estimate below is pessimistic for CJK and emoji but still optimistic for ASCII
|
||||
that tokenises densely -- source code, minified JSON, hashes, command output all run
|
||||
nearer two or three characters per token than four -- so a result could be admitted at
|
||||
well under its real cost and then land in the current tool exchange, which the window
|
||||
is not allowed to evict. A tokenizer-backed caller passes its own counter, and the
|
||||
GGUF path is one, so the check that already computes the budget exactly can now spend
|
||||
it exactly too.
|
||||
"""
|
||||
if counter is not None:
|
||||
try:
|
||||
return int(counter(text)) + _TOOL_MESSAGE_FRAMING_TOKENS
|
||||
except Exception:
|
||||
logger.debug("conversation search: exact result count failed", exc_info = True)
|
||||
return _conversation_search_tokens(text) + _TOOL_MESSAGE_FRAMING_TOKENS
|
||||
|
||||
|
||||
def _conversation_search_tokens(text: str) -> int:
|
||||
"""A deliberately pessimistic size for a search result, in tokens.
|
||||
|
||||
The shared estimator charges four characters per token, which is about right for
|
||||
English and badly wrong for text that tokenises densely: CJK and emoji run closer to
|
||||
one token per character, so a result could be accepted at a quarter of its real cost
|
||||
and then land in the current tool exchange, which the window cannot evict. No exact
|
||||
counter is reachable from here, the provider loop having no tokenizer at all, so
|
||||
non-ASCII characters are charged one token each and the rest at the usual rate.
|
||||
"""
|
||||
dense = sum(1 for char in text if ord(char) > 127)
|
||||
return max(1, dense + (len(text) - dense) // 4)
|
||||
|
||||
|
||||
def _rendered_conversation_search(found) -> str:
|
||||
"""The tool result exactly as the model would receive it."""
|
||||
text, sources = found
|
||||
if sources:
|
||||
import json as _json
|
||||
return text + RAG_SOURCES_SENTINEL + _json.dumps(sources, ensure_ascii = False)
|
||||
return text
|
||||
|
||||
|
||||
def _search_knowledge_base_with_budget(
|
||||
arguments: dict,
|
||||
rag_scope: dict | None,
|
||||
timeout: int | None,
|
||||
cancel_event = None,
|
||||
search_fn = None,
|
||||
) -> str:
|
||||
"""Admission-controlled RAG search.
|
||||
|
||||
``search_fn`` swaps in a different search over the same capacity-of-one slot, so
|
||||
archive lookups queue behind document lookups instead of racing for the embedder."""
|
||||
search_fn = search_fn or _search_knowledge_base
|
||||
if cancel_event is not None and cancel_event.is_set():
|
||||
return "Error: knowledge base search cancelled."
|
||||
deadline = time.monotonic() + timeout if timeout is not None else None
|
||||
|
|
@ -10179,7 +10374,7 @@ def _search_knowledge_base_with_budget(
|
|||
|
||||
if timeout is None and cancel_event is None:
|
||||
try:
|
||||
return _search_knowledge_base(arguments, rag_scope)
|
||||
return search_fn(arguments, rag_scope)
|
||||
finally:
|
||||
release_slot()
|
||||
|
||||
|
|
@ -10187,7 +10382,7 @@ def _search_knowledge_base_with_budget(
|
|||
|
||||
def search() -> None:
|
||||
try:
|
||||
result.put((True, _search_knowledge_base(arguments, rag_scope)))
|
||||
result.put((True, search_fn(arguments, rag_scope)))
|
||||
except BaseException as exc:
|
||||
result.put((False, exc))
|
||||
finally:
|
||||
|
|
@ -10330,6 +10525,146 @@ def _last_user_text(conversation: list[dict]) -> str:
|
|||
return ""
|
||||
|
||||
|
||||
def build_synthetic_search_exchange(
|
||||
*,
|
||||
tool_name: str,
|
||||
call_prefix: str,
|
||||
status_label: str,
|
||||
query: str,
|
||||
text: str,
|
||||
sources: list[dict],
|
||||
) -> dict:
|
||||
"""Render a retrieval the loop never asked for as a normal tool exchange.
|
||||
|
||||
Returns ``{"events": [...], "messages": [...]}``: the messages are what the model
|
||||
reads, the events what the UI draws, so a forced retrieval shows up as an ordinary
|
||||
tool card with working citations instead of appearing from nowhere.
|
||||
"""
|
||||
import json as _json
|
||||
import uuid as _uuid
|
||||
|
||||
call_id = call_prefix + _uuid.uuid4().hex[:12]
|
||||
args = {"query": query}
|
||||
full_result = text + RAG_SOURCES_SENTINEL + _json.dumps(sources, ensure_ascii = False)
|
||||
events = [
|
||||
{"type": "status", "text": f"{status_label}: {query[:60]}"},
|
||||
{
|
||||
"type": "tool_start",
|
||||
"tool_name": tool_name,
|
||||
"tool_call_id": call_id,
|
||||
"arguments": args,
|
||||
},
|
||||
{
|
||||
"type": "tool_end",
|
||||
"tool_name": tool_name,
|
||||
"tool_call_id": call_id,
|
||||
"result": full_result,
|
||||
},
|
||||
{"type": "status", "text": ""},
|
||||
]
|
||||
messages = [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": call_id,
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": tool_name,
|
||||
"arguments": _json.dumps(args, ensure_ascii = False),
|
||||
},
|
||||
}
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "tool",
|
||||
"name": tool_name,
|
||||
"tool_call_id": call_id,
|
||||
"content": text,
|
||||
},
|
||||
]
|
||||
return {"events": events, "messages": messages}
|
||||
|
||||
|
||||
_RECALL_BLOCK = (
|
||||
"<recalled_conversation>\n"
|
||||
"This conversation was compacted and earlier turns were removed from your context. "
|
||||
"These are the most relevant earlier turns, retrieved verbatim from this chat.\n"
|
||||
"{text}\n"
|
||||
"</recalled_conversation>\n\n"
|
||||
)
|
||||
|
||||
|
||||
def build_conversation_recall(
|
||||
conversation: list[dict],
|
||||
thread_id: str | None,
|
||||
*,
|
||||
style: str = "tool",
|
||||
top_k: int | None = None,
|
||||
branch_messages: list[dict] | None = None,
|
||||
) -> dict | None:
|
||||
"""Retrieve the archived turns most relevant to the latest user message.
|
||||
|
||||
Deliberately NOT gated on ``rag_scope``: compaction happens whether or not document
|
||||
RAG is on, and these turns are the conversation's own, not uploaded files.
|
||||
|
||||
Forcing this one retrieval is the whole feature. Given only a search tool, a 35B on
|
||||
MRCR v2 declined on 56% of rows, scoring 0.099 when it skipped against 0.461 when it
|
||||
searched; forcing the lookup on the evicting turn took tool-only 0.258 to 0.604, and
|
||||
the model then called the tool on 0% of rows, so it costs nothing on the common path.
|
||||
|
||||
``style="tool"`` renders a tool exchange, for the tool loop which already carries a
|
||||
tools array. ``style="inline"`` prefixes the latest user message instead, for the
|
||||
plain path: forging tool_calls without a tools array breaks strict templates.
|
||||
"""
|
||||
if not thread_id:
|
||||
return None
|
||||
try:
|
||||
from core.rag import conversation_archive
|
||||
except Exception:
|
||||
return None
|
||||
if not conversation_archive.enabled():
|
||||
return None
|
||||
|
||||
# The BRANCH's latest user turn, not the loop conversation's: a later tool-loop
|
||||
# iteration can end with an internal user-role re-prompt (the plan-without-action
|
||||
# nudge), and searching for that controller instruction defeats the forced retrieval.
|
||||
# branch_messages is what the client sent, so its last user turn is the real request.
|
||||
query = _last_user_text(branch_messages or conversation) or _last_user_text(conversation)
|
||||
if not query:
|
||||
return None
|
||||
try:
|
||||
found = conversation_archive.recall(
|
||||
thread_id, query, top_k = top_k, branch_messages = branch_messages
|
||||
)
|
||||
except Exception:
|
||||
logger.warning("Conversation recall failed", exc_info = True)
|
||||
return None
|
||||
if not found:
|
||||
return None
|
||||
text, sources = found
|
||||
|
||||
if style == "inline":
|
||||
return {
|
||||
"events": [],
|
||||
"messages": [],
|
||||
"prefix": _RECALL_BLOCK.format(text = text),
|
||||
"sources": len(sources),
|
||||
}
|
||||
built = build_synthetic_search_exchange(
|
||||
tool_name = "search_conversation",
|
||||
call_prefix = "conv_recall_",
|
||||
status_label = "Recalling earlier conversation",
|
||||
query = query,
|
||||
text = text,
|
||||
sources = sources,
|
||||
)
|
||||
built["sources"] = len(sources)
|
||||
logger.info("Conversation recall: %d earlier passage(s) for %r", len(sources), query[:80])
|
||||
return built
|
||||
|
||||
|
||||
def build_rag_autoinject(conversation: list[dict], rag_scope: dict | None) -> dict | None:
|
||||
"""Pre-retrieve the latest user turn; if a hit clears the cosine floor return
|
||||
``{"events": [...], "messages": [...]}`` to splice into the loop, else ``None``.
|
||||
|
|
@ -10431,52 +10766,16 @@ def build_rag_autoinject(conversation: list[dict], rag_scope: dict | None) -> di
|
|||
if text is None:
|
||||
return None
|
||||
|
||||
import json as _json
|
||||
import uuid as _uuid
|
||||
|
||||
call_id = "rag_auto_" + _uuid.uuid4().hex[:12]
|
||||
args = {"query": query}
|
||||
full_result = text + RAG_SOURCES_SENTINEL + _json.dumps(sources, ensure_ascii = False)
|
||||
events = [
|
||||
{"type": "status", "text": f"Searching documents: {query[:60]}"},
|
||||
{
|
||||
"type": "tool_start",
|
||||
"tool_name": "search_knowledge_base",
|
||||
"tool_call_id": call_id,
|
||||
"arguments": args,
|
||||
},
|
||||
{
|
||||
"type": "tool_end",
|
||||
"tool_name": "search_knowledge_base",
|
||||
"tool_call_id": call_id,
|
||||
"result": full_result,
|
||||
},
|
||||
{"type": "status", "text": ""},
|
||||
]
|
||||
messages = [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": call_id,
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "search_knowledge_base",
|
||||
"arguments": _json.dumps(args, ensure_ascii = False),
|
||||
},
|
||||
}
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "tool",
|
||||
"name": "search_knowledge_base",
|
||||
"tool_call_id": call_id,
|
||||
"content": text,
|
||||
},
|
||||
]
|
||||
built = build_synthetic_search_exchange(
|
||||
tool_name = "search_knowledge_base",
|
||||
call_prefix = "rag_auto_",
|
||||
status_label = "Searching documents",
|
||||
query = query,
|
||||
text = text,
|
||||
sources = sources,
|
||||
)
|
||||
logger.info("RAG auto-inject: %d passage(s) for %r", len(sources), query[:80])
|
||||
return {"events": events, "messages": messages}
|
||||
return built
|
||||
|
||||
|
||||
_MAX_PAGE_CHARS = 16000 # cap fetched page text (after HTML-to-MD conversion)
|
||||
|
|
|
|||
|
|
@ -24,6 +24,24 @@ RRF_K = int(os.environ.get("RAG_RRF_K", "60"))
|
|||
THREAD_WHOLE_DOC = os.environ.get("RAG_THREAD_WHOLE_DOC", "1") == "1"
|
||||
WHOLE_DOC_MAX_TOKENS = int(os.environ.get("RAG_WHOLE_DOC_MAX_TOKENS", "6000"))
|
||||
|
||||
# Conversation archive: turns evicted by the rolling context window go to a per-thread
|
||||
# searchable scope and the relevant ones are recalled on the turn that evicted them. Off,
|
||||
# evicted turns are simply dropped again, and the recall reserve is not taken. Only applies
|
||||
# once the window evicts, which is itself opt-in per request via
|
||||
# context_overflow="truncate_oldest".
|
||||
#
|
||||
# It does NOT turn the rolling window back into what it was before: the compaction headroom
|
||||
# and the sticky boundary belong to the window, not to the archive, and have their own knob
|
||||
# (ROLLING_COMPACTION_HEADROOM_RATIO). Gating them here instead would make a host without
|
||||
# sqlite-vec silently compact differently.
|
||||
CONVERSATION_ARCHIVE = os.environ.get("RAG_CONVERSATION_ARCHIVE", "1") == "1"
|
||||
CONVERSATION_ARCHIVE_TOP_K = int(os.environ.get("RAG_CONVERSATION_ARCHIVE_TOP_K", "4"))
|
||||
# Room held back during the fit for the turns recalled straight after it. Sized to
|
||||
# CONVERSATION_ARCHIVE_TOP_K * CHUNK_TOKENS with slack for the wrapper text.
|
||||
CONVERSATION_RECALL_RESERVE_TOKENS = int(
|
||||
os.environ.get("RAG_CONVERSATION_RECALL_RESERVE_TOKENS", "2048")
|
||||
)
|
||||
|
||||
UPLOAD_EXTS = {".pdf", ".txt", ".md", ".markdown", ".docx", ".html", ".htm"}
|
||||
# Reject uploads larger than this, so one pathological file can't drive unbounded parse
|
||||
# + vision work at ingest. 0 disables the cap. Default 200 MB.
|
||||
|
|
|
|||
1190
studio/backend/core/rag/conversation_archive.py
Normal file
1190
studio/backend/core/rag/conversation_archive.py
Normal file
File diff suppressed because it is too large
Load diff
|
|
@ -35,6 +35,20 @@ def project_scope(project_id: str) -> str:
|
|||
return f"project_{project_id}"
|
||||
|
||||
|
||||
CONVERSATION_ARCHIVE_PREFIX = "convarchive_"
|
||||
|
||||
|
||||
def conversation_archive_scope(thread_id: str) -> str:
|
||||
"""Scope holding the turns a thread's rolling context window has evicted.
|
||||
|
||||
Deliberately NOT ``thread_scope``: with ``config.THREAD_WHOLE_DOC`` on, that scope is
|
||||
rendered in full into every request, so archiving turns there would re-inject the
|
||||
history and undo the compaction. A separate scope also keeps the archive out of the
|
||||
attachments UI and the citation panel.
|
||||
"""
|
||||
return f"{CONVERSATION_ARCHIVE_PREFIX}{thread_id}"
|
||||
|
||||
|
||||
def _scopes(scope) -> list[str]:
|
||||
"""Search helpers accept one scope or several (e.g. project + thread)."""
|
||||
return [scope] if isinstance(scope, str) else list(scope)
|
||||
|
|
@ -130,13 +144,14 @@ def create_document(
|
|||
embedding_model: str | None = None,
|
||||
linked_folder_id: str | None = None,
|
||||
linked_relative_path: str | None = None,
|
||||
archive_messages: int | None = None,
|
||||
commit: bool = True,
|
||||
) -> str:
|
||||
document_id = document_id or str(uuid.uuid4())
|
||||
conn.execute(
|
||||
"INSERT INTO documents(id, scope, kb_id, thread_id, project_id, filename, sha256, "
|
||||
"status, stored_path, created_at, embedding_model, linked_folder_id, "
|
||||
"linked_relative_path) VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?)",
|
||||
"linked_relative_path, archive_messages) VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?)",
|
||||
(
|
||||
document_id,
|
||||
scope,
|
||||
|
|
@ -151,6 +166,7 @@ def create_document(
|
|||
embedding_model,
|
||||
linked_folder_id,
|
||||
linked_relative_path,
|
||||
archive_messages,
|
||||
),
|
||||
)
|
||||
if commit:
|
||||
|
|
@ -197,12 +213,17 @@ def list_documents(conn: sqlite3.Connection, scope: str) -> list[dict]:
|
|||
|
||||
|
||||
def list_all_documents(conn: sqlite3.Connection) -> list[dict]:
|
||||
"""Every uploaded document across all scopes (KBs, threads, projects)."""
|
||||
"""Every uploaded document across all scopes (KBs, threads, projects).
|
||||
|
||||
Archived conversation turns are excluded: nobody uploaded them, so listing them would
|
||||
show a chat's own history back as files the user never added.
|
||||
"""
|
||||
rows = conn.execute(
|
||||
"SELECT id, scope, kb_id, thread_id, project_id, filename, sha256, status, error, "
|
||||
"num_chunks, stored_path, created_at, linked_folder_id "
|
||||
"FROM documents d WHERE NOT EXISTS "
|
||||
"(SELECT 1 FROM linked_folder_retired_scopes r WHERE r.scope=d.scope) "
|
||||
"AND d.scope NOT LIKE 'convarchive#_%' ESCAPE '#' "
|
||||
"ORDER BY created_at DESC"
|
||||
).fetchall()
|
||||
return [dict(r) for r in rows]
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ from storage import rag_db
|
|||
from . import config, retrieval
|
||||
from .store import (
|
||||
all_chunks_for_scope,
|
||||
conversation_archive_scope,
|
||||
kb_scope,
|
||||
project_scope,
|
||||
scope_token_estimate,
|
||||
|
|
@ -53,9 +54,16 @@ def _resolve_scope(
|
|||
scope_kb_id: str | None,
|
||||
scope_thread_id: str | None,
|
||||
scope_project_id: str | None = None,
|
||||
scope_conversation_id: str | None = None,
|
||||
) -> str | list[str] | None:
|
||||
"""KB (an explicit pick) is exclusive; project and thread scopes combine so a
|
||||
project chat also retrieves from its own attached documents."""
|
||||
project chat also retrieves from its own attached documents.
|
||||
|
||||
The conversation archive is exclusive too, and takes precedence: it is a different
|
||||
corpus (this chat's evicted turns), so sharing a top-K would have old turns and
|
||||
document passages crowd each other out."""
|
||||
if scope_conversation_id:
|
||||
return conversation_archive_scope(scope_conversation_id)
|
||||
if scope_kb_id:
|
||||
return kb_scope(scope_kb_id)
|
||||
scopes = []
|
||||
|
|
@ -126,6 +134,7 @@ def search_knowledge_base_with_sources(
|
|||
scope_kb_id: str | None = None,
|
||||
scope_thread_id: str | None = None,
|
||||
scope_project_id: str | None = None,
|
||||
scope_conversation_id: str | None = None,
|
||||
top_k: int | None = None,
|
||||
min_score: float = 0.0,
|
||||
model_name: str | None = None,
|
||||
|
|
@ -135,7 +144,7 @@ def search_knowledge_base_with_sources(
|
|||
rendered ``<chunk>`` block's ``id``."""
|
||||
if not query or not query.strip():
|
||||
return "Error: query is empty.", []
|
||||
scope = _resolve_scope(scope_kb_id, scope_thread_id, scope_project_id)
|
||||
scope = _resolve_scope(scope_kb_id, scope_thread_id, scope_project_id, scope_conversation_id)
|
||||
if scope is None:
|
||||
return "No documents are attached to this chat.", []
|
||||
|
||||
|
|
|
|||
|
|
@ -707,6 +707,8 @@ async def delete_threads(
|
|||
):
|
||||
from starlette.concurrency import run_in_threadpool
|
||||
|
||||
# Before the rows go, so a thread id that comes back in the gap is cut here.
|
||||
cutoff = _archive_cutoff()
|
||||
deleted_research_run_ids = await run_in_threadpool(delete_chat_threads, payload.ids)
|
||||
_cancel_research_runs(request, deleted_research_run_ids)
|
||||
_cancel_active_generations(payload.ids)
|
||||
|
|
@ -715,9 +717,43 @@ async def delete_threads(
|
|||
# In a worker: right after an upgrade this also runs the legacy move, and a
|
||||
# cross-filesystem copy on the event loop stops every other request.
|
||||
removed, kept = await _remove_sandboxes(payload.ids, payload.delete_files)
|
||||
# Archived turns are keyed by thread id and unreferenced once the thread is gone, so
|
||||
# drop them rather than leaking a scope per deleted chat.
|
||||
await run_in_threadpool(_remove_conversation_archives, payload.ids, cutoff = cutoff)
|
||||
return {"status": "deleted", "sandboxes_removed": removed, "sandboxes_kept": kept}
|
||||
|
||||
|
||||
def _archive_cutoff() -> str:
|
||||
"""The instant a delete was accepted, as an ISO-8601 UTC string. See below."""
|
||||
from datetime import datetime, timezone
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
|
||||
|
||||
def _remove_conversation_archives(thread_ids, *, cutoff: "str | None" = None) -> None:
|
||||
"""Drop each deleted thread's archived turns. Never raises."""
|
||||
try:
|
||||
from core.rag import conversation_archive
|
||||
except Exception:
|
||||
return
|
||||
for thread_id in thread_ids or []:
|
||||
# As next to the sandbox removal: the rows went first and the sandbox pass ran in
|
||||
# between, so another tab can have recreated this id and already be archiving
|
||||
# turns under it. That chat is alive, and its memory is not this delete's to take
|
||||
# -- but the conversation the user DID delete is, and skipping the scope kept it
|
||||
# too, recallable under the live id with nothing left to sweep it. So cut at the
|
||||
# instant the delete was accepted instead of skipping: everything archived before
|
||||
# it belongs to the deleted conversation, everything after to the new one.
|
||||
recreated = get_chat_thread(str(thread_id)) is not None
|
||||
if recreated and not cutoff:
|
||||
continue
|
||||
try:
|
||||
conversation_archive.delete_for_thread(
|
||||
str(thread_id), created_before = cutoff if recreated else None
|
||||
)
|
||||
except Exception:
|
||||
logger.warning("Could not remove the conversation archive for %s", thread_id)
|
||||
|
||||
|
||||
async def _remove_sandboxes(thread_ids, delete_files: bool) -> "tuple[int, list[str]]":
|
||||
"""Drop each thread's sandbox off the event loop. Never raises.
|
||||
|
||||
|
|
@ -1008,6 +1044,7 @@ async def delete_project(
|
|||
# Rows first, files last: a member chat can still be running a tool in the
|
||||
# workspace, and its cwd disappearing mid-call either kills the call or
|
||||
# leaves what it writes next in a directory no project owns.
|
||||
cutoff = _archive_cutoff()
|
||||
try:
|
||||
project = await run_in_threadpool(
|
||||
lambda: delete_chat_project(project_id, delete_files = False)
|
||||
|
|
@ -1041,6 +1078,8 @@ async def delete_project(
|
|||
# By run id: the rows are gone by now, so there is nothing left to look up.
|
||||
_cancel_research_runs(request, list(project.get("activeResearchRunIds") or []))
|
||||
_cancel_active_generations(member_ids)
|
||||
# The project's chats go with it, so their archives have to as well.
|
||||
await run_in_threadpool(_remove_conversation_archives, member_ids, cutoff = cutoff)
|
||||
if project.get("sandboxPath"):
|
||||
from core.inference.tools import (
|
||||
finish_workspace_delete_when_idle,
|
||||
|
|
@ -1298,6 +1337,7 @@ async def clear_history(
|
|||
):
|
||||
from starlette.concurrency import run_in_threadpool
|
||||
|
||||
cutoff = _archive_cutoff()
|
||||
# Admission is already closed in the frontend. Include its pending and legacy ids in the
|
||||
# transaction's fence so a delayed POST cannot recreate a chat after this returns.
|
||||
thread_ids = (
|
||||
|
|
@ -1330,6 +1370,11 @@ async def clear_history(
|
|||
_cancel_active_generations(late)
|
||||
# By id: the rows went with the threads, so nothing can look them up now.
|
||||
_cancel_research_runs(request, cleared_runs)
|
||||
# Same archive cleanup as DELETE /threads. Without it "Clear all chats" leaves every
|
||||
# conversation searchable in rag.db, and a reused thread id reads the old archive.
|
||||
await run_in_threadpool(
|
||||
_remove_conversation_archives, list(dict.fromkeys(thread_ids + cleared)), cutoff = cutoff
|
||||
)
|
||||
# "Clear all chats" is the common bulk delete, so it has to clean up the
|
||||
# same folders DELETE /threads does; otherwise every sandbox is stranded.
|
||||
# delete_files matches DELETE /threads: off by default, since the files are
|
||||
|
|
|
|||
|
|
@ -682,9 +682,9 @@ def _accumulate_context_truncation(current: Optional[dict], event: dict) -> dict
|
|||
if current is None:
|
||||
return incoming
|
||||
combined = {**current, **incoming}
|
||||
combined["dropped_messages"] = int(current.get("dropped_messages") or 0) + int(
|
||||
incoming.get("dropped_messages") or 0
|
||||
)
|
||||
for counter in ("dropped_messages", "archived_messages", "recalled_chunks"):
|
||||
if counter in current or counter in incoming:
|
||||
combined[counter] = int(current.get(counter) or 0) + int(incoming.get(counter) or 0)
|
||||
if current.get("prompt_tokens_before") is not None:
|
||||
combined["prompt_tokens_before"] = current["prompt_tokens_before"]
|
||||
return combined
|
||||
|
|
@ -3426,6 +3426,17 @@ _RAG_GROUNDING_NUDGE = (
|
|||
)
|
||||
|
||||
|
||||
def _thread_has_conversation_archive(thread_id) -> bool:
|
||||
"""Whether the rolling window has archived anything for this thread yet."""
|
||||
if not thread_id:
|
||||
return False
|
||||
try:
|
||||
from core.rag import conversation_archive
|
||||
return conversation_archive.has_archive(str(thread_id))
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
async def _select_request_tools(
|
||||
payload: ChatCompletionRequest, *, tools_on: bool, mcp_allowed: bool
|
||||
) -> list[dict]:
|
||||
|
|
@ -3455,6 +3466,20 @@ async def _select_request_tools(
|
|||
# Drop the RAG tool without a scope: nothing to search over.
|
||||
if not payload.rag_scope:
|
||||
tools = [t for t in tools if t["function"]["name"] != "search_knowledge_base"]
|
||||
# Same rule for the conversation archive: offered only once this thread has had turns
|
||||
# evicted, so a short chat never sees the extra schema. On the first compaction the
|
||||
# tool is still absent (the archive is written mid-request) and the forced recall
|
||||
# covers that turn. getattr, because this helper also serves the token-count request
|
||||
# model, which carries no thread_id.
|
||||
# Follows the ARCHIVE, not the caller's allowlist: Studio always sends an explicit
|
||||
# enabled_tools array and has no reason to name an internal tool it shows no pill for,
|
||||
# so the filter above removed search_conversation and neither it nor the compaction
|
||||
# nudge gated on it ever reached a Studio chat. It is read-only and always-safe, so it
|
||||
# is added on that condition rather than requested.
|
||||
has_archive = _thread_has_conversation_archive(getattr(payload, "thread_id", None))
|
||||
tools = [t for t in tools if t["function"]["name"] != "search_conversation"]
|
||||
if has_archive and tools_on:
|
||||
tools = tools + [t for t in ALL_TOOLS if t["function"]["name"] == "search_conversation"]
|
||||
# Built-ins only, so this runs before the MCP append: an MCP tool's
|
||||
# description is the server's to write, and Full access says nothing about
|
||||
# how that server runs.
|
||||
|
|
@ -3465,6 +3490,28 @@ async def _select_request_tools(
|
|||
return tools
|
||||
|
||||
|
||||
_COMPACTED_SESSION_NUDGE = (
|
||||
"This conversation is long, so older turns have been removed from your context. "
|
||||
"The relevant ones are retrieved and shown to you automatically when that happens. "
|
||||
"If the user refers to something you cannot see, call search_conversation before "
|
||||
"answering. Never tell the user you have no record of an earlier turn, and never "
|
||||
"assume the conversation began where your visible context begins."
|
||||
)
|
||||
|
||||
|
||||
def _apply_compaction_nudge(nudge: str, tools: list[dict]) -> str:
|
||||
"""Append the compacted-session nudge when the conversation-archive tool is active.
|
||||
|
||||
Gated on the tool rather than separate state, so it appears exactly when there is an
|
||||
archive to search and stays a no-op for chats that never compacted."""
|
||||
tool_names = {(t.get("function") or {}).get("name") for t in (tools or [])}
|
||||
if "search_conversation" not in tool_names:
|
||||
return nudge
|
||||
if not nudge:
|
||||
return _COMPACTED_SESSION_NUDGE
|
||||
return nudge + " " + _COMPACTED_SESSION_NUDGE
|
||||
|
||||
|
||||
def _apply_rag_nudge(nudge: str, tools: list[dict], *, rag_scope) -> str:
|
||||
"""Append the RAG grounding nudge to ``nudge`` when the knowledge-base tool
|
||||
is active (search_knowledge_base present and a retrieval scope is set). The
|
||||
|
|
@ -14219,6 +14266,7 @@ async def openai_chat_completions(
|
|||
|
||||
# Nudge the model to ground in attached documents instead of memory.
|
||||
_nudge = _apply_rag_nudge(_nudge, tools_to_use, rag_scope = payload.rag_scope)
|
||||
_nudge = _apply_compaction_nudge(_nudge, tools_to_use)
|
||||
|
||||
if _nudge:
|
||||
# Append nudge to system prompt (preserve user's prompt)
|
||||
|
|
@ -14982,6 +15030,7 @@ async def openai_chat_completions(
|
|||
seed = _seed,
|
||||
perf_callback = _gguf_perf_callback,
|
||||
context_overflow = _rolling_context_policy(payload),
|
||||
thread_id = payload.thread_id,
|
||||
)
|
||||
|
||||
_gguf_sentinel = object()
|
||||
|
|
@ -15483,8 +15532,8 @@ async def openai_chat_completions(
|
|||
or int(_choice_context_truncation.get("dropped_messages") or 0)
|
||||
>= int(_context_truncation.get("dropped_messages") or 0)
|
||||
):
|
||||
# Choices share the original prompt. Keep the shortest
|
||||
# choice's cumulative fit instead of summing choices.
|
||||
# Choices share one prompt: report the choice that dropped
|
||||
# the most, rather than summing across choices.
|
||||
_context_truncation = _choice_context_truncation
|
||||
|
||||
reasoning_text, visible_text = _extract_responses_reasoning(
|
||||
|
|
@ -15778,6 +15827,7 @@ async def openai_chat_completions(
|
|||
|
||||
# RAG nudge, mirroring the GGUF path.
|
||||
_sf_nudge = _apply_rag_nudge(_sf_nudge, _sf_tools_to_use, rag_scope = payload.rag_scope)
|
||||
_sf_nudge = _apply_compaction_nudge(_sf_nudge, _sf_tools_to_use)
|
||||
|
||||
_sf_system_prompt = system_prompt
|
||||
if _sf_nudge:
|
||||
|
|
@ -19650,7 +19700,9 @@ _STUDIO_ANTHROPIC_TOOL_ALIASES = {
|
|||
# this channel has no way to present, so an omitted permission_mode ("ask") only
|
||||
# asks then. render_html is excluded because a networked canvas prompts in auto,
|
||||
# and this channel invokes the loop without confirm; auto/ask reject, off/full run.
|
||||
_ANTHROPIC_UNPROMPTED_SAFE_TOOLS = frozenset({"web_search", "search_knowledge_base"})
|
||||
_ANTHROPIC_UNPROMPTED_SAFE_TOOLS = frozenset(
|
||||
{"web_search", "search_knowledge_base", "search_conversation"}
|
||||
)
|
||||
|
||||
|
||||
def _anthropic_selects_server_tools(
|
||||
|
|
|
|||
|
|
@ -255,6 +255,10 @@ def _ensure_schema(conn: sqlite3.Connection) -> None:
|
|||
conn.execute("ALTER TABLE documents ADD COLUMN linked_folder_id TEXT")
|
||||
if "linked_relative_path" not in cols:
|
||||
conn.execute("ALTER TABLE documents ADD COLUMN linked_relative_path TEXT")
|
||||
# How many messages an archived turn was rendered from. NULL for everything else and
|
||||
# for older archives, which fall back to counting role labels in the rendered text.
|
||||
if "archive_messages" not in cols:
|
||||
conn.execute("ALTER TABLE documents ADD COLUMN archive_messages INTEGER")
|
||||
# After the ALTER that adds the column on an older database. Partial, so it holds only
|
||||
# folder-owned rows and is empty with nothing linked, which keeps the lexical fast-path
|
||||
# gate an index probe rather than a scan of documents.
|
||||
|
|
|
|||
|
|
@ -3355,6 +3355,22 @@ def count_forks_for_message(thread_id: str, message_id: str) -> int:
|
|||
conn.close()
|
||||
|
||||
|
||||
def chat_thread_has_messages(thread_id: str) -> bool:
|
||||
"""Whether this thread has any saved message. Existence only, no rows hydrated.
|
||||
|
||||
A temporary (incognito) chat is never written here, so this is what tells a thread
|
||||
whose turns can be archived from one whose turns must not be.
|
||||
"""
|
||||
conn = get_connection()
|
||||
try:
|
||||
row = conn.execute(
|
||||
"SELECT 1 FROM chat_messages WHERE thread_id = ? LIMIT 1", (thread_id,)
|
||||
).fetchone()
|
||||
return row is not None
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def fork_counts_for_thread(thread_id: str) -> dict[str, int]:
|
||||
"""Fork counts for every message of one thread, keyed by message id."""
|
||||
conn = get_connection()
|
||||
|
|
|
|||
|
|
@ -2132,6 +2132,19 @@ class TestAnthropicMessagesToolRouting:
|
|||
_drive(anthropic_messages(payload, request = None, current_subject = "t"))
|
||||
assert backend.calls[0][0] == "tools"
|
||||
|
||||
# Reading this conversation's own archive is as read-only as the other two, and
|
||||
# is_potentially_unsafe_tool_call says so, so selecting it must not trip the gate.
|
||||
# Adding the schema to ALL_TOOLS without adding the name here made the Anthropic
|
||||
# selector pick it and the pre-switch guard reject the whole request with the
|
||||
# terminal/python message, on auto and on the omitted default alike.
|
||||
for extra in ({"permission_mode": "auto"}, {}):
|
||||
backend = _mock_backend(monkeypatch)
|
||||
payload = _basic_payload(
|
||||
enable_tools = True, enabled_tools = ["search_conversation"], **extra
|
||||
)
|
||||
_drive(anthropic_messages(payload, request = None, current_subject = "t"))
|
||||
assert backend.calls[0][0] == "tools"
|
||||
|
||||
# But auto or an omitted mode that would run a local tool (terminal/python,
|
||||
# via a bare Anthropic tool type or enabled_tools) is rejected, since that
|
||||
# tool could need the gate this channel lacks.
|
||||
|
|
|
|||
|
|
@ -37,7 +37,13 @@ from routes.inference import (
|
|||
_truncate_middle_messages,
|
||||
_truncate_oldest_messages,
|
||||
)
|
||||
from core.inference.context_window import fit_rolling_context, messages_have_media
|
||||
from core.inference import context_window
|
||||
from core.inference.context_window import (
|
||||
evicted_messages,
|
||||
fit_rolling_context,
|
||||
group_turns,
|
||||
messages_have_media,
|
||||
)
|
||||
from models.inference import ChatCompletion
|
||||
import routes.inference as routes_mod
|
||||
|
||||
|
|
@ -234,9 +240,8 @@ def test_rolling_media_detection_covers_image_and_audio_parts():
|
|||
assert messages_have_media(
|
||||
[{"role": "user", "content": [{"type": "input_audio", "input_audio": {}}]}]
|
||||
)
|
||||
# llama.cpp's own part type, written by `_inject_video_part` into the very list this
|
||||
# fit is handed. Missed, a video prompt ran the preflight that video tokens are not
|
||||
# counted by, so it could be certified as fitting when it does not.
|
||||
# llama.cpp's own part type; missing it would send a video prompt through a
|
||||
# preflight that does not count its tokens.
|
||||
assert messages_have_media(
|
||||
[{"role": "user", "content": [{"type": "input_video", "input_video": {"data": "AAAA"}}]}]
|
||||
)
|
||||
|
|
@ -270,7 +275,18 @@ def test_rolling_truncation_can_drop_assistant_after_instruction(instruction_rol
|
|||
assert new == [instruction, latest]
|
||||
|
||||
|
||||
def test_rolling_fit_recounts_until_the_real_template_fits():
|
||||
@pytest.fixture
|
||||
def no_compaction_headroom(monkeypatch):
|
||||
"""Pin the compaction headroom to zero.
|
||||
|
||||
For tests about the MINIMUM eviction needed to fit, the headroom is noise: it drops
|
||||
more than necessary, so an exact count would assert the headroom's value rather than
|
||||
the fit's behaviour. Tests about the headroom set it explicitly.
|
||||
"""
|
||||
monkeypatch.setattr(context_window, "_COMPACTION_HEADROOM_RATIO", 0.0)
|
||||
|
||||
|
||||
def test_rolling_fit_recounts_until_the_real_template_fits(no_compaction_headroom):
|
||||
messages = [
|
||||
{"role": "system", "content": "system"},
|
||||
{"role": "user", "content": "one" * 40},
|
||||
|
|
@ -314,9 +330,132 @@ def test_rolling_fit_never_clips_an_irreducible_latest_turn():
|
|||
|
||||
assert fitted is messages
|
||||
assert fitted == messages
|
||||
# Unchanged messages, but not a silent None: the fit says WHY it gave up, so the
|
||||
# user hears the single message is the problem rather than the history.
|
||||
assert info is not None and info["fits"] is False
|
||||
assert info["dropped_messages"] == 0
|
||||
assert info["latest_turn_tokens"] > info["context_length"]
|
||||
|
||||
|
||||
def _length_counter(candidate):
|
||||
return sum(len(str(message.get("content", ""))) for message in candidate)
|
||||
|
||||
|
||||
def test_evicted_messages_returns_dropped_turns_in_original_order(no_compaction_headroom):
|
||||
messages = [
|
||||
{"role": "system", "content": "system"},
|
||||
{"role": "user", "content": "one" * 40},
|
||||
{"role": "assistant", "content": "answer" * 40},
|
||||
{"role": "user", "content": "two" * 40},
|
||||
{"role": "assistant", "content": "answer" * 40},
|
||||
{"role": "user", "content": "latest"},
|
||||
]
|
||||
|
||||
fitted, info = fit_rolling_context(
|
||||
messages,
|
||||
context_length = 500,
|
||||
max_tokens = 100,
|
||||
count_tokens = _length_counter,
|
||||
)
|
||||
gone = evicted_messages(messages, fitted)
|
||||
|
||||
assert info is not None
|
||||
assert len(gone) == info["dropped_messages"]
|
||||
assert gone == [messages[1], messages[2]]
|
||||
|
||||
|
||||
def test_evicted_messages_uses_identity_not_equality():
|
||||
"""Two byte-identical turns must not collapse into one.
|
||||
|
||||
An equality diff reports BOTH copies as evicted when only the older one was, so
|
||||
downstream acts on a turn the model can still see.
|
||||
"""
|
||||
first = {"role": "user", "content": "same question"}
|
||||
second = {"role": "user", "content": "same question"}
|
||||
before = [first, {"role": "assistant", "content": "reply"}, second]
|
||||
after = [second]
|
||||
|
||||
gone = evicted_messages(before, after)
|
||||
|
||||
assert len(gone) == 2
|
||||
assert gone[0] is first
|
||||
assert all(message is not second for message in gone)
|
||||
|
||||
|
||||
def test_group_turns_matches_the_unit_truncation_drops():
|
||||
messages = [
|
||||
{"role": "system", "content": "system"},
|
||||
{"role": "user", "content": "ask"},
|
||||
{"role": "assistant", "content": None, "tool_calls": [{"id": "c1"}]},
|
||||
{"role": "tool", "tool_call_id": "c1", "content": "result"},
|
||||
{"role": "assistant", "content": "answer"},
|
||||
{"role": "user", "content": "latest"},
|
||||
]
|
||||
|
||||
groups = group_turns(messages)
|
||||
|
||||
assert [[message["role"] for message in group] for group in groups] == [
|
||||
["system"],
|
||||
["user"],
|
||||
["assistant", "tool", "assistant"],
|
||||
["user"],
|
||||
]
|
||||
|
||||
|
||||
def test_reserve_tokens_does_not_trim_a_prompt_that_already_fits():
|
||||
"""The reserve must never be what causes eviction.
|
||||
|
||||
A conversation inside the window comes back untouched even when the reserve would
|
||||
not fit alongside it, or recall would start evicting chats nowhere near the limit.
|
||||
"""
|
||||
messages = [
|
||||
{"role": "system", "content": "system"},
|
||||
{"role": "user", "content": "short"},
|
||||
{"role": "assistant", "content": "short answer"},
|
||||
{"role": "user", "content": "latest"},
|
||||
]
|
||||
|
||||
fitted, info = fit_rolling_context(
|
||||
messages,
|
||||
context_length = 500,
|
||||
max_tokens = 100,
|
||||
count_tokens = _length_counter,
|
||||
reserve_tokens = 380,
|
||||
)
|
||||
|
||||
assert fitted is messages
|
||||
assert info is None
|
||||
|
||||
|
||||
def test_reserve_tokens_trims_further_once_trimming_is_needed(no_compaction_headroom):
|
||||
messages = [
|
||||
{"role": "system", "content": "system"},
|
||||
{"role": "user", "content": "one" * 40},
|
||||
{"role": "assistant", "content": "answer" * 40},
|
||||
{"role": "user", "content": "two" * 40},
|
||||
{"role": "assistant", "content": "answer" * 40},
|
||||
{"role": "user", "content": "latest"},
|
||||
]
|
||||
|
||||
_, plain = fit_rolling_context(
|
||||
messages,
|
||||
context_length = 500,
|
||||
max_tokens = 100,
|
||||
count_tokens = _length_counter,
|
||||
)
|
||||
_, reserved = fit_rolling_context(
|
||||
messages,
|
||||
context_length = 500,
|
||||
max_tokens = 100,
|
||||
count_tokens = _length_counter,
|
||||
reserve_tokens = 200,
|
||||
)
|
||||
|
||||
assert plain is not None and reserved is not None
|
||||
assert reserved["dropped_messages"] > plain["dropped_messages"]
|
||||
assert reserved["prompt_tokens_after"] < plain["prompt_tokens_after"]
|
||||
|
||||
|
||||
def test_rolling_fit_keeps_original_when_protected_messages_still_do_not_fit():
|
||||
latest = {"role": "user", "content": "latest" * 200}
|
||||
messages = [
|
||||
|
|
@ -335,7 +474,114 @@ def test_rolling_fit_keeps_original_when_protected_messages_still_do_not_fit():
|
|||
|
||||
assert fitted is messages
|
||||
assert fitted == messages
|
||||
assert info is None
|
||||
assert info is not None and info["fits"] is False
|
||||
# The partial eviction is deliberately NOT applied: the request fails either way,
|
||||
# and dropping turns off a doomed request loses them for nothing.
|
||||
assert info["dropped_messages"] == 0
|
||||
assert info["prompt_tokens_after"] == info["prompt_tokens_before"]
|
||||
|
||||
|
||||
def test_an_irreducible_fit_says_WHOSE_turn_does_not_fit():
|
||||
"""A tool loop refits with the tool result appended.
|
||||
|
||||
The turn that will not fit is then output the user never wrote and cannot edit, so
|
||||
"shorten this message" has no remedy. The role is what tells the two apart.
|
||||
"""
|
||||
user_turn = [
|
||||
{"role": "system", "content": "system"},
|
||||
{"role": "user", "content": "what does this file contain"},
|
||||
{"role": "assistant", "content": "reading it"},
|
||||
{"role": "tool", "tool_call_id": "c1", "content": "output" * 200},
|
||||
]
|
||||
_, info = fit_rolling_context(
|
||||
user_turn,
|
||||
context_length = 500,
|
||||
max_tokens = 100,
|
||||
count_tokens = _length_counter,
|
||||
)
|
||||
assert info["fits"] is False
|
||||
assert info["latest_turn_role"] == "tool"
|
||||
|
||||
# And an ordinary overflowing user message still says so.
|
||||
_, info = fit_rolling_context(
|
||||
[
|
||||
{"role": "system", "content": "system"},
|
||||
{"role": "user", "content": "latest" * 200},
|
||||
],
|
||||
context_length = 500,
|
||||
max_tokens = 100,
|
||||
count_tokens = _length_counter,
|
||||
)
|
||||
assert info["latest_turn_role"] == "user"
|
||||
|
||||
|
||||
def test_an_irreducible_fit_survives_a_template_that_refuses_a_lone_tool_result():
|
||||
"""The diagnosis is produced exactly where a tool loop is most likely to be.
|
||||
|
||||
Strict templates refuse to render a tool result on its own, so counting that slice
|
||||
threw out of the fit and the caller fell back to the untrimmed request, telling the
|
||||
client nothing on the one path this diagnosis exists for.
|
||||
"""
|
||||
|
||||
def strict_counter(messages):
|
||||
if len(messages) == 1 and messages[0].get("role") == "tool":
|
||||
raise RuntimeError("a tool result must follow an assistant tool call")
|
||||
return _length_counter(messages)
|
||||
|
||||
messages = [
|
||||
{"role": "user", "content": "read it"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"tool_calls": [{"id": "c1", "function": {"name": "python", "arguments": "{}"}}],
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "c1", "content": "output" * 500},
|
||||
]
|
||||
_, info = fit_rolling_context(
|
||||
messages,
|
||||
context_length = 500,
|
||||
max_tokens = 100,
|
||||
count_tokens = strict_counter,
|
||||
)
|
||||
|
||||
assert info is not None and info["fits"] is False
|
||||
assert info["latest_turn_role"] == "tool"
|
||||
# Estimated rather than counted: an approximation beats no diagnosis at all.
|
||||
assert info["latest_turn_tokens"] > 0
|
||||
|
||||
|
||||
def test_an_irreducible_fit_says_whether_the_message_or_the_history_is_at_fault():
|
||||
"""The two numbers that make the error actionable.
|
||||
|
||||
llama-server's error reports the WHOLE conversation's size and advises shortening
|
||||
it, which cannot work when the latest turn alone is over the window.
|
||||
"""
|
||||
huge_message = [
|
||||
{"role": "system", "content": "system"},
|
||||
{"role": "user", "content": "latest" * 200},
|
||||
]
|
||||
_, info = fit_rolling_context(
|
||||
huge_message,
|
||||
context_length = 500,
|
||||
max_tokens = 100,
|
||||
count_tokens = _length_counter,
|
||||
)
|
||||
assert info["fits"] is False
|
||||
assert info["latest_turn_tokens"] > info["context_length"]
|
||||
|
||||
# A conversation that fits reports nothing, not a fits:False dict a caller could
|
||||
# mistake for a failure.
|
||||
small = [
|
||||
{"role": "system", "content": "system"},
|
||||
{"role": "user", "content": "hello"},
|
||||
]
|
||||
_, none_info = fit_rolling_context(
|
||||
small,
|
||||
context_length = 500,
|
||||
max_tokens = 100,
|
||||
count_tokens = _length_counter,
|
||||
)
|
||||
assert none_info is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -596,3 +842,178 @@ def test_reasoning_clip_alone_prevents_middle_eviction():
|
|||
assert body["max_tokens"] == max(
|
||||
1024, int(8192 * (1.0 - routes_mod._OVERFLOW_PROMPT_TARGET_FRACTION))
|
||||
)
|
||||
|
||||
|
||||
def test_compaction_headroom_does_not_trim_a_prompt_that_already_fits():
|
||||
"""Same rule as the reserve: headroom must never be what causes eviction.
|
||||
|
||||
The headroom makes a compaction take a chunk out in one go; charging it up front
|
||||
would evict from chats that comfortably fit today.
|
||||
"""
|
||||
messages = [
|
||||
{"role": "system", "content": "system"},
|
||||
{"role": "user", "content": "short"},
|
||||
{"role": "assistant", "content": "short answer"},
|
||||
{"role": "user", "content": "latest"},
|
||||
]
|
||||
|
||||
fitted, info = fit_rolling_context(
|
||||
messages,
|
||||
context_length = 500,
|
||||
max_tokens = 100,
|
||||
count_tokens = _length_counter,
|
||||
)
|
||||
|
||||
assert fitted is messages
|
||||
assert info is None
|
||||
|
||||
|
||||
def test_compaction_leaves_headroom_below_the_budget():
|
||||
"""A compaction lands clear of the budget, not flush against it.
|
||||
|
||||
Trimming to the brim makes the boundary creep every turn: the client re-sends the
|
||||
whole transcript, so an exactly fitted prompt is over again on the next turn.
|
||||
"""
|
||||
messages = [{"role": "system", "content": "system"}]
|
||||
for index in range(12):
|
||||
messages.append({"role": "user", "content": f"question {index} " * 20})
|
||||
messages.append({"role": "assistant", "content": f"answer {index} " * 20})
|
||||
messages.append({"role": "user", "content": "latest"})
|
||||
|
||||
_, info = fit_rolling_context(
|
||||
messages,
|
||||
context_length = 800,
|
||||
max_tokens = 100,
|
||||
count_tokens = _length_counter,
|
||||
)
|
||||
|
||||
assert info is not None
|
||||
prompt_target = 800 - min(100, 800 // 4)
|
||||
assert info["prompt_tokens_after"] <= prompt_target
|
||||
# The point of the change: comfortably under, not just under.
|
||||
assert info["prompt_tokens_after"] < prompt_target * 0.9
|
||||
|
||||
|
||||
def _long_thread(turns: int = 40):
|
||||
"""A thread several times its window, in turns big enough to be evicted as units."""
|
||||
messages = [{"role": "system", "content": "system"}]
|
||||
for index in range(turns):
|
||||
messages.append({"role": "user", "content": f"question {index} " * 200})
|
||||
messages.append({"role": "assistant", "content": f"answer {index} " * 200})
|
||||
return messages
|
||||
|
||||
|
||||
def _fit_with_appended(
|
||||
base,
|
||||
appended,
|
||||
sticky = 0,
|
||||
):
|
||||
messages = list(base)
|
||||
for index in range(appended):
|
||||
messages.append({"role": "user", "content": f"follow up {index} " * 20})
|
||||
messages.append({"role": "assistant", "content": f"reply {index} " * 20})
|
||||
messages.append({"role": "user", "content": "latest"})
|
||||
_, info = fit_rolling_context(
|
||||
messages,
|
||||
context_length = 8000,
|
||||
max_tokens = 512,
|
||||
count_tokens = _length_counter,
|
||||
sticky_dropped = sticky,
|
||||
)
|
||||
return info
|
||||
|
||||
|
||||
def test_sticky_boundary_holds_still_while_short_turns_are_appended():
|
||||
"""After a compaction, ordinary turns do not push the boundary again.
|
||||
|
||||
The notice depends on this, and it is why the boundary is read back rather than
|
||||
recomputed: the client re-sends the whole transcript, so a recomputed "keep the
|
||||
newest N tokens" slides forward and every reply reports a fresh compaction.
|
||||
"""
|
||||
base = _long_thread()
|
||||
first = _fit_with_appended(base, 0)
|
||||
assert first is not None and first["dropped_messages"] > 0
|
||||
|
||||
boundary = first["dropped_messages"]
|
||||
# Not "forever": the appended turns consume the headroom and the next test pins
|
||||
# down that it eventually moves. A handful of turns, against a baseline that moved
|
||||
# on nearly every one.
|
||||
for appended in range(1, 6):
|
||||
info = _fit_with_appended(base, appended, sticky = boundary)
|
||||
assert info is not None
|
||||
assert (
|
||||
info["dropped_messages"] == boundary
|
||||
), f"the boundary moved after {appended} appended turns"
|
||||
|
||||
|
||||
def test_sticky_boundary_moves_again_once_the_headroom_is_used_up():
|
||||
"""It holds still, but it does not hold forever: enough new turns re-compact."""
|
||||
base = _long_thread()
|
||||
boundary = _fit_with_appended(base, 0)["dropped_messages"]
|
||||
|
||||
moved = None
|
||||
for appended in range(1, 60):
|
||||
info = _fit_with_appended(base, appended, sticky = boundary)
|
||||
if info["dropped_messages"] > boundary:
|
||||
moved = appended
|
||||
break
|
||||
|
||||
assert moved is not None, "the boundary never moved, so the window would overflow"
|
||||
assert moved > 4, f"the boundary moved again after only {moved} turns"
|
||||
|
||||
|
||||
def test_sticky_boundary_never_causes_eviction_on_a_thread_that_fits():
|
||||
"""A stale boundary from a longer branch must not evict a conversation that fits.
|
||||
|
||||
After a rollback the saved boundary describes a branch that no longer exists. The
|
||||
fit may reapply it, but never report a compaction on a prompt that already fits.
|
||||
"""
|
||||
messages = [
|
||||
{"role": "system", "content": "system"},
|
||||
{"role": "user", "content": "short"},
|
||||
{"role": "assistant", "content": "short answer"},
|
||||
{"role": "user", "content": "latest"},
|
||||
]
|
||||
|
||||
fitted, info = fit_rolling_context(
|
||||
messages,
|
||||
context_length = 4000,
|
||||
max_tokens = 100,
|
||||
count_tokens = _length_counter,
|
||||
sticky_dropped = 40,
|
||||
)
|
||||
|
||||
assert fitted is messages
|
||||
assert info is None
|
||||
|
||||
|
||||
def test_the_compaction_headroom_needs_a_boundary_to_be_worth_it():
|
||||
"""Cutting deeper than needed buys quiet turns only if the cut is remembered.
|
||||
|
||||
An incognito chat, an API request with no persisted thread, or a request whose turns
|
||||
are not saved gets neither the boundary back nor a recall of what went, so there the
|
||||
headroom is simply less history than plain eviction would have kept, on every
|
||||
overflow, and turning the archive off did not restore the old behaviour.
|
||||
"""
|
||||
messages = []
|
||||
for index in range(20):
|
||||
messages.append({"role": "user", "content": f"q{index} " + "u" * 80})
|
||||
messages.append({"role": "assistant", "content": f"a{index} " + "a" * 80})
|
||||
messages.append({"role": "user", "content": "latest"})
|
||||
|
||||
def _fit(keeps_boundary):
|
||||
return fit_rolling_context(
|
||||
list(messages),
|
||||
context_length = 2000,
|
||||
max_tokens = 200,
|
||||
count_tokens = _length_counter,
|
||||
keeps_boundary = keeps_boundary,
|
||||
)
|
||||
|
||||
plain, plain_info = _fit(False)
|
||||
sticky, sticky_info = _fit(True)
|
||||
|
||||
assert plain_info["fits"] and sticky_info["fits"]
|
||||
# The one that can restore its boundary is the one that pays for headroom.
|
||||
assert sticky_info["dropped_messages"] > plain_info["dropped_messages"]
|
||||
assert len(plain) > len(sticky)
|
||||
|
|
|
|||
1925
studio/backend/tests/test_conversation_archive.py
Normal file
1925
studio/backend/tests/test_conversation_archive.py
Normal file
File diff suppressed because it is too large
Load diff
1532
studio/backend/tests/test_conversation_recall_injection.py
Normal file
1532
studio/backend/tests/test_conversation_recall_injection.py
Normal file
File diff suppressed because it is too large
Load diff
|
|
@ -0,0 +1,193 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""`search_conversation` under a safetensors model.
|
||||
|
||||
The tool is advertised by thread, not by backend: a GGUF-compacted chat keeps its
|
||||
archive across a model switch and `_select_request_tools` is shared, so this loop offers
|
||||
the tool too and needs both guards the GGUF loop applies.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import threading
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from core.inference.safetensors_agentic import run_safetensors_tool_loop # noqa: E402
|
||||
from core.inference.tool_call_parser import RAG_MAX_SEARCHES_PER_TURN # noqa: E402
|
||||
|
||||
MESSAGES = [
|
||||
{"role": "user", "content": "what was the code from earlier"},
|
||||
{"role": "assistant", "content": "let me look"},
|
||||
{"role": "user", "content": "please"},
|
||||
]
|
||||
|
||||
|
||||
def _searching_model(calls_wanted):
|
||||
"""A model that issues one distinct conversation search per iteration."""
|
||||
state = {"n": 0}
|
||||
|
||||
def single_turn(messages):
|
||||
state["n"] += 1
|
||||
if state["n"] > calls_wanted:
|
||||
yield "done"
|
||||
return
|
||||
call = (
|
||||
'<tool_call>{"name":"search_conversation","arguments":'
|
||||
'{"query":"the code %d"}}</tool_call>' % state["n"]
|
||||
)
|
||||
buffer = ""
|
||||
for char in call:
|
||||
buffer += char
|
||||
yield buffer
|
||||
|
||||
return single_turn
|
||||
|
||||
|
||||
def _run(calls_wanted, execute_tool):
|
||||
return list(
|
||||
run_safetensors_tool_loop(
|
||||
single_turn = _searching_model(calls_wanted),
|
||||
messages = list(MESSAGES),
|
||||
tools = [{"type": "function", "function": {"name": "search_conversation"}}],
|
||||
execute_tool = execute_tool,
|
||||
cancel_event = threading.Event(),
|
||||
max_tool_iterations = calls_wanted + 1,
|
||||
thread_id = "t-sf",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def test_the_loop_hands_its_search_the_active_branch():
|
||||
"""Without it the search falls back to the whole stored DAG, siblings included."""
|
||||
seen = []
|
||||
|
||||
def execute_tool(name, arguments, **kwargs):
|
||||
seen.append(kwargs.get("conversation_branch"))
|
||||
return "an earlier turn"
|
||||
|
||||
_run(1, execute_tool)
|
||||
|
||||
assert seen == [MESSAGES]
|
||||
|
||||
|
||||
def test_conversation_searches_share_the_per_turn_cap():
|
||||
"""Paraphrased re-searches slip past the exact-args duplicate guard.
|
||||
|
||||
Each appends passages into the protected current tool exchange, so uncapped the turn
|
||||
can only end in a context-length error.
|
||||
"""
|
||||
executed = []
|
||||
|
||||
def execute_tool(name, arguments, **kwargs):
|
||||
executed.append(arguments.get("query"))
|
||||
return "an earlier turn"
|
||||
|
||||
_run(RAG_MAX_SEARCHES_PER_TURN + 2, execute_tool)
|
||||
|
||||
assert len(executed) == RAG_MAX_SEARCHES_PER_TURN
|
||||
|
||||
|
||||
def test_the_budget_charges_token_dense_text_at_its_real_rate():
|
||||
"""Four characters per token is about right for English and half the truth for CJK.
|
||||
|
||||
The result side already prices non-ASCII at a token per character; the spend side did
|
||||
not, so a CJK chat reported roughly twice the room it had. This path runs no rolling
|
||||
fit, so nothing downstream recovers: the tool exchange it sized lands in the next
|
||||
prompt and takes it past the window.
|
||||
"""
|
||||
from core.inference.context_window import (
|
||||
estimate_messages_tokens,
|
||||
estimate_messages_tokens_dense,
|
||||
prompt_budget,
|
||||
)
|
||||
|
||||
dense_messages = [
|
||||
{"role": "user", "content": "設定を確認してください。" * 40},
|
||||
{"role": "assistant", "content": "承知しました。" * 40},
|
||||
{"role": "user", "content": "先ほどのコードは何でしたか"},
|
||||
]
|
||||
seen = {}
|
||||
tools = [{"type": "function", "function": {"name": "search_conversation"}}]
|
||||
|
||||
list(
|
||||
run_safetensors_tool_loop(
|
||||
single_turn = _searching_model(1),
|
||||
messages = list(dense_messages),
|
||||
tools = tools,
|
||||
execute_tool = lambda name, arguments, **kwargs: (
|
||||
seen.update(kwargs) or "an earlier turn"
|
||||
),
|
||||
cancel_event = threading.Event(),
|
||||
max_tool_iterations = 2,
|
||||
thread_id = "t-sf",
|
||||
context_length = 4096,
|
||||
max_tokens = 512,
|
||||
)
|
||||
)
|
||||
|
||||
budget = seen.get("conversation_budget_tokens")
|
||||
assert budget is not None
|
||||
# The flat estimate is the one that overstates the room. Charged densely, the budget
|
||||
# has to be at most what is left after the real cost of what is already there.
|
||||
assert budget <= prompt_budget(4096, 512) - estimate_messages_tokens_dense(dense_messages)
|
||||
assert estimate_messages_tokens_dense(dense_messages) > estimate_messages_tokens(dense_messages)
|
||||
|
||||
|
||||
def test_the_dense_estimate_matches_the_flat_one_on_plain_ascii():
|
||||
"""It corrects a known undercount, it does not make every English chat pessimistic."""
|
||||
from core.inference.context_window import (
|
||||
estimate_messages_tokens,
|
||||
estimate_messages_tokens_dense,
|
||||
)
|
||||
assert estimate_messages_tokens_dense(MESSAGES) == estimate_messages_tokens(MESSAGES)
|
||||
|
||||
|
||||
def test_the_loop_budgets_its_search_against_this_models_context():
|
||||
"""Without it the clamp in the tool is skipped and top_k 8 lands unbudgeted.
|
||||
|
||||
Roughly 4K tokens appended into the current tool exchange, which the rolling window
|
||||
protects and cannot evict.
|
||||
"""
|
||||
from core.inference.context_window import estimate_messages_tokens, prompt_budget
|
||||
|
||||
seen = {}
|
||||
tools = [{"type": "function", "function": {"name": "search_conversation"}}]
|
||||
|
||||
def execute_tool(name, arguments, **kwargs):
|
||||
seen.update(kwargs)
|
||||
return "an earlier turn"
|
||||
|
||||
list(
|
||||
run_safetensors_tool_loop(
|
||||
single_turn = _searching_model(1),
|
||||
messages = list(MESSAGES),
|
||||
tools = tools,
|
||||
execute_tool = execute_tool,
|
||||
cancel_event = threading.Event(),
|
||||
max_tool_iterations = 2,
|
||||
thread_id = "t-sf",
|
||||
context_length = 4096,
|
||||
max_tokens = 512,
|
||||
)
|
||||
)
|
||||
|
||||
budget = seen.get("conversation_budget_tokens")
|
||||
assert budget is not None
|
||||
# Everything already in the prompt is charged: the messages and the catalogue.
|
||||
assert budget <= prompt_budget(4096, 512) - estimate_messages_tokens(tools)
|
||||
assert budget < prompt_budget(4096, 512) - estimate_messages_tokens(MESSAGES)
|
||||
|
||||
|
||||
def test_the_loop_omits_the_budget_when_the_context_is_unknown():
|
||||
"""An absent budget must not become a budget of zero, which would refuse every search."""
|
||||
seen = {}
|
||||
|
||||
def execute_tool(name, arguments, **kwargs):
|
||||
seen.update(kwargs)
|
||||
return "an earlier turn"
|
||||
|
||||
_run(1, execute_tool)
|
||||
|
||||
assert "conversation_budget_tokens" not in seen
|
||||
|
|
@ -3697,7 +3697,12 @@ def test_tool_loop_refits_each_preflight_path_after_context_shrinking_respawn(mo
|
|||
],
|
||||
payloads,
|
||||
)
|
||||
backend._effective_context_length = 100
|
||||
# Sized so each window overflows by roughly one turn-group. Compaction trims a
|
||||
# headroom margin BELOW the budget and the turn-picking estimator is coarser than
|
||||
# the exact count, so single-group steps would evict the whole history in one pass
|
||||
# and leave the second preflight nothing to refit. The property under test is that
|
||||
# BOTH preflight paths refit against the window they were given.
|
||||
backend._effective_context_length = 2000
|
||||
monkeypatch.setattr(
|
||||
backend,
|
||||
"count_chat_tokens",
|
||||
|
|
@ -3707,17 +3712,17 @@ def test_tool_loop_refits_each_preflight_path_after_context_shrinking_respawn(mo
|
|||
)
|
||||
|
||||
def fake_respawn():
|
||||
backend._effective_context_length = 60
|
||||
backend._effective_context_length = 1000
|
||||
return True
|
||||
|
||||
monkeypatch.setattr(backend, "_respawn_if_dead", fake_respawn)
|
||||
events = list(
|
||||
backend.generate_chat_completion_with_tools(
|
||||
messages = [
|
||||
{"role": "user", "content": "u" * 25},
|
||||
{"role": "assistant", "content": "a" * 25},
|
||||
{"role": "user", "content": "u" * 25},
|
||||
{"role": "assistant", "content": "a" * 25},
|
||||
{"role": "user", "content": "u" * 400},
|
||||
{"role": "assistant", "content": "a" * 400},
|
||||
{"role": "user", "content": "u" * 400},
|
||||
{"role": "assistant", "content": "a" * 400},
|
||||
{"role": "user", "content": "final"},
|
||||
],
|
||||
tools = [{"type": "function", "function": {"name": "python"}}],
|
||||
|
|
@ -3728,8 +3733,8 @@ def test_tool_loop_refits_each_preflight_path_after_context_shrinking_respawn(mo
|
|||
|
||||
notices = [event for event in events if event.get("type") == "context_truncated"]
|
||||
assert [notice["dropped_messages"] for notice in notices] == [2, 2]
|
||||
assert [notice["context_length"] for notice in notices] == [100, 60]
|
||||
assert [payload["max_tokens"] for payload in payloads] == [100, 60]
|
||||
assert [notice["context_length"] for notice in notices] == [2000, 1000]
|
||||
assert [payload["max_tokens"] for payload in payloads] == [2000, 1000]
|
||||
assert len(payloads[0]["messages"]) == 3
|
||||
assert len(payloads[1]["messages"]) == 1
|
||||
|
||||
|
|
@ -3897,6 +3902,119 @@ def test_rolling_preflight_counts_the_sanitized_payload(monkeypatch):
|
|||
assert counted[0] != messages
|
||||
|
||||
|
||||
def test_a_respawn_refit_archives_what_it_evicts(monkeypatch):
|
||||
"""The respawn refits run against a smaller replacement window.
|
||||
|
||||
They evict more of the conversation, and without archiving there those turns are
|
||||
gone for good: unlike the ordinary preflight, nothing else sees them.
|
||||
"""
|
||||
import httpx
|
||||
from core.inference import llama_cpp
|
||||
|
||||
archived: list = []
|
||||
|
||||
def fake_archive(conversation, before, **kwargs):
|
||||
archived.append(llama_cpp.evicted_messages(before, conversation))
|
||||
return {"conversation": conversation, "events": [], "counts": {}, "recalled": False}
|
||||
|
||||
monkeypatch.setattr(llama_cpp, "_archive_and_recall", fake_archive)
|
||||
|
||||
payloads: list[dict] = []
|
||||
backend = _make_backend(
|
||||
monkeypatch,
|
||||
[httpx.ConnectError("server is down"), [_sse({"content": "OK"}), _done()]],
|
||||
payloads,
|
||||
)
|
||||
backend._effective_context_length = 2000
|
||||
monkeypatch.setattr(
|
||||
backend,
|
||||
"count_chat_tokens",
|
||||
lambda candidate, *_args, **_kwargs: sum(
|
||||
len(str(message.get("content", ""))) for message in candidate
|
||||
),
|
||||
)
|
||||
|
||||
def fake_respawn():
|
||||
backend._effective_context_length = 1000
|
||||
return True
|
||||
|
||||
monkeypatch.setattr(backend, "_respawn_if_dead", fake_respawn)
|
||||
list(
|
||||
backend.generate_chat_completion_with_tools(
|
||||
messages = [
|
||||
{"role": "user", "content": "u" * 400},
|
||||
{"role": "assistant", "content": "a" * 400},
|
||||
{"role": "user", "content": "u" * 400},
|
||||
{"role": "assistant", "content": "a" * 400},
|
||||
{"role": "user", "content": "final"},
|
||||
],
|
||||
tools = [{"type": "function", "function": {"name": "python"}}],
|
||||
max_tool_iterations = 1,
|
||||
context_overflow = "truncate_oldest",
|
||||
thread_id = "t-respawn-archive",
|
||||
)
|
||||
)
|
||||
|
||||
# More than one archiving pass, and the respawn's own evictions are among them.
|
||||
assert len(archived) >= 2
|
||||
assert any(batch for batch in archived[1:])
|
||||
|
||||
|
||||
def test_the_respawn_retry_keeps_the_thread(monkeypatch):
|
||||
"""The retry refits for the replacement window, so it can evict more.
|
||||
|
||||
Without the thread those extra turns are archived nowhere and no reserve or boundary
|
||||
applies, on the one path that deliberately compacts a second time.
|
||||
"""
|
||||
import httpx
|
||||
from core.inference import llama_cpp
|
||||
|
||||
payloads: list[dict] = []
|
||||
backend = _make_backend(
|
||||
monkeypatch,
|
||||
[httpx.ConnectError("server is down"), [_sse({"content": "OK"}), _done()]],
|
||||
payloads,
|
||||
)
|
||||
backend._effective_context_length = 2000
|
||||
monkeypatch.setattr(
|
||||
backend,
|
||||
"count_chat_tokens",
|
||||
lambda candidate, *_args, **_kwargs: sum(
|
||||
len(str(message.get("content", ""))) for message in candidate
|
||||
),
|
||||
)
|
||||
|
||||
def fake_respawn():
|
||||
backend._effective_context_length = 1000
|
||||
return True
|
||||
|
||||
monkeypatch.setattr(backend, "_respawn_if_dead", fake_respawn)
|
||||
seen: list = []
|
||||
monkeypatch.setattr(
|
||||
llama_cpp,
|
||||
"_conversation_recall_reserve",
|
||||
lambda thread_id: seen.append(thread_id) or 0,
|
||||
)
|
||||
|
||||
list(
|
||||
backend.generate_chat_completion(
|
||||
messages = [
|
||||
{"role": "user", "content": "u" * 400},
|
||||
{"role": "assistant", "content": "a" * 400},
|
||||
{"role": "user", "content": "u" * 400},
|
||||
{"role": "assistant", "content": "a" * 400},
|
||||
{"role": "user", "content": "final"},
|
||||
],
|
||||
context_overflow = "truncate_oldest",
|
||||
thread_id = "t-respawn",
|
||||
)
|
||||
)
|
||||
|
||||
# Both fits, the original and the one the retry runs, know which thread they are on.
|
||||
assert len(seen) == 2
|
||||
assert seen == ["t-respawn", "t-respawn"]
|
||||
|
||||
|
||||
def test_rolling_respawn_retry_refits_when_the_effective_context_changes(monkeypatch):
|
||||
"""A smaller replacement window can evict more without repeating the first eviction."""
|
||||
import httpx
|
||||
|
|
@ -3907,7 +4025,12 @@ def test_rolling_respawn_retry_refits_when_the_effective_context_changes(monkeyp
|
|||
[httpx.ConnectError("server is down"), [_sse({"content": "OK"}), _done()]],
|
||||
payloads,
|
||||
)
|
||||
backend._effective_context_length = 100
|
||||
# Sized so each window overflows by roughly one turn-group. Compaction trims a
|
||||
# headroom margin BELOW the budget and the turn-picking estimator is coarser than
|
||||
# the exact count, so single-group steps would evict the whole history in one pass
|
||||
# and leave the second preflight nothing to refit. The property under test is that
|
||||
# BOTH preflight paths refit against the window they were given.
|
||||
backend._effective_context_length = 2000
|
||||
monkeypatch.setattr(
|
||||
backend,
|
||||
"count_chat_tokens",
|
||||
|
|
@ -3917,17 +4040,17 @@ def test_rolling_respawn_retry_refits_when_the_effective_context_changes(monkeyp
|
|||
)
|
||||
|
||||
def fake_respawn():
|
||||
backend._effective_context_length = 60
|
||||
backend._effective_context_length = 1000
|
||||
return True
|
||||
|
||||
monkeypatch.setattr(backend, "_respawn_if_dead", fake_respawn)
|
||||
events = list(
|
||||
backend.generate_chat_completion(
|
||||
messages = [
|
||||
{"role": "user", "content": "u" * 25},
|
||||
{"role": "assistant", "content": "a" * 25},
|
||||
{"role": "user", "content": "u" * 25},
|
||||
{"role": "assistant", "content": "a" * 25},
|
||||
{"role": "user", "content": "u" * 400},
|
||||
{"role": "assistant", "content": "a" * 400},
|
||||
{"role": "user", "content": "u" * 400},
|
||||
{"role": "assistant", "content": "a" * 400},
|
||||
{"role": "user", "content": "final"},
|
||||
],
|
||||
context_overflow = "truncate_oldest",
|
||||
|
|
@ -3940,8 +4063,8 @@ def test_rolling_respawn_retry_refits_when_the_effective_context_changes(monkeyp
|
|||
if isinstance(event, dict) and event.get("type") == "context_truncated"
|
||||
]
|
||||
assert [notice["dropped_messages"] for notice in notices] == [2, 2]
|
||||
assert [notice["context_length"] for notice in notices] == [100, 60]
|
||||
assert [payload["max_tokens"] for payload in payloads] == [100, 60]
|
||||
assert [notice["context_length"] for notice in notices] == [2000, 1000]
|
||||
assert [payload["max_tokens"] for payload in payloads] == [2000, 1000]
|
||||
assert len(payloads[0]["messages"]) == 3
|
||||
assert len(payloads[1]["messages"]) == 1
|
||||
|
||||
|
|
@ -5575,3 +5698,313 @@ def test_parallel_disabled_suppresses_provisional_for_reused_index(monkeypatch):
|
|||
if e.get("tool_call_id") == "call_term"
|
||||
and e.get("type") in {"tool_start", "tool_args", "tool_end"}
|
||||
]
|
||||
|
||||
|
||||
def test_conversation_search_budget_counts_the_tool_catalogue(monkeypatch):
|
||||
"""The estimator sees the messages only; the tools array is prompt too.
|
||||
|
||||
A large (MCP) catalogue can be thousands of tokens, so a budget ignoring it reports
|
||||
room the request lacks, into a tool exchange the next iteration cannot evict.
|
||||
"""
|
||||
payloads: list[dict] = []
|
||||
backend = _make_backend(
|
||||
monkeypatch,
|
||||
[
|
||||
[
|
||||
_sse(
|
||||
{
|
||||
"tool_calls": [
|
||||
{
|
||||
"index": 0,
|
||||
"id": "call_s",
|
||||
"function": {
|
||||
"name": "search_conversation",
|
||||
"arguments": '{"query":"the code"}',
|
||||
},
|
||||
}
|
||||
]
|
||||
}
|
||||
),
|
||||
_finish("tool_calls"),
|
||||
_done(),
|
||||
],
|
||||
[_sse({"content": "It was 5150."}), _finish("stop"), _done()],
|
||||
],
|
||||
payloads,
|
||||
)
|
||||
backend._effective_context_length = 4096
|
||||
# What llama-server would really return: the messages, plus a catalogue that on its
|
||||
# own fills most of the window. The estimator counts the messages and nothing else.
|
||||
monkeypatch.setattr(
|
||||
backend,
|
||||
"count_chat_tokens",
|
||||
lambda candidate, *_a, **_k: 2800
|
||||
+ sum(len(str(message.get("content", ""))) for message in candidate) // 10,
|
||||
)
|
||||
|
||||
seen = {}
|
||||
|
||||
def execute_tool(name, arguments, **kwargs):
|
||||
seen.update(kwargs)
|
||||
return "an earlier turn"
|
||||
|
||||
monkeypatch.setattr("core.inference.tools.execute_tool", execute_tool)
|
||||
|
||||
list(
|
||||
backend.generate_chat_completion_with_tools(
|
||||
messages = [
|
||||
{"role": "user", "content": "u" * 2000},
|
||||
{"role": "assistant", "content": "a" * 2000},
|
||||
{"role": "user", "content": "u" * 2000},
|
||||
{"role": "assistant", "content": "a" * 2000},
|
||||
{"role": "user", "content": "what was the code"},
|
||||
],
|
||||
tools = [{"type": "function", "function": {"name": "search_conversation"}}],
|
||||
max_tokens = 512,
|
||||
context_overflow = "truncate_oldest",
|
||||
)
|
||||
)
|
||||
|
||||
from core.inference.context_window import prompt_budget
|
||||
|
||||
budget = seen.get("conversation_budget_tokens")
|
||||
assert budget is not None
|
||||
# 2,800 of the 3,584-token budget is catalogue and framing the estimator cannot see,
|
||||
# so what is left is hundreds of tokens, not the thousands it would have claimed.
|
||||
assert 0 <= budget < 1000
|
||||
|
||||
|
||||
def test_a_long_tool_run_reports_a_boundary_in_the_requests_own_terms(monkeypatch):
|
||||
"""dropped_messages is summed by the client, and it counts THIS request's messages.
|
||||
|
||||
A tool loop refits every iteration, so a long agent run also counts the tool
|
||||
exchanges it created, which the next request's transcript lacks. Re-applying that
|
||||
total advances the boundary past the turns actually evicted, so the boundary is
|
||||
carried separately, measured against the messages the request was sent with.
|
||||
"""
|
||||
calls = 6
|
||||
streams = []
|
||||
for index in range(calls):
|
||||
streams.append(
|
||||
[
|
||||
_sse(
|
||||
{
|
||||
"tool_calls": [
|
||||
{
|
||||
"index": 0,
|
||||
"id": f"c{index}",
|
||||
"function": {
|
||||
"name": "python",
|
||||
"arguments": '{"code": "step %d"}' % index,
|
||||
},
|
||||
}
|
||||
]
|
||||
}
|
||||
),
|
||||
_finish("tool_calls"),
|
||||
_done(),
|
||||
]
|
||||
)
|
||||
streams.append([_sse({"content": "done."}), _finish("stop"), _done()])
|
||||
|
||||
payloads: list[dict] = []
|
||||
backend = _make_backend(monkeypatch, streams, payloads)
|
||||
backend._effective_context_length = 4000
|
||||
monkeypatch.setattr(
|
||||
backend,
|
||||
"count_chat_tokens",
|
||||
lambda candidate, *_a, **_k: sum(
|
||||
len(str(message.get("content", ""))) for message in candidate
|
||||
)
|
||||
// 4,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"core.inference.tools.execute_tool", lambda name, arguments, **_k: "R" * 3200
|
||||
)
|
||||
|
||||
branch = [
|
||||
# Studio always prepends one and a fit never evicts it, so counting it as the
|
||||
# front of the branch reported zero on every compaction.
|
||||
{"role": "system", "content": "you are helpful"},
|
||||
{"role": "user", "content": "u" * 1200},
|
||||
{"role": "assistant", "content": "a" * 1200},
|
||||
{"role": "user", "content": "u2" * 600},
|
||||
{"role": "assistant", "content": "a2" * 600},
|
||||
{"role": "user", "content": "keep going"},
|
||||
]
|
||||
events = list(
|
||||
backend.generate_chat_completion_with_tools(
|
||||
messages = branch,
|
||||
tools = [{"type": "function", "function": {"name": "python"}}],
|
||||
max_tokens = 400,
|
||||
max_tool_iterations = calls + 1,
|
||||
context_overflow = "truncate_oldest",
|
||||
)
|
||||
)
|
||||
|
||||
notices = [
|
||||
event for event in events if event.get("type") == "context_truncated" and event.get("fits")
|
||||
]
|
||||
assert len(notices) > 1, "the fixture must refit more than once"
|
||||
# Summed, this passes the number of evictable messages the branch ever had.
|
||||
assert sum(notice["dropped_messages"] for notice in notices) > len(branch)
|
||||
# The boundary does not: it says where the branch was cut, so it never passes what the
|
||||
# branch had to give (4; the system prompt and the latest turn are neither evictable
|
||||
# nor counted) and it only ever moves forward.
|
||||
boundaries = [notice["boundary_messages"] for notice in notices]
|
||||
assert max(boundaries) == 4
|
||||
assert boundaries == sorted(boundaries)
|
||||
|
||||
|
||||
def test_conversation_search_budget_is_exact_when_nothing_was_truncated(monkeypatch):
|
||||
"""`fit_rolling_context` returns None when it drops nothing.
|
||||
|
||||
A prompt that simply FITS, after a context-length increase or on a shorter branch,
|
||||
therefore left the budget to a character estimate that cannot see the template's own
|
||||
framing. It reported room the request did not have, the recall appended a passage too
|
||||
large for the real window, and the next iteration could not evict it again because the
|
||||
current tool exchange is protected.
|
||||
"""
|
||||
payloads: list[dict] = []
|
||||
backend = _make_backend(
|
||||
monkeypatch,
|
||||
[
|
||||
[
|
||||
_sse(
|
||||
{
|
||||
"tool_calls": [
|
||||
{
|
||||
"index": 0,
|
||||
"id": "call_s",
|
||||
"function": {
|
||||
"name": "search_conversation",
|
||||
"arguments": '{"query":"the code"}',
|
||||
},
|
||||
}
|
||||
]
|
||||
}
|
||||
),
|
||||
_finish("tool_calls"),
|
||||
_done(),
|
||||
],
|
||||
[_sse({"content": "It was 5150."}), _finish("stop"), _done()],
|
||||
],
|
||||
payloads,
|
||||
)
|
||||
backend._effective_context_length = 4096
|
||||
# Most of the window is catalogue and template framing, which no character estimate
|
||||
# can see. The messages themselves are short, so the fit drops nothing at all.
|
||||
monkeypatch.setattr(
|
||||
backend,
|
||||
"count_chat_tokens",
|
||||
lambda candidate, *_a, **_k: 2800
|
||||
+ sum(len(str(message.get("content", ""))) for message in candidate) // 10,
|
||||
)
|
||||
|
||||
seen = {}
|
||||
|
||||
def execute_tool(name, arguments, **kwargs):
|
||||
seen.update(kwargs)
|
||||
return "an earlier turn"
|
||||
|
||||
monkeypatch.setattr("core.inference.tools.execute_tool", execute_tool)
|
||||
|
||||
list(
|
||||
backend.generate_chat_completion_with_tools(
|
||||
messages = [{"role": "user", "content": "what was the code"}],
|
||||
tools = [{"type": "function", "function": {"name": "search_conversation"}}],
|
||||
max_tokens = 512,
|
||||
context_overflow = "truncate_oldest",
|
||||
)
|
||||
)
|
||||
|
||||
budget = seen.get("conversation_budget_tokens")
|
||||
assert budget is not None
|
||||
# 3,584 of budget against a real prompt of roughly 2,800: hundreds of tokens of room,
|
||||
# not the thousands the estimate claimed from a handful of short messages.
|
||||
assert 0 <= budget < 1000, budget
|
||||
|
||||
|
||||
def test_the_exact_recall_budget_is_recomputed_after_an_intervening_tool(monkeypatch):
|
||||
"""The exact count is absolute, so caching it for the request goes stale.
|
||||
|
||||
The loop appends the assistant call and the tool result of every intervening tool to
|
||||
the conversation, so a figure taken before them understates the prompt by exactly
|
||||
those exchanges and hands the search room that is already spent.
|
||||
"""
|
||||
payloads: list[dict] = []
|
||||
backend = _make_backend(
|
||||
monkeypatch,
|
||||
[
|
||||
[
|
||||
_sse(
|
||||
{
|
||||
"tool_calls": [
|
||||
{
|
||||
"index": 0,
|
||||
"id": "call_t",
|
||||
"function": {"name": "terminal", "arguments": '{"command":"ls"}'},
|
||||
}
|
||||
]
|
||||
}
|
||||
),
|
||||
_finish("tool_calls"),
|
||||
_done(),
|
||||
],
|
||||
[
|
||||
_sse(
|
||||
{
|
||||
"tool_calls": [
|
||||
{
|
||||
"index": 0,
|
||||
"id": "call_s",
|
||||
"function": {
|
||||
"name": "search_conversation",
|
||||
"arguments": '{"query":"the code"}',
|
||||
},
|
||||
}
|
||||
]
|
||||
}
|
||||
),
|
||||
_finish("tool_calls"),
|
||||
_done(),
|
||||
],
|
||||
[_sse({"content": "It was 5150."}), _finish("stop"), _done()],
|
||||
],
|
||||
payloads,
|
||||
)
|
||||
backend._effective_context_length = 4096
|
||||
monkeypatch.setattr(
|
||||
backend,
|
||||
"count_chat_tokens",
|
||||
lambda candidate, *_a, **_k: 1000
|
||||
+ sum(len(str(message.get("content", ""))) for message in candidate) // 10,
|
||||
)
|
||||
|
||||
budgets: list = []
|
||||
|
||||
def execute_tool(name, arguments, **kwargs):
|
||||
if name == "search_conversation":
|
||||
budgets.append(kwargs.get("conversation_budget_tokens"))
|
||||
return "an earlier turn"
|
||||
# A big result, which the loop appends before the search runs.
|
||||
return "x" * 12000
|
||||
|
||||
monkeypatch.setattr("core.inference.tools.execute_tool", execute_tool)
|
||||
|
||||
list(
|
||||
backend.generate_chat_completion_with_tools(
|
||||
messages = [{"role": "user", "content": "what was the code"}],
|
||||
tools = [
|
||||
{"type": "function", "function": {"name": "terminal"}},
|
||||
{"type": "function", "function": {"name": "search_conversation"}},
|
||||
],
|
||||
max_tokens = 512,
|
||||
context_overflow = "truncate_oldest",
|
||||
)
|
||||
)
|
||||
|
||||
assert budgets and budgets[0] is not None
|
||||
# The 12,000-character tool result is roughly 1,200 tokens of the 3,584-token budget,
|
||||
# and the count taken before it cannot see them.
|
||||
assert budgets[0] < 1400, budgets
|
||||
|
|
|
|||
|
|
@ -218,6 +218,56 @@ def test_structured_call_executes_and_continues(executed):
|
|||
assert follow_up[-1]["content"] == "RESULT<web_search>"
|
||||
|
||||
|
||||
def test_a_conversation_search_here_gets_the_active_branch(executed):
|
||||
"""The provider loops share the local paths' tool catalogue.
|
||||
|
||||
So search_conversation is advertised here once a thread has an archive, and needs the
|
||||
branch for the same reason: the stored rows are the whole DAG, Retry included.
|
||||
"""
|
||||
branch = [
|
||||
{"role": "user", "content": "what was the code"},
|
||||
{"role": "assistant", "content": "let me look"},
|
||||
{"role": "user", "content": "please"},
|
||||
]
|
||||
transport = FakeTransport(
|
||||
[
|
||||
[
|
||||
_sse(
|
||||
{
|
||||
"tool_calls": [
|
||||
{
|
||||
"index": 0,
|
||||
"id": "call_c",
|
||||
"function": {
|
||||
"name": "search_conversation",
|
||||
"arguments": '{"query":"the code"}',
|
||||
},
|
||||
}
|
||||
]
|
||||
}
|
||||
),
|
||||
_sse(finish = "tool_calls"),
|
||||
_DONE,
|
||||
],
|
||||
[_sse({"content": "It was 5150."}), _sse(finish = "stop"), _DONE],
|
||||
]
|
||||
)
|
||||
|
||||
_run(transport, tools = [_tool("search_conversation")], messages = branch)
|
||||
|
||||
assert [call["name"] for call in executed] == ["search_conversation"]
|
||||
assert executed[0]["conversation_branch"] == branch
|
||||
# And a budget, or the tool's clamp is skipped and a model-chosen top_k of 8 appends
|
||||
# roughly 4K tokens to a prompt this loop replays. Studio cannot measure an external
|
||||
# model's window, so the cap is one ordinary recall's worth.
|
||||
from core.rag import config as rag_config
|
||||
|
||||
assert (
|
||||
executed[0]["conversation_budget_tokens"]
|
||||
== rag_config.CHUNK_TOKENS * rag_config.CONVERSATION_ARCHIVE_TOP_K
|
||||
)
|
||||
|
||||
|
||||
def test_streamed_tool_name_fragments_are_not_concatenated(executed):
|
||||
"""llama-server re-sends the whole name as it grows: web -> web_search."""
|
||||
transport = FakeTransport(
|
||||
|
|
|
|||
|
|
@ -0,0 +1,61 @@
|
|||
// SPDX-License-Identifier: AGPL-3.0-only
|
||||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import { ArchiveIcon } from "lucide-react";
|
||||
import type { FC } from "react";
|
||||
|
||||
import type { ContextTruncation } from "@/features/chat/utils/context-truncation";
|
||||
|
||||
/**
|
||||
* Tells the user their conversation was compacted, on the turn where it STARTED.
|
||||
*
|
||||
* Deliberately NOT a message: it renders inside the assistant message's container from
|
||||
* `metadata.custom.contextTruncation`, so it is never sent to the model, editable, or
|
||||
* exported as content, yet unlike a toast it survives a reload and stays attached to the
|
||||
* turn it describes.
|
||||
*
|
||||
* Rendered once per COMPACTION, gated by the caller, not once per compacted turn: a
|
||||
* thread past its window refits on every request, so per-turn would mean a notice on
|
||||
* every reply forever. The caller shows this only when the eviction boundary moved.
|
||||
*/
|
||||
export const CompactionNotice: FC<{ truncation: ContextTruncation }> = ({
|
||||
truncation,
|
||||
}) => {
|
||||
if (!truncation?.fits || !truncation.dropped_messages) return null;
|
||||
|
||||
const dropped = truncation.dropped_messages;
|
||||
const archived = truncation.archived_messages ?? 0;
|
||||
const recalled = truncation.recalled_chunks ?? 0;
|
||||
|
||||
const detail = archived
|
||||
? "They are saved and searchable, and the parts relevant to each question are brought back automatically."
|
||||
: "The full conversation is still visible and saved here.";
|
||||
|
||||
return (
|
||||
<div
|
||||
className="aui-compaction-notice mb-3 flex items-start gap-2 rounded-lg border border-border/60 bg-muted/40 px-3 py-2 text-ui-13 text-muted-foreground"
|
||||
data-testid="compaction-notice"
|
||||
data-dropped={dropped}
|
||||
data-archived={archived}
|
||||
data-recalled={recalled}
|
||||
>
|
||||
<ArchiveIcon className="mt-0.5 size-3.5 shrink-0" aria-hidden />
|
||||
<div className="min-w-0">
|
||||
<span className="font-medium text-foreground/80">
|
||||
This conversation got long, so it was compacted.
|
||||
</span>{" "}
|
||||
<span>
|
||||
Older messages were dropped from the model's context to make room. {detail}
|
||||
</span>
|
||||
<span>
|
||||
{" "}
|
||||
({dropped} {dropped === 1 ? "message" : "messages"} dropped here
|
||||
{recalled > 0
|
||||
? `, ${recalled} earlier ${recalled === 1 ? "passage" : "passages"} recalled`
|
||||
: ""}
|
||||
.)
|
||||
</span>
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
};
|
||||
|
|
@ -9,6 +9,11 @@ import {
|
|||
GeneratedImageOverlayProvider,
|
||||
useGeneratedImageOverlay,
|
||||
} from "@/components/assistant-ui/generated-image-overlay-context";
|
||||
import { CompactionNotice } from "@/components/assistant-ui/compaction-notice";
|
||||
import {
|
||||
compactionBoundary,
|
||||
type ContextTruncation,
|
||||
} from "@/features/chat/utils/context-truncation";
|
||||
import { downloadImagePart } from "@/components/assistant-ui/image";
|
||||
import { MarkdownText } from "@/components/assistant-ui/markdown-text";
|
||||
import { MessageHtmlArtifacts } from "@/components/assistant-ui/message-html-artifacts";
|
||||
|
|
@ -6792,6 +6797,42 @@ const AssistantMessage: FC = () => {
|
|||
? custom.researchRunId
|
||||
: null;
|
||||
});
|
||||
// Persisted on the assistant turn that compacted, so the notice survives a reload.
|
||||
const contextTruncation = useAuiState(({ message }) => {
|
||||
const custom = (
|
||||
message.metadata as
|
||||
| { custom?: { contextTruncation?: unknown } }
|
||||
| undefined
|
||||
)?.custom;
|
||||
const value = custom?.contextTruncation;
|
||||
return value && typeof value === "object"
|
||||
? (value as ContextTruncation)
|
||||
: null;
|
||||
});
|
||||
// Once a thread outgrows the window every request runs the fit, so "this turn
|
||||
// compacted" is true of every later reply and would put a notice on all of them. What
|
||||
// matters is when MORE of the conversation fell out of view: the eviction boundary
|
||||
// rising above the last turn that reported one. Between moves the model sees the same
|
||||
// history, so there is nothing new to say.
|
||||
const showsNotice = useAuiState(({ thread }) => {
|
||||
let previousDropped = 0;
|
||||
for (const message of thread.messages) {
|
||||
if (message.role !== "assistant") continue;
|
||||
const value = (
|
||||
message.metadata as
|
||||
| { custom?: { contextTruncation?: unknown } }
|
||||
| undefined
|
||||
)?.custom?.contextTruncation as ContextTruncation | undefined;
|
||||
const dropped = compactionBoundary(value);
|
||||
if (dropped > previousDropped) {
|
||||
if (message.id === messageId) return true;
|
||||
previousDropped = dropped;
|
||||
} else if (message.id === messageId) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
return false;
|
||||
});
|
||||
const incognito = useChatRuntimeStore((s) => s.incognito);
|
||||
|
||||
// Use global store for editing state to ensure a single source of truth
|
||||
|
|
@ -6866,6 +6907,9 @@ const AssistantMessage: FC = () => {
|
|||
onBlur={focusReveal.onBlur}
|
||||
>
|
||||
<div className="aui-assistant-message-content wrap-break-word min-w-0 text-[#0d0d0d] dark:text-foreground leading-relaxed">
|
||||
{contextTruncation && showsNotice && !isEditing && (
|
||||
<CompactionNotice truncation={contextTruncation} />
|
||||
)}
|
||||
{isEditing ? (
|
||||
<div className="flex flex-col gap-2 w-full">
|
||||
<textarea
|
||||
|
|
|
|||
|
|
@ -5663,16 +5663,29 @@ export function createOpenAIStreamAdapter(
|
|||
chunk.context_truncated,
|
||||
);
|
||||
const activeThreadId = useChatRuntimeStore.getState().activeThreadId;
|
||||
// fits:false means the fitter could NOT make the request fit and returned
|
||||
// the original messages. Toasting "older turns were removed" is untrue
|
||||
// there, and burns the once-per-thread flag so a real one is silent.
|
||||
const reallyCompacted =
|
||||
chunk.context_truncated.fits === true &&
|
||||
(chunk.context_truncated.dropped_messages ?? 0) > 0;
|
||||
if (
|
||||
reallyCompacted &&
|
||||
resolvedThreadId &&
|
||||
activeThreadId === resolvedThreadId &&
|
||||
!rollingContextNoticeThreads.has(resolvedThreadId)
|
||||
) {
|
||||
rollingContextNoticeThreads.add(resolvedThreadId);
|
||||
toast.info("Older turns omitted from model context", {
|
||||
description:
|
||||
"The full conversation is still visible and saved. " +
|
||||
"Studio removed complete older turns from this request so the chat can continue.",
|
||||
// Once per thread per page load; the persistent record is the notice on
|
||||
// the assistant turn that compacted.
|
||||
const archived = contextTruncation?.archived_messages ?? 0;
|
||||
toast.info("This conversation was compacted", {
|
||||
description: archived
|
||||
? "It got long, so older turns were removed from the model's " +
|
||||
"context. They are saved and searchable, and relevant parts are " +
|
||||
"brought back automatically."
|
||||
: "The full conversation is still visible and saved. " +
|
||||
"Studio removed complete older turns from this request so the chat can continue.",
|
||||
duration: 8000,
|
||||
});
|
||||
}
|
||||
|
|
@ -6764,14 +6777,45 @@ export function createOpenAIStreamAdapter(
|
|||
duration: 8000,
|
||||
});
|
||||
} else if (isContextLimitError(msg)) {
|
||||
// llama-server runs with --no-context-shift, returning a hard
|
||||
// error instead of silently dropping old KV-cache turns. Point
|
||||
// the user at the control that raises the ceiling.
|
||||
// `fits: false` means everything evictable was evicted and the request STILL
|
||||
// does not fit, so the message just sent is the problem and the usual advice
|
||||
// is wrong: the history is already gone. Say which part is too long.
|
||||
const irreducible =
|
||||
contextTruncation?.fits === false ? contextTruncation : null;
|
||||
// Against prompt_target, not context_length: the fit reserves up to a quarter
|
||||
// of the window for the reply, so a 3,500-token message cannot fit a 4,096
|
||||
// context, and comparing with the raw window would blame the conversation and
|
||||
// send the user to a new chat that fails identically.
|
||||
const budget =
|
||||
irreducible?.prompt_target ?? irreducible?.context_length ?? 0;
|
||||
const oneTurnIsTheProblem =
|
||||
irreducible != null && (irreducible.latest_turn_tokens ?? 0) > budget;
|
||||
// Whose turn it is decides the advice: in a tool loop the offending turn is
|
||||
// often output the user never wrote and cannot edit, so "shorten this
|
||||
// message" names the wrong thing and offers no remedy.
|
||||
const userCanShortenIt =
|
||||
(irreducible?.latest_turn_role ?? "user") === "user";
|
||||
const tooLong =
|
||||
`${irreducible?.latest_turn_tokens?.toLocaleString()} tokens on its own, ` +
|
||||
`against the ${budget.toLocaleString()} tokens this ` +
|
||||
`${irreducible?.context_length?.toLocaleString()}-token window leaves for the prompt. ` +
|
||||
"The earlier turns were already removed and it still does not fit, so " +
|
||||
"shortening the conversation will not help. ";
|
||||
toast.error("Context limit reached", {
|
||||
description:
|
||||
"The conversation has filled the model's context window. " +
|
||||
'Increase "Context Length" in the chat Settings panel (⚙ in the top-right), ' +
|
||||
"or start a new chat.",
|
||||
description: oneTurnIsTheProblem
|
||||
? userCanShortenIt
|
||||
? `This message is ${tooLong}` +
|
||||
'Shorten this message, or raise "Context Length" ' +
|
||||
"in the chat Settings panel (⚙ in the top-right)."
|
||||
: `The last tool result is ${tooLong}` +
|
||||
'Raise "Context Length" in the chat Settings panel (⚙ in the top-right), ' +
|
||||
"or ask for less output from that tool."
|
||||
: // llama-server runs with --no-context-shift, returning a hard
|
||||
// error instead of silently dropping old KV-cache turns. Point
|
||||
// the user at the control that raises the ceiling.
|
||||
"The conversation has filled the model's context window. " +
|
||||
'Increase "Context Length" in the chat Settings panel (⚙ in the top-right), ' +
|
||||
"or start a new chat.",
|
||||
duration: 8000,
|
||||
});
|
||||
} else {
|
||||
|
|
|
|||
|
|
@ -698,5 +698,29 @@ export interface OpenAIChatChunk {
|
|||
prompt_tokens_after?: number;
|
||||
context_length?: number;
|
||||
fits: boolean;
|
||||
// Present when the evicted turns were archived and searched. Counts only, never
|
||||
// message text: this rides an SSE chunk that reaches the client.
|
||||
archived_messages?: number;
|
||||
recalled_chunks?: number;
|
||||
// Present only when `fits` is false: the floor the conversation cannot go below, and
|
||||
// how much of it is the message just sent. Together they say whether the history or
|
||||
// that one message is the problem, i.e. whether "shorten the conversation" helps.
|
||||
irreducible_tokens?: number;
|
||||
latest_turn_tokens?: number;
|
||||
// Where the compaction boundary sits in the messages THIS request was sent with.
|
||||
// Absolute, unlike dropped_messages, so re-sending it after a turn that refit several
|
||||
// times cannot advance the boundary past the turns actually evicted.
|
||||
boundary_messages?: number;
|
||||
// The text the boundary landed ON, so the count can be re-derived by position. A count
|
||||
// is only valid against the transcript it was counted on, and deleting an already
|
||||
// evicted prompt shortens that transcript; without the anchor the replayed count then
|
||||
// evicts live turns instead. Carried through untouched, like boundary_messages.
|
||||
boundary_anchor?: string;
|
||||
// Whose message that is: in a tool loop the last one is often a tool result rather
|
||||
// than anything the user typed.
|
||||
latest_turn_role?: string;
|
||||
// The prompt's share of the window (context_length minus the reply reserve), which is
|
||||
// what one turn must fit inside. Not re-derived here: the formula lives in the fit.
|
||||
prompt_target?: number;
|
||||
};
|
||||
}
|
||||
|
|
|
|||
|
|
@ -7,13 +7,33 @@ export type ContextTruncation = NonNullable<
|
|||
OpenAIChatChunk["context_truncated"]
|
||||
>;
|
||||
|
||||
function spreadSum(
|
||||
key: "archived_messages" | "recalled_chunks",
|
||||
a: number | undefined,
|
||||
b: number | undefined,
|
||||
): Record<string, number> {
|
||||
if (a === undefined && b === undefined) return {};
|
||||
return { [key]: (a ?? 0) + (b ?? 0) };
|
||||
}
|
||||
|
||||
export function compactionBoundary(
|
||||
truncation: ContextTruncation | undefined,
|
||||
): number {
|
||||
if (!truncation?.fits) return 0;
|
||||
// boundary_messages is where the boundary sits in the saved transcript.
|
||||
// dropped_messages accumulates what each fit removed, so a tool-heavy turn reports far
|
||||
// more than the boundary moved and a later real advance looks like none. Fallback only,
|
||||
// for turns saved before the boundary was recorded.
|
||||
return truncation.boundary_messages ?? truncation.dropped_messages ?? 0;
|
||||
}
|
||||
|
||||
export function mergeContextTruncation(
|
||||
current: ContextTruncation | undefined,
|
||||
incoming: ContextTruncation,
|
||||
): ContextTruncation {
|
||||
if (!current) return incoming;
|
||||
|
||||
return {
|
||||
const merged = {
|
||||
...current,
|
||||
...incoming,
|
||||
dropped_messages: current.dropped_messages + incoming.dropped_messages,
|
||||
|
|
@ -21,5 +41,25 @@ export function mergeContextTruncation(
|
|||
current.prompt_tokens_before ?? incoming.prompt_tokens_before,
|
||||
prompt_tokens_after:
|
||||
incoming.prompt_tokens_after ?? current.prompt_tokens_after,
|
||||
// A turn can compact more than once (the tool loop refits per iteration), so these
|
||||
// accumulate rather than taking the last chunk's value. Spread conditionally so a
|
||||
// plain rolling-window response keeps its old shape, with no archive keys set to
|
||||
// undefined.
|
||||
...spreadSum("archived_messages", current.archived_messages, incoming.archived_messages),
|
||||
...spreadSum("recalled_chunks", current.recalled_chunks, incoming.recalled_chunks),
|
||||
};
|
||||
|
||||
// boundary_messages needs no rule: it is absolute, so the spread above already keeps
|
||||
// the latest fit's value. Summing it is the bug it exists to fix. boundary_anchor rides
|
||||
// along with it for the same reason, and the two must come from the SAME fit.
|
||||
|
||||
// The irreducible diagnosis describes ONE fit that gave up, so an earlier failure
|
||||
// followed by a later success would otherwise leave those numbers on a result that fit.
|
||||
// Deleted rather than spread as undefined, which would put both keys on every ordinary
|
||||
// response; delete on an absent key is a no-op.
|
||||
if (incoming.fits) {
|
||||
delete merged.irreducible_tokens;
|
||||
delete merged.latest_turn_tokens;
|
||||
}
|
||||
return merged;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -4,7 +4,10 @@
|
|||
import assert from "node:assert/strict";
|
||||
import { readFileSync } from "node:fs";
|
||||
import test from "node:test";
|
||||
import { mergeContextTruncation } from "../src/features/chat/utils/context-truncation.ts";
|
||||
import {
|
||||
compactionBoundary,
|
||||
mergeContextTruncation,
|
||||
} from "../src/features/chat/utils/context-truncation.ts";
|
||||
|
||||
const adapter = readFileSync(
|
||||
new URL("../src/features/chat/api/chat-adapter.ts", import.meta.url),
|
||||
|
|
@ -18,7 +21,7 @@ const transport = readFileSync(
|
|||
test("local chat opts into the rolling context policy", () => {
|
||||
assert.match(adapter, /isGguf === true/);
|
||||
assert.match(adapter, /context_overflow:\s*"truncate_oldest"/);
|
||||
assert.match(adapter, /Older turns omitted from model context/);
|
||||
assert.match(adapter, /This conversation was compacted/);
|
||||
});
|
||||
|
||||
test("the transport preserves standard chunks with context metadata", () => {
|
||||
|
|
@ -27,6 +30,40 @@ test("the transport preserves standard chunks with context metadata", () => {
|
|||
assert.match(adapter, /contextTruncation = mergeContextTruncation\(/);
|
||||
});
|
||||
|
||||
test("the compaction notice follows the boundary, not the accumulated drops", () => {
|
||||
// A tool-heavy turn reports 12 drops while the boundary moved to 4. Recording 12 as
|
||||
// the high-water mark means the next two real advances, to 8 and to 10, are silent.
|
||||
assert.equal(
|
||||
compactionBoundary({ dropped_messages: 12, boundary_messages: 4, fits: true }),
|
||||
4,
|
||||
);
|
||||
// Turns saved before the boundary existed still report something.
|
||||
assert.equal(compactionBoundary({ dropped_messages: 6, fits: true }), 6);
|
||||
// A fit that gave up moved no boundary at all.
|
||||
assert.equal(
|
||||
compactionBoundary({ dropped_messages: 0, boundary_messages: 0, fits: false }),
|
||||
0,
|
||||
);
|
||||
assert.equal(compactionBoundary(undefined), 0);
|
||||
});
|
||||
|
||||
test("the compaction boundary takes the latest value, never the sum", () => {
|
||||
// dropped_messages counts what each fit removed in front of it, this turn's tool
|
||||
// messages included, so summing it and re-applying the total advances the boundary
|
||||
// past the turns actually evicted. The boundary is carried separately and absolutely.
|
||||
const combined = mergeContextTruncation(
|
||||
mergeContextTruncation(undefined, {
|
||||
dropped_messages: 4,
|
||||
boundary_messages: 4,
|
||||
fits: true,
|
||||
}),
|
||||
{ dropped_messages: 4, boundary_messages: 4, fits: true },
|
||||
);
|
||||
|
||||
assert.equal(combined.dropped_messages, 8);
|
||||
assert.equal(combined.boundary_messages, 4);
|
||||
});
|
||||
|
||||
test("tool-loop truncation metadata accumulates across stream events", () => {
|
||||
const first = mergeContextTruncation(undefined, {
|
||||
dropped_messages: 2,
|
||||
|
|
@ -51,3 +88,233 @@ test("tool-loop truncation metadata accumulates across stream events", () => {
|
|||
fits: true,
|
||||
});
|
||||
});
|
||||
|
||||
|
||||
test("compaction counts accumulate and stay absent on a plain rolling window", () => {
|
||||
// A plain rolling-window response must keep exactly the shape it had before the
|
||||
// conversation archive existed, rather than carrying archive keys set to undefined.
|
||||
const plain = mergeContextTruncation(
|
||||
{ dropped_messages: 1, fits: true },
|
||||
{ dropped_messages: 2, fits: true },
|
||||
);
|
||||
assert.ok(!("archived_messages" in plain));
|
||||
assert.ok(!("recalled_chunks" in plain));
|
||||
|
||||
const archived = mergeContextTruncation(
|
||||
{ dropped_messages: 1, fits: true, archived_messages: 2, recalled_chunks: 4 },
|
||||
{ dropped_messages: 2, fits: true, archived_messages: 3, recalled_chunks: 1 },
|
||||
);
|
||||
assert.equal(archived.archived_messages, 5);
|
||||
assert.equal(archived.recalled_chunks, 5);
|
||||
});
|
||||
|
||||
test("the compaction notice renders from persisted metadata, not from a message", () => {
|
||||
const notice = readFileSync(
|
||||
new URL("../src/components/assistant-ui/compaction-notice.tsx", import.meta.url),
|
||||
"utf8",
|
||||
);
|
||||
const thread = readFileSync(
|
||||
new URL("../src/components/assistant-ui/thread.tsx", import.meta.url),
|
||||
"utf8",
|
||||
);
|
||||
// Read off metadata.custom so it can never become part of the conversation.
|
||||
assert.match(thread, /custom\?\.contextTruncation/);
|
||||
assert.match(thread, /<CompactionNotice truncation=\{contextTruncation\}/);
|
||||
assert.match(notice, /This conversation got long, so it was compacted/);
|
||||
});
|
||||
|
||||
test("the compaction notice is gated on the eviction boundary MOVING", () => {
|
||||
const thread = readFileSync(
|
||||
new URL("../src/components/assistant-ui/thread.tsx", import.meta.url),
|
||||
"utf8",
|
||||
);
|
||||
// Every request after the window fills runs the fit, so "this turn compacted" puts a
|
||||
// notice on every reply. The trigger is dropped_messages rising above the last turn
|
||||
// that reported it: more of the conversation actually leaving the context.
|
||||
assert.match(thread, /const showsNotice = useAuiState/);
|
||||
assert.match(thread, /contextTruncation && showsNotice && !isEditing/);
|
||||
assert.match(thread, /dropped > previousDropped/);
|
||||
// Walked in order, not against the preceding message: turns between two moves report
|
||||
// the same count and must not reset the baseline.
|
||||
assert.match(thread, /for \(const message of thread\.messages\)/);
|
||||
});
|
||||
|
||||
// The gate is a pure function of the thread's persisted truncation counts, so it can be
|
||||
// evaluated directly on the sequences the server actually produces.
|
||||
const noticeTurns = (dropped: (number | null)[]): number[] => {
|
||||
const shown: number[] = [];
|
||||
let previousDropped = 0;
|
||||
dropped.forEach((value, index) => {
|
||||
const d = value ?? 0;
|
||||
if (d > previousDropped) {
|
||||
shown.push(index);
|
||||
previousDropped = d;
|
||||
}
|
||||
});
|
||||
return shown;
|
||||
};
|
||||
|
||||
test("one notice per compaction, and silence on the turns in between", () => {
|
||||
// A compaction, a stretch of turns whose boundary does not move, then another.
|
||||
assert.deepStrictEqual(
|
||||
noticeTurns([0, 0, 52, 52, 52, 52, 52, 62, 62, 62, 74]),
|
||||
[2, 7, 10],
|
||||
);
|
||||
// The uncompacted case stays silent throughout.
|
||||
assert.deepStrictEqual(noticeTurns([0, 0, 0]), []);
|
||||
// A single compaction that never moves again is reported exactly once.
|
||||
assert.deepStrictEqual(noticeTurns([36, 36, 36]), [0]);
|
||||
});
|
||||
|
||||
test("a boundary that goes BACKWARDS does not re-announce", () => {
|
||||
// A rollback leaves a shorter branch needing less eviction. Less is missing than
|
||||
// before, so there is nothing to say and the baseline must not be dragged down.
|
||||
assert.deepStrictEqual(noticeTurns([52, 20, 20, 20]), [0]);
|
||||
});
|
||||
|
||||
/** The source of one function, by brace matching from its declaration. */
|
||||
const functionBody = (source: string, name: string): string => {
|
||||
const start = source.indexOf(`function ${name}(`);
|
||||
if (start < 0) return "";
|
||||
const open = source.indexOf("{", start);
|
||||
let depth = 0;
|
||||
for (let index = open; index < source.length; index += 1) {
|
||||
if (source[index] === "{") depth += 1;
|
||||
else if (source[index] === "}") {
|
||||
depth -= 1;
|
||||
if (depth === 0) return source.slice(start, index + 1);
|
||||
}
|
||||
}
|
||||
return "";
|
||||
};
|
||||
|
||||
test("the notice is a NOTICE, never part of the conversation", () => {
|
||||
const thread = readFileSync(
|
||||
new URL("../src/components/assistant-ui/thread.tsx", import.meta.url),
|
||||
"utf8",
|
||||
);
|
||||
const adapter = readFileSync(
|
||||
new URL("../src/features/chat/api/chat-adapter.ts", import.meta.url),
|
||||
"utf8",
|
||||
);
|
||||
const exporter = readFileSync(
|
||||
new URL("../src/features/chat/utils/conversation-markdown-export.ts", import.meta.url),
|
||||
"utf8",
|
||||
);
|
||||
|
||||
// 1. A sibling of the rendered content parts, not one of them: inside
|
||||
// MessagePrimitive.Parts everything that walks parts would pick it up.
|
||||
const noticeAt = thread.indexOf("<CompactionNotice");
|
||||
const partsAt = thread.indexOf("<MessagePrimitive.Parts", noticeAt);
|
||||
assert.ok(noticeAt > 0 && partsAt > noticeAt);
|
||||
assert.ok(
|
||||
!/<MessagePrimitive\.Parts[^>]*>[\s\S]*<CompactionNotice/.test(thread),
|
||||
"the notice must not be rendered inside the message's content parts",
|
||||
);
|
||||
|
||||
// 2. Nothing that builds a request may read the key it renders from. Bounded to the
|
||||
// function bodies: slicing to end-of-file also catches the streaming handler,
|
||||
// which reads contextTruncation legitimately on the way IN.
|
||||
for (const name of ["toOpenAIMessages", "serializeAssistantReplayMessages"]) {
|
||||
const body = functionBody(adapter, name);
|
||||
assert.ok(body.length > 0, `${name} not found`);
|
||||
assert.ok(
|
||||
!body.includes("contextTruncation"),
|
||||
`${name} must never read contextTruncation`,
|
||||
);
|
||||
}
|
||||
|
||||
// 3. Nor may the user-facing export, which is the other way text leaves a thread.
|
||||
assert.ok(!exporter.includes("contextTruncation"));
|
||||
assert.ok(!exporter.includes("compacted"));
|
||||
|
||||
// 4. Suppressed while editing, so it cannot be saved back as message text.
|
||||
assert.match(thread, /contextTruncation && showsNotice && !isEditing/);
|
||||
});
|
||||
|
||||
test("an irreducible fit reports a diagnosis, and it is dropped once something fits", () => {
|
||||
// A fit that gave up carries the numbers that say WHICH part is too long.
|
||||
const failed = mergeContextTruncation(undefined, {
|
||||
dropped_messages: 0,
|
||||
fits: false,
|
||||
prompt_tokens_before: 10290,
|
||||
prompt_tokens_after: 10290,
|
||||
context_length: 4096,
|
||||
irreducible_tokens: 5050,
|
||||
latest_turn_tokens: 5000,
|
||||
});
|
||||
assert.equal(failed.fits, false);
|
||||
assert.equal(failed.latest_turn_tokens, 5000);
|
||||
|
||||
// The loop refits per iteration, and an iteration that DOES fit must not carry the
|
||||
// earlier failure's numbers forward, where they describe nothing.
|
||||
const recovered = mergeContextTruncation(failed, {
|
||||
dropped_messages: 12,
|
||||
fits: true,
|
||||
prompt_tokens_after: 3000,
|
||||
context_length: 4096,
|
||||
});
|
||||
assert.equal(recovered.fits, true);
|
||||
assert.ok(!("irreducible_tokens" in recovered));
|
||||
assert.ok(!("latest_turn_tokens" in recovered));
|
||||
|
||||
// And an ordinary response never grows the keys at all, not even set to undefined.
|
||||
const plain = mergeContextTruncation(
|
||||
{ dropped_messages: 1, fits: true },
|
||||
{ dropped_messages: 2, fits: true },
|
||||
);
|
||||
assert.ok(!("irreducible_tokens" in plain));
|
||||
assert.ok(!("latest_turn_tokens" in plain));
|
||||
});
|
||||
|
||||
test("the too-long advice depends on WHICH part does not fit", () => {
|
||||
const adapterSource = readFileSync(
|
||||
new URL("../src/features/chat/api/chat-adapter.ts", import.meta.url),
|
||||
"utf8",
|
||||
);
|
||||
// Telling someone to shorten the conversation is a dead end when the history has
|
||||
// already been evicted and the single message is what overflows.
|
||||
assert.match(adapterSource, /contextTruncation\?\.fits === false/);
|
||||
assert.match(adapterSource, /shortening the conversation will not help/);
|
||||
assert.match(adapterSource, /latest_turn_tokens/);
|
||||
});
|
||||
|
||||
test("a fits:false diagnosis is not a compaction", () => {
|
||||
const source = readFileSync(
|
||||
new URL("../src/features/chat/api/chat-adapter.ts", import.meta.url),
|
||||
"utf8",
|
||||
);
|
||||
// The fitter returned the ORIGINAL messages with dropped_messages 0, so "older turns
|
||||
// were removed" is untrue, and toasting it burns the once-per-thread flag.
|
||||
assert.match(source, /const reallyCompacted =/);
|
||||
assert.match(source, /context_truncated\.fits === true/);
|
||||
assert.match(source, /dropped_messages \?\? 0\) > 0/);
|
||||
});
|
||||
|
||||
test("the advice depends on WHOSE turn does not fit", () => {
|
||||
const source = readFileSync(
|
||||
new URL("../src/features/chat/api/chat-adapter.ts", import.meta.url),
|
||||
"utf8",
|
||||
);
|
||||
// A tool loop refits with the tool result appended, so the offending turn is often
|
||||
// output the user never wrote and cannot edit, leaving no remedy.
|
||||
assert.match(source, /latest_turn_role/);
|
||||
assert.match(source, /const userCanShortenIt =/);
|
||||
assert.match(source, /The last tool result is/);
|
||||
// The user-authored case keeps its advice, and an older server that sends no role
|
||||
// still gets it (the default is "user").
|
||||
assert.match(source, /latest_turn_role \?\? "user"/);
|
||||
assert.match(source, /Shorten this message/);
|
||||
});
|
||||
|
||||
test("the too-long check uses the prompt budget, not the raw window", () => {
|
||||
const source = readFileSync(
|
||||
new URL("../src/features/chat/api/chat-adapter.ts", import.meta.url),
|
||||
"utf8",
|
||||
);
|
||||
// The fit reserves up to a quarter of the window for the reply, so a 3,500-token
|
||||
// message cannot fit a 4,096-token context. The raw window would blame the
|
||||
// conversation and send the user to a new chat that fails identically.
|
||||
assert.match(source, /irreducible\?\.prompt_target \?\? irreducible\?\.context_length/);
|
||||
assert.match(source, /latest_turn_tokens \?\? 0\) > budget/);
|
||||
});
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue