mirror of
https://github.com/arsvendg/Stirling-PDF.git
synced 2026-07-14 10:34:06 +02:00
Add ability for Stirling engine to reason across large documents (#6314)
# Description of Changes Adds storage in the database for full document content alongside the RAG content (and changes the service to `DocumentService` instead of `RagService`). Then adds a generic capability that should be usable by any agent (currently just used by the Question Agent) which allows the agent to pull out the full contents of the doc, chunks it into various sections that will fit in the context window, and then processes them in parallel to create an intermediate result, and then processes the intermediate result into a final answer. It will re-chunk as many times as necessary to get the content small enough for the actual answer to be analysed (I've tested on PDFs ~3500 pages long, which is well above the context limit and requires maybe 3 rounds of compression to get an answer). The new full doc analysis stuff is heavier than the RAG lookup so both remain. The agents should use RAG for targeted info and the chunked reasoner for info that requires reading the full doc.
This commit is contained in:
James Brunton
@@ -0,0 +1,511 @@
|
||||
"""Tests for ``ChunkedReasoner``: chunking, fan-out, and synthesis wiring.
|
||||
|
||||
LLM calls are stubbed at the agent boundary; the runtime fixture supplies a
|
||||
``test`` model so construction succeeds without provider credentials.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from dataclasses import dataclass
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from pydantic import BaseModel
|
||||
|
||||
from stirling.agents.shared.chunked_reasoner import ChunkedReasoner, ChunkNotes
|
||||
from stirling.contracts import WholeDocSliceDone
|
||||
from stirling.contracts.documents import Page
|
||||
from stirling.services import reset_progress_emitter, set_progress_emitter
|
||||
from stirling.services.runtime import AppRuntime
|
||||
|
||||
|
||||
@dataclass
|
||||
class _StubAgentResult[T]:
|
||||
output: T
|
||||
|
||||
|
||||
class _Answer(BaseModel):
|
||||
answer: str
|
||||
|
||||
|
||||
def _page(n: int, text: str) -> Page:
|
||||
return Page(page_number=n, text=text, char_count=len(text))
|
||||
|
||||
|
||||
# Slicing logic
|
||||
|
||||
|
||||
class TestSlicePages:
|
||||
"""``_slice_pages`` is pure: no model calls, no I/O."""
|
||||
|
||||
def test_groups_consecutive_pages_under_budget(self, runtime: AppRuntime) -> None:
|
||||
reasoner = ChunkedReasoner(runtime, chars_per_slice=20)
|
||||
pages = [_page(1, "abc"), _page(2, "defgh"), _page(3, "ij")]
|
||||
slices = reasoner._slice_pages(pages)
|
||||
|
||||
assert len(slices) == 1
|
||||
assert [p.page_number for p in slices[0]] == [1, 2, 3]
|
||||
|
||||
def test_starts_new_slice_when_budget_exceeded(self, runtime: AppRuntime) -> None:
|
||||
reasoner = ChunkedReasoner(runtime, chars_per_slice=10)
|
||||
pages = [_page(1, "a" * 6), _page(2, "b" * 6), _page(3, "c" * 6)]
|
||||
slices = reasoner._slice_pages(pages)
|
||||
|
||||
# Each page is 6 chars, budget 10 -> two pages would be 12 (over),
|
||||
# so the slicer breaks after each page.
|
||||
assert [[p.page_number for p in s] for s in slices] == [[1], [2], [3]]
|
||||
|
||||
def test_oversized_page_becomes_its_own_slice(self, runtime: AppRuntime) -> None:
|
||||
"""A single page larger than the budget is never split. The reasoner
|
||||
accepts that this slice exceeds the budget rather than breaking page
|
||||
boundaries."""
|
||||
reasoner = ChunkedReasoner(runtime, chars_per_slice=10)
|
||||
pages = [_page(1, "small"), _page(2, "x" * 100), _page(3, "tiny")]
|
||||
slices = reasoner._slice_pages(pages)
|
||||
|
||||
assert [[p.page_number for p in s] for s in slices] == [[1], [2], [3]]
|
||||
|
||||
def test_preserves_input_order(self, runtime: AppRuntime) -> None:
|
||||
reasoner = ChunkedReasoner(runtime, chars_per_slice=1000)
|
||||
pages = [_page(i, f"page-{i}") for i in range(1, 11)]
|
||||
slices = reasoner._slice_pages(pages)
|
||||
|
||||
flattened = [p.page_number for s in slices for p in s]
|
||||
assert flattened == list(range(1, 11))
|
||||
|
||||
|
||||
# End-to-end orchestration
|
||||
|
||||
|
||||
class TestReason:
|
||||
@pytest.mark.anyio
|
||||
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."""
|
||||
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_answer = _Answer(answer="final answer")
|
||||
|
||||
with (
|
||||
patch.object(reasoner, "_extract_chunk", AsyncMock(return_value=(canned_notes, 0.0))) as chunk_mock,
|
||||
patch.object(reasoner, "_synthesise", AsyncMock(return_value=canned_answer)) as synth_mock,
|
||||
):
|
||||
result = await reasoner.reason(
|
||||
pages=pages,
|
||||
question="summarise this",
|
||||
answer_prompt="answer the question from the notes",
|
||||
answer_type=_Answer,
|
||||
)
|
||||
|
||||
assert result == canned_answer
|
||||
assert chunk_mock.await_count == 1
|
||||
synth_mock.assert_awaited_once()
|
||||
synth_args = synth_mock.await_args
|
||||
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 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."""
|
||||
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_answer = _Answer(answer="ok")
|
||||
|
||||
with (
|
||||
patch.object(reasoner, "_extract_chunk", AsyncMock(return_value=(canned_notes, 0.0))) as chunk_mock,
|
||||
patch.object(reasoner, "_synthesise", AsyncMock(return_value=canned_answer)),
|
||||
):
|
||||
await reasoner.reason(
|
||||
pages=pages,
|
||||
question="aggregate",
|
||||
answer_prompt="answer",
|
||||
answer_type=_Answer,
|
||||
)
|
||||
|
||||
# 5 pages, budget 10, each page 8 chars -> 5 slices -> 5 chunk calls.
|
||||
assert chunk_mock.await_count == 5
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_skips_first_round_chunks_that_raise_and_continues(self, runtime: AppRuntime) -> None:
|
||||
"""First-round chunks have no fallback notes, so a failure is dropped
|
||||
rather than preserving anything; the surviving notes still flow into
|
||||
synthesis."""
|
||||
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]
|
||||
|
||||
async def _chunk(*_args: object, **_kwargs: object) -> tuple[ChunkNotes, float]:
|
||||
value = async_results.pop(0)
|
||||
if isinstance(value, BaseException):
|
||||
raise value
|
||||
return value, 0.0
|
||||
|
||||
canned_answer = _Answer(answer="resilient")
|
||||
|
||||
with (
|
||||
patch.object(reasoner, "_extract_chunk", AsyncMock(side_effect=_chunk)),
|
||||
patch.object(reasoner, "_synthesise", AsyncMock(return_value=canned_answer)) as synth_mock,
|
||||
):
|
||||
result = await reasoner.reason(
|
||||
pages=pages,
|
||||
question="aggregate",
|
||||
answer_prompt="answer",
|
||||
answer_type=_Answer,
|
||||
)
|
||||
|
||||
assert result == canned_answer
|
||||
synth_args = synth_mock.await_args
|
||||
assert synth_args is not None
|
||||
_, notes_arg, _, _ = synth_args.args
|
||||
assert len(notes_arg) == 2
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_raises_when_every_first_round_chunk_fails(self, runtime: AppRuntime) -> None:
|
||||
reasoner = ChunkedReasoner(runtime, chars_per_slice=10)
|
||||
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, "_synthesise", AsyncMock()) as synth_mock,
|
||||
pytest.raises(RuntimeError, match="no notes"),
|
||||
):
|
||||
await reasoner.reason(
|
||||
pages=pages,
|
||||
question="anything",
|
||||
answer_prompt="answer",
|
||||
answer_type=_Answer,
|
||||
)
|
||||
|
||||
synth_mock.assert_not_awaited()
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_rejects_empty_pages(self, runtime: AppRuntime) -> None:
|
||||
reasoner = ChunkedReasoner(runtime)
|
||||
with pytest.raises(ValueError, match="at least one page"):
|
||||
await reasoner.reason(
|
||||
pages=[],
|
||||
question="x",
|
||||
answer_prompt="y",
|
||||
answer_type=_Answer,
|
||||
)
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_progress_events_carry_monotonic_completion_counter(self, runtime: AppRuntime) -> None:
|
||||
reasoner = ChunkedReasoner(runtime, chars_per_slice=10)
|
||||
pages = [_page(i, "x" * 8) for i in range(1, 4)]
|
||||
|
||||
# Each chunk's worker awaits a different release event; we release them in
|
||||
# reverse order so completion order is the inverse of slice order.
|
||||
release_events = [asyncio.Event() for _ in pages]
|
||||
next_call_index = 0
|
||||
|
||||
async def _gated_worker(*_args: object, **_kwargs: object) -> _StubAgentResult[ChunkNotes]:
|
||||
nonlocal next_call_index
|
||||
this_call = next_call_index
|
||||
next_call_index += 1
|
||||
await release_events[this_call].wait()
|
||||
return _StubAgentResult(output=ChunkNotes(pages=[this_call + 1], summary=f"slice-{this_call}"))
|
||||
|
||||
emitted: list[WholeDocSliceDone] = []
|
||||
|
||||
async def _capture_emitter(event: object) -> None:
|
||||
if isinstance(event, WholeDocSliceDone):
|
||||
emitted.append(event)
|
||||
|
||||
async def _release_in_reverse() -> None:
|
||||
# Wait briefly so all three worker tasks are blocked on their events
|
||||
# before we start releasing them.
|
||||
await asyncio.sleep(0)
|
||||
for ev in reversed(release_events):
|
||||
ev.set()
|
||||
# Yield so the just-released worker can run to completion before
|
||||
# we release the next one — keeps ordering deterministic.
|
||||
await asyncio.sleep(0)
|
||||
await asyncio.sleep(0)
|
||||
|
||||
token = set_progress_emitter(_capture_emitter)
|
||||
try:
|
||||
with patch.object(reasoner._extractor, "run", AsyncMock(side_effect=_gated_worker)):
|
||||
gather_task = asyncio.create_task(reasoner.gather_notes(pages, "anything"))
|
||||
await _release_in_reverse()
|
||||
notes = await gather_task
|
||||
finally:
|
||||
reset_progress_emitter(token)
|
||||
|
||||
assert len(notes) == 3
|
||||
assert [event.completed for event in emitted] == [1, 2, 3]
|
||||
assert all(event.total == 3 for event in emitted)
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_worker_timeout_is_terminal_for_the_chunk(self, runtime: AppRuntime) -> None:
|
||||
reasoner = ChunkedReasoner(runtime, chars_per_slice=10, worker_timeout_seconds=0.05)
|
||||
pages = [_page(1, "x" * 8), _page(2, "y" * 8)]
|
||||
attempts = 0
|
||||
|
||||
async def _hang_forever(*_args: object, **_kwargs: object) -> _StubAgentResult[ChunkNotes]:
|
||||
nonlocal attempts
|
||||
attempts += 1
|
||||
await asyncio.sleep(10)
|
||||
return _StubAgentResult(output=ChunkNotes(pages=[0], summary="never"))
|
||||
|
||||
with (
|
||||
patch.object(reasoner._extractor, "run", AsyncMock(side_effect=_hang_forever)),
|
||||
patch.object(reasoner, "_synthesise", AsyncMock()) as synth_mock,
|
||||
pytest.raises(RuntimeError, match="no notes"),
|
||||
):
|
||||
await reasoner.reason(
|
||||
pages=pages,
|
||||
question="anything",
|
||||
answer_prompt="answer",
|
||||
answer_type=_Answer,
|
||||
)
|
||||
|
||||
# One attempt per slice; no retry path.
|
||||
assert attempts == len(pages)
|
||||
synth_mock.assert_not_awaited()
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_worker_timeout_drops_stalled_chunks(self, runtime: AppRuntime) -> None:
|
||||
"""A worker that exceeds ``worker_timeout_seconds`` is abandoned, not awaited.
|
||||
|
||||
Without this guard one stuck upstream call would pin gather_notes to its
|
||||
provider HTTP timeout (~10 minutes), starving the orchestrator request.
|
||||
"""
|
||||
reasoner = ChunkedReasoner(runtime, chars_per_slice=10, worker_timeout_seconds=0.05)
|
||||
pages = [_page(i, "x" * 8) for i in range(1, 4)]
|
||||
|
||||
async def _hang(*_args: object, **_kwargs: object) -> _StubAgentResult[ChunkNotes]:
|
||||
await asyncio.sleep(10)
|
||||
return _StubAgentResult(output=ChunkNotes(pages=[0], summary="never"))
|
||||
|
||||
with (
|
||||
patch.object(reasoner._extractor, "run", AsyncMock(side_effect=_hang)),
|
||||
patch.object(reasoner, "_synthesise", AsyncMock()) as synth_mock,
|
||||
pytest.raises(RuntimeError, match="no notes"),
|
||||
):
|
||||
await reasoner.reason(
|
||||
pages=pages,
|
||||
question="anything",
|
||||
answer_prompt="answer",
|
||||
answer_type=_Answer,
|
||||
)
|
||||
|
||||
synth_mock.assert_not_awaited()
|
||||
|
||||
|
||||
# Prompt construction
|
||||
|
||||
|
||||
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."""
|
||||
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?")
|
||||
|
||||
assert "what is on page two?" in prompt
|
||||
assert "[Page 2]" in prompt
|
||||
assert "[Page 3]" in prompt
|
||||
assert "page two body" in prompt
|
||||
|
||||
def test_format_notes_groups_by_page_label(self) -> None:
|
||||
notes = [
|
||||
ChunkNotes(pages=[1], summary="single", facts=["f-1"]),
|
||||
ChunkNotes(pages=[2, 3, 4], summary="range", relevant_excerpts=["quote-1"]),
|
||||
]
|
||||
rendered = ChunkedReasoner.format_notes(notes)
|
||||
|
||||
assert "[Notes from page 1]" in rendered
|
||||
assert "[Notes from pages 2-4]" in rendered
|
||||
assert "f-1" in rendered
|
||||
assert "quote-1" in rendered
|
||||
|
||||
|
||||
# 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.
|
||||
|
||||
|
||||
class TestCompression:
|
||||
@pytest.mark.anyio
|
||||
async def test_no_compression_when_under_budget(self, runtime: AppRuntime) -> None:
|
||||
"""First-round notes that already fit the budget result in zero
|
||||
compression rounds: the only extractor calls are one per slice."""
|
||||
from stirling.agents.shared.chunked_reasoner import _ExtractedNotes
|
||||
|
||||
reasoner = ChunkedReasoner(runtime, chars_per_slice=200, notes_char_budget=10_000)
|
||||
pages = [_page(i, "x" * 150) for i in range(1, 5)] # each page exceeds slice budget alone -> 4 slices
|
||||
|
||||
canned = _ExtractedNotes(summary="ok")
|
||||
with patch.object(
|
||||
reasoner._extractor,
|
||||
"run",
|
||||
AsyncMock(return_value=_StubAgentResult(output=canned)),
|
||||
) as ext_mock:
|
||||
notes = await reasoner.gather_notes(pages, "anything")
|
||||
|
||||
assert ext_mock.await_count == 4
|
||||
assert len(notes) == 4
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_runs_compression_when_over_budget(self, runtime: AppRuntime) -> None:
|
||||
"""When first-round notes overflow the budget, the loop regroups them
|
||||
and runs the extractor again. Output is shorter than input; pages from
|
||||
every input slice survive in the consolidated notes."""
|
||||
from stirling.agents.shared.chunked_reasoner import _ExtractedNotes
|
||||
|
||||
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 ~80 chars rendered. 4 * 80 = 320 chars, over the 200 budget.
|
||||
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"))
|
||||
|
||||
with patch.object(reasoner._extractor, "run", AsyncMock(side_effect=_stub)) as ext_mock:
|
||||
notes = await reasoner.gather_notes(pages, "anything")
|
||||
|
||||
# 4 first-round + 2 compression-round calls = 6 total.
|
||||
assert ext_mock.await_count == 6
|
||||
# Compressed from 4 notes to 2.
|
||||
assert len(notes) == 2
|
||||
# Pages from every original slice are preserved through compression.
|
||||
assert sorted({p for n in notes for p in n.pages}) == [1, 2, 3, 4]
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_compression_preserves_input_notes_when_a_group_fails(self, runtime: AppRuntime) -> None:
|
||||
"""A compression chunk that raises has its input notes carried forward
|
||||
rather than dropped, so page coverage isn't silently lost. The
|
||||
succeeding chunk is replaced by its consolidated note.
|
||||
|
||||
Budget is sized so the post-round survivors (2 preserved + 1
|
||||
consolidated) fit, leaving a single compression round as the
|
||||
observable interaction."""
|
||||
from stirling.agents.shared.chunked_reasoner import _ExtractedNotes
|
||||
|
||||
reasoner = ChunkedReasoner(runtime, chars_per_slice=200, notes_char_budget=300)
|
||||
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:
|
||||
return _StubAgentResult(output=_ExtractedNotes(summary="x" * 60))
|
||||
if call_count == 5:
|
||||
# The first compression call (covering 2 of the round-1 notes) fails.
|
||||
raise RuntimeError("compression group fails")
|
||||
return _StubAgentResult(output=_ExtractedNotes(summary="ok"))
|
||||
|
||||
with patch.object(reasoner._extractor, "run", AsyncMock(side_effect=_stub)):
|
||||
notes = await reasoner.gather_notes(pages, "anything")
|
||||
|
||||
# 2 preserved round-1 notes + 1 consolidated note = 3 notes total. Pages
|
||||
# from every original slice are still covered (preservation worked).
|
||||
assert len(notes) == 3
|
||||
assert sorted({p for n in notes for p in n.pages}) == [1, 2, 3, 4]
|
||||
|
||||
consolidated = [n for n in notes if n.summary == "ok"]
|
||||
preserved = [n for n in notes if n.summary.startswith("x")]
|
||||
assert len(consolidated) == 1
|
||||
assert len(preserved) == 2
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_compression_bails_when_every_group_fails(self, runtime: AppRuntime) -> None:
|
||||
"""If every chunk in a compression round fails, every input note is
|
||||
preserved (none consolidated). The loop exits rather than retrying
|
||||
the same shape forever."""
|
||||
from stirling.agents.shared.chunked_reasoner import _ExtractedNotes
|
||||
|
||||
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:
|
||||
return _StubAgentResult(output=_ExtractedNotes(summary="x" * 60))
|
||||
raise RuntimeError("compression always fails")
|
||||
|
||||
with patch.object(reasoner._extractor, "run", AsyncMock(side_effect=_stub)):
|
||||
notes = await reasoner.gather_notes(pages, "anything")
|
||||
|
||||
# All 4 round-1 notes preserved through the bailed compression round.
|
||||
assert len(notes) == 4
|
||||
assert sorted({p for n in notes for p in n.pages}) == [1, 2, 3, 4]
|
||||
|
||||
|
||||
class TestGroupNotesForCompression:
|
||||
"""``_group_notes_for_compression`` is pure and packs by rendered char count."""
|
||||
|
||||
def test_packs_consecutive_notes_under_budget(self, runtime: AppRuntime) -> None:
|
||||
reasoner = ChunkedReasoner(runtime, chars_per_slice=10_000)
|
||||
notes = [ChunkNotes(pages=[i], summary=f"s-{i}") for i in range(1, 5)]
|
||||
|
||||
groups = reasoner._group_notes_for_compression(notes)
|
||||
|
||||
assert len(groups) == 1
|
||||
assert [n.pages[0] for n in groups[0]] == [1, 2, 3, 4]
|
||||
|
||||
def test_starts_new_group_when_budget_exceeded(self, runtime: AppRuntime) -> None:
|
||||
"""Each note already exceeds the per-group budget, so each becomes its
|
||||
own group; this matches how slice-pages handles oversize pages."""
|
||||
reasoner = ChunkedReasoner(runtime, chars_per_slice=5)
|
||||
notes = [ChunkNotes(pages=[i], summary=f"slice-{i}-with-prose-to-fill-space") for i in range(1, 5)]
|
||||
|
||||
groups = reasoner._group_notes_for_compression(notes)
|
||||
|
||||
assert [[n.pages[0] for n in g] for g in groups] == [[1], [2], [3], [4]]
|
||||
|
||||
|
||||
class TestExtractChunk:
|
||||
@pytest.mark.anyio
|
||||
async def test_pages_are_unioned_for_compression_chunks(self, runtime: AppRuntime) -> None:
|
||||
"""A compression chunk's resulting note carries the union of input pages.
|
||||
The model output schema doesn't include pages, so the wrapper is the
|
||||
single source of truth."""
|
||||
from stirling.agents.shared.chunked_reasoner import _ExtractedNotes
|
||||
|
||||
reasoner = ChunkedReasoner(runtime)
|
||||
group = [
|
||||
ChunkNotes(pages=[1, 2], summary="a"),
|
||||
ChunkNotes(pages=[3], summary="b"),
|
||||
ChunkNotes(pages=[4, 5], summary="c"),
|
||||
]
|
||||
chunk = reasoner._chunk_from_notes(group)
|
||||
canned = _ExtractedNotes(summary="merged", facts=["x"], relevant_excerpts=["y"])
|
||||
|
||||
with patch.object(
|
||||
reasoner._extractor,
|
||||
"run",
|
||||
AsyncMock(return_value=_StubAgentResult(output=canned)),
|
||||
):
|
||||
note, _ = await reasoner._extract_chunk(chunk, "anything")
|
||||
|
||||
assert note.pages == [1, 2, 3, 4, 5]
|
||||
assert note.summary == "merged"
|
||||
assert note.facts == ["x"]
|
||||
assert note.relevant_excerpts == ["y"]
|
||||
@@ -0,0 +1,220 @@
|
||||
"""Tests for ``WholeDocReaderCapability``: tool dispatch, multi-file iteration,
|
||||
budget enforcement, and graceful handling of missing pages.
|
||||
|
||||
The map-phase LLM call is patched at the reasoner boundary so tests don't hit
|
||||
any model.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import replace
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from stirling.agents.shared import ChunkedReasoner, ChunkNotes, WholeDocReaderCapability
|
||||
from stirling.contracts import AiFile, PageText
|
||||
from stirling.documents import Document, DocumentService, SqliteVecStore
|
||||
from stirling.models import FileId
|
||||
from stirling.services.runtime import AppRuntime
|
||||
|
||||
|
||||
class StubEmbedder:
|
||||
"""Deterministic embeddings so tests don't need a real provider."""
|
||||
|
||||
def __init__(self, dim: int = 8) -> None:
|
||||
self._dim = dim
|
||||
|
||||
async def embed_query(self, text: str) -> list[float]:
|
||||
h = hash(text) % 1000
|
||||
return [(h + i) / 1000.0 for i in range(self._dim)]
|
||||
|
||||
async def embed_documents(self, texts: list[str]) -> list[list[float]]:
|
||||
return [await self.embed_query(t) for t in texts]
|
||||
|
||||
def chunk_and_prepare(
|
||||
self,
|
||||
text: str,
|
||||
source: str = "",
|
||||
base_metadata: dict[str, str] | None = None,
|
||||
) -> list[Document]:
|
||||
from stirling.documents.chunker import chunk_text
|
||||
|
||||
chunks = chunk_text(text, 100, 10)
|
||||
docs: list[Document] = []
|
||||
for i, chunk in enumerate(chunks):
|
||||
meta = dict(base_metadata) if base_metadata else {}
|
||||
meta["source"] = source
|
||||
meta["chunk_index"] = str(i)
|
||||
doc_id = f"{source}:chunk:{i}" if source else f"chunk:{i}"
|
||||
docs.append(Document(id=doc_id, text=chunk, metadata=meta))
|
||||
return docs
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def runtime_with_stub_docs(runtime: AppRuntime) -> AppRuntime:
|
||||
stub = DocumentService(
|
||||
embedder=StubEmbedder(), # type: ignore[arg-type]
|
||||
store=SqliteVecStore.ephemeral(),
|
||||
default_top_k=runtime.settings.rag_default_top_k,
|
||||
)
|
||||
return replace(runtime, documents=stub)
|
||||
|
||||
|
||||
def _ai_file(file_id: str, name: str) -> AiFile:
|
||||
return AiFile(id=FileId(file_id), name=name)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_read_full_document_returns_formatted_notes_for_single_file(
|
||||
runtime_with_stub_docs: AppRuntime,
|
||||
) -> None:
|
||||
"""The tool reads the file's stored pages, calls the reasoner's map phase,
|
||||
and returns the formatted notes prefixed by the file name."""
|
||||
pages = [
|
||||
PageText(page_number=1, text="Chapter one prose."),
|
||||
PageText(page_number=2, text="Chapter two prose."),
|
||||
]
|
||||
await runtime_with_stub_docs.documents.ingest(FileId("doc-id"), pages, source="doc.pdf")
|
||||
|
||||
reasoner = ChunkedReasoner(runtime_with_stub_docs)
|
||||
canned_notes = [ChunkNotes(pages=[1, 2], summary="overview", facts=["fact-A"])]
|
||||
with patch.object(reasoner, "gather_notes", AsyncMock(return_value=canned_notes)) as gather_mock:
|
||||
capability = WholeDocReaderCapability(
|
||||
runtime=runtime_with_stub_docs,
|
||||
files=[_ai_file("doc-id", "doc.pdf")],
|
||||
reasoner=reasoner,
|
||||
)
|
||||
result = await capability._read_full_document("what is in the document?")
|
||||
|
||||
gather_mock.assert_awaited_once()
|
||||
call = gather_mock.await_args
|
||||
assert call is not None
|
||||
pages_arg = call.args[0]
|
||||
assert [p.page_number for p in pages_arg] == [1, 2]
|
||||
assert "=== doc.pdf ===" in result
|
||||
assert "fact-A" in result
|
||||
assert "[Notes from pages 1-2]" in result
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_read_full_document_iterates_multiple_files(runtime_with_stub_docs: AppRuntime) -> None:
|
||||
"""Multi-file requests run the map phase per file and return one section
|
||||
per file in the formatted output."""
|
||||
for cid, source in (("doc-a", "a.pdf"), ("doc-b", "b.pdf")):
|
||||
await runtime_with_stub_docs.documents.ingest(
|
||||
FileId(cid),
|
||||
[PageText(page_number=1, text=f"contents of {cid}")],
|
||||
source=source,
|
||||
)
|
||||
|
||||
reasoner = ChunkedReasoner(runtime_with_stub_docs)
|
||||
notes_by_call = [
|
||||
[ChunkNotes(pages=[1], summary="a-summary")],
|
||||
[ChunkNotes(pages=[1], summary="b-summary")],
|
||||
]
|
||||
|
||||
async def _gather(*_args: object, **_kwargs: object) -> list[ChunkNotes]:
|
||||
return notes_by_call.pop(0)
|
||||
|
||||
with patch.object(reasoner, "gather_notes", AsyncMock(side_effect=_gather)) as gather_mock:
|
||||
capability = WholeDocReaderCapability(
|
||||
runtime=runtime_with_stub_docs,
|
||||
files=[_ai_file("doc-a", "a.pdf"), _ai_file("doc-b", "b.pdf")],
|
||||
reasoner=reasoner,
|
||||
)
|
||||
result = await capability._read_full_document("compare them")
|
||||
|
||||
assert gather_mock.await_count == 2
|
||||
assert "=== a.pdf ===" in result
|
||||
assert "a-summary" in result
|
||||
assert "=== b.pdf ===" in result
|
||||
assert "b-summary" in result
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_read_full_document_skips_files_without_pages(runtime_with_stub_docs: AppRuntime) -> None:
|
||||
"""Files with no stored pages are quietly skipped; the tool still runs
|
||||
the map phase for files that do have pages."""
|
||||
await runtime_with_stub_docs.documents.ingest(
|
||||
FileId("present"),
|
||||
[PageText(page_number=1, text="real content")],
|
||||
source="present.pdf",
|
||||
)
|
||||
# 'missing' is never ingested -> read_pages returns [].
|
||||
|
||||
reasoner = ChunkedReasoner(runtime_with_stub_docs)
|
||||
canned = [ChunkNotes(pages=[1], summary="present summary")]
|
||||
with patch.object(reasoner, "gather_notes", AsyncMock(return_value=canned)) as gather_mock:
|
||||
capability = WholeDocReaderCapability(
|
||||
runtime=runtime_with_stub_docs,
|
||||
files=[_ai_file("missing", "missing.pdf"), _ai_file("present", "present.pdf")],
|
||||
reasoner=reasoner,
|
||||
)
|
||||
result = await capability._read_full_document("anything")
|
||||
|
||||
gather_mock.assert_awaited_once()
|
||||
assert "=== present.pdf ===" in result
|
||||
assert "missing.pdf" not in result
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_read_full_document_returns_empty_message_when_no_pages_anywhere(
|
||||
runtime_with_stub_docs: AppRuntime,
|
||||
) -> None:
|
||||
reasoner = ChunkedReasoner(runtime_with_stub_docs)
|
||||
with patch.object(reasoner, "gather_notes", AsyncMock()) as gather_mock:
|
||||
capability = WholeDocReaderCapability(
|
||||
runtime=runtime_with_stub_docs,
|
||||
files=[_ai_file("nope", "nope.pdf")],
|
||||
reasoner=reasoner,
|
||||
)
|
||||
result = await capability._read_full_document("anything")
|
||||
|
||||
gather_mock.assert_not_awaited()
|
||||
assert result == "Could not read any document content."
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_read_full_document_budget_hides_tool_when_exhausted(
|
||||
runtime_with_stub_docs: AppRuntime,
|
||||
) -> None:
|
||||
"""The prepare callback returns None once max_reads is reached so the
|
||||
agent can no longer call the tool on subsequent turns. Mirrors
|
||||
RagCapability's per-run budget."""
|
||||
await runtime_with_stub_docs.documents.ingest(
|
||||
FileId("doc-id"),
|
||||
[PageText(page_number=1, text="content")],
|
||||
source="doc.pdf",
|
||||
)
|
||||
reasoner = ChunkedReasoner(runtime_with_stub_docs)
|
||||
with patch.object(reasoner, "gather_notes", AsyncMock(return_value=[ChunkNotes(pages=[1], summary="s")])):
|
||||
capability = WholeDocReaderCapability(
|
||||
runtime=runtime_with_stub_docs,
|
||||
files=[_ai_file("doc-id", "doc.pdf")],
|
||||
reasoner=reasoner,
|
||||
max_reads=1,
|
||||
)
|
||||
sentinel: object = object()
|
||||
|
||||
# Budget intact -> prepare returns the tool.
|
||||
assert await capability._prepare_read_full_document(None, sentinel) is sentinel # type: ignore[arg-type]
|
||||
|
||||
# Spend the budget.
|
||||
await capability._read_full_document("anything")
|
||||
|
||||
# Budget spent -> prepare returns None.
|
||||
assert await capability._prepare_read_full_document(None, sentinel) is None # type: ignore[arg-type]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_instructions_mention_attached_files(runtime_with_stub_docs: AppRuntime) -> None:
|
||||
capability = WholeDocReaderCapability(
|
||||
runtime=runtime_with_stub_docs,
|
||||
files=[_ai_file("doc-a", "alpha.pdf"), _ai_file("doc-b", "beta.pdf")],
|
||||
)
|
||||
text = capability.instructions
|
||||
|
||||
assert "alpha.pdf" in text
|
||||
assert "beta.pdf" in text
|
||||
assert "read_full_document" in text
|
||||
Reference in New Issue
Block a user