mirror of
https://github.com/usestrix/strix.git
synced 2026-08-21 10:48:59 +02:00
feat(context): model-aware conversation compaction for long scans
Restore cumulative history compaction (removed in the SDK migration) so a long scan no longer replays an ever-growing transcript until it overflows the model's context window and the run fails. - strix/llm/context_budget.py: resolve the model's real input/output token limits from LiteLLM metadata (128k gpt-4o, 272k gpt-5, 1M claude-sonnet-4, 131k deepseek), with a large configurable fallback for unmapped models and a chars/4 token-count fallback. - strix/llm/compaction.py: provider-agnostic compaction via litellm. Keeps a security-focused structured summary (objective, findings, credentials, payloads, URLs/paths, work state, dead ends, next move), keeps the most recent turns by token budget, and snaps the summary boundary so no tool call is separated from its result. Both triggers: proactive before each run and reactive compact-and-retry on a real context-overflow error. - Wire both triggers into the agent run loop next to the existing image recovery; add replace_session_items() with restore-on-failure. Env-tunable via STRIX_CONTEXT_* (auto-compact on by default).
This commit is contained in:
Ahmed Allam
@@ -22,6 +22,7 @@ from strix.core.sessions import (
|
||||
open_agent_session,
|
||||
strip_all_images_from_session,
|
||||
)
|
||||
from strix.llm.compaction import is_context_overflow, maybe_compact
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -40,6 +41,41 @@ logger = logging.getLogger(__name__)
|
||||
StreamEventSink = Callable[[str, Any], None]
|
||||
|
||||
_INPUT_REJECTION_CODES = frozenset({400, 404, 422})
|
||||
_MAX_COMPACTIONS_PER_CYCLE = 2
|
||||
|
||||
|
||||
def _run_config_model(run_config: RunConfig) -> str | None:
|
||||
return run_config.model if isinstance(run_config.model, str) else None
|
||||
|
||||
|
||||
def _agent_instructions(agent: Any) -> str:
|
||||
instructions = getattr(agent, "instructions", None)
|
||||
return instructions if isinstance(instructions, str) else ""
|
||||
|
||||
|
||||
def _agent_tools_text(agent: Any) -> str:
|
||||
parts: list[str] = []
|
||||
for tool in getattr(agent, "tools", []) or []:
|
||||
name = getattr(tool, "name", "")
|
||||
description = getattr(tool, "description", "") or ""
|
||||
schema = getattr(tool, "params_json_schema", "") or ""
|
||||
parts.append(f"{name} {description} {schema}")
|
||||
return "\n".join(parts)
|
||||
|
||||
|
||||
async def _compact_session(
|
||||
agent: Any, session: Session, run_config: RunConfig, *, force: bool
|
||||
) -> bool:
|
||||
model = _run_config_model(run_config)
|
||||
if session is None or model is None:
|
||||
return False
|
||||
return await maybe_compact(
|
||||
session,
|
||||
model=model,
|
||||
instructions=_agent_instructions(agent),
|
||||
tools_text=_agent_tools_text(agent),
|
||||
force=force,
|
||||
)
|
||||
|
||||
|
||||
async def run_agent_loop(
|
||||
@@ -350,6 +386,7 @@ async def _run_cycle( # noqa: PLR0912, PLR0915
|
||||
hooks: RunHooks[dict[str, Any]] | None,
|
||||
) -> RunResultBase | None:
|
||||
image_strips = 0
|
||||
compactions = 0
|
||||
while True:
|
||||
try:
|
||||
await coordinator.mark_running(agent_id)
|
||||
@@ -360,6 +397,10 @@ async def _run_cycle( # noqa: PLR0912, PLR0915
|
||||
await enforce_image_budget(session, max_images)
|
||||
except Exception:
|
||||
logger.exception("image-budget enforcement failed for %s", agent_id)
|
||||
try:
|
||||
await _compact_session(agent, session, run_config, force=False)
|
||||
except Exception:
|
||||
logger.exception("proactive compaction failed for %s", agent_id)
|
||||
stream = Runner.run_streamed(
|
||||
agent,
|
||||
input=input_data,
|
||||
@@ -428,6 +469,25 @@ async def _run_cycle( # noqa: PLR0912, PLR0915
|
||||
)
|
||||
input_data = []
|
||||
continue
|
||||
if (
|
||||
compactions < _MAX_COMPACTIONS_PER_CYCLE
|
||||
and session is not None
|
||||
and is_context_overflow(exc)
|
||||
):
|
||||
try:
|
||||
compacted = await _compact_session(agent, session, run_config, force=True)
|
||||
except Exception:
|
||||
logger.exception("overflow compaction recovery failed for %s", agent_id)
|
||||
compacted = False
|
||||
if compacted:
|
||||
compactions += 1
|
||||
logger.info(
|
||||
"Compacted %s session after context overflow; retrying (%d)",
|
||||
agent_id,
|
||||
compactions,
|
||||
)
|
||||
input_data = []
|
||||
continue
|
||||
if not interactive:
|
||||
raise
|
||||
if isinstance(exc, MaxTurnsExceeded):
|
||||
|
||||
@@ -92,6 +92,39 @@ async def _rewrite_session(
|
||||
return True
|
||||
|
||||
|
||||
async def replace_session_items(
|
||||
session: Session,
|
||||
new_items: list[Any],
|
||||
*,
|
||||
expected_len: int | None = None,
|
||||
) -> bool:
|
||||
"""Overwrite the session's items, restoring the originals on failure.
|
||||
|
||||
When ``expected_len`` is given, the rewrite is skipped if the session no
|
||||
longer has that many items (a concurrent writer changed it), so a slow
|
||||
compaction summary can't clobber newer turns.
|
||||
"""
|
||||
async with session_write_lock(session):
|
||||
original = list(await session.get_items())
|
||||
if expected_len is not None and len(original) != expected_len:
|
||||
logger.warning(
|
||||
"skipping session rewrite: expected %d items, found %d",
|
||||
expected_len,
|
||||
len(original),
|
||||
)
|
||||
return False
|
||||
rebuilt = cast("list[TResponseInputItem]", new_items)
|
||||
await session.clear_session()
|
||||
try:
|
||||
await session.add_items(rebuilt)
|
||||
except Exception:
|
||||
logger.exception("session rewrite failed; restoring original items")
|
||||
await session.clear_session()
|
||||
await session.add_items(original)
|
||||
raise
|
||||
return True
|
||||
|
||||
|
||||
async def strip_all_images_from_session(session: Session) -> bool:
|
||||
"""Replace every image tool output with a text placeholder (rejection recovery)."""
|
||||
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
"""LLM-facing context management: model-aware budgets and history compaction."""
|
||||
@@ -0,0 +1,285 @@
|
||||
"""Provider-agnostic conversation compaction.
|
||||
|
||||
When an agent's session grows past the model's usable context window, older
|
||||
turns are summarised into a single checkpoint while the most recent turns are
|
||||
kept verbatim. This runs for every LiteLLM provider (not just OpenAI), keeps a
|
||||
security-focused structured summary, and preserves tool-call/tool-result
|
||||
pairing so the trimmed history is still valid provider input.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import litellm
|
||||
|
||||
from strix.config import load_settings
|
||||
from strix.core.sessions import replace_session_items, session_write_lock
|
||||
from strix.llm.context_budget import context_window, count_tokens, output_limit
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from agents.memory import Session
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_CHECKPOINT_TAG = "<conversation-checkpoint>"
|
||||
_TOOL_OUTPUT_MAX_CHARS = 2_000
|
||||
_MIN_ITEMS_TO_COMPACT = 6
|
||||
|
||||
# Substrings that identify a context-window-overflow error across providers.
|
||||
# Deliberately excludes rate-limit/throttle wording, which must not trigger
|
||||
# compaction.
|
||||
_OVERFLOW_MARKERS = (
|
||||
"context length",
|
||||
"context window",
|
||||
"maximum context",
|
||||
"context_length_exceeded",
|
||||
"too many tokens",
|
||||
"reduce the length",
|
||||
"input is too long",
|
||||
"prompt is too long",
|
||||
"exceeds the maximum",
|
||||
"string too long",
|
||||
)
|
||||
|
||||
|
||||
def is_context_overflow(exc: BaseException) -> bool:
|
||||
"""Whether ``exc`` looks like a model context-window-overflow error."""
|
||||
overflow_error = getattr(litellm, "ContextWindowExceededError", None)
|
||||
if overflow_error is not None and isinstance(exc, overflow_error):
|
||||
return True
|
||||
message = str(exc).lower()
|
||||
return any(marker in message for marker in _OVERFLOW_MARKERS)
|
||||
|
||||
|
||||
_SUMMARY_INSTRUCTIONS = """\
|
||||
You are compacting the earlier part of an autonomous security-testing agent's \
|
||||
conversation so it fits the model context window. Produce a dense, factual \
|
||||
record that lets the agent continue with no loss of important state. Preserve \
|
||||
exact values: URLs, endpoints, file paths, parameters, payloads, credentials \
|
||||
and tokens, software versions, and error messages. Do not invent anything and \
|
||||
do not describe this compaction process.
|
||||
|
||||
Return Markdown with exactly these sections:
|
||||
|
||||
## Objective
|
||||
The overall goal and target scope.
|
||||
|
||||
## Important Details
|
||||
Discovered vulnerabilities and attack vectors, scan/tool findings, \
|
||||
credentials and auth material, system architecture and weak points, and any \
|
||||
exact identifiers worth keeping (URLs, paths, params, payloads, versions).
|
||||
|
||||
## Work State
|
||||
- Completed: what has been verified or finished.
|
||||
- Active: what is in progress right now.
|
||||
- Blocked: anything stuck and why.
|
||||
|
||||
## Failed Attempts & Dead Ends
|
||||
Approaches already tried that did not work, so they are not repeated.
|
||||
|
||||
## Next Move
|
||||
The concrete next step(s) the agent intended to take.
|
||||
|
||||
## Relevant Files
|
||||
Files/notes/reports created or modified and their purpose."""
|
||||
|
||||
|
||||
def _content_text(content: Any) -> str:
|
||||
if isinstance(content, str):
|
||||
return content
|
||||
if isinstance(content, list):
|
||||
parts: list[str] = []
|
||||
for block in content:
|
||||
if not isinstance(block, dict):
|
||||
continue
|
||||
text = block.get("text")
|
||||
if isinstance(text, str):
|
||||
parts.append(text)
|
||||
elif block.get("type") in {"input_image", "image_url", "output_image"}:
|
||||
parts.append("[image]")
|
||||
return "\n".join(parts)
|
||||
return ""
|
||||
|
||||
|
||||
def _truncate(text: str, limit: int) -> str:
|
||||
return text if len(text) <= limit else f"{text[:limit]}\n[truncated]"
|
||||
|
||||
|
||||
def _serialize_item(item: Any) -> str:
|
||||
if not isinstance(item, dict):
|
||||
return str(item)
|
||||
item_type = item.get("type")
|
||||
role = item.get("role")
|
||||
if item_type == "function_call":
|
||||
args = _truncate(str(item.get("arguments", "")), _TOOL_OUTPUT_MAX_CHARS)
|
||||
return f"[tool_call {item.get('name', '?')}] {args}"
|
||||
if item_type == "function_call_output":
|
||||
output = item.get("output")
|
||||
text = output if isinstance(output, str) else _content_text(output)
|
||||
return f"[tool_result] {_truncate(text, _TOOL_OUTPUT_MAX_CHARS)}"
|
||||
if item_type == "reasoning":
|
||||
return ""
|
||||
if role or item_type == "message":
|
||||
return f"[{role or 'assistant'}] {_content_text(item.get('content'))}".strip()
|
||||
return ""
|
||||
|
||||
|
||||
def _serialize_items(items: list[Any]) -> str:
|
||||
return "\n".join(s for s in (_serialize_item(item) for item in items) if s)
|
||||
|
||||
|
||||
def _is_tool_call(item: Any) -> bool:
|
||||
return isinstance(item, dict) and item.get("type") == "function_call"
|
||||
|
||||
|
||||
def _is_tool_output(item: Any) -> bool:
|
||||
return isinstance(item, dict) and item.get("type") == "function_call_output"
|
||||
|
||||
|
||||
def _open_calls_at(items: list[Any]) -> list[int]:
|
||||
"""Prefix count of tool calls still awaiting their result at each index.
|
||||
|
||||
``result[i]`` is the number of open calls *before* index ``i``. A split is
|
||||
only safe (won't orphan a tool result from its call) where this is zero.
|
||||
"""
|
||||
balance = [0] * (len(items) + 1)
|
||||
for i, item in enumerate(items):
|
||||
delta = 1 if _is_tool_call(item) else -1 if _is_tool_output(item) else 0
|
||||
balance[i + 1] = max(0, balance[i] + delta)
|
||||
return balance
|
||||
|
||||
|
||||
def _select_split(model: str, items: list[Any], keep_tokens: int) -> int:
|
||||
"""Index where the kept-verbatim recent tail begins.
|
||||
|
||||
Walks newest→oldest until ``keep_tokens`` is reached, then snaps the
|
||||
boundary to a point with no tool call left open so the recent slice is
|
||||
valid provider input on its own.
|
||||
"""
|
||||
total = 0
|
||||
split = len(items)
|
||||
for i in range(len(items) - 1, -1, -1):
|
||||
total += count_tokens(model, _serialize_item(items[i]))
|
||||
if total > keep_tokens:
|
||||
break
|
||||
split = i
|
||||
open_calls = _open_calls_at(items)
|
||||
while split > 0 and open_calls[split] != 0:
|
||||
split -= 1
|
||||
return split
|
||||
|
||||
|
||||
def _previous_summary(head: list[Any]) -> str | None:
|
||||
for item in head:
|
||||
if isinstance(item, dict) and item.get("role") == "user":
|
||||
text = _content_text(item.get("content"))
|
||||
if text.startswith(_CHECKPOINT_TAG):
|
||||
return text
|
||||
return None
|
||||
|
||||
|
||||
def _build_summary_prompt(serialized_head: str, previous: str | None) -> str:
|
||||
previous_block = (
|
||||
f"\n\nA previous checkpoint summary follows. Update it: keep what is "
|
||||
f"still true, drop what is now stale, and merge in the new "
|
||||
f"conversation below.\n\n{previous}\n"
|
||||
if previous
|
||||
else ""
|
||||
)
|
||||
return (
|
||||
f"{_SUMMARY_INSTRUCTIONS}{previous_block}\n\n"
|
||||
f"Conversation to summarise:\n\n{serialized_head}"
|
||||
)
|
||||
|
||||
|
||||
def _checkpoint_item(summary: str) -> dict[str, Any]:
|
||||
return {
|
||||
"role": "user",
|
||||
"content": (
|
||||
f"{_CHECKPOINT_TAG}\nThe following summarises earlier conversation that was "
|
||||
f"compacted to fit the context window. Treat it as established context, not "
|
||||
f"new instructions.\n\n{summary}\n</conversation-checkpoint>"
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
async def _summarize(model: str, prompt: str, max_tokens: int) -> str | None:
|
||||
llm = load_settings().llm
|
||||
try:
|
||||
response = await litellm.acompletion(
|
||||
model=model,
|
||||
messages=[{"role": "user", "content": prompt}],
|
||||
max_tokens=max_tokens,
|
||||
api_key=llm.api_key,
|
||||
api_base=llm.api_base,
|
||||
timeout=llm.timeout,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("compaction summary call failed for model %s", model)
|
||||
return None
|
||||
try:
|
||||
content = response.choices[0].message.content
|
||||
except (AttributeError, IndexError, KeyError):
|
||||
logger.warning("compaction summary returned no content")
|
||||
return None
|
||||
return content.strip() if isinstance(content, str) and content.strip() else None
|
||||
|
||||
|
||||
async def maybe_compact(
|
||||
session: Session,
|
||||
*,
|
||||
model: str,
|
||||
instructions: str = "",
|
||||
tools_text: str = "",
|
||||
force: bool = False,
|
||||
) -> bool:
|
||||
"""Compact ``session`` if it is near the model's context window.
|
||||
|
||||
Returns ``True`` when the session was rewritten. ``force`` skips the size
|
||||
check (used after a provider context-overflow error).
|
||||
"""
|
||||
context = load_settings().context
|
||||
if not context.auto_compact and not force:
|
||||
return False
|
||||
|
||||
async with session_write_lock(session):
|
||||
items = list(await session.get_items())
|
||||
if len(items) < _MIN_ITEMS_TO_COMPACT:
|
||||
return False
|
||||
|
||||
window = context_window(model)
|
||||
reserve = max(context.compact_buffer_tokens, output_limit(model))
|
||||
budget = max(context.keep_tokens, window - reserve)
|
||||
used = count_tokens(model, "\n".join((instructions, tools_text, _serialize_items(items))))
|
||||
if not force and used <= budget:
|
||||
return False
|
||||
|
||||
split = _select_split(model, items, context.keep_tokens)
|
||||
head, recent = items[:split], items[split:]
|
||||
if not head:
|
||||
return False
|
||||
|
||||
summary = await _summarize(
|
||||
model,
|
||||
_build_summary_prompt(_serialize_items(head), _previous_summary(head)),
|
||||
context.summary_max_tokens,
|
||||
)
|
||||
if summary is None:
|
||||
return False
|
||||
|
||||
new_items = [_checkpoint_item(summary), *recent]
|
||||
rewritten = await replace_session_items(session, new_items, expected_len=len(items))
|
||||
if rewritten:
|
||||
logger.info(
|
||||
"compacted %s: %d items (~%d tok) -> %d items (summary + %d recent)",
|
||||
model,
|
||||
len(items),
|
||||
used,
|
||||
len(new_items),
|
||||
len(recent),
|
||||
)
|
||||
return rewritten
|
||||
@@ -0,0 +1,77 @@
|
||||
"""Model-aware token budgets.
|
||||
|
||||
The context window and output cap vary widely by model (128k for gpt-4o, 272k
|
||||
for gpt-5, 1M for claude-sonnet-4, 131k for deepseek). We resolve them from
|
||||
LiteLLM's model metadata so compaction triggers at the right point for the
|
||||
selected model instead of a fixed guess, falling back to a large configurable
|
||||
default for models LiteLLM doesn't map.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from functools import lru_cache
|
||||
from typing import Any
|
||||
|
||||
import litellm
|
||||
|
||||
from strix.config import load_settings
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# LiteLLM keys models without the routing prefix users type (``openai/``,
|
||||
# ``litellm/``, ``ollama/`` ...). Strip a leading provider segment on lookup.
|
||||
_STRIPPABLE_PREFIXES = ("openai/", "litellm/", "any-llm/", "ollama/", "ollama_chat/")
|
||||
|
||||
_DEFAULT_OUTPUT_TOKENS = 8_192
|
||||
|
||||
|
||||
def _lookup_key(model: str) -> str:
|
||||
for prefix in _STRIPPABLE_PREFIXES:
|
||||
if model.startswith(prefix):
|
||||
return model[len(prefix) :]
|
||||
return model
|
||||
|
||||
|
||||
def _safe_get_model_info(model: str) -> dict[str, Any] | None:
|
||||
try:
|
||||
return dict(litellm.get_model_info(model))
|
||||
except Exception: # noqa: BLE001 - unmapped models raise; caller falls back.
|
||||
return None
|
||||
|
||||
|
||||
@lru_cache(maxsize=128)
|
||||
def _model_info(model: str) -> dict[str, int]:
|
||||
for candidate in (model, _lookup_key(model)):
|
||||
info = _safe_get_model_info(candidate)
|
||||
if info is not None:
|
||||
return {
|
||||
"max_input_tokens": int(
|
||||
info.get("max_input_tokens") or info.get("max_tokens") or 0
|
||||
),
|
||||
"max_output_tokens": int(info.get("max_output_tokens") or 0),
|
||||
}
|
||||
logger.debug("No LiteLLM model info for %r; using configured fallbacks", model)
|
||||
return {"max_input_tokens": 0, "max_output_tokens": 0}
|
||||
|
||||
|
||||
def context_window(model: str) -> int:
|
||||
"""Input token capacity for ``model`` (configured fallback when unmapped)."""
|
||||
resolved = _model_info(model)["max_input_tokens"]
|
||||
return resolved or load_settings().context.fallback_context_tokens
|
||||
|
||||
|
||||
def output_limit(model: str) -> int:
|
||||
"""Max output tokens for ``model`` (a conservative default when unmapped)."""
|
||||
return _model_info(model)["max_output_tokens"] or _DEFAULT_OUTPUT_TOKENS
|
||||
|
||||
|
||||
def count_tokens(model: str, text: str) -> int:
|
||||
"""Token count for ``text`` under ``model`` (chars/4 fallback)."""
|
||||
if not text:
|
||||
return 0
|
||||
try:
|
||||
return int(litellm.token_counter(model=_lookup_key(model), text=text))
|
||||
except Exception: # noqa: BLE001 - tokenizer may be unavailable for some models.
|
||||
return len(text) // 4
|
||||
@@ -0,0 +1,186 @@
|
||||
"""Tests for provider-agnostic conversation compaction."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import pytest
|
||||
|
||||
from strix.config import ContextSettings
|
||||
from strix.llm import compaction
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from agents.memory.session_settings import SessionSettings
|
||||
|
||||
|
||||
class FakeSession:
|
||||
"""Minimal in-memory Session for exercising compaction."""
|
||||
|
||||
session_id = "fake"
|
||||
session_settings: SessionSettings | None = None
|
||||
|
||||
def __init__(self, items: list[Any]) -> None:
|
||||
self._items = list(items)
|
||||
|
||||
async def get_items(self, limit: int | None = None) -> list[Any]:
|
||||
return list(self._items) if limit is None else list(self._items[-limit:])
|
||||
|
||||
async def add_items(self, items: list[Any]) -> None:
|
||||
self._items.extend(items)
|
||||
|
||||
async def clear_session(self) -> None:
|
||||
self._items = []
|
||||
|
||||
async def pop_item(self) -> Any:
|
||||
return self._items.pop() if self._items else None
|
||||
|
||||
|
||||
def _user(text: str) -> dict[str, Any]:
|
||||
return {"role": "user", "content": text}
|
||||
|
||||
|
||||
def _assistant(text: str) -> dict[str, Any]:
|
||||
return {"role": "assistant", "content": text}
|
||||
|
||||
|
||||
def _call(call_id: str, name: str = "exec_command") -> dict[str, Any]:
|
||||
return {"type": "function_call", "call_id": call_id, "name": name, "arguments": "{}"}
|
||||
|
||||
|
||||
def _output(call_id: str, text: str = "done") -> dict[str, Any]:
|
||||
return {"type": "function_call_output", "call_id": call_id, "output": text}
|
||||
|
||||
|
||||
def _turns(n: int) -> list[dict[str, Any]]:
|
||||
items: list[dict[str, Any]] = []
|
||||
for i in range(n):
|
||||
items += [
|
||||
_user(f"task {i}"),
|
||||
_call(f"c{i}"),
|
||||
_output(f"c{i}", f"result {i}"),
|
||||
_assistant(f"ok {i}"),
|
||||
]
|
||||
return items
|
||||
|
||||
|
||||
def _has_orphan_tool_output(items: list[Any]) -> bool:
|
||||
call_ids = {i["call_id"] for i in items if compaction._is_tool_call(i)}
|
||||
return any(i["call_id"] not in call_ids for i in items if compaction._is_tool_output(i))
|
||||
|
||||
|
||||
def test_is_context_overflow_matches_and_excludes() -> None:
|
||||
assert compaction.is_context_overflow(
|
||||
RuntimeError("This model's maximum context length is 8192")
|
||||
)
|
||||
assert compaction.is_context_overflow(ValueError("input is too long for the model"))
|
||||
assert not compaction.is_context_overflow(RuntimeError("rate limit exceeded, retry later"))
|
||||
|
||||
|
||||
def test_select_split_never_orphans_tool_output(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(compaction, "count_tokens", lambda _m, t: len(t))
|
||||
items = _turns(10)
|
||||
split = compaction._select_split("m", items, keep_tokens=25)
|
||||
recent = items[split:]
|
||||
assert recent # something is kept
|
||||
assert not _has_orphan_tool_output(recent)
|
||||
|
||||
|
||||
def test_select_split_handles_parallel_calls(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(compaction, "count_tokens", lambda _m, _t: 1)
|
||||
# Two parallel calls then their two outputs.
|
||||
items = [
|
||||
_user("start"),
|
||||
_call("a"),
|
||||
_call("b"),
|
||||
_output("a"),
|
||||
_output("b"),
|
||||
_assistant("done"),
|
||||
]
|
||||
# keep_tokens picks a boundary that would land between the calls/outputs.
|
||||
split = compaction._select_split("m", items, keep_tokens=2)
|
||||
assert not _has_orphan_tool_output(items[split:])
|
||||
|
||||
|
||||
def _patch_budget(monkeypatch: pytest.MonkeyPatch, *, keep_tokens: int, window: int) -> None:
|
||||
monkeypatch.setattr(compaction, "count_tokens", lambda _m, t: len(t))
|
||||
monkeypatch.setattr(compaction, "context_window", lambda _m: window)
|
||||
monkeypatch.setattr(compaction, "output_limit", lambda _m: 0)
|
||||
context = ContextSettings()
|
||||
context.keep_tokens = keep_tokens
|
||||
context.compact_buffer_tokens = 0
|
||||
context.summary_max_tokens = 64
|
||||
context.auto_compact = True
|
||||
settings = SimpleNamespace(
|
||||
context=context,
|
||||
llm=SimpleNamespace(api_key=None, api_base=None, timeout=1),
|
||||
)
|
||||
monkeypatch.setattr(compaction, "load_settings", lambda: settings)
|
||||
|
||||
|
||||
def _patch_summary(monkeypatch: pytest.MonkeyPatch, text: str) -> None:
|
||||
async def fake_acompletion(**_kwargs: Any) -> Any:
|
||||
message = SimpleNamespace(content=text)
|
||||
return SimpleNamespace(choices=[SimpleNamespace(message=message)])
|
||||
|
||||
monkeypatch.setattr("strix.llm.compaction.litellm.acompletion", fake_acompletion)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_maybe_compact_noop_when_within_budget(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
_patch_budget(monkeypatch, keep_tokens=50, window=1_000_000)
|
||||
session = FakeSession(_turns(10))
|
||||
before = await session.get_items()
|
||||
|
||||
assert await compaction.maybe_compact(session, model="m") is False
|
||||
assert await session.get_items() == before
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_maybe_compact_rewrites_and_keeps_pairs(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
_patch_budget(monkeypatch, keep_tokens=30, window=50)
|
||||
_patch_summary(monkeypatch, "SUMMARY BODY")
|
||||
session = FakeSession(_turns(12))
|
||||
|
||||
assert await compaction.maybe_compact(session, model="m") is True
|
||||
|
||||
items = await session.get_items()
|
||||
assert items[0]["role"] == "user"
|
||||
assert items[0]["content"].startswith(compaction._CHECKPOINT_TAG)
|
||||
assert "SUMMARY BODY" in items[0]["content"]
|
||||
assert len(items) < len(_turns(12))
|
||||
assert not _has_orphan_tool_output(items)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_maybe_compact_updates_previous_summary(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
_patch_budget(monkeypatch, keep_tokens=30, window=50)
|
||||
captured: dict[str, str] = {}
|
||||
|
||||
async def fake_acompletion(**kwargs: Any) -> Any:
|
||||
captured["prompt"] = kwargs["messages"][0]["content"]
|
||||
return SimpleNamespace(choices=[SimpleNamespace(message=SimpleNamespace(content="NEW"))])
|
||||
|
||||
monkeypatch.setattr("strix.llm.compaction.litellm.acompletion", fake_acompletion)
|
||||
|
||||
prior = compaction._checkpoint_item("OLD SUMMARY TEXT")
|
||||
session = FakeSession([prior, *_turns(12)])
|
||||
|
||||
assert await compaction.maybe_compact(session, model="m") is True
|
||||
assert "OLD SUMMARY TEXT" in captured["prompt"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_maybe_compact_skips_when_summary_fails(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
_patch_budget(monkeypatch, keep_tokens=30, window=50)
|
||||
|
||||
async def fake_acompletion(**_kwargs: Any) -> Any:
|
||||
raise RuntimeError("boom")
|
||||
|
||||
monkeypatch.setattr("strix.llm.compaction.litellm.acompletion", fake_acompletion)
|
||||
session = FakeSession(_turns(12))
|
||||
before = await session.get_items()
|
||||
|
||||
assert await compaction.maybe_compact(session, model="m") is False
|
||||
assert await session.get_items() == before
|
||||
@@ -0,0 +1,46 @@
|
||||
"""Tests for model-aware token budgets."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from strix.config import load_settings
|
||||
from strix.llm import context_budget
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import pytest
|
||||
|
||||
|
||||
def test_context_window_known_model() -> None:
|
||||
# gpt-4o is mapped by LiteLLM at 128k input tokens.
|
||||
assert context_budget.context_window("gpt-4o") == 128_000
|
||||
|
||||
|
||||
def test_context_window_strips_provider_prefix() -> None:
|
||||
assert context_budget.context_window("openai/gpt-4o") == 128_000
|
||||
|
||||
|
||||
def test_context_window_unmapped_uses_fallback(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
context_budget._model_info.cache_clear()
|
||||
|
||||
def _raise(_model: str) -> dict[str, int]:
|
||||
raise ValueError("This model isn't mapped yet.")
|
||||
|
||||
monkeypatch.setattr("strix.llm.context_budget.litellm.get_model_info", _raise)
|
||||
expected = load_settings().context.fallback_context_tokens
|
||||
assert context_budget.context_window("totally-made-up-model") == expected
|
||||
context_budget._model_info.cache_clear()
|
||||
|
||||
|
||||
def test_count_tokens_fallback_on_error(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
def _raise(**_kwargs: object) -> int:
|
||||
raise RuntimeError("no tokenizer")
|
||||
|
||||
monkeypatch.setattr("strix.llm.context_budget.litellm.token_counter", _raise)
|
||||
text = "x" * 400
|
||||
assert context_budget.count_tokens("weird-model", text) == 100
|
||||
|
||||
|
||||
def test_count_tokens_empty_is_zero() -> None:
|
||||
assert context_budget.count_tokens("gpt-4o", "") == 0
|
||||
Reference in New Issue
Block a user