arsvendg/Stirling-PDF
1
0
Fork 0
mirror of https://github.com/arsvendg/Stirling-PDF.git synced 2026-07-14 10:34:06 +02:00
Code Issues Actions 4 Packages Projects Releases 1 Wiki Activity

Flesh out RAG system (#6197)

Browse Source 
# Description of Changes
Flesh out the RAG system and connect it to the PDF Question Agent so it
can respond to questions about PDFs of an extremely large size.

I'd expect lots more work will need to be done to finish off the RAG
system to really be what we need, but this should be a reasonable start
which will let us connect it to tools and have the ingestion mostly
handled automatically. I'm leaving file deletion and proper file ID
management to be done in a future PR. We also need to consider whether
all tools should retrieve content exclusively via RAG, or whether it's
beneficial to have tools sometimes fetch the direct content and other
times fetch it from RAG.

A diagram of the expected interaction is as follows:

```mermaid
sequenceDiagram
    autonumber
    actor U as User
    participant FE as Frontend<br/>(ChatPanel)
    participant J as Java<br/>(AiWorkflowService)
    participant O as Engine:<br/>OrchestratorAgent
    participant QA as Engine:<br/>PdfQuestionAgent
    participant RAG as Engine:<br/>RagService + SqliteVecStore
    participant V as VoyageAI<br/>(embeddings)
    participant L as LLM<br/>(Claude / etc.)

    U->>FE: types "Summarise this PDF"<br/>(PDF already uploaded)
    FE->>J: POST /api/v1/ai/orchestrate/stream<br/>multipart: fileInputs[], userMessage
    Note over J: ByteHashFileIdStrategy<br/>id = sha256(bytes)[:16]
    J->>O: POST /api/v1/orchestrator<br/>{ files:[{id,name}], userMessage }

    O->>L: route via fast model
    L-->>O: delegate_pdf_question
    O->>QA: PdfQuestionRequest

    loop for each file
        QA->>RAG: has_collection(file.id)
        RAG-->>QA: false
    end
    QA-->>O: NeedIngestResponse(files_to_ingest)
    O-->>J: { outcome:"need_ingest", filesToIngest:[...] }

    Note over J: onNeedIngest
    loop per file
        J->>J: PDFBox: extract page text
        J->>O: POST /api/v1/rag/documents<br/>(long-running timeout)
        O->>RAG: chunk + stage documents
        O->>V: embed_documents (batches of 256)
        V-->>O: embeddings
        O->>RAG: add_documents
        O-->>J: { chunks_indexed: N }
    end

    Note over J: retry with resumeWith=pdf_question
    J->>O: POST /api/v1/orchestrator
    Note over O: fast-path to PdfQuestionAgent

    O->>QA: PdfQuestionRequest
    Note over QA: build RagCapability<br/>pinned to file IDs
    QA->>L: run(prompt) with search_knowledge tool

    loop up to max_searches
        L->>QA: search_knowledge(query)
        QA->>V: embed_query
        V-->>QA: query vector
        QA->>RAG: search(vector, collections=[file.id])
        RAG-->>QA: top-k chunks
        QA-->>L: formatted chunks
    end

    Note over QA: once budget spent,<br/>prepare() hides the tool
    L-->>QA: PdfQuestionAnswerResponse
    QA-->>O: answer
    O-->>J: { outcome:"answer", answer, evidence }
    J-->>FE: SSE "result"
    FE->>U: assistant bubble
```
This commit is contained in:
James Brunton
2026-05-01 14:11:54 +01:00
committed by GitHub
parent 5605062153
commit 5541dd666c
48 changed files with 1067 additions and 534 deletions
Download Patch File Download Diff File Expand all files Collapse all files
engine/src/stirling/agents/orchestrator.py
+9 -8
View File
@@ -18,10 +18,11 @@ from stirling.contracts import (
OrchestratorRequest,
OrchestratorResponse,
PdfEditResponse,
PdfQuestionResponse,
PdfQuestionOrchestrateResponse,
SupportedCapability,
UnsupportedCapabilityResponse,
format_conversation_history,
format_file_names,
)
from stirling.contracts.pdf_edit import EditPlanResponse
from stirling.services import AppRuntime
@@ -78,12 +79,13 @@ class OrchestratorAgent:
"You are the top-level orchestrator. "
"Choose exactly one output function that best handles the request. "
"Use delegate_pdf_edit for requested modifications of single or multiple PDFs. "
"Use delegate_pdf_question for questions about PDF contents. "
"Use delegate_pdf_question for questions about the contents of the attached PDFs. "
"Use delegate_user_spec for requests to create or define an agent spec. "
"Use delegate_pdf_review when the user wants the PDF returned with review"
" comments attached — anything like 'review this', 'annotate with comments',"
" 'leave feedback on the PDF'. "
"Use unsupported_capability only when none of the other outputs fit."
"Use unsupported_capability when the user asks about the assistant itself "
"or when none of the other outputs fit; supply a helpful message."
),
model_settings=runtime.fast_model_settings,
)
@@ -91,7 +93,7 @@ class OrchestratorAgent:
async def handle(self, request: OrchestratorRequest) -> OrchestratorResponse:
logger.info(
"[orchestrator] handle: files=%s resume_with=%s artifacts=%s msg=%r",
request.file_names,
[file.name for file in request.files],
request.resume_with,
[type(a).__name__ for a in request.artifacts],
request.user_message,
@@ -137,10 +139,10 @@ class OrchestratorAgent:
async def _run_pdf_edit(self, request: OrchestratorRequest) -> PdfEditResponse:
return await PdfEditAgent(self.runtime).orchestrate(request)
async def delegate_pdf_question(self, ctx: RunContext[OrchestratorDeps]) -> PdfQuestionResponse:
async def delegate_pdf_question(self, ctx: RunContext[OrchestratorDeps]) -> PdfQuestionOrchestrateResponse:
return await self._run_pdf_question(ctx.deps.request)
async def _run_pdf_question(self, request: OrchestratorRequest) -> PdfQuestionResponse:
async def _run_pdf_question(self, request: OrchestratorRequest) -> PdfQuestionOrchestrateResponse:
return await PdfQuestionAgent(self.runtime).orchestrate(request)
async def delegate_user_spec(self, ctx: RunContext[OrchestratorDeps]) -> AgentDraftWorkflowResponse:
@@ -165,12 +167,11 @@ class OrchestratorAgent:
def _build_prompt(self, request: OrchestratorRequest) -> str:
artifact_summary = self._describe_artifacts(request)
file_names = ", ".join(request.file_names) if request.file_names else "Unknown files"
history = format_conversation_history(request.conversation_history)
return (
f"Conversation history:\n{history}\n"
f"User message: {request.user_message}\n"
f"Files: {file_names}\n"
f"Files: {format_file_names(request.files)}\n"
f"Available artifacts:\n{artifact_summary}"
)
engine/src/stirling/agents/pdf_edit.py
+6 -8
View File
@@ -22,6 +22,7 @@ from stirling.contracts import (
SupportedCapability,
ToolOperationStep,
format_conversation_history,
format_file_names,
)
from stirling.logging import Pretty
from stirling.models import OPERATIONS, ApiModel, ParamToolModel, ToolEndpoint
@@ -116,7 +117,6 @@ class PdfEditParameterSelector:
) -> str:
operation_id = operation_plan[operation_index]
operation_list = ", ".join(operation.name for operation in operation_plan)
file_names = ", ".join(request.file_names) if request.file_names else "No file names were provided."
generated_steps_text = (
"\n".join(
f"- Step {step_index + 1}: {step.model_dump_json()}" for step_index, step in enumerate(generated_steps)
@@ -127,7 +127,7 @@ class PdfEditParameterSelector:
return (
f"Conversation history:\n{format_conversation_history(request.conversation_history)}\n"
f"User request: {request.user_message}\n"
f"Files: {file_names}\n"
f"Files: {format_file_names(request.files)}\n"
f"Operation plan: {operation_list}\n"
f"Selected operation index: {operation_index + 1} of {len(operation_plan)}\n"
f"Selected operation: {operation_id.name}\n"
@@ -153,7 +153,7 @@ class PdfEditAgent:
return await self.handle(
PdfEditRequest(
user_message=request.user_message,
file_names=request.file_names,
files=request.files,
conversation_history=request.conversation_history,
page_text=extracted_text.files if extracted_text is not None else [],
)
@@ -166,7 +166,7 @@ class PdfEditAgent:
async def handle(self, request: PdfEditRequest, allow_need_content: bool = True) -> PdfEditResponse:
logger.info(
"[pdf-edit] handle: files=%s has_text=%s allow_need_content=%s msg=%r",
request.file_names,
[file.name for file in request.files],
has_page_text(request.page_text),
allow_need_content,
request.user_message,
@@ -225,11 +225,10 @@ class PdfEditAgent:
)
def _build_selection_prompt(self, request: PdfEditRequest) -> str:
file_names = ", ".join(request.file_names) if request.file_names else "No file names were provided."
return (
f"Conversation history:\n{format_conversation_history(request.conversation_history)}\n"
f"User request: {request.user_message}\n"
f"Files: {file_names}\n"
f"Files: {format_file_names(request.files)}\n"
f"Supported operations: {self._supported_operations_prompt()}\n"
f"Extracted page text:\n{format_page_text(request.page_text)}"
)
@@ -243,8 +242,7 @@ class PdfEditAgent:
request: PdfEditRequest,
) -> NeedContentResponse:
files = selection.files or [
NeedContentFileRequest(file_name=file_name, content_types=[PdfContentType.PAGE_TEXT])
for file_name in request.file_names
NeedContentFileRequest(file=file, content_types=[PdfContentType.PAGE_TEXT]) for file in request.files
]
return NeedContentResponse(
resume_with=SupportedCapability.PDF_EDIT,
engine/src/stirling/agents/pdf_questions.py
+100 -69
View File
@@ -1,32 +1,67 @@
from __future__ import annotations
import logging
from pydantic_ai import Agent
from pydantic_ai.output import NativeOutput
from stirling.agents._page_text import (
format_page_text,
get_extracted_text_artifact,
has_page_text,
)
from stirling.agents.math_presentation import MathIntentClassifier, extract_math_verdict
from stirling.contracts import (
AiFile,
EditPlanResponse,
NeedContentFileRequest,
NeedContentResponse,
NeedIngestResponse,
OrchestratorRequest,
PdfContentType,
PdfQuestionAnswerResponse,
PdfQuestionNotFoundResponse,
PdfQuestionOrchestrateResponse,
PdfQuestionRequest,
PdfQuestionResponse,
PdfQuestionTerminalResponse,
SupportedCapability,
ToolOperationStep,
Verdict,
format_conversation_history,
format_file_names,
)
from stirling.models.agent_tool_models import AgentToolId, MathAuditorAgentParams
from stirling.rag import RagCapability
from stirling.services import AppRuntime
logger = logging.getLogger(__name__)
PDF_QUESTION_SYSTEM_PROMPT = (
"You answer questions about PDF documents by retrieving relevant content with the "
"search_knowledge tool. Use it before answering. Do not guess or use outside knowledge.\n"
"\n"
"The search_knowledge tool has a finite call budget per run. When it is no longer "
"available, answer from what you have already retrieved.\n"
"\n"
"Guidelines:\n"
"- Make targeted search_knowledge calls. Typically one or two is enough.\n"
"- Answer from the retrieved text. If the retrieved content doesn't support a confident "
"answer, return not_found.\n"
"- For questions that would require reading the entire document end-to-end (e.g. "
"'what's the shortest chapter', 'how many X are there'), return not_found.\n"
"- Include a short list of evidence snippets (with page numbers where available) drawn "
"from what search_knowledge returned.\n"
"\n"
"Writing the not_found reason:\n"
"- The reason is shown directly to the end user, so write it in plain, friendly "
"language. One or two short sentences.\n"
"- NEVER mention 'RAG', 'retrieval', 'chunks', 'search results', 'targeted search', "
"'search_knowledge', or other implementation details.\n"
"- Be honest about the actual limitation. For questions that require full-document "
"analysis (shortest chapter, word counts, etc.), explain that the document is too "
"long to analyse end-to-end: you can only look up specific passages, and that's "
"not enough to compare every part of the document against every other.\n"
"- For questions where the answer just isn't in the document, say so directly: "
"'I couldn't find that information in the document.'\n"
"- Do not make it sound like you're choosing not to answer. Be clear that it's "
"a genuine constraint."
)
_MATH_SYNTH_SYSTEM_PROMPT = (
"You are given a math-audit Verdict (structured JSON) and the user's "
"original question. Answer the question in plain prose using only "
@@ -41,26 +76,6 @@ _MATH_SYNTH_SYSTEM_PROMPT = (
class PdfQuestionAgent:
def __init__(self, runtime: AppRuntime) -> None:
self.runtime = runtime
rag = runtime.rag_capability
self.agent = Agent(
model=runtime.smart_model,
output_type=NativeOutput(
[
PdfQuestionAnswerResponse,
PdfQuestionNotFoundResponse,
]
),
system_prompt=(
"Answer questions about PDFs using only the extracted page text provided in the prompt. "
"Do not guess or use outside knowledge. "
"If the answer is not supported by the provided text, return not_found. "
"When answering, include a short list of evidence snippets with their page numbers. "
"Reply in the SAME LANGUAGE as the question."
),
instructions=rag.instructions,
toolsets=[rag.toolset],
model_settings=runtime.smart_model_settings,
)
self._math_synth_agent: Agent[None, str] = Agent(
model=runtime.fast_model,
output_type=str,
@@ -70,30 +85,31 @@ class PdfQuestionAgent:
self._math_intent_classifier = MathIntentClassifier(runtime)
async def handle(self, request: PdfQuestionRequest) -> PdfQuestionResponse:
if not has_page_text(request.page_text):
return NeedContentResponse(
logger.info(
"[pdf-question] handle: files=%s question=%r",
[file.name for file in request.files],
request.question,
)
missing = await self._find_missing_files(request.files)
if missing:
logger.info("[pdf-question] missing ingestions: %s", [file.name for file in missing])
return NeedIngestResponse(
resume_with=SupportedCapability.PDF_QUESTION,
reason="No extracted PDF page text was provided, so the question cannot be answered yet.",
files=[
NeedContentFileRequest(
file_name=file_name,
content_types=[PdfContentType.PAGE_TEXT],
)
for file_name in request.file_names
],
max_pages=self.runtime.settings.max_pages,
max_characters=self.runtime.settings.max_characters,
reason="Some files have not been ingested into RAG yet.",
files_to_ingest=missing,
content_types=[PdfContentType.PAGE_TEXT],
)
return await self._run_answer_agent(request)
async def orchestrate(self, request: OrchestratorRequest) -> PdfQuestionResponse:
async def orchestrate(self, request: OrchestratorRequest) -> PdfQuestionOrchestrateResponse:
"""Entry point for the orchestrator delegate.
Decides math intent locally via a small classifier LLM (language-agnostic).
On a math first turn, embeds an :class:`EditPlanResponse` in the answer
response; on the resume turn, digests the captured :class:`Verdict` into
a localised prose answer. Non-math first turns fall through to the
text-grounded :meth:`handle` pipeline.
On a math first turn, returns an :class:`EditPlanResponse` (``outcome=PLAN``)
with ``resume_with=PDF_QUESTION`` so the caller runs the math specialist
and re-invokes the orchestrator. On the resume turn, the captured
:class:`Verdict` is digested into a localised prose answer. Non-math
first turns fall through to the text-grounded :meth:`handle` pipeline.
"""
verdict = extract_math_verdict(request)
if verdict is not None:
@@ -104,36 +120,53 @@ class PdfQuestionAgent:
return PdfQuestionAnswerResponse(answer=answer, evidence=[])
if await self._math_intent_classifier.classify(request.user_message):
# First turn — ask the caller to run the math specialist and come back.
# The plan rides on the answer response as a nullable member; ``answer``
# is empty on this turn and the caller resumes once the plan is run.
return PdfQuestionAnswerResponse(
answer="",
evidence=[],
edit_plan=EditPlanResponse(
summary="",
steps=[
ToolOperationStep(
tool=AgentToolId.MATH_AUDITOR_AGENT,
parameters=MathAuditorAgentParams(),
)
],
resume_with=SupportedCapability.PDF_QUESTION,
),
# First turn — emit a one-step plan calling the math specialist,
# with resume_with set so the caller comes back with the verdict
# in artifacts (handled by the resume branch above).
return EditPlanResponse(
summary="",
steps=[
ToolOperationStep(
tool=AgentToolId.MATH_AUDITOR_AGENT,
parameters=MathAuditorAgentParams(),
)
],
resume_with=SupportedCapability.PDF_QUESTION,
)
extracted_text = get_extracted_text_artifact(request)
return await self.handle(
PdfQuestionRequest(
question=request.user_message,
file_names=request.file_names,
page_text=extracted_text.files if extracted_text is not None else [],
files=request.files,
conversation_history=request.conversation_history,
)
)
async def _run_answer_agent(self, request: PdfQuestionRequest) -> PdfQuestionResponse:
result = await self.agent.run(self._build_prompt(request))
async def _find_missing_files(self, files: list[AiFile]) -> list[AiFile]:
missing: list[AiFile] = []
for file in files:
if not await self.runtime.rag_service.has_collection(file.id):
missing.append(file)
return missing
async def _run_answer_agent(self, request: PdfQuestionRequest) -> PdfQuestionTerminalResponse:
rag = RagCapability(
rag_service=self.runtime.rag_service,
collections=[file.id for file in request.files],
top_k=self.runtime.settings.rag_default_top_k,
max_searches=self.runtime.settings.rag_max_searches,
)
agent = Agent(
model=self.runtime.smart_model,
output_type=NativeOutput([PdfQuestionAnswerResponse, PdfQuestionNotFoundResponse]),
system_prompt=PDF_QUESTION_SYSTEM_PROMPT,
instructions=rag.instructions,
toolsets=[rag.toolset],
model_settings=self.runtime.smart_model_settings,
)
prompt = self._build_prompt(request)
logger.debug("[pdf-question] prompt:\n%s", prompt)
result = await agent.run(prompt)
return result.output
async def _synthesise_math_answer(self, user_message: str, verdict: Verdict) -> str:
@@ -146,12 +179,10 @@ class PdfQuestionAgent:
return result.output
def _build_prompt(self, request: PdfQuestionRequest) -> str:
file_names = ", ".join(request.file_names) if request.file_names else "Unknown files"
pages = format_page_text(request.page_text, empty="")
history = format_conversation_history(request.conversation_history)
return (
f"Conversation history:\n{history}\n"
f"Files: {file_names}\n"
f"Files: {format_file_names(request.files)}\n"
f"Question: {request.question}\n"
f"Extracted page text:\n{pages}"
"Use search_knowledge to retrieve the relevant content, then answer."
)
engine/src/stirling/api/app.py
+8 -1
View File
@@ -7,7 +7,13 @@ from fastapi import Depends, FastAPI
from pydantic_ai import Agent
from pydantic_ai.models.instrumented import InstrumentationSettings
from stirling.agents import ExecutionPlanningAgent, OrchestratorAgent, PdfEditAgent, PdfQuestionAgent, UserSpecAgent
from stirling.agents import (
ExecutionPlanningAgent,
OrchestratorAgent,
PdfEditAgent,
PdfQuestionAgent,
UserSpecAgent,
)
from stirling.agents.ledger import MathAuditorAgent
from stirling.agents.pdf_comment import PdfCommentAgent
from stirling.api.middleware import UserIdMiddleware
@@ -51,6 +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()
if tracer_provider:
tracer_provider.shutdown()
engine/src/stirling/api/dependencies.py
+7 -5
View File
@@ -2,7 +2,13 @@ from __future__ import annotations
from fastapi import Request
from stirling.agents import ExecutionPlanningAgent, OrchestratorAgent, PdfEditAgent, PdfQuestionAgent, UserSpecAgent
from stirling.agents import (
ExecutionPlanningAgent,
OrchestratorAgent,
PdfEditAgent,
PdfQuestionAgent,
UserSpecAgent,
)
from stirling.agents.ledger import MathAuditorAgent
from stirling.agents.pdf_comment import PdfCommentAgent
from stirling.rag import RagService
@@ -37,10 +43,6 @@ def get_rag_service(request: Request) -> RagService:
return request.app.state.runtime.rag_service
def get_rag_embedding_model(request: Request) -> str:
return request.app.state.runtime.settings.rag_embedding_model
def get_math_auditor_agent(request: Request) -> MathAuditorAgent:
return request.app.state.math_auditor_agent
engine/src/stirling/api/routes/rag.py
+47 -62
View File
@@ -4,75 +4,60 @@ from typing import Annotated
from fastapi import APIRouter, Depends
from stirling.api.dependencies import get_rag_embedding_model, get_rag_service
from stirling.api.dependencies import get_rag_service
from stirling.contracts import (
RagCollectionsResponse,
RagDeleteCollectionResponse,
RagIndexRequest,
RagIndexResponse,
RagSearchRequest,
RagSearchResponse,
RagSearchResultItem,
RagStatusResponse,
DeleteDocumentResponse,
IngestDocumentRequest,
IngestDocumentResponse,
PdfContentType,
)
from stirling.rag import RagService
from stirling.models import FileId
from stirling.rag import Document, RagService
router = APIRouter(prefix="/api/v1/rag", tags=["rag"])
@router.get("/status", response_model=RagStatusResponse)
async def rag_status(
@router.post("/documents", response_model=IngestDocumentResponse)
async def ingest_document(
request: IngestDocumentRequest,
rag: Annotated[RagService, Depends(get_rag_service)],
embedding_model: Annotated[str, Depends(get_rag_embedding_model)],
) -> RagStatusResponse:
collections = await rag.list_collections()
return RagStatusResponse(embedding_model=embedding_model, collections=collections)
) -> 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.post("/index", response_model=RagIndexResponse)
async def rag_index(
request: RagIndexRequest,
@router.delete("/documents/{document_id}", response_model=DeleteDocumentResponse)
async def delete_document(
document_id: FileId,
rag: Annotated[RagService, Depends(get_rag_service)],
) -> RagIndexResponse:
count = await rag.index_text(
collection=request.collection,
text=request.text,
source=request.source,
metadata=request.metadata,
)
return RagIndexResponse(collection=request.collection, chunks_indexed=count)
@router.post("/search", response_model=RagSearchResponse)
async def rag_search(
request: RagSearchRequest,
rag: Annotated[RagService, Depends(get_rag_service)],
) -> RagSearchResponse:
results = await rag.search(query=request.query, collection=request.collection, top_k=request.top_k)
items = [
RagSearchResultItem(
text=r.document.text,
source=r.document.metadata.get("source", ""),
chunk_id=r.document.metadata.get("chunk_index", ""),
score=r.score,
)
for r in results
]
return RagSearchResponse(query=request.query, results=items)
@router.get("/collections", response_model=RagCollectionsResponse)
async def rag_collections(
rag: Annotated[RagService, Depends(get_rag_service)],
) -> RagCollectionsResponse:
collections = await rag.list_collections()
return RagCollectionsResponse(collections=collections)
@router.delete("/collections/{name}", response_model=RagDeleteCollectionResponse)
async def rag_delete_collection(
name: str,
rag: Annotated[RagService, Depends(get_rag_service)],
) -> RagDeleteCollectionResponse:
await rag.delete_collection(name)
return RagDeleteCollectionResponse(status="deleted", collection=name)
) -> 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)
engine/src/stirling/config/settings.py
+1
View File
@@ -36,6 +36,7 @@ class AppSettings(BaseSettings):
rag_chunk_size: int = Field(validation_alias="STIRLING_RAG_CHUNK_SIZE")
rag_chunk_overlap: int = Field(validation_alias="STIRLING_RAG_CHUNK_OVERLAP")
rag_default_top_k: int = Field(validation_alias="STIRLING_RAG_TOP_K")
rag_max_searches: int = Field(validation_alias="STIRLING_RAG_MAX_SEARCHES")
max_pages: int = Field(validation_alias="STIRLING_MAX_PAGES")
max_characters: int = Field(validation_alias="STIRLING_MAX_CHARACTERS")
engine/src/stirling/contracts/__init__.py
+16 -18
View File
@@ -10,12 +10,14 @@ from .agent_drafts import (
from .agent_specs import AgentSpec, AgentSpecStep, AiToolAgentStep
from .comments import CommentSpec
from .common import (
AiFile,
ArtifactKind,
ConversationMessage,
ExtractedFileText,
MathAuditorToolReportArtifact,
NeedContentFileRequest,
NeedContentResponse,
NeedIngestResponse,
PdfContentType,
PdfTextSelection,
StepKind,
@@ -24,6 +26,7 @@ from .common import (
ToolReportArtifact,
WorkflowOutcome,
format_conversation_history,
format_file_names,
)
from .execution import (
AgentExecutionRequest,
@@ -71,24 +74,20 @@ from .pdf_edit import (
from .pdf_questions import (
PdfQuestionAnswerResponse,
PdfQuestionNotFoundResponse,
PdfQuestionOrchestrateResponse,
PdfQuestionRequest,
PdfQuestionResponse,
PdfQuestionTerminalResponse,
)
from .rag import (
MAX_INDEX_TEXT_LENGTH,
RagCollectionsResponse,
RagDeleteCollectionResponse,
RagIndexRequest,
RagIndexResponse,
RagSearchRequest,
RagSearchResponse,
RagSearchResultItem,
RagStatusResponse,
DeleteDocumentResponse,
IngestDocumentRequest,
IngestDocumentResponse,
IngestedPageText,
)
__all__ = [
"MAX_INDEX_TEXT_LENGTH",
"AiFile",
"AgentDraft",
"AgentDraftRequest",
"AgentDraftResponse",
@@ -105,6 +104,7 @@ __all__ = [
"CommentSpec",
"CompletedExecutionAction",
"ConversationMessage",
"DeleteDocumentResponse",
"Discrepancy",
"DiscrepancyKind",
"EditCannotDoResponse",
@@ -119,10 +119,15 @@ __all__ = [
"FolioManifest",
"FolioType",
"format_conversation_history",
"format_file_names",
"HealthResponse",
"IngestDocumentRequest",
"IngestDocumentResponse",
"IngestedPageText",
"MathAuditorToolReportArtifact",
"NeedContentFileRequest",
"NeedContentResponse",
"NeedIngestResponse",
"NextExecutionAction",
"OrchestratorRequest",
"OrchestratorResponse",
@@ -136,18 +141,11 @@ __all__ = [
"PdfEditTerminalResponse",
"PdfQuestionAnswerResponse",
"PdfQuestionNotFoundResponse",
"PdfQuestionOrchestrateResponse",
"PdfQuestionRequest",
"PdfQuestionResponse",
"PdfQuestionTerminalResponse",
"PdfTextSelection",
"RagCollectionsResponse",
"RagDeleteCollectionResponse",
"RagIndexRequest",
"RagIndexResponse",
"RagSearchRequest",
"RagSearchResponse",
"RagSearchResultItem",
"RagStatusResponse",
"Requisition",
"Severity",
"StepKind",
engine/src/stirling/contracts/common.py
+35 -2
View File
@@ -6,7 +6,7 @@ from typing import Literal, assert_never
from pydantic import Field, model_validator
from stirling.contracts.ledger import Verdict
from stirling.models import OPERATIONS, ApiModel, ToolEndpoint
from stirling.models import OPERATIONS, ApiModel, FileId, ToolEndpoint
from stirling.models.agent_tool_models import AGENT_OPERATIONS, AgentToolId, AnyParamModel, AnyToolId
@@ -50,6 +50,7 @@ class WorkflowOutcome(StrEnum):
ANSWER = "answer"
NEED_CONTENT = "need_content"
NEED_INGEST = "need_ingest"
NOT_FOUND = "not_found"
PLAN = "plan"
NEED_CLARIFICATION = "need_clarification"
@@ -94,12 +95,30 @@ class ConversationMessage(ApiModel):
content: str
class AiFile(ApiModel):
"""A file the user has supplied, identified by both a stable id and a display name.
The id is opaque to the engine: Java generates it (content hash, file path, UUID, etc.)
and the engine uses it as the RAG collection key for any agent that indexes content.
The name is used in user-facing prompts and responses.
"""
id: FileId = Field(min_length=1)
name: str = Field(min_length=1)
def format_conversation_history(conversation_history: list[ConversationMessage]) -> str:
if not conversation_history:
return "None"
return "\n".join(f"- {message.role}: {message.content}" for message in conversation_history)
def format_file_names(files: list[AiFile]) -> str:
if not files:
return "No file names were provided."
return ", ".join(file.name for file in files)
class PdfTextSelection(ApiModel):
page_number: int | None = None
text: str
@@ -111,7 +130,7 @@ class ExtractedFileText(ApiModel):
class NeedContentFileRequest(ApiModel):
file_name: str
file: AiFile
page_numbers: list[int] = Field(default_factory=list)
content_types: list[PdfContentType]
@@ -146,6 +165,20 @@ class MathAuditorToolReportArtifact(ApiModel):
ToolReportArtifact = MathAuditorToolReportArtifact
class NeedIngestResponse(ApiModel):
"""Signal that the listed files must be ingested into RAG before the agent can continue.
Java's handling: for each file, extract the requested content types, POST to
``/api/v1/rag/documents`` keyed by ``file.id``, then retry the original request.
"""
outcome: Literal[WorkflowOutcome.NEED_INGEST] = WorkflowOutcome.NEED_INGEST
resume_with: SupportedCapability
reason: str
files_to_ingest: list[AiFile]
content_types: list[PdfContentType] = Field(default_factory=list)
class ToolOperationStep(ApiModel):
kind: Literal[StepKind.TOOL] = StepKind.TOOL
tool: AnyToolId
engine/src/stirling/contracts/orchestrator.py
+4 -1
View File
@@ -8,10 +8,12 @@ from stirling.models import ApiModel
from .agent_drafts import AgentDraftResponse
from .common import (
AiFile,
ArtifactKind,
ConversationMessage,
ExtractedFileText,
NeedContentResponse,
NeedIngestResponse,
SupportedCapability,
ToolReportArtifact,
WorkflowOutcome,
@@ -31,7 +33,7 @@ WorkflowArtifact = Annotated[ExtractedTextArtifact | ToolReportArtifact, Field(d
class OrchestratorRequest(ApiModel):
user_message: str
file_names: list[str]
files: list[AiFile] = Field(default_factory=list)
conversation_history: list[ConversationMessage] = Field(default_factory=list)
artifacts: list[WorkflowArtifact] = Field(default_factory=list)
resume_with: SupportedCapability | None = None
@@ -47,6 +49,7 @@ type OrchestratorResponse = Annotated[
PdfEditTerminalResponse
| PdfQuestionTerminalResponse
| NeedContentResponse
| NeedIngestResponse
| AgentDraftResponse
| NextExecutionAction
| UnsupportedCapabilityResponse,
engine/src/stirling/contracts/pdf_edit.py
+2 -1
View File
@@ -7,6 +7,7 @@ from pydantic import Field
from stirling.models import ApiModel
from .common import (
AiFile,
ConversationMessage,
ExtractedFileText,
NeedContentResponse,
@@ -18,7 +19,7 @@ from .common import (
class PdfEditRequest(ApiModel):
user_message: str
file_names: list[str] = Field(default_factory=list)
files: list[AiFile] = Field(default_factory=list)
conversation_history: list[ConversationMessage] = Field(default_factory=list)
page_text: list[ExtractedFileText] = Field(default_factory=list)
engine/src/stirling/contracts/pdf_questions.py
+12 -13
View File
@@ -7,9 +7,10 @@ from pydantic import Field
from stirling.models import ApiModel
from .common import (
AiFile,
ConversationMessage,
ExtractedFileText,
NeedContentResponse,
NeedIngestResponse,
WorkflowOutcome,
)
from .pdf_edit import EditPlanResponse
@@ -17,8 +18,7 @@ from .pdf_edit import EditPlanResponse
class PdfQuestionRequest(ApiModel):
question: str
page_text: list[ExtractedFileText] = Field(default_factory=list)
file_names: list[str]
files: list[AiFile] = Field(default_factory=list)
conversation_history: list[ConversationMessage] = Field(default_factory=list)
@@ -26,15 +26,6 @@ class PdfQuestionAnswerResponse(ApiModel):
outcome: Literal[WorkflowOutcome.ANSWER] = WorkflowOutcome.ANSWER
answer: str
evidence: list[ExtractedFileText] = Field(default_factory=list)
edit_plan: EditPlanResponse | None = Field(
default=None,
description=(
"Optional plan the caller must run before the answer is final. When"
" populated, ``answer`` is empty on this turn — the caller executes"
" the plan and re-invokes the orchestrator with ``resume_with`` set"
" to PDF_QUESTION; the real answer arrives on the resume turn."
),
)
class PdfQuestionNotFoundResponse(ApiModel):
@@ -44,6 +35,14 @@ class PdfQuestionNotFoundResponse(ApiModel):
type PdfQuestionTerminalResponse = PdfQuestionAnswerResponse | PdfQuestionNotFoundResponse
type PdfQuestionResponse = Annotated[
PdfQuestionTerminalResponse | NeedContentResponse,
PdfQuestionTerminalResponse | NeedIngestResponse,
Field(discriminator="outcome"),
]
# ``orchestrate`` may also emit an ``EditPlanResponse`` on the math-routing
# first turn (``outcome=PLAN`` with ``resume_with=PDF_QUESTION``). It's not in
# ``PdfQuestionTerminalResponse`` because that alias would otherwise duplicate
# the PLAN branch already provided by ``PdfEditTerminalResponse`` in the
# top-level :class:`OrchestratorResponse` discriminated union.
type PdfQuestionOrchestrateResponse = PdfQuestionResponse | EditPlanResponse
engine/src/stirling/contracts/rag.py
+24 -36
View File
@@ -4,48 +4,36 @@ from pydantic import Field
from stirling.models import ApiModel
MAX_INDEX_TEXT_LENGTH = 1_000_000 # 1MB text limit per index request
from .common import FileId
class RagStatusResponse(ApiModel):
embedding_model: str
collections: list[str]
class IngestedPageText(ApiModel):
page_number: int = Field(ge=1)
text: str
class RagIndexRequest(ApiModel):
collection: str = Field(min_length=1)
text: str = Field(max_length=MAX_INDEX_TEXT_LENGTH)
source: str = ""
metadata: dict[str, str] = Field(default_factory=dict)
class IngestDocumentRequest(ApiModel):
"""Replace-ingest a document's content into RAG under the given document_id.
Each content-type field is optional; the endpoint replaces the document's entire
stored content with whatever is provided. To add a content type later, call again
with all content types the document should have (incremental-add-without-replace
will be a separate endpoint if/when we need it).
``source`` is a human-readable label (typically the original filename) that flows
into chunk metadata so search results are readable when document_id is a hash.
"""
document_id: FileId = Field(min_length=1)
source: str = Field(min_length=1)
page_text: list[IngestedPageText] | None = None
class RagIndexResponse(ApiModel):
collection: str
class IngestDocumentResponse(ApiModel):
document_id: FileId
chunks_indexed: int
class RagSearchRequest(ApiModel):
query: str
collection: str | None = Field(default=None, min_length=1)
top_k: int = 5
class RagSearchResultItem(ApiModel):
text: str
source: str
chunk_id: str
score: float
class RagSearchResponse(ApiModel):
query: str
results: list[RagSearchResultItem]
class RagCollectionsResponse(ApiModel):
collections: list[str]
class RagDeleteCollectionResponse(ApiModel):
status: str
collection: str
class DeleteDocumentResponse(ApiModel):
document_id: FileId
deleted: bool
engine/src/stirling/models/__init__.py
+2 -1
View File
@@ -1,9 +1,10 @@
from . import tool_models
from .base import ApiModel
from .base import ApiModel, FileId
from .tool_models import OPERATIONS, ParamToolModel, ToolEndpoint
__all__ = [
"ApiModel",
"FileId",
"OPERATIONS",
"ParamToolModel",
"ToolEndpoint",
engine/src/stirling/models/base.py
+7
View File
@@ -1,8 +1,15 @@
from __future__ import annotations
from typing import NewType
from pydantic import BaseModel, ConfigDict
from pydantic.alias_generators import to_camel
# Stable, opaque identifier for a file supplied by the caller. Owned by the caller's
# ID strategy (content hash, filesystem path, etc.) and used as the RAG collection key
# throughout the engine.
FileId = NewType("FileId", str)
class ApiModel(BaseModel):
model_config = ConfigDict(
engine/src/stirling/rag/capability.py
+45 -4
View File
@@ -1,11 +1,16 @@
from __future__ import annotations
import logging
from collections.abc import Awaitable, Callable
from pydantic_ai import FunctionToolset
from pydantic_ai import FunctionToolset, RunContext, ToolDefinition
from pydantic_ai.toolsets import AbstractToolset
from stirling.models import FileId
from stirling.rag.service import RagService
from stirling.rag.store import SearchResult
logger = logging.getLogger(__name__)
class RagCapability:
@@ -22,19 +27,29 @@ class RagCapability:
When no collections are pinned, the instructions are generated dynamically at
run time so the agent sees the current list of collections in the store.
Lifecycle: a ``RagCapability`` instance is intended to live for the duration of a
single agent run.
"""
def __init__(
self,
rag_service: RagService,
collections: list[str] | None = None,
collections: list[FileId] | None = None,
top_k: int = 5,
max_searches: int = 5,
) -> None:
self._rag_service = rag_service
self._collections = collections
self._top_k = top_k
self._max_searches = max_searches
self._search_count = 0
toolset: FunctionToolset[None] = FunctionToolset()
toolset.add_function(self._search_knowledge, name="search_knowledge")
toolset.add_function(
self._search_knowledge,
name="search_knowledge",
prepare=self._prepare_search_knowledge,
)
self._toolset = toolset
@property
@@ -48,7 +63,7 @@ class RagCapability:
return self._toolset
@staticmethod
def _static_instructions_text(collections: list[str]) -> str:
def _static_instructions_text(collections: list[FileId]) -> str:
collection_desc = f"collections: {', '.join(collections)}"
return (
"You have access to a knowledge base search tool called 'search_knowledge'. "
@@ -73,6 +88,18 @@ class RagCapability:
"You do not have to use it if the answer is already clear from the provided text."
)
async def _prepare_search_knowledge(
self,
ctx: RunContext[None],
tool_def: ToolDefinition,
) -> ToolDefinition | None:
"""Remove the search tool from the agent's toolset once the per-run search
budget is exhausted. The agent then has no choice but to answer from what it
has already retrieved, which prevents runaway search loops."""
if self._search_count >= self._max_searches:
return None
return tool_def
async def _search_knowledge(self, query: str, max_results: int | None = None) -> str:
"""Search the knowledge base for information relevant to the query.
@@ -83,6 +110,7 @@ class RagCapability:
Returns:
Formatted text with the most relevant knowledge base excerpts.
"""
self._search_count += 1
k = max_results if max_results is not None else self._top_k
if self._collections:
all_results = []
@@ -95,8 +123,21 @@ class RagCapability:
results = await self._rag_service.search(query, top_k=k)
if not results:
logger.info("[rag] search_knowledge query=%r -> 0 results", query)
return "No relevant results found in the knowledge base."
formatted = self._format_results(results)
logger.info(
"[rag] search_knowledge query=%r -> %d results, %d chars",
query,
len(results),
len(formatted),
)
logger.debug("[rag] search_knowledge query=%r returned:\n%s", query, formatted)
return formatted
@staticmethod
def _format_results(results: list[SearchResult]) -> str:
sections = []
for i, result in enumerate(results, 1):
source = result.document.metadata.get("source", "unknown")
engine/src/stirling/rag/embedder.py
+25 -4
View File
@@ -5,14 +5,27 @@ from pydantic_ai import Embedder
from stirling.rag.chunker import chunk_text
from stirling.rag.store import Document
# Keep each upstream embed request under every major provider's per-call limit while
# still batching large enough that a book-sized document ingests in a reasonable number
# of round trips. VoyageAI caps at 1000, OpenAI at 2048, Cohere at 96; 256 is a good
# default for Voyage/OpenAI. Cohere users should pass a lower value via construction.
DEFAULT_EMBED_BATCH_SIZE = 256
class EmbeddingService:
"""Wraps Pydantic AI's Embedder to provide document chunking and embedding."""
def __init__(self, model_name: str, chunk_size: int = 512, chunk_overlap: int = 64) -> None:
def __init__(
self,
model_name: str,
chunk_size: int = 512,
chunk_overlap: int = 64,
embed_batch_size: int = DEFAULT_EMBED_BATCH_SIZE,
) -> None:
self._embedder = Embedder(model_name)
self._chunk_size = chunk_size
self._chunk_overlap = chunk_overlap
self._embed_batch_size = embed_batch_size
async def embed_query(self, text: str) -> list[float]:
"""Embed a search query, optimised for retrieval."""
@@ -20,11 +33,19 @@ class EmbeddingService:
return list(result.embeddings[0])
async def embed_documents(self, texts: list[str]) -> list[list[float]]:
"""Embed multiple document texts for indexing."""
"""Embed multiple document texts for indexing.
Splits the input into batches of ``embed_batch_size`` so callers can hand us
any number of chunks without hitting provider per-request limits.
"""
if not texts:
return []
result = await self._embedder.embed_documents(texts)
return [list(emb) for emb in result.embeddings]
all_embeddings: list[list[float]] = []
for start in range(0, len(texts), self._embed_batch_size):
batch = texts[start : start + self._embed_batch_size]
result = await self._embedder.embed_documents(batch)
all_embeddings.extend(list(emb) for emb in result.embeddings)
return all_embeddings
def chunk_and_prepare(
self,
engine/src/stirling/rag/pgvector_store.py
+4
View File
@@ -131,3 +131,7 @@ class PgVectorStore(VectorStore):
)
row = await cur.fetchone()
return row is not None
async def close(self) -> None:
# Connections are opened and closed per call, so nothing persistent to release.
return None
engine/src/stirling/rag/service.py
+28 -6
View File
@@ -2,6 +2,7 @@ from __future__ import annotations
import logging
from stirling.models import FileId
from stirling.rag.embedder import EmbeddingService
from stirling.rag.store import Document, SearchResult, VectorStore
@@ -18,7 +19,7 @@ class RagService:
async def index_text(
self,
collection: str,
collection: FileId,
text: str,
source: str = "",
metadata: dict[str, str] | None = None,
@@ -31,7 +32,7 @@ class RagService:
await self._store.add_documents(collection, documents, embeddings)
return len(documents)
async def index_documents(self, collection: str, documents: list[Document]) -> int:
async def index_documents(self, collection: FileId, documents: list[Document]) -> int:
"""Embed and store pre-chunked documents. Returns the number stored."""
if not documents:
return 0
@@ -39,10 +40,23 @@ class RagService:
await self._store.add_documents(collection, documents, embeddings)
return len(documents)
def chunk_text(
self,
text: str,
source: str = "",
base_metadata: dict[str, str] | None = None,
) -> list[Document]:
"""Chunk text into Document objects ready for indexing. Does NOT embed.
Exposed so callers that ingest many chunks can accumulate them across calls
and then pass the full batch to ``index_documents`` for a single embedding pass.
"""
return self._embedder.chunk_and_prepare(text, source=source, base_metadata=base_metadata)
async def search(
self,
query: str,
collection: str | None = None,
collection: FileId | None = None,
top_k: int | None = None,
) -> list[SearchResult]:
"""Embed query and search across one or all collections.
@@ -71,10 +85,18 @@ class RagService:
all_results.sort(key=lambda r: r.score, reverse=True)
return all_results[:k]
async def delete_collection(self, collection: str) -> None:
async def delete_collection(self, collection: FileId) -> None:
"""Remove a collection and all its documents."""
await self._store.delete_collection(collection)
async def list_collections(self) -> list[str]:
async def has_collection(self, collection: FileId) -> bool:
"""Check whether a collection exists."""
return await self._store.has_collection(collection)
async def list_collections(self) -> list[FileId]:
"""List all available collections."""
return await self._store.list_collections()
return [FileId(name) for name in await self._store.list_collections()]
async def close(self) -> None:
"""Release the underlying vector store's resources."""
await self._store.close()
engine/src/stirling/rag/sqlite_vec_store.py
+16
View File
@@ -225,3 +225,19 @@ class SqliteVecStore(VectorStore):
def _sync_has_collection(self, collection: str) -> bool:
row = self._conn.execute("SELECT 1 FROM collections WHERE name = ?", (collection,)).fetchone()
return row is not None
async def close(self) -> None:
async with self._lock:
await asyncio.to_thread(self._sync_close)
def _sync_close(self) -> None:
"""Checkpoint the WAL into the main database file and close the connection so
the .db-shm and .db-wal files are cleaned up on graceful shutdown."""
if self._db_path is not None:
try:
self._conn.execute("PRAGMA wal_checkpoint(TRUNCATE)")
self._conn.commit()
except sqlite3.Error:
# Best effort: if checkpointing fails we still want to close the connection.
pass
self._conn.close()
engine/src/stirling/rag/store.py
+4
View File
@@ -57,3 +57,7 @@ class VectorStore(ABC):
@abstractmethod
async def has_collection(self, collection: str) -> bool:
"""Check whether a collection exists."""
@abstractmethod
async def close(self) -> None:
"""Release any resources held by the store (connections, handles, etc.)."""