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.

engine/src/stirling/api/app.py

@@ -19,13 +19,13 @@ from stirling.agents.pdf_comment import PdfCommentAgent
 from stirling.api.middleware import UserIdMiddleware
 from stirling.api.routes import (
     agent_draft_router,
+    document_router,
     execution_router,
     ledger_router,
     orchestrator_router,
     pdf_comments_router,
     pdf_edit_router,
     pdf_question_router,
-    rag_router,
 )
 from stirling.config import AppSettings, load_settings
 from stirling.contracts import HealthResponse
@@ -57,7 +57,7 @@ async def lifespan(fast_api: FastAPI):
     if tracer_provider:
         Agent.instrument_all(InstrumentationSettings(tracer_provider=tracer_provider))
     yield
-    await runtime.rag_service.close()
+    await runtime.documents.close()
     if tracer_provider:
         tracer_provider.shutdown()
 
@@ -69,7 +69,7 @@ app.include_router(pdf_edit_router)
 app.include_router(pdf_question_router)
 app.include_router(agent_draft_router)
 app.include_router(execution_router)
-app.include_router(rag_router)
+app.include_router(document_router)
 app.include_router(ledger_router)
 app.include_router(pdf_comments_router)
 

engine/src/stirling/api/dependencies.py

@@ -11,7 +11,7 @@ from stirling.agents import (
 )
 from stirling.agents.ledger import MathAuditorAgent
 from stirling.agents.pdf_comment import PdfCommentAgent
-from stirling.rag import RagService
+from stirling.documents import DocumentService
 from stirling.services import AppRuntime
 
 
@@ -39,8 +39,8 @@ def get_execution_planning_agent(request: Request) -> ExecutionPlanningAgent:
     return request.app.state.execution_planning_agent
 
 
-def get_rag_service(request: Request) -> RagService:
-    return request.app.state.runtime.rag_service
+def get_document_service(request: Request) -> DocumentService:
+    return request.app.state.runtime.documents
 
 
 def get_math_auditor_agent(request: Request) -> MathAuditorAgent:

engine/src/stirling/api/routes/__init__.py

@@ -1,19 +1,19 @@
 from .agent_drafts import router as agent_draft_router
+from .documents import router as document_router
 from .execution import router as execution_router
 from .ledger import router as ledger_router
 from .orchestrator import router as orchestrator_router
 from .pdf_comments import router as pdf_comments_router
 from .pdf_edit import router as pdf_edit_router
 from .pdf_questions import router as pdf_question_router
-from .rag import router as rag_router
 
 __all__ = [
     "agent_draft_router",
+    "document_router",
     "execution_router",
     "ledger_router",
     "orchestrator_router",
     "pdf_comments_router",
     "pdf_edit_router",
     "pdf_question_router",
-    "rag_router",
 ]

engine/src/stirling/api/routes/documents.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+from typing import Annotated
+
+from fastapi import APIRouter, Depends
+
+from stirling.api.dependencies import get_document_service
+from stirling.contracts import (
+    DeleteDocumentResponse,
+    IngestDocumentRequest,
+    IngestDocumentResponse,
+)
+from stirling.documents import DocumentService
+from stirling.models import FileId
+
+router = APIRouter(prefix="/api/v1/documents", tags=["documents"])
+
+
+@router.post("", response_model=IngestDocumentResponse)
+async def ingest_document(
+    request: IngestDocumentRequest,
+    documents: Annotated[DocumentService, Depends(get_document_service)],
+) -> IngestDocumentResponse:
+    """Replace-ingest a document's content under ``document_id``.
+
+    Stores both representations in one shot:
+      * embedded chunks for RAG search,
+      * ordered page text for whole-document reading.
+    Any previously-stored content for this document is removed first.
+    """
+    pages = request.page_text or []
+    chunks_indexed = await documents.ingest(
+        collection=request.document_id,
+        pages=pages,
+        source=request.source,
+    )
+    return IngestDocumentResponse(document_id=request.document_id, chunks_indexed=chunks_indexed)
+
+
+@router.delete("/{document_id}", response_model=DeleteDocumentResponse)
+async def delete_document(
+    document_id: FileId,
+    documents: Annotated[DocumentService, Depends(get_document_service)],
+) -> DeleteDocumentResponse:
+    """Remove a document's content. Idempotent."""
+    existed = await documents.has_collection(document_id)
+    if existed:
+        await documents.delete_collection(document_id)
+    return DeleteDocumentResponse(document_id=document_id, deleted=existed)

engine/src/stirling/api/routes/orchestrator.py

@@ -1,19 +1,170 @@
 from __future__ import annotations
 
-from typing import Annotated
+import asyncio
+import json
+import logging
+from collections.abc import AsyncIterator
+from dataclasses import dataclass
+from typing import Annotated, assert_never
 
 from fastapi import APIRouter, Depends
+from fastapi.responses import StreamingResponse
 
 from stirling.agents import OrchestratorAgent
 from stirling.api.dependencies import get_orchestrator_agent
-from stirling.contracts import OrchestratorRequest, OrchestratorResponse
+from stirling.contracts import OrchestratorRequest, OrchestratorResponse, ProgressEvent
+from stirling.services import reset_progress_emitter, set_progress_emitter
+
+logger = logging.getLogger(__name__)
+
+# Cadence for keep-alive heartbeats on the streaming endpoint. Java forwards
+# them to the frontend as SSE comments; their job is to make every layer of
+# the connection visibly alive at this rhythm so disconnects surface within a
+# bounded window instead of waiting for the next progress event.
+HEARTBEAT_INTERVAL_SECONDS = 10.0
 
 router = APIRouter(prefix="/api/v1/orchestrator", tags=["orchestrator"])
 
 
-@router.post("", response_model=OrchestratorResponse)
+@router.post("")
 async def orchestrate(
     request: OrchestratorRequest,
     agent: Annotated[OrchestratorAgent, Depends(get_orchestrator_agent)],
-) -> OrchestratorResponse:
-    return await agent.handle(request)
+) -> StreamingResponse:
+    """Run the orchestrator and stream NDJSON events.
+
+    Each output line is a JSON object with an ``event`` field. ``progress``
+    events arrive whenever an inner agent reports work (e.g. each
+    chunked-reasoner slice completing); the final ``result`` event carries the
+    typed orchestrator response. ``error`` events surface failures without
+    breaking the connection. ``heartbeat`` events fire on a fixed cadence to
+    keep idle connections visibly alive so disconnects propagate.
+
+    The stream itself is the liveness signal: as long as events flow, work is
+    alive. Java consumes this with a long total timeout and treats line
+    arrival as forward progress.
+    """
+    return StreamingResponse(
+        _OrchestratorStream(
+            agent=agent,
+            request=request,
+            heartbeat_interval_seconds=HEARTBEAT_INTERVAL_SECONDS,
+        ).iterate(),
+        media_type="application/x-ndjson",
+    )
+
+
+@dataclass(frozen=True, slots=True)
+class _ProgressFrame:
+    event: ProgressEvent
+
+
+@dataclass(frozen=True, slots=True)
+class _ResultFrame:
+    response: OrchestratorResponse
+
+
+@dataclass(frozen=True, slots=True)
+class _ErrorFrame:
+    message: str
+
+
+@dataclass(frozen=True, slots=True)
+class _HeartbeatFrame:
+    """No payload: a heartbeat exists only to push bytes through the pipe.
+
+    Without periodic traffic, a slow workflow phase (e.g. all extractor
+    workers busy on long calls) leaves the engine writer, Java's SSE
+    forwarder, and the frontend's fetch all silently waiting. A closed
+    connection at any layer wouldn't surface until the next real event,
+    which could be many tens of seconds away. Heartbeats bound that window
+    to :data:`HEARTBEAT_INTERVAL_SECONDS`.
+    """
+
+
+type _StreamFrame = _ProgressFrame | _ResultFrame | _ErrorFrame | _HeartbeatFrame
+
+
+def _serialize_frame(frame: _StreamFrame) -> bytes:
+    """Render a frame as one NDJSON line."""
+    match frame:
+        case _ProgressFrame(event=event):
+            body = {"event": "progress", **event.model_dump(mode="json")}
+        case _ResultFrame(response=response):
+            body = {"event": "result", "response": response.model_dump(mode="json")}
+        case _ErrorFrame(message=message):
+            body = {"event": "error", "message": message}
+        case _HeartbeatFrame():
+            body = {"event": "heartbeat"}
+        case _:
+            assert_never(frame)
+    return (json.dumps(body) + "\n").encode("utf-8")
+
+
+class _OrchestratorStream:
+    """Drives one streaming orchestrator request.
+
+    Owns the per-request queue and pumps progress events through it; the agent
+    runs as a child task so its emissions and the streaming response interleave.
+    A heartbeat task pushes keep-alive messages onto the same queue at a fixed
+    cadence so the connection stays visibly alive between progress events.
+    """
+
+    def __init__(
+        self,
+        *,
+        agent: OrchestratorAgent,
+        request: OrchestratorRequest,
+        heartbeat_interval_seconds: float,
+    ) -> None:
+        self._agent = agent
+        self._request = request
+        self._heartbeat_interval_seconds = heartbeat_interval_seconds
+        self._queue: asyncio.Queue[_StreamFrame | None] = asyncio.Queue()
+
+    async def iterate(self) -> AsyncIterator[bytes]:
+        token = set_progress_emitter(self._emit_progress)
+        agent_task = asyncio.create_task(self._run_agent())
+        heartbeat_task = asyncio.create_task(self._emit_heartbeats())
+        try:
+            while True:
+                frame = await self._queue.get()
+                if frame is None:
+                    break
+                yield _serialize_frame(frame)
+        finally:
+            reset_progress_emitter(token)
+            await self._cancel_task(heartbeat_task)
+            await self._cancel_task(agent_task)
+
+    async def _emit_progress(self, event: ProgressEvent) -> None:
+        await self._queue.put(_ProgressFrame(event=event))
+
+    async def _emit_heartbeats(self) -> None:
+        while True:
+            await asyncio.sleep(self._heartbeat_interval_seconds)
+            await self._queue.put(_HeartbeatFrame())
+
+    async def _run_agent(self) -> None:
+        try:
+            response = await self._agent.handle(self._request)
+            await self._queue.put(_ResultFrame(response=response))
+        except asyncio.CancelledError:
+            raise
+        except Exception as exc:
+            logger.exception("orchestrator stream failed")
+            await self._queue.put(_ErrorFrame(message=str(exc)))
+        finally:
+            await self._queue.put(None)
+
+    @staticmethod
+    async def _cancel_task(task: asyncio.Task[None]) -> None:
+        if task.done():
+            return
+        task.cancel()
+        try:
+            await task
+        except asyncio.CancelledError:
+            pass
+        except Exception:
+            logger.exception("background task failed during cancellation", exc_info=True)

