Add document context for edit agent (#6152)

# Description of Changes
Adds the ability for the Edit agent to request the content of the
document before it decides which parameters it needs. This makes it able
to process requests like `Split the document after the page containing
the "My Section" section`, allowing for document context-based requests
for all[^1] tools.

I had to make a few changes elsewhere to make this work, including:
- Moving the requesting of content out of the Question Agent and into a
common location
- Added specific API docs for the Split param because the generic ones
were not specific enough for the AI to be able to reliably perform the
correct operation
- Fixed an issue in the tool models generator which caused the Redact
params to only be half-generated (causing Pydantic to crash when the AI
tried to run Redact)
- Added missing logging to a bunch of tools and hooked it up properly so
it'll print to stderr
- Made the limits for the max pages/chars to extract from PDFs
configurable via env var

[^1]: Many of the tools can't actually do anything useful with the
context at this stage, but will just need the tool API to be extended
with new features like page-specific operations to be automatically able
to do smart operations without needing to change the Edit agent itself.

engine/tests/conftest.py

@@ -30,6 +30,8 @@ def build_app_settings() -> AppSettings:
         rag_chunk_size=512,
         rag_chunk_overlap=64,
         rag_default_top_k=5,
+        max_pages=200,
+        max_characters=200_000,
         posthog_enabled=False,
         posthog_api_key="",
         posthog_host="https://eu.i.posthog.com",

engine/tests/test_pdf_edit_agent.py

@@ -5,11 +5,18 @@ from dataclasses import dataclass
 import pytest
 
 from stirling.agents import PdfEditAgent, PdfEditParameterSelector, PdfEditPlanSelection
+from stirling.agents.pdf_edit import PdfEditPlanOutput
 from stirling.contracts import (
     EditCannotDoResponse,
     EditClarificationRequest,
     EditPlanResponse,
+    ExtractedFileText,
+    NeedContentFileRequest,
+    NeedContentResponse,
+    PdfContentType,
     PdfEditRequest,
+    PdfTextSelection,
+    SupportedCapability,
     ToolOperationStep,
 )
 from stirling.models.tool_models import Angle, FlattenParams, RotatePdfParams, ToolEndpoint
@@ -52,7 +59,7 @@ class StubPdfEditAgent(PdfEditAgent):
     def __init__(
         self,
         runtime: AppRuntime,
-        selection: PdfEditPlanSelection | EditClarificationRequest | EditCannotDoResponse,
+        selection: PdfEditPlanOutput,
         parameter_selector: RecordingParameterSelector | PdfEditParameterSelector | None = None,
     ) -> None:
         super().__init__(runtime)
@@ -63,7 +70,8 @@ class StubPdfEditAgent(PdfEditAgent):
     async def _select_plan(
         self,
         request: PdfEditRequest,
-    ) -> PdfEditPlanSelection | EditClarificationRequest | EditCannotDoResponse:
+        allow_need_content: bool = True,
+    ) -> PdfEditPlanOutput:
         return self.selection
 
 
