mirror of
https://github.com/usestrix/strix.git
synced 2026-10-12 03:47:56 +00:00
fix(compaction): chain any input filter already set on the run config
The image scrub for text-only models also sets call_model_input_filter; replacing it dropped the scrub.
This commit is contained in:
parent
94adabd67e
commit
f565cece23
2 changed files with 29 additions and 3 deletions
|
|
@ -4,6 +4,7 @@ from __future__ import annotations
|
|||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import inspect
|
||||
import logging
|
||||
import uuid
|
||||
from collections.abc import Callable
|
||||
|
|
@ -147,22 +148,27 @@ def _with_compaction_check(run_config: RunConfig) -> RunConfig:
|
|||
model = _run_config_model(run_config)
|
||||
if model is None:
|
||||
return run_config
|
||||
inner = run_config.call_model_input_filter
|
||||
first_call = True
|
||||
|
||||
async def _check(data: CallModelData[Any]) -> ModelInputData:
|
||||
nonlocal first_call
|
||||
model_data = data.model_data
|
||||
if inner is not None:
|
||||
result = inner(data)
|
||||
model_data = await result if inspect.isawaitable(result) else result
|
||||
# Skip the first call: the cycle compacts right before the run, and the run's
|
||||
# new input is only persisted after this filter, so stopping here would lose it.
|
||||
if not first_call and await asyncio.to_thread(
|
||||
needs_compaction,
|
||||
model,
|
||||
data.model_data.instructions or "",
|
||||
model_data.instructions or "",
|
||||
_agent_tools_text(data.agent),
|
||||
data.model_data.input,
|
||||
model_data.input,
|
||||
):
|
||||
raise CompactionNeededError(model)
|
||||
first_call = False
|
||||
return data.model_data
|
||||
return model_data
|
||||
|
||||
return replace(run_config, call_model_input_filter=_check)
|
||||
|
||||
|
|
|
|||
|
|
@ -1553,6 +1553,26 @@ async def test_compaction_check_stops_run_after_first_call_when_over_budget(
|
|||
await guard(data)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_compaction_check_chains_an_existing_input_filter(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
seen: list[Any] = []
|
||||
monkeypatch.setattr(execution, "needs_compaction", lambda *args: seen.append(args[3]))
|
||||
filtered = ModelInputData(input=[{"role": "user", "content": "scrubbed"}], instructions="sys")
|
||||
run_config = RunConfig(model="m", call_model_input_filter=lambda _data: filtered)
|
||||
check: Any = execution._with_compaction_check(run_config).call_model_input_filter
|
||||
data = CallModelData(
|
||||
model_data=ModelInputData(input=[], instructions="sys"),
|
||||
agent=MagicMock(tools=[]),
|
||||
context=None,
|
||||
)
|
||||
|
||||
assert await check(data) is filtered
|
||||
assert await check(data) is filtered
|
||||
assert seen == [filtered.input]
|
||||
|
||||
|
||||
class _BudgetStream:
|
||||
def __init__(self, exc: BaseException | None) -> None:
|
||||
self.run_loop_exception = exc
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue