arsvendg/Stirling-PDF
1
0
Fork 0
mirror of https://github.com/arsvendg/Stirling-PDF.git synced 2026-07-16 03:20:46 +02:00
Code Issues Actions 3 Packages Projects Releases 1 Wiki Activity

feat(ai): add Contradiction Agent on a new ChunkedMapper primitive (#6369)

Browse Source 
## Summary

Adds a new AI specialist that finds **textual contradictions** across
one or more PDFs — conflicting claims, recommendations, points of view,
contested facts — built entirely in Python on top of the new
`DocumentService` + `ChunkedReasoner` stack from #6314. Replaces the
closed #6304, which was started before #6314 landed and therefore
over-engineered (Java orchestrator, two-round handshake, resume
artifact, discriminated-union lift).

Two commits:
1. **`refactor(engine): extract ChunkedMapper[T] from ChunkedReasoner`**
— pure refactor, public API of ChunkedReasoner unchanged. New
`ChunkedMapper[T: BaseModel]` is a generic parallel-chunk primitive
(slicing, semaphore, time-bounded extraction, cancellation drain,
progress events) that's now a peer to ChunkedReasoner rather than locked
inside it. The compression loop stays on ChunkedReasoner where it
belongs.
2. **`feat(ai): add Contradiction Agent on ChunkedMapper`** — the agent
itself, plus integrations into `PdfReviewAgent` and `PdfQuestionAgent`.

## Architecture

- **Python-only.** No Java code. No `AgentToolId.CONTRADICTION_AGENT`.
No dedicated HTTP endpoint. No resume artifact, no discriminated-union
lift in `contracts/common.py`. Detector runs inside the Python engine
and the Python engine alone.
- **Review path** (`PdfReviewAgent`): a new
`ContradictionIntentClassifier` fires on contradiction-flavoured
prompts; agent runs detection synchronously and emits a single
`EditPlanResponse(steps=[ADD_COMMENTS])`. Single-turn flow — no resume.
- **Question path** (`PdfQuestionAgent`): a new
`ContradictionCapability` joins `RagCapability` and
`WholeDocReaderCapability` in the smart-model toolset, exposing
`find_contradictions(query)`. The smart model picks it from the toolset
alongside `search_knowledge` and `read_full_document`.

## Inside `ContradictionDetector.detect()`

1. `DocumentService.read_pages(file_id)` → ordered `list[Page]`.
2. `ChunkedMapper[_ExtractedClaims].map_pages(...)` — char-budgeted
multi-page slicing; each slice runs the claim-extractor LLM in parallel
under a semaphore.
3. Page-traceability: the extractor returns `_ExtractedClaim.page`
(which `[Page N]` marker the claim came from). The wrapper validates
`page ∈ chunk.pages`; if not, mechanical fallback searches the chunk's
page text for the verbatim quote and reassigns. If still no match, drop
the claim.
4. `Claim.anchor_quality: Literal[\"verbatim\", \"paraphrased\"]` is set
by a substring check against the declared page's text. Verbatim quotes
feed `anchor_text` for snap-to-quote add-comments placement; paraphrased
ones fall back to margin geometry.
5. Subject canonicalisation: ONE fast-model LLM call collapses synonyms
across the document. Fails open to lexical bucketing.
6. Pre-filters: drop identical-quote pairs; drop same-page same-polarity
paraphrases.
7. Per-bucket pair detection in parallel (separate semaphore, cap 5).
Buckets > 12 claims chunk into windows of 12 with overlap 2; pairs
deduped across overlapping windows by frozen `(i, j)` index pair.
8. Summary fast-model call with fallback string on error.

## Prompt-injection hardening

Every prompt that interpolates user-supplied or PDF-extracted text wraps
content in `<user_message>` / `<verdict>` / `<content>` tags with an
explicit SECURITY preamble instructing the model to treat tagged content
as data only.

## Limitations

- **Combined math + contradiction intent**: when both intent classifiers
fire on the same prompt, contradiction takes precedence and the math
intent is silently dropped. Documented in the Review module docstring
and pinned by
`test_review_integration.py::test_contradiction_precedence_over_math`.
- **Cross-window contradiction reach**: within a subject bucket, pairs
more than ~10 claim indices apart in the same chunked window may be
missed by the overlap-2 strategy. Documented in `test_detector.py`.
Acceptable for v1.

## Settings (engine/src/stirling/config/settings.py)

```python
contradiction_detect_concurrency = 5     # per-bucket detector semaphore
contradiction_bucket_chunk_size = 12     # max claims per detector call
contradiction_bucket_chunk_overlap = 2   # overlap for >threshold buckets
```

`chars_per_slice` and extraction concurrency are reused from the
existing `chunked_reasoner_*` settings.

## Test plan

- [x] `uv run pytest tests/ -v` — **245/245 pass** (210 pre-existing +
35 new)
- [x] `uv run ruff check src/ tests/` — clean
- [x] `uv run pyright src/stirling/agents/contradiction/
src/stirling/contracts/contradiction.py` — 0 errors
- [x] `./gradlew :proprietary:test` — green; no Java was touched, but
verified untouched
- [x] Page-traceability tests cover: valid page kept, hallucinated page
dropped, mechanical-reassign on misattribution, anchor-quality verbatim
vs paraphrased
- [x] Review integration: ADD_COMMENTS plan with two paired CommentSpecs
per contradiction; NeedIngestResponse precheck; precedence vs math
intent pinned
- [x] Question integration: all three capabilities wired into
smart-model toolset; `find_contradictions` returns formatted report text
- [x] ChunkedMapper standalone: slicing, multi-chunk ordering, worker
failures, timeouts, cancellation drain, semaphore saturation
- [x] ChunkedReasoner regression: all pre-existing tests pass unchanged
after the internal split

## Relationship to closed #6304

#6304 was closed in favour of this PR. The closed PR predated #6314 and
modelled the agent as a Java-orchestrated two-round examine/deliberate
flow with its own HTTP endpoint and a discriminated-union resume
artifact. With #6314 making full ordered page text available to the
engine via `DocumentService.read_pages`, none of that is needed. Net
effect: drop ~600 lines of Java, drop the two-round handshake, drop the
`ToolReportArtifact` lift, while ending up with a more scalable agent
(chunk-based instead of page-based extraction; tested to
ChunkedReasoner-equivalent scale).
This commit is contained in:
ConnorYoh
2026-05-22 13:23:52 +00:00
committed by GitHub
parent 0a50e765b7
commit 017c8d59fa
27 changed files with 4401 additions and 268 deletions
Download Patch File Download Diff File Expand all files Collapse all files
engine/tests/agents/test_chunked_mapper.py
+356
View File
@@ -0,0 +1,356 @@
"""Tests for the generic ``ChunkedMapper`` primitive.
The mapper is the per-chunk fan-out machinery extracted from
``ChunkedReasoner``: char-budgeted slicing, parallel scheduling under a
semaphore, time-bounded extraction with cancellation, progress events, and
worker-failure tolerance. These tests drive it with a stubbed
``Agent[None, T]`` so the model boundary stays patched out.
"""
from __future__ import annotations
import asyncio
from dataclasses import dataclass
from unittest.mock import AsyncMock, patch
import pytest
from pydantic import BaseModel
from pydantic_ai import Agent
from stirling.agents.shared.chunked_mapper import ChunkedMapper
from stirling.contracts.documents import Page
from stirling.services.runtime import AppRuntime
@dataclass
class _StubAgentResult[T]:
output: T
class _Extracted(BaseModel):
"""Tiny per-chunk extractor payload used by these tests."""
label: str
def _page(n: int, text: str) -> Page:
return Page(page_number=n, text=text, char_count=len(text))
def _build_mapper(
runtime: AppRuntime,
*,
chars_per_slice: int | None = None,
concurrency: int | None = None,
worker_timeout_seconds: float | None = None,
) -> ChunkedMapper[_Extracted]:
"""Build a mapper wrapping a real ``Agent`` whose ``.run`` is patched per test."""
extractor: Agent[None, _Extracted] = Agent(
model=runtime.fast_model,
output_type=_Extracted,
model_settings=runtime.fast_model_settings,
)
return ChunkedMapper(
runtime,
extractor=extractor,
chars_per_slice=chars_per_slice,
concurrency=concurrency,
worker_timeout_seconds=worker_timeout_seconds,
)
class TestSlicePages:
"""The static helper is pure: no I/O, no scheduling."""
def test_single_slice_when_under_budget(self) -> None:
pages = [_page(1, "abc"), _page(2, "def"), _page(3, "gh")]
slices = ChunkedMapper.slice_pages(pages, chars_per_slice=20)
assert [[p.page_number for p in s] for s in slices] == [[1, 2, 3]]
def test_starts_new_slice_when_budget_exceeded(self) -> None:
pages = [_page(1, "a" * 6), _page(2, "b" * 6), _page(3, "c" * 6)]
slices = ChunkedMapper.slice_pages(pages, chars_per_slice=10)
# 6 + 6 > 10 → break after each page
assert [[p.page_number for p in s] for s in slices] == [[1], [2], [3]]
def test_oversized_page_is_its_own_slice(self) -> None:
"""Page boundaries are never broken: an oversize page becomes its own slice."""
pages = [_page(1, "small"), _page(2, "x" * 100), _page(3, "tiny")]
slices = ChunkedMapper.slice_pages(pages, chars_per_slice=10)
assert [[p.page_number for p in s] for s in slices] == [[1], [2], [3]]
def test_rejects_non_positive_budget(self) -> None:
with pytest.raises(ValueError, match="chars_per_slice"):
ChunkedMapper.slice_pages([_page(1, "x")], chars_per_slice=0)
class TestFormatChunkContent:
def test_renders_page_markers(self) -> None:
rendered = ChunkedMapper.format_chunk_content([_page(2, "two"), _page(3, "three")])
assert "[Page 2]\ntwo" in rendered
assert "[Page 3]\nthree" in rendered
# Blank-line separator between pages
assert "two\n\n[Page 3]" in rendered
class TestMapPages:
@pytest.mark.anyio
async def test_single_chunk_returns_single_output(self, runtime: AppRuntime) -> None:
mapper = _build_mapper(runtime, chars_per_slice=1000)
pages = [_page(1, "alpha"), _page(2, "beta"), _page(3, "gamma")]
canned = _Extracted(label="one")
with patch.object(
mapper._extractor,
"run",
AsyncMock(return_value=_StubAgentResult(output=canned)),
) as run_mock:
outputs = await mapper.map_pages(pages, "what")
assert run_mock.await_count == 1
assert len(outputs) == 1
assert outputs[0].pages == [1, 2, 3]
assert outputs[0].output == canned
assert outputs[0].label == "pages=1-3"
@pytest.mark.anyio
async def test_multi_chunk_outputs_are_in_document_order(self, runtime: AppRuntime) -> None:
"""Outputs are sorted by first covered page regardless of completion order."""
mapper = _build_mapper(runtime, chars_per_slice=10, concurrency=3)
pages = [_page(i, "x" * 8) for i in range(1, 4)]
# Each chunk's worker awaits a release event; we release in reverse
# order so completion order is the inverse of slice order.
release = [asyncio.Event() for _ in pages]
call_index = 0
async def _gated(*_args: object, **_kwargs: object) -> _StubAgentResult[_Extracted]:
nonlocal call_index
mine = call_index
call_index += 1
await release[mine].wait()
return _StubAgentResult(output=_Extracted(label=f"slice-{mine + 1}"))
async def _release_in_reverse() -> None:
await asyncio.sleep(0)
for ev in reversed(release):
ev.set()
await asyncio.sleep(0)
await asyncio.sleep(0)
with patch.object(mapper._extractor, "run", AsyncMock(side_effect=_gated)):
task = asyncio.create_task(mapper.map_pages(pages, "anything"))
await _release_in_reverse()
outputs = await task
assert [o.pages for o in outputs] == [[1], [2], [3]]
@pytest.mark.anyio
async def test_worker_failure_drops_only_that_chunk(self, runtime: AppRuntime) -> None:
mapper = _build_mapper(runtime, chars_per_slice=10)
pages = [_page(i, "x" * 8) for i in range(1, 4)]
results: list[_Extracted | BaseException] = [
_Extracted(label="a"),
RuntimeError("boom"),
_Extracted(label="c"),
]
async def _stub(*_args: object, **_kwargs: object) -> _StubAgentResult[_Extracted]:
value = results.pop(0)
if isinstance(value, BaseException):
raise value
return _StubAgentResult(output=value)
with patch.object(mapper._extractor, "run", AsyncMock(side_effect=_stub)):
outputs = await mapper.map_pages(pages, "anything")
assert len(outputs) == 2
assert {o.output.label for o in outputs} == {"a", "c"}
@pytest.mark.anyio
async def test_worker_timeout_drops_only_that_chunk(self, runtime: AppRuntime) -> None:
mapper = _build_mapper(runtime, chars_per_slice=10, worker_timeout_seconds=0.05)
pages = [_page(i, "x" * 8) for i in range(1, 4)]
async def _stub(*_args: object, **_kwargs: object) -> _StubAgentResult[_Extracted]:
# Page 2 hangs forever; pages 1 and 3 return immediately.
prompt = _args[0]
assert isinstance(prompt, str)
if "[Page 2]" in prompt:
await asyncio.sleep(10)
return _StubAgentResult(output=_Extracted(label="ok"))
with patch.object(mapper._extractor, "run", AsyncMock(side_effect=_stub)):
outputs = await mapper.map_pages(pages, "anything")
covered = sorted({p for o in outputs for p in o.pages})
assert covered == [1, 3]
@pytest.mark.anyio
async def test_outer_cancellation_drains_pending_tasks(self, runtime: AppRuntime) -> None:
"""Cancellation propagating in from upstream cancels per-chunk model
calls rather than letting them keep billing tokens."""
mapper = _build_mapper(runtime, chars_per_slice=10, concurrency=5)
pages = [_page(i, "x" * 8) for i in range(1, 5)]
cancellations = 0
async def _hang(*_args: object, **_kwargs: object) -> _StubAgentResult[_Extracted]:
nonlocal cancellations
try:
await asyncio.sleep(60)
except asyncio.CancelledError:
cancellations += 1
raise
return _StubAgentResult(output=_Extracted(label="never"))
with patch.object(mapper._extractor, "run", AsyncMock(side_effect=_hang)):
task = asyncio.create_task(mapper.map_pages(pages, "anything"))
# Yield once so all four workers are blocked on their sleep.
await asyncio.sleep(0)
await asyncio.sleep(0)
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
assert cancellations == len(pages)
@pytest.mark.anyio
async def test_semaphore_caps_concurrency(self, runtime: AppRuntime) -> None:
"""At most ``concurrency`` workers run at once; with strictly more work
items than slots the observed max is exactly the configured cap."""
concurrency = 2
mapper = _build_mapper(runtime, chars_per_slice=10, concurrency=concurrency)
pages = [_page(i, "x" * 8) for i in range(1, 6)] # 5 items > 2 slots
active = 0
peak = 0
async def _track(*_args: object, **_kwargs: object) -> _StubAgentResult[_Extracted]:
nonlocal active, peak
active += 1
peak = max(peak, active)
# Yield enough times that other waiters get a chance to enter.
for _ in range(5):
await asyncio.sleep(0)
active -= 1
return _StubAgentResult(output=_Extracted(label="ok"))
with patch.object(mapper._extractor, "run", AsyncMock(side_effect=_track)):
outputs = await mapper.map_pages(pages, "anything")
assert len(outputs) == 5
assert peak == concurrency
@pytest.mark.anyio
async def test_rejects_empty_pages(self, runtime: AppRuntime) -> None:
mapper = _build_mapper(runtime)
with pytest.raises(ValueError, match="at least one page"):
await mapper.map_pages([], "anything")
class TestSummaryCounts:
"""``summary_counts`` callback feeds the WholeDocSliceDone event's
excerpts/facts counters from the consumer's extractor output shape
without the mapper itself duck-typing fields on ``T``."""
@pytest.mark.anyio
async def test_default_callback_emits_zero_counts(self, runtime: AppRuntime) -> None:
"""No callback supplied → events emit ``excerpts=0 facts=0``."""
from stirling.contracts import WholeDocSliceDone
from stirling.services import reset_progress_emitter, set_progress_emitter
mapper = _build_mapper(runtime, chars_per_slice=1000)
pages = [_page(1, "small")]
canned = _Extracted(label="ok")
emitted: list[WholeDocSliceDone] = []
async def _emit(event: object) -> None:
if isinstance(event, WholeDocSliceDone):
emitted.append(event)
token = set_progress_emitter(_emit)
try:
with patch.object(
mapper._extractor,
"run",
AsyncMock(return_value=_StubAgentResult(output=canned)),
):
await mapper.map_pages(pages, "q")
finally:
reset_progress_emitter(token)
assert len(emitted) == 1
assert emitted[0].excerpts == 0
assert emitted[0].facts == 0
@pytest.mark.anyio
async def test_user_callback_drives_counts(self, runtime: AppRuntime) -> None:
"""A supplied callback receives each chunk's typed output and its
returned tuple is what the event carries."""
from stirling.contracts import WholeDocSliceDone
from stirling.services import reset_progress_emitter, set_progress_emitter
captured: list[_Extracted] = []
def _counts(output: _Extracted) -> tuple[int, int]:
captured.append(output)
return (3, 7)
extractor: Agent[None, _Extracted] = Agent(
model=runtime.fast_model,
output_type=_Extracted,
model_settings=runtime.fast_model_settings,
)
mapper: ChunkedMapper[_Extracted] = ChunkedMapper(
runtime,
extractor=extractor,
chars_per_slice=1000,
summary_counts=_counts,
)
canned = _Extracted(label="ok")
emitted: list[WholeDocSliceDone] = []
async def _emit(event: object) -> None:
if isinstance(event, WholeDocSliceDone):
emitted.append(event)
token = set_progress_emitter(_emit)
try:
with patch.object(
mapper._extractor,
"run",
AsyncMock(return_value=_StubAgentResult(output=canned)),
):
await mapper.map_pages([_page(1, "small")], "q")
finally:
reset_progress_emitter(token)
assert len(captured) == 1
assert captured[0].label == "ok"
assert emitted[0].excerpts == 3
assert emitted[0].facts == 7
class TestChunkOutputShape:
@pytest.mark.anyio
async def test_single_page_label(self, runtime: AppRuntime) -> None:
mapper = _build_mapper(runtime, chars_per_slice=5)
pages = [_page(7, "x" * 6)] # one oversize page → one slice
canned = _Extracted(label="solo")
with patch.object(
mapper._extractor,
"run",
AsyncMock(return_value=_StubAgentResult(output=canned)),
):
outputs = await mapper.map_pages(pages, "q")
assert outputs[0].label == "pages=7"
engine/tests/agents/test_chunked_reasoner.py
+89 -18
View File
@@ -8,11 +8,13 @@ from __future__ import annotations
import asyncio
from dataclasses import dataclass
from typing import Any
from unittest.mock import AsyncMock, patch
import pytest
from pydantic import BaseModel
from stirling.agents.shared.chunked_mapper import _ChunkExtraction
from stirling.agents.shared.chunked_reasoner import ChunkedReasoner, ChunkNotes
from stirling.contracts import WholeDocSliceDone
from stirling.contracts.documents import Page
@@ -83,14 +85,20 @@ class TestReason:
async def test_runs_one_chunk_per_slice_and_synthesises(self, runtime: AppRuntime) -> None:
"""Three small pages with a generous budget produce one chunk and one extractor call;
the synthesis stage receives notes from all chunks and returns the final answer."""
from stirling.agents.shared.chunked_reasoner import _ExtractedNotes
reasoner = ChunkedReasoner(runtime, chars_per_slice=1000)
pages = [_page(1, "alpha"), _page(2, "beta"), _page(3, "gamma")]
canned_notes = ChunkNotes(pages=[1, 2, 3], summary="all three pages", facts=["fact-1"])
canned_extracted = _ExtractedNotes(summary="all three pages", facts=["fact-1"])
canned_answer = _Answer(answer="final answer")
with (
patch.object(reasoner, "_extract_chunk", AsyncMock(return_value=(canned_notes, 0.0))) as chunk_mock,
patch.object(
reasoner._mapper,
"_extract_chunk",
AsyncMock(return_value=_ChunkExtraction(output=canned_extracted, duration_seconds=0.0)),
) as chunk_mock,
patch.object(reasoner, "_synthesise", AsyncMock(return_value=canned_answer)) as synth_mock,
):
result = await reasoner.reason(
@@ -107,20 +115,29 @@ class TestReason:
assert synth_args is not None
# _synthesise(question, notes, answer_prompt, answer_type)
_, notes_arg, _, type_arg = synth_args.args
assert notes_arg == [canned_notes]
assert len(notes_arg) == 1
assert notes_arg[0].pages == [1, 2, 3]
assert notes_arg[0].summary == "all three pages"
assert notes_arg[0].facts == ["fact-1"]
assert type_arg is _Answer
@pytest.mark.anyio
async def test_fans_out_when_pages_exceed_slice_budget(self, runtime: AppRuntime) -> None:
"""Pages that don't fit into a single slice produce one extractor call per slice."""
from stirling.agents.shared.chunked_reasoner import _ExtractedNotes
reasoner = ChunkedReasoner(runtime, chars_per_slice=10)
pages = [_page(i, "x" * 8) for i in range(1, 6)]
canned_notes = ChunkNotes(pages=[0], summary="placeholder")
canned_extracted = _ExtractedNotes(summary="placeholder")
canned_answer = _Answer(answer="ok")
with (
patch.object(reasoner, "_extract_chunk", AsyncMock(return_value=(canned_notes, 0.0))) as chunk_mock,
patch.object(
reasoner._mapper,
"_extract_chunk",
AsyncMock(return_value=_ChunkExtraction(output=canned_extracted, duration_seconds=0.0)),
) as chunk_mock,
patch.object(reasoner, "_synthesise", AsyncMock(return_value=canned_answer)),
):
await reasoner.reason(
@@ -138,22 +155,24 @@ class TestReason:
"""First-round chunks have no fallback notes, so a failure is dropped
rather than preserving anything; the surviving notes still flow into
synthesis."""
from stirling.agents.shared.chunked_reasoner import _ExtractedNotes
reasoner = ChunkedReasoner(runtime, chars_per_slice=10)
pages = [_page(i, "x" * 8) for i in range(1, 4)]
good = ChunkNotes(pages=[1], summary="ok")
async_results = [good, RuntimeError("chunk boom"), good]
good = _ExtractedNotes(summary="ok")
async_results: list[_ExtractedNotes | BaseException] = [good, RuntimeError("chunk boom"), good]
async def _chunk(*_args: object, **_kwargs: object) -> tuple[ChunkNotes, float]:
async def _chunk(*_args: object, **_kwargs: object) -> _ChunkExtraction[_ExtractedNotes]:
value = async_results.pop(0)
if isinstance(value, BaseException):
raise value
return value, 0.0
return _ChunkExtraction(output=value, duration_seconds=0.0)
canned_answer = _Answer(answer="resilient")
with (
patch.object(reasoner, "_extract_chunk", AsyncMock(side_effect=_chunk)),
patch.object(reasoner._mapper, "_extract_chunk", AsyncMock(side_effect=_chunk)),
patch.object(reasoner, "_synthesise", AsyncMock(return_value=canned_answer)) as synth_mock,
):
result = await reasoner.reason(
@@ -175,7 +194,7 @@ class TestReason:
pages = [_page(i, "x" * 8) for i in range(1, 3)]
with (
patch.object(reasoner, "_extract_chunk", AsyncMock(side_effect=RuntimeError("boom"))),
patch.object(reasoner._mapper, "_extract_chunk", AsyncMock(side_effect=RuntimeError("boom"))),
patch.object(reasoner, "_synthesise", AsyncMock()) as synth_mock,
pytest.raises(RuntimeError, match="no notes"),
):
@@ -310,9 +329,13 @@ class TestPromptConstruction:
def test_extraction_prompt_includes_question_and_page_markers(self, runtime: AppRuntime) -> None:
"""A first-round chunk's content carries ``[Page N]`` markers; the
extraction prompt prepends the user question."""
from stirling.agents.shared.chunked_mapper import ChunkedMapper
reasoner = ChunkedReasoner(runtime)
chunk = reasoner._chunk_from_pages([_page(2, "page two body"), _page(3, "page three body")])
prompt = reasoner._build_extraction_prompt(chunk.content, "what is on page two?")
# Render chunk content through the mapper's public helper — the
# first-round chunk shape lives in ChunkedMapper.
content = ChunkedMapper.format_chunk_content([_page(2, "page two body"), _page(3, "page three body")])
prompt = reasoner._build_extraction_prompt(content, "what is on page two?")
assert "what is on page two?" in prompt
assert "[Page 2]" in prompt
@@ -334,10 +357,11 @@ class TestPromptConstruction:
# Hierarchical compression
#
# The compression loop is part of ``_run_chunks`` and isn't exposed
# directly, so these tests drive it end-to-end via ``gather_notes`` with a
# stubbed extractor that controls per-call output (and per-call failure
# patterns) by counting calls.
# The compression loop is part of ``_compress_until_fits`` /
# ``_run_compression_round`` and isn't exposed directly, so these tests
# drive it end-to-end via ``gather_notes`` with a stubbed extractor that
# controls per-call output (and per-call failure patterns) by counting
# calls.
class TestCompression:
@@ -503,9 +527,56 @@ class TestExtractChunk:
"run",
AsyncMock(return_value=_StubAgentResult(output=canned)),
):
note, _ = await reasoner._extract_chunk(chunk, "anything")
extraction = await reasoner._extract_compression_chunk(chunk, "compress these")
note = extraction.output
assert note.pages == [1, 2, 3, 4, 5]
assert note.summary == "merged"
assert note.facts == ["x"]
assert note.relevant_excerpts == ["y"]
assert extraction.duration_seconds >= 0
@pytest.mark.anyio
async def test_compression_rounds_receive_user_question_through_gather_notes(self, runtime: AppRuntime) -> None:
"""Regression — every extractor call (first round AND every
compression round) MUST carry the same user question. The pre-fix
bug passed ``""`` to the compression-round prompt builder, so the
model consolidated notes against different relevance criteria
than it extracted them under. Flagged by Aikido on PR #6369;
pinned end-to-end here by capturing every prompt the extractor
sees while ``gather_notes`` forces a compression round through a
tight notes budget.
"""
from stirling.agents.shared.chunked_reasoner import _ExtractedNotes
# Small notes budget forces a compression round; small slice
# budget produces multiple first-round chunks that overflow it.
reasoner = ChunkedReasoner(runtime, chars_per_slice=200, notes_char_budget=200)
pages = [_page(i, "x" * 150) for i in range(1, 5)]
call_count = 0
async def _stub(*_args: object, **_kwargs: object) -> _StubAgentResult[object]:
nonlocal call_count
call_count += 1
if call_count <= 4:
# Round 1: each note ~60 chars rendered. 4 * 80 = 320 chars,
# over the 200 budget so a compression round must fire.
return _StubAgentResult(output=_ExtractedNotes(summary="x" * 60))
# Round 2: smaller note so the post-round set fits the budget.
return _StubAgentResult(output=_ExtractedNotes(summary="ok"))
seen_prompts: list[str] = []
async def _capture(prompt: str, *_a: Any, **_kw: Any) -> Any:
seen_prompts.append(prompt)
return await _stub()
with patch.object(reasoner._extractor, "run", side_effect=_capture):
await reasoner.gather_notes(pages, "what is the deadline?")
# At least four first-round calls plus the compression-round
# calls — every single one must carry the user question.
assert len(seen_prompts) >= 5
for prompt in seen_prompts:
assert "what is the deadline?" in prompt