mirror of
https://github.com/arsvendg/Stirling-PDF.git
synced 2026-07-16 03:20:46 +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:
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user