wehub-resource-sync
241 lines
8.8 KiB
Python
241 lines
8.8 KiB
Python
"""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, OwnerId, PrincipalId
|
|
from stirling.services.runtime import AppRuntime
|
|
|
|
OWNER = OwnerId("test-user")
|
|
PRINCIPALS = [PrincipalId("test-user")]
|
|
|
|
|
|
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", owner_id=OWNER, read_principals=PRINCIPALS, expires_at=None
|
|
)
|
|
|
|
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")],
|
|
principals=PRINCIPALS,
|
|
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,
|
|
owner_id=OWNER,
|
|
read_principals=PRINCIPALS,
|
|
expires_at=None,
|
|
)
|
|
|
|
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")],
|
|
principals=PRINCIPALS,
|
|
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",
|
|
owner_id=OWNER,
|
|
read_principals=PRINCIPALS,
|
|
expires_at=None,
|
|
)
|
|
# '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")],
|
|
principals=PRINCIPALS,
|
|
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")],
|
|
principals=PRINCIPALS,
|
|
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",
|
|
owner_id=OWNER,
|
|
read_principals=PRINCIPALS,
|
|
expires_at=None,
|
|
)
|
|
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")],
|
|
principals=PRINCIPALS,
|
|
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")],
|
|
principals=PRINCIPALS,
|
|
)
|
|
text = capability.instructions
|
|
|
|
assert "alpha.pdf" in text
|
|
assert "beta.pdf" in text
|
|
assert "read_full_document" in text
|