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
2026-05-14 13:19:38 +00:00
committed by GitHub
parent 8abe734f0b
commit 672e81d286
55 changed files with 3327 additions and 578 deletions
@@ -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
+4
View File
@@ -31,6 +31,10 @@ def build_app_settings() -> AppSettings:
rag_chunk_overlap=64,
rag_default_top_k=5,
rag_max_searches=5,
chunked_reasoner_chars_per_slice=16_000,
chunked_reasoner_concurrency=10,
chunked_reasoner_notes_char_budget=250_000,
chunked_reasoner_worker_timeout_seconds=60.0,
max_pages=200,
max_characters=200_000,
posthog_enabled=False,
+4
View File
@@ -43,6 +43,10 @@ class StubSettingsProvider:
rag_chunk_overlap=64,
rag_default_top_k=5,
rag_max_searches=5,
chunked_reasoner_chars_per_slice=16_000,
chunked_reasoner_concurrency=10,
chunked_reasoner_worker_timeout_seconds=60.0,
chunked_reasoner_notes_char_budget=250_000,
max_pages=100,
max_characters=100_000,
posthog_enabled=False,
@@ -2,14 +2,15 @@ from __future__ import annotations
import pytest
from stirling.contracts import PageText
from stirling.documents.chunker import chunk_text
from stirling.documents.rag_capability import RagCapability
from stirling.documents.service import DocumentService
from stirling.documents.sqlite_vec_store import SqliteVecStore
from stirling.documents.store import Document, SearchResult
from stirling.models import FileId
from stirling.rag.capability import RagCapability
from stirling.rag.chunker import chunk_text
from stirling.rag.service import RagService
from stirling.rag.sqlite_vec_store import SqliteVecStore
from stirling.rag.store import Document, SearchResult
# ── chunk_text ──────────────────────────────────────────────────────────
# chunk_text
class TestChunkText:
@@ -26,7 +27,6 @@ class TestChunkText:
def test_splits_on_paragraph_boundaries(self) -> None:
text = "First paragraph.\n\nSecond paragraph.\n\nThird paragraph."
chunks = chunk_text(text, chunk_size=30, overlap=0)
# Each paragraph fits in 30 chars, so they should be split
assert len(chunks) >= 2
assert "First paragraph." in chunks[0]
@@ -35,22 +35,19 @@ class TestChunkText:
chunks = chunk_text(text, chunk_size=100, overlap=10)
assert len(chunks) > 1
for chunk in chunks:
# Chunks may slightly exceed due to sentence boundary snapping
assert len(chunk) <= 200 # generous upper bound
assert len(chunk) <= 200
def test_overlap_produces_shared_content(self) -> None:
sentences = [f"Sentence number {i}." for i in range(20)]
text = " ".join(sentences)
chunks = chunk_text(text, chunk_size=100, overlap=30)
if len(chunks) >= 2:
# After word-boundary snapping, the second chunk should share
# some content with the tail of the first chunk
words_in_first_tail = chunks[0].split()[-3:] # last 3 words
words_in_first_tail = chunks[0].split()[-3:]
overlap_text = " ".join(words_in_first_tail)
assert overlap_text in chunks[1], f"Expected overlap '{overlap_text}' in chunk[1]: '{chunks[1][:80]}...'"
# ── SqliteVecStore ──────────────────────────────────────────────────────
# SqliteVecStore
class TestSqliteVecStore:
@@ -59,12 +56,12 @@ class TestSqliteVecStore:
@pytest.mark.anyio
async def test_add_and_search(self) -> None:
store = SqliteVecStore.ephemeral()
await store.ensure_collection("test-col", "test.pdf")
docs = [
Document(id="1", text="Python is a programming language", metadata={"source": "test"}),
Document(id="2", text="Java is another programming language", metadata={"source": "test"}),
Document(id="3", text="The weather today is sunny", metadata={"source": "test"}),
]
# Simple 3-dimensional embeddings for testing
embeddings = [
[1.0, 0.0, 0.0],
[0.9, 0.1, 0.0],
@@ -72,17 +69,16 @@ class TestSqliteVecStore:
]
await store.add_documents("test-col", docs, embeddings)
# Search with a query close to the programming-related docs
results = await store.search("test-col", [1.0, 0.05, 0.0], top_k=2)
assert len(results) == 2
assert isinstance(results[0], SearchResult)
# The closest should be doc "1" (exact match on first dimension)
assert results[0].document.id == "1"
assert results[0].score > 0.5
@pytest.mark.anyio
async def test_list_and_has_collection(self) -> None:
store = SqliteVecStore.ephemeral()
await store.ensure_collection("my-collection", "test.pdf")
docs = [Document(id="1", text="test", metadata={})]
await store.add_documents("my-collection", docs, [[1.0, 0.0]])
@@ -94,6 +90,7 @@ class TestSqliteVecStore:
@pytest.mark.anyio
async def test_delete_collection(self) -> None:
store = SqliteVecStore.ephemeral()
await store.ensure_collection("to-delete", "test.pdf")
docs = [Document(id="1", text="test", metadata={})]
await store.add_documents("to-delete", docs, [[1.0]])
@@ -104,6 +101,7 @@ class TestSqliteVecStore:
@pytest.mark.anyio
async def test_search_empty_collection(self) -> None:
store = SqliteVecStore.ephemeral()
await store.ensure_collection("empty-test", "test.pdf")
docs = [Document(id="1", text="test", metadata={})]
await store.add_documents("empty-test", docs, [[1.0, 0.0]])
results = await store.search("empty-test", [1.0, 0.0], top_k=5)
@@ -117,7 +115,7 @@ class TestSqliteVecStore:
await store.add_documents("bad", docs, [[1.0], [2.0]])
# ── RagService (with stub embedder) ────────────────────────────────────
# DocumentService (with stub embedder)
class StubEmbeddingService:
@@ -127,7 +125,6 @@ class StubEmbeddingService:
self._dim = dim
async def embed_query(self, text: str) -> list[float]:
# Deterministic embedding based on hash of text
h = hash(text) % 1000
return [(h + i) / 1000.0 for i in range(self._dim)]
@@ -140,8 +137,6 @@ class StubEmbeddingService:
source: str = "",
base_metadata: dict[str, str] | None = None,
) -> list[Document]:
from stirling.rag.chunker import chunk_text
chunks = chunk_text(text, 100, 10)
docs = []
for i, chunk in enumerate(chunks):
@@ -154,53 +149,113 @@ class StubEmbeddingService:
@pytest.fixture
def rag_service() -> RagService:
"""Each RagService test gets its own fresh ephemeral store to avoid dimension conflicts."""
def documents() -> DocumentService:
"""Each DocumentService test gets its own fresh ephemeral store to avoid dimension conflicts."""
store = SqliteVecStore.ephemeral()
return RagService(embedder=StubEmbeddingService(), store=store, default_top_k=3) # type: ignore[arg-type]
return DocumentService(embedder=StubEmbeddingService(), store=store, default_top_k=3) # type: ignore[arg-type]
class TestRagService:
def _pages(text: str) -> list[PageText]:
return [PageText(page_number=1, text=text)]
class TestDocumentService:
@pytest.mark.anyio
async def test_index_and_search(self, rag_service: RagService) -> None:
async def test_ingest_and_search(self, documents: DocumentService) -> None:
text = "Python is great for data science. It has many libraries like pandas and numpy."
count = await rag_service.index_text(FileId("docs"), text, source="guide.pdf")
count = await documents.ingest(FileId("docs"), _pages(text), source="guide.pdf")
assert count > 0
results = await rag_service.search("Python libraries", collection=FileId("docs"))
results = await documents.search("Python libraries", collection=FileId("docs"))
assert len(results) > 0
assert results[0].document.text # non-empty text
assert results[0].document.text
@pytest.mark.anyio
async def test_index_empty_text_returns_zero(self, rag_service: RagService) -> None:
count = await rag_service.index_text(FileId("docs"), "", source="empty.pdf")
async def test_ingest_empty_text_returns_zero_chunks(self, documents: DocumentService) -> None:
count = await documents.ingest(FileId("docs"), _pages(""), source="empty.pdf")
assert count == 0
@pytest.mark.anyio
async def test_search_nonexistent_collection_returns_empty(self, rag_service: RagService) -> None:
results = await rag_service.search("anything", collection=FileId("nonexistent"))
async def test_search_nonexistent_collection_returns_empty(self, documents: DocumentService) -> None:
results = await documents.search("anything", collection=FileId("nonexistent"))
assert results == []
@pytest.mark.anyio
async def test_search_all_collections(self, rag_service: RagService) -> None:
await rag_service.index_text(FileId("col-a"), "Machine learning overview.", source="ml.pdf")
await rag_service.index_text(FileId("col-b"), "Deep learning with neural networks.", source="dl.pdf")
async def test_search_all_collections(self, documents: DocumentService) -> None:
await documents.ingest(FileId("col-a"), _pages("Machine learning overview."), source="ml.pdf")
await documents.ingest(FileId("col-b"), _pages("Deep learning with neural networks."), source="dl.pdf")
results = await rag_service.search("neural networks")
results = await documents.search("neural networks")
assert len(results) > 0
@pytest.mark.anyio
async def test_delete_collection(self, rag_service: RagService) -> None:
await rag_service.index_text(FileId("temp"), "Temporary data.", source="tmp.pdf")
collections = await rag_service.list_collections()
async def test_delete_collection(self, documents: DocumentService) -> None:
await documents.ingest(FileId("temp"), _pages("Temporary data."), source="tmp.pdf")
collections = await documents.list_collections()
assert "temp" in collections
await rag_service.delete_collection(FileId("temp"))
collections = await rag_service.list_collections()
await documents.delete_collection(FileId("temp"))
collections = await documents.list_collections()
assert "temp" not in collections
@pytest.mark.anyio
async def test_ingest_stores_pages_in_order(self, documents: DocumentService) -> None:
pages = [
PageText(page_number=1, text="First page text."),
PageText(page_number=2, text="Second page text."),
PageText(page_number=3, text="Third page text."),
]
await documents.ingest(FileId("ordered"), pages, source="ordered.pdf")
# ── RagCapability ──────────────────────────────────────────────────────
stored = await documents.read_pages(FileId("ordered"))
assert [p.page_number for p in stored] == [1, 2, 3]
assert stored[0].text == "First page text."
assert stored[0].char_count == len("First page text.")
@pytest.mark.anyio
async def test_read_pages_with_range(self, documents: DocumentService) -> None:
from stirling.contracts import PageRange
pages = [PageText(page_number=i, text=f"page {i}") for i in range(1, 6)]
await documents.ingest(FileId("ranged"), pages, source="r.pdf")
subset = await documents.read_pages(FileId("ranged"), PageRange(start=2, end=4))
assert [p.page_number for p in subset] == [2, 3, 4]
@pytest.mark.anyio
async def test_ingest_replaces_previous_pages(self, documents: DocumentService) -> None:
await documents.ingest(
FileId("doc"),
[PageText(page_number=1, text="old"), PageText(page_number=2, text="old2")],
source="v1.pdf",
)
await documents.ingest(
FileId("doc"),
[PageText(page_number=1, text="new")],
source="v2.pdf",
)
stored = await documents.read_pages(FileId("doc"))
assert [p.page_number for p in stored] == [1]
assert stored[0].text == "new"
@pytest.mark.anyio
async def test_ingest_keeps_blank_pages_in_page_store(self, documents: DocumentService) -> None:
"""Blank pages are skipped for embedding but retained in the page store
so page numbering stays continuous when reading back."""
pages = [
PageText(page_number=1, text="Real text on page 1."),
PageText(page_number=2, text=" "),
PageText(page_number=3, text="Real text on page 3."),
]
await documents.ingest(FileId("with-blanks"), pages, source="blanks.pdf")
stored = await documents.read_pages(FileId("with-blanks"))
assert [p.page_number for p in stored] == [1, 2, 3]
assert stored[1].text.strip() == ""
# RagCapability
async def _invoke_search_knowledge(capability: RagCapability, query: str, max_results: int = 5) -> str:
@@ -210,27 +265,27 @@ async def _invoke_search_knowledge(capability: RagCapability, query: str, max_re
toolset = capability.toolset
assert isinstance(toolset, FunctionToolset)
tool = toolset.tools["search_knowledge"]
return await tool.function(query=query, max_results=max_results) # type: ignore[call-arg] — pyright can't infer the generic tool function's kwargs
return await tool.function(query=query, max_results=max_results) # type: ignore[call-arg]
class TestRagCapability:
def test_instructions_static_when_collections_pinned(self, rag_service: RagService) -> None:
cap = RagCapability(rag_service, collections=[FileId("docs"), FileId("manuals")])
def test_instructions_static_when_collections_pinned(self, documents: DocumentService) -> None:
cap = RagCapability(documents, collections=[FileId("docs"), FileId("manuals")])
instructions = cap.instructions
assert isinstance(instructions, str)
assert "docs, manuals" in instructions
assert "search_knowledge" in instructions
def test_instructions_dynamic_when_no_collections(self, rag_service: RagService) -> None:
cap = RagCapability(rag_service)
def test_instructions_dynamic_when_no_collections(self, documents: DocumentService) -> None:
cap = RagCapability(documents)
instructions = cap.instructions
assert callable(instructions)
@pytest.mark.anyio
async def test_dynamic_instructions_list_available_collections(self, rag_service: RagService) -> None:
await rag_service.index_text(FileId("col-a"), "Alpha content.", source="a.pdf")
await rag_service.index_text(FileId("col-b"), "Beta content.", source="b.pdf")
cap = RagCapability(rag_service)
async def test_dynamic_instructions_list_available_collections(self, documents: DocumentService) -> None:
await documents.ingest(FileId("col-a"), _pages("Alpha content."), source="a.pdf")
await documents.ingest(FileId("col-b"), _pages("Beta content."), source="b.pdf")
cap = RagCapability(documents)
instructions_fn = cap.instructions
assert callable(instructions_fn)
text = await instructions_fn()
@@ -238,23 +293,23 @@ class TestRagCapability:
assert "col-b" in text
@pytest.mark.anyio
async def test_dynamic_instructions_when_store_empty(self, rag_service: RagService) -> None:
cap = RagCapability(rag_service)
async def test_dynamic_instructions_when_store_empty(self, documents: DocumentService) -> None:
cap = RagCapability(documents)
instructions_fn = cap.instructions
assert callable(instructions_fn)
text = await instructions_fn()
assert "empty" in text.lower()
@pytest.mark.anyio
async def test_search_knowledge_returns_no_results_message_when_empty(self, rag_service: RagService) -> None:
cap = RagCapability(rag_service)
async def test_search_knowledge_returns_no_results_message_when_empty(self, documents: DocumentService) -> None:
cap = RagCapability(documents)
output = await _invoke_search_knowledge(cap, "anything")
assert output == "No relevant results found in the knowledge base."
@pytest.mark.anyio
async def test_search_knowledge_formats_results_with_source_and_score(self, rag_service: RagService) -> None:
await rag_service.index_text(FileId("docs"), "Python is a programming language.", source="guide.pdf")
cap = RagCapability(rag_service)
async def test_search_knowledge_formats_results_with_source_and_score(self, documents: DocumentService) -> None:
await documents.ingest(FileId("docs"), _pages("Python is a programming language."), source="guide.pdf")
cap = RagCapability(documents)
output = await _invoke_search_knowledge(cap, "Python")
assert "[Result 1" in output
assert "source: guide.pdf" in output
@@ -262,43 +317,39 @@ class TestRagCapability:
assert "relevance:" in output
@pytest.mark.anyio
async def test_search_knowledge_restricts_to_pinned_collections(self, rag_service: RagService) -> None:
await rag_service.index_text(FileId("pinned"), "Pinned collection content.", source="pinned.pdf")
await rag_service.index_text(FileId("other"), "Content in another collection.", source="other.pdf")
async def test_search_knowledge_restricts_to_pinned_collections(self, documents: DocumentService) -> None:
await documents.ingest(FileId("pinned"), _pages("Pinned collection content."), source="pinned.pdf")
await documents.ingest(FileId("other"), _pages("Content in another collection."), source="other.pdf")
cap = RagCapability(rag_service, collections=[FileId("pinned")])
cap = RagCapability(documents, collections=[FileId("pinned")])
output = await _invoke_search_knowledge(cap, "content")
assert "pinned.pdf" in output
assert "other.pdf" not in output
@pytest.mark.anyio
async def test_search_knowledge_respects_max_results(self, rag_service: RagService) -> None:
async def test_search_knowledge_respects_max_results(self, documents: DocumentService) -> None:
paragraphs = "\n\n".join(f"Paragraph {i} about topic." for i in range(10))
await rag_service.index_text(FileId("bulk"), paragraphs, source="bulk.pdf")
await documents.ingest(FileId("bulk"), _pages(paragraphs), source="bulk.pdf")
cap = RagCapability(rag_service)
cap = RagCapability(documents)
output = await _invoke_search_knowledge(cap, "topic", max_results=2)
# Only two results requested, so only Result 1 and Result 2 should appear
assert "[Result 1" in output
assert "[Result 2" in output
assert "[Result 3" not in output
@pytest.mark.anyio
async def test_search_knowledge_tool_is_hidden_after_budget_exhausted(self, rag_service: RagService) -> None:
async def test_search_knowledge_tool_is_hidden_after_budget_exhausted(self, documents: DocumentService) -> None:
"""The prepare callback must return None once max_searches has been reached
so the agent can no longer call the tool on subsequent turns."""
await rag_service.index_text(FileId("docs"), "Some content.", source="x.pdf")
cap = RagCapability(rag_service, max_searches=2)
await documents.ingest(FileId("docs"), _pages("Some content."), source="x.pdf")
cap = RagCapability(documents, max_searches=2)
tool_def = _dummy_tool_def()
# Budget intact: prepare returns the tool definition.
assert await cap._prepare_search_knowledge(None, tool_def) is tool_def # type: ignore[arg-type]
# Use the budget.
await _invoke_search_knowledge(cap, "content")
await _invoke_search_knowledge(cap, "content")
# Budget spent: prepare returns None, removing the tool from the agent's next turn.
assert await cap._prepare_search_knowledge(None, tool_def) is None # type: ignore[arg-type]
@@ -6,9 +6,9 @@ import pytest
from fastapi.testclient import TestClient
from stirling.api import app
from stirling.api.dependencies import get_rag_service
from stirling.api.dependencies import get_document_service
from stirling.documents import Document, DocumentService, SqliteVecStore
from stirling.models import FileId
from stirling.rag import Document, RagService, SqliteVecStore
class StubEmbedder:
@@ -30,7 +30,7 @@ class StubEmbedder:
source: str = "",
base_metadata: dict[str, str] | None = None,
) -> list[Document]:
from stirling.rag.chunker import chunk_text
from stirling.documents.chunker import chunk_text
chunks = chunk_text(text, 100, 10)
docs = []
@@ -43,8 +43,8 @@ class StubEmbedder:
return docs
def _build_service() -> RagService:
return RagService(
def _build_service() -> DocumentService:
return DocumentService(
embedder=StubEmbedder(), # type: ignore[arg-type]
store=SqliteVecStore.ephemeral(),
default_top_k=3,
@@ -52,25 +52,25 @@ def _build_service() -> RagService:
@pytest.fixture
def service() -> RagService:
def service() -> DocumentService:
return _build_service()
@pytest.fixture
def client(service: RagService) -> Iterator[TestClient]:
app.dependency_overrides[get_rag_service] = lambda: service
def client(service: DocumentService) -> Iterator[TestClient]:
app.dependency_overrides[get_document_service] = lambda: service
try:
yield TestClient(app)
finally:
app.dependency_overrides.pop(get_rag_service, None)
app.dependency_overrides.pop(get_document_service, None)
# ── POST /documents ─────────────────────────────────────────────────────
def test_ingest_document_indexes_page_text(client: TestClient, service: RagService) -> None:
def test_ingest_document_indexes_page_text(client: TestClient, service: DocumentService) -> None:
response = client.post(
"/api/v1/rag/documents",
"/api/v1/documents",
json={
"documentId": "doc-123",
"source": "report.pdf",
@@ -87,9 +87,9 @@ def test_ingest_document_indexes_page_text(client: TestClient, service: RagServi
@pytest.mark.anyio
async def test_ingest_document_replaces_existing_content(client: TestClient, service: RagService) -> None:
async def test_ingest_document_replaces_existing_content(client: TestClient, service: DocumentService) -> None:
client.post(
"/api/v1/rag/documents",
"/api/v1/documents",
json={
"documentId": "replace-me",
"source": "replace-me.pdf",
@@ -98,7 +98,7 @@ async def test_ingest_document_replaces_existing_content(client: TestClient, ser
)
# Second ingest with different content should replace the first entirely
response = client.post(
"/api/v1/rag/documents",
"/api/v1/documents",
json={
"documentId": "replace-me",
"source": "replace-me.pdf",
@@ -115,7 +115,7 @@ async def test_ingest_document_replaces_existing_content(client: TestClient, ser
def test_ingest_document_skips_empty_pages(client: TestClient) -> None:
response = client.post(
"/api/v1/rag/documents",
"/api/v1/documents",
json={
"documentId": "mixed",
"source": "mixed.pdf",
@@ -130,14 +130,14 @@ def test_ingest_document_skips_empty_pages(client: TestClient) -> None:
def test_ingest_document_with_no_content_returns_zero(client: TestClient) -> None:
response = client.post("/api/v1/rag/documents", json={"documentId": "empty", "source": "empty.pdf"})
response = client.post("/api/v1/documents", json={"documentId": "empty", "source": "empty.pdf"})
assert response.status_code == 200
assert response.json()["chunksIndexed"] == 0
def test_ingest_document_rejects_empty_id(client: TestClient) -> None:
response = client.post(
"/api/v1/rag/documents",
"/api/v1/documents",
json={"documentId": "", "source": "x.pdf", "pageText": [{"pageNumber": 1, "text": "something"}]},
)
assert response.status_code == 422
@@ -145,7 +145,7 @@ def test_ingest_document_rejects_empty_id(client: TestClient) -> None:
def test_ingest_document_rejects_missing_source(client: TestClient) -> None:
response = client.post(
"/api/v1/rag/documents",
"/api/v1/documents",
json={"documentId": "doc-1", "pageText": [{"pageNumber": 1, "text": "something"}]},
)
assert response.status_code == 422
@@ -153,7 +153,7 @@ def test_ingest_document_rejects_missing_source(client: TestClient) -> None:
def test_ingest_document_rejects_empty_source(client: TestClient) -> None:
response = client.post(
"/api/v1/rag/documents",
"/api/v1/documents",
json={"documentId": "doc-1", "source": "", "pageText": [{"pageNumber": 1, "text": "something"}]},
)
assert response.status_code == 422
@@ -161,7 +161,7 @@ def test_ingest_document_rejects_empty_source(client: TestClient) -> None:
def test_ingest_document_rejects_non_positive_page_number(client: TestClient) -> None:
response = client.post(
"/api/v1/rag/documents",
"/api/v1/documents",
json={
"documentId": "bad-page",
"source": "bad-page.pdf",
@@ -176,30 +176,30 @@ def test_ingest_document_rejects_non_positive_page_number(client: TestClient) ->
def test_delete_document_reports_deleted_true_when_existed(client: TestClient) -> None:
client.post(
"/api/v1/rag/documents",
"/api/v1/documents",
json={
"documentId": "to-delete",
"source": "to-delete.pdf",
"pageText": [{"pageNumber": 1, "text": "Text."}],
},
)
response = client.delete("/api/v1/rag/documents/to-delete")
response = client.delete("/api/v1/documents/to-delete")
assert response.status_code == 200
assert response.json() == {"documentId": "to-delete", "deleted": True}
def test_delete_document_is_idempotent(client: TestClient) -> None:
response = client.delete("/api/v1/rag/documents/never-existed")
response = client.delete("/api/v1/documents/never-existed")
assert response.status_code == 200
assert response.json() == {"documentId": "never-existed", "deleted": False}
@pytest.mark.anyio
async def test_delete_document_removes_collection(client: TestClient, service: RagService) -> None:
async def test_delete_document_removes_collection(client: TestClient, service: DocumentService) -> None:
client.post(
"/api/v1/rag/documents",
"/api/v1/documents",
json={"documentId": "gone", "source": "gone.pdf", "pageText": [{"pageNumber": 1, "text": "Text."}]},
)
assert await service.has_collection(FileId("gone"))
client.delete("/api/v1/rag/documents/gone")
client.delete("/api/v1/documents/gone")
assert not await service.has_collection(FileId("gone"))
+15 -14
View File
@@ -9,6 +9,7 @@ from stirling.contracts import (
AiFile,
ExtractedFileText,
NeedIngestResponse,
PageText,
PdfContentType,
PdfQuestionAnswerResponse,
PdfQuestionNotFoundResponse,
@@ -17,8 +18,8 @@ from stirling.contracts import (
PdfTextSelection,
SupportedCapability,
)
from stirling.documents import Document, DocumentService, SqliteVecStore
from stirling.models import FileId
from stirling.rag import Document, RagService, SqliteVecStore
from stirling.services.runtime import AppRuntime
@@ -41,7 +42,7 @@ class StubEmbedder:
source: str = "",
base_metadata: dict[str, str] | None = None,
) -> list[Document]:
from stirling.rag.chunker import chunk_text
from stirling.documents.chunker import chunk_text
chunks = chunk_text(text, 100, 10)
docs: list[Document] = []
@@ -65,13 +66,13 @@ class StubPdfQuestionAgent(PdfQuestionAgent):
@pytest.fixture
def runtime_with_stub_rag(runtime: AppRuntime) -> AppRuntime:
"""A runtime whose RAG service uses a stub embedder + ephemeral store."""
stub = RagService(
"""A runtime whose document service uses a stub embedder + ephemeral store."""
stub = DocumentService(
embedder=StubEmbedder(), # type: ignore[arg-type]
store=SqliteVecStore.ephemeral(),
default_top_k=runtime.settings.rag_default_top_k,
)
return replace(runtime, rag_service=stub)
return replace(runtime, documents=stub)
@pytest.mark.anyio
@@ -89,9 +90,9 @@ async def test_requests_ingest_when_file_missing_from_rag(runtime_with_stub_rag:
@pytest.mark.anyio
async def test_reports_only_missing_files(runtime_with_stub_rag: AppRuntime) -> None:
await runtime_with_stub_rag.rag_service.index_text(
collection=FileId("present-id"),
text="Invoice total: 120.00.",
await runtime_with_stub_rag.documents.ingest(
FileId("present-id"),
[PageText(page_number=1, text="Invoice total: 120.00.")],
source="present.pdf",
)
agent = PdfQuestionAgent(runtime_with_stub_rag)
@@ -106,9 +107,9 @@ async def test_reports_only_missing_files(runtime_with_stub_rag: AppRuntime) ->
@pytest.mark.anyio
async def test_returns_grounded_answer_when_all_files_ingested(runtime_with_stub_rag: AppRuntime) -> None:
await runtime_with_stub_rag.rag_service.index_text(
collection=FileId("invoice-id"),
text="Invoice total: 120.00.",
await runtime_with_stub_rag.documents.ingest(
FileId("invoice-id"),
[PageText(page_number=1, text="Invoice total: 120.00.")],
source="invoice.pdf",
)
agent = StubPdfQuestionAgent(
@@ -137,9 +138,9 @@ async def test_returns_grounded_answer_when_all_files_ingested(runtime_with_stub
@pytest.mark.anyio
async def test_returns_not_found_when_answer_not_in_doc(runtime_with_stub_rag: AppRuntime) -> None:
await runtime_with_stub_rag.rag_service.index_text(
collection=FileId("shipping-id"),
text="This page contains only a shipping address.",
await runtime_with_stub_rag.documents.ingest(
FileId("shipping-id"),
[PageText(page_number=1, text="This page contains only a shipping address.")],
source="shipping.pdf",
)
agent = StubPdfQuestionAgent(
+115 -5
View File
@@ -1,3 +1,7 @@
import asyncio
import json
from unittest.mock import patch
from conftest import build_app_settings
from fastapi.testclient import TestClient
@@ -25,8 +29,11 @@ from stirling.contracts import (
PdfQuestionNotFoundResponse,
PdfQuestionRequest,
SupportedCapability,
WholeDocReadStarted,
WholeDocSliceDone,
)
from stirling.models.tool_models import Angle, RotatePdfParams
from stirling.services import emit_progress
class StubOrchestratorAgent:
@@ -40,6 +47,34 @@ class StubOrchestratorAgent:
)
class StubProgressOrchestratorAgent:
"""Orchestrator stub that emits two progress events before returning.
Used to verify the streaming endpoint plumbs the ContextVar emitter through
to deep callees and forwards events as NDJSON in order.
"""
async def handle(self, request: OrchestratorRequest) -> NeedContentResponse:
await emit_progress(WholeDocReadStarted(question="x", pages=10, slices=2))
await emit_progress(
WholeDocSliceDone(
completed=1,
total=2,
pages="pages=1-5",
duration_ms=42,
excerpts=2,
facts=3,
)
)
return NeedContentResponse(
resume_with=SupportedCapability.PDF_QUESTION,
reason=request.user_message,
files=[],
max_pages=1,
max_characters=1000,
)
class StubPdfEditAgent:
async def handle(self, request: PdfEditRequest) -> EditCannotDoResponse:
return EditCannotDoResponse(reason=request.user_message)
@@ -87,14 +122,89 @@ def test_health_route() -> None:
assert response.json()["status"] == "ok"
def test_orchestrator_route() -> None:
response = client.post(
def test_orchestrator_route_streams_result_only_when_no_progress() -> None:
"""The orchestrator endpoint always streams NDJSON. An agent that emits no
progress events still produces a single ``result`` frame with the typed
response body."""
with client.stream(
"POST",
"/api/v1/orchestrator",
json={"userMessage": "route this", "files": [{"id": "test-id", "name": "test.pdf"}]},
)
) as response:
assert response.status_code == 200
events = [json.loads(line) for line in response.iter_lines() if line]
assert response.status_code == 200
assert response.json()["outcome"] == "need_content"
assert [e["event"] for e in events] == ["result"]
body = events[0]["response"]
assert body["outcome"] == "need_content"
def test_orchestrator_route_streams_progress_then_result() -> None:
"""When an agent emits progress via the ContextVar emitter, those frames
arrive on the wire before the final result frame."""
app.dependency_overrides[get_orchestrator_agent] = lambda: StubProgressOrchestratorAgent()
try:
with client.stream(
"POST",
"/api/v1/orchestrator",
json={"userMessage": "stream this", "files": [{"id": "test-id", "name": "test.pdf"}]},
) as response:
assert response.status_code == 200
events = [json.loads(line) for line in response.iter_lines() if line]
finally:
app.dependency_overrides[get_orchestrator_agent] = lambda: StubOrchestratorAgent()
progress = [e for e in events if e["event"] == "progress"]
results = [e for e in events if e["event"] == "result"]
assert [p["phase"] for p in progress] == ["whole_doc_read_started", "whole_doc_slice_done"]
assert progress[1]["completed"] == 1
assert len(results) == 1
response = results[0]["response"]
assert response["outcome"] == "need_content"
# Wire format must be camelCase: Java's Jackson deserializer expects camelCase
# field names. ``maxPages`` here doubles as a regression guard against the
# snake_case bug that surfaced as "need_ingest without listing any files to ingest".
assert "maxPages" in response
assert "max_pages" not in response
def test_orchestrator_route_emits_heartbeats_while_agent_is_busy() -> None:
"""While the agent is in flight, the streaming endpoint emits heartbeat
frames at the configured cadence so each layer of the connection stays
visibly alive and disconnects propagate within bounded latency."""
class _SlowAgent:
async def handle(self, _request: OrchestratorRequest) -> NeedContentResponse:
# Sleep long enough for several heartbeats at the patched cadence.
await asyncio.sleep(0.2)
return NeedContentResponse(
resume_with=SupportedCapability.PDF_QUESTION,
reason="ok",
files=[],
max_pages=1,
max_characters=1000,
)
app.dependency_overrides[get_orchestrator_agent] = lambda: _SlowAgent()
try:
with patch("stirling.api.routes.orchestrator.HEARTBEAT_INTERVAL_SECONDS", 0.03):
with client.stream(
"POST",
"/api/v1/orchestrator",
json={"userMessage": "wait", "files": [{"id": "test-id", "name": "test.pdf"}]},
) as response:
assert response.status_code == 200
events = [json.loads(line) for line in response.iter_lines() if line]
finally:
app.dependency_overrides[get_orchestrator_agent] = lambda: StubOrchestratorAgent()
heartbeats = [e for e in events if e["event"] == "heartbeat"]
results = [e for e in events if e["event"] == "result"]
# At least a couple of heartbeats fired during the 0.2s agent sleep at 0.03s cadence.
assert len(heartbeats) >= 2
# The result still arrives after the agent finishes.
assert len(results) == 1
assert results[0]["response"]["outcome"] == "need_content"
def test_pdf_edit_route() -> None:
+4
View File
@@ -92,6 +92,10 @@ def test_app_settings_accepts_model_configuration() -> None:
rag_chunk_overlap=64,
rag_default_top_k=5,
rag_max_searches=5,
chunked_reasoner_chars_per_slice=16_000,
chunked_reasoner_concurrency=10,
chunked_reasoner_worker_timeout_seconds=60.0,
chunked_reasoner_notes_char_budget=250_000,
max_pages=200,
max_characters=200_000,
posthog_enabled=False,