@@ -153,3 +161,117 @@ async def test_pdf_edit_agent_returns_cannot_do_without_partial_plan(runtime: Ap
     response = await agent.handle(PdfEditRequest(user_message="Read this scan and summarize it."))
 
     assert isinstance(response, EditCannotDoResponse)
+
+
+@pytest.mark.anyio
+async def test_pdf_edit_agent_returns_need_content_without_building_plan(runtime: AppRuntime) -> None:
+    parameter_selector = RecordingParameterSelector()
+    agent = StubPdfEditAgent(
+        runtime,
+        NeedContentResponse(
+            resume_with=SupportedCapability.PDF_EDIT,
+            reason="Need page text to locate the NEW PAGE markers.",
+            files=[],
+            max_pages=0,
+            max_characters=0,
+        ),
+        parameter_selector=parameter_selector,
+    )
+
+    response = await agent.handle(
+        PdfEditRequest(
+            user_message="Split after every page that says 'NEW PAGE'.",
+            file_names=["report.pdf"],
+        )
+    )
+
+    assert isinstance(response, NeedContentResponse)
+    assert response.resume_with == SupportedCapability.PDF_EDIT
+    assert response.files == [NeedContentFileRequest(file_name="report.pdf", content_types=[PdfContentType.PAGE_TEXT])]
+    assert response.max_pages == runtime.settings.max_pages
+    assert response.max_characters == runtime.settings.max_characters
+    assert parameter_selector.calls == []
+
+
+@pytest.mark.anyio
+async def test_pdf_edit_agent_builds_selection_agent_matching_content_availability(runtime: AppRuntime) -> None:
+    from stirling.agents.pdf_edit import PdfEditSelectionAgent
+
+    agent = PdfEditAgent(runtime)
+    captured: list[bool] = []
+
+    def record(*, allow_need_content: bool) -> PdfEditSelectionAgent:
+        captured.append(allow_need_content)
+        raise _StopSelectionError()
+
+    agent._build_selection_agent = record
+
+    with pytest.raises(_StopSelectionError):
+        await agent._select_plan(PdfEditRequest(user_message="Rotate."))
+    with pytest.raises(_StopSelectionError):
+        await agent._select_plan(
+            PdfEditRequest(
+                user_message="Rotate.",
+                page_text=[
+                    ExtractedFileText(
+                        file_name="report.pdf",
+                        pages=[PdfTextSelection(page_number=1, text="content")],
+                    )
+                ],
+            )
+        )
+    with pytest.raises(_StopSelectionError):
+        await agent._select_plan(PdfEditRequest(user_message="Rotate."), allow_need_content=False)
+
+    assert captured == [True, False, False]
+
+
+@pytest.mark.anyio
+async def test_pdf_edit_selection_agent_excludes_need_content_from_schema_when_not_allowed(
+    runtime: AppRuntime,
+) -> None:
+    from stirling.agents.pdf_edit import PdfEditSelectionAgent
+
+    can_request = PdfEditSelectionAgent(runtime, "base", allow_need_content=True)
+    cannot_request = PdfEditSelectionAgent(runtime, "base", allow_need_content=False)
+
+    assert NeedContentResponse in _agent_output_types(can_request)
+    assert NeedContentResponse not in _agent_output_types(cannot_request)
+
+
+def _agent_output_types(agent: object) -> list[type]:
+    native = getattr(getattr(agent, "agent"), "output_type")
+    return list(getattr(native, "outputs", []))
+
+
+class _StopSelectionError(Exception):
+    pass
+
+
+@pytest.mark.anyio
+async def test_pdf_edit_agent_passes_page_text_to_parameter_selector(runtime: AppRuntime) -> None:
+    parameter_selector = RecordingParameterSelector()
+    agent = StubPdfEditAgent(
+        runtime,
+        PdfEditPlanSelection(
+            operations=[ToolEndpoint.ROTATE_PDF],
+            summary="Rotate the PDF.",
+        ),
+        parameter_selector=parameter_selector,
+    )
+
+    page_text = [
+        ExtractedFileText(
+            file_name="report.pdf",
+            pages=[PdfTextSelection(page_number=1, text="NEW PAGE")],
+        )
+    ]
+    await agent.handle(
+        PdfEditRequest(
+            user_message="Rotate clockwise.",
+            file_names=["report.pdf"],
+            page_text=page_text,
+        )
+    )
+
+    assert parameter_selector.calls[0].request.page_text == page_text

engine/tests/test_pdf_question_agent.py

@@ -5,8 +5,8 @@ import pytest
 from stirling.agents import PdfQuestionAgent
 from stirling.contracts import (
     ExtractedFileText,
+    NeedContentResponse,
     PdfQuestionAnswerResponse,
-    PdfQuestionNeedContentResponse,
     PdfQuestionNotFoundResponse,
     PdfQuestionRequest,
     PdfTextSelection,
@@ -41,7 +41,7 @@ async def test_pdf_question_agent_requires_extracted_text(runtime: AppRuntime) -
         PdfQuestionRequest(question="What is the total?", page_text=[], file_names=["test.pdf"])
     )
 
-    assert isinstance(response, PdfQuestionNeedContentResponse)
+    assert isinstance(response, NeedContentResponse)
 
 
 @pytest.mark.anyio

engine/tests/test_stirling_api.py

@@ -19,18 +19,25 @@ from stirling.contracts import (
     AgentRevisionResponse,
     CannotContinueExecutionAction,
     EditCannotDoResponse,
+    NeedContentResponse,
     OrchestratorRequest,
     PdfEditRequest,
-    PdfQuestionNeedContentResponse,
     PdfQuestionNotFoundResponse,
     PdfQuestionRequest,
+    SupportedCapability,
 )
 from stirling.models.tool_models import Angle, RotatePdfParams
 
 
 class StubOrchestratorAgent:
-    async def handle(self, request: OrchestratorRequest) -> PdfQuestionNeedContentResponse:
-        return PdfQuestionNeedContentResponse(reason=request.user_message, files=[], max_pages=1, max_characters=1000)
+    async def handle(self, request: OrchestratorRequest) -> NeedContentResponse:
+        return NeedContentResponse(
+            resume_with=SupportedCapability.PDF_QUESTION,
+            reason=request.user_message,
+            files=[],
+            max_pages=1,
+            max_characters=1000,
+        )
 
 
 class StubPdfEditAgent:

engine/tests/test_stirling_contracts.py

@@ -89,6 +89,8 @@ def test_app_settings_accepts_model_configuration() -> None:
         rag_chunk_size=512,
         rag_chunk_overlap=64,
         rag_default_top_k=5,
+        max_pages=200,
+        max_characters=200_000,
         posthog_enabled=False,
         posthog_api_key="",
         posthog_host="https://eu.i.posthog.com",