engine/src/stirling/api/routes/rag.py

@@ -1,63 +0,0 @@
-from __future__ import annotations
-
-from typing import Annotated
-
-from fastapi import APIRouter, Depends
-
-from stirling.api.dependencies import get_rag_service
-from stirling.contracts import (
-    DeleteDocumentResponse,
-    IngestDocumentRequest,
-    IngestDocumentResponse,
-    PdfContentType,
-)
-from stirling.models import FileId
-from stirling.rag import Document, RagService
-
-router = APIRouter(prefix="/api/v1/rag", tags=["rag"])
-
-
-@router.post("/documents", response_model=IngestDocumentResponse)
-async def ingest_document(
-    request: IngestDocumentRequest,
-    rag: Annotated[RagService, Depends(get_rag_service)],
-) -> IngestDocumentResponse:
-    """Replace-ingest a document's content under ``document_id``.
-
-    Any previously-stored content for this document is removed and the
-    provided content replaces it wholesale. All pages are chunked up front
-    and then embedded in a single batched call so large documents (e.g. a
-    500-page book) don't fan out into hundreds of embedding requests.
-    """
-    await rag.delete_collection(request.document_id)
-
-    chunks: list[Document] = []
-    if request.page_text:
-        for page in request.page_text:
-            if not page.text.strip():
-                continue
-            chunks.extend(
-                rag.chunk_text(
-                    text=page.text,
-                    source=f"{request.source}:page:{page.page_number}",
-                    base_metadata={
-                        "page_number": str(page.page_number),
-                        "content_type": PdfContentType.PAGE_TEXT.value,
-                    },
-                )
-            )
-
-    indexed = await rag.index_documents(request.document_id, chunks) if chunks else 0
-    return IngestDocumentResponse(document_id=request.document_id, chunks_indexed=indexed)
-
-
-@router.delete("/documents/{document_id}", response_model=DeleteDocumentResponse)
-async def delete_document(
-    document_id: FileId,
-    rag: Annotated[RagService, Depends(get_rag_service)],
-) -> DeleteDocumentResponse:
-    """Remove a document's content from RAG. Idempotent."""
-    existed = await rag.has_collection(document_id)
-    if existed:
-        await rag.delete_collection(document_id)
-    return DeleteDocumentResponse(document_id=document_id, deleted=existed)