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/src/stirling/agents/_page_text.py

@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+from stirling.contracts import ExtractedFileText
+
+
+def has_page_text(page_text: list[ExtractedFileText]) -> bool:
+    return any(selection.text.strip() for file_text in page_text for selection in file_text.pages)
+
+
+def format_page_text(page_text: list[ExtractedFileText], empty: str = "None") -> str:
+    if not has_page_text(page_text):
+        return empty
+    sections = [
+        f"[File: {file_text.file_name}, Page {selection.page_number or '?'}]\n{selection.text}"
+        for file_text in page_text
+        for selection in file_text.pages
+    ]
+    return "\n\n".join(sections)

engine/src/stirling/agents/orchestrator.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import logging
 from dataclasses import dataclass
 from typing import assert_never
 
@@ -29,6 +30,8 @@ from stirling.contracts.pdf_edit import EditPlanResponse
 from stirling.models.agent_tool_models import AgentToolId, MathAuditorAgentParams
 from stirling.services import AppRuntime
 
+logger = logging.getLogger(__name__)
+
 
 @dataclass(frozen=True)
 class OrchestratorDeps:
@@ -86,12 +89,20 @@ 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,
+            request.resume_with,
+            [type(a).__name__ for a in request.artifacts],
+            request.user_message,
+        )
         if request.resume_with is not None:
             return await self._resume(request, request.resume_with)
         result = await self.agent.run(
             self._build_prompt(request),
             deps=OrchestratorDeps(runtime=self.runtime, request=request),
         )
+        logger.info("[orchestrator] routed -> %s", type(result.output).__name__)
         return result.output
 
     async def _resume(self, request: OrchestratorRequest, capability: SupportedCapability) -> OrchestratorResponse:
@@ -117,11 +128,13 @@ class OrchestratorAgent:
         return await self._run_pdf_edit(ctx.deps.request)
 
     async def _run_pdf_edit(self, request: OrchestratorRequest) -> PdfEditResponse:
+        extracted_text = self._get_extracted_text_artifact(request)
         return await PdfEditAgent(self.runtime).handle(
             PdfEditRequest(
                 user_message=request.user_message,
                 file_names=request.file_names,
                 conversation_history=request.conversation_history,
+                page_text=extracted_text.files if extracted_text is not None else [],
             )
         )
 

engine/src/stirling/agents/pdf_edit.py

@@ -1,23 +1,33 @@
 from __future__ import annotations
 
-from typing import Literal
+import logging
+from typing import Literal, overload
 
 from pydantic import Field
 from pydantic_ai import Agent
 from pydantic_ai.output import NativeOutput
 
+from stirling.agents._page_text import format_page_text, has_page_text
 from stirling.contracts import (
     EditCannotDoResponse,
     EditClarificationRequest,
     EditPlanResponse,
+    NeedContentFileRequest,
+    NeedContentResponse,
+    PdfContentType,
     PdfEditRequest,
     PdfEditResponse,
+    PdfEditTerminalResponse,
+    SupportedCapability,
     ToolOperationStep,
     format_conversation_history,
 )
+from stirling.logging import Pretty
 from stirling.models import OPERATIONS, ApiModel, ParamToolModel, ToolEndpoint
 from stirling.services import AppRuntime
 
+logger = logging.getLogger(__name__)
+
 
 class PdfEditPlanSelection(ApiModel):
     outcome: Literal["plan"] = "plan"
@@ -26,6 +36,39 @@ class PdfEditPlanSelection(ApiModel):
     rationale: str | None = None
 
 
+type PdfEditPlanOutput = PdfEditPlanSelection | EditClarificationRequest | EditCannotDoResponse | NeedContentResponse
+
+
+class PdfEditSelectionAgent:
+    def __init__(self, runtime: AppRuntime, base_system_prompt: str, *, allow_need_content: bool) -> None:
+        self.runtime = runtime
+        output_types: list[type[PdfEditPlanOutput]] = [
+            PdfEditPlanSelection,
+            EditClarificationRequest,
+            EditCannotDoResponse,
+        ]
+        system_prompt = base_system_prompt
+        if allow_need_content:
+            output_types.append(NeedContentResponse)
+            system_prompt += (
+                " Return need_content when planning a correct answer requires inspecting the actual PDF "
+                "page text (e.g. 'split after every page that says NEW PAGE', "
+                "'rotate pages that mention draft')."
+            )
+        self.agent = Agent(
+            model=runtime.smart_model,
+            output_type=NativeOutput(output_types),
+            system_prompt=system_prompt,
+            model_settings=runtime.smart_model_settings,
+        )
+
+    async def select(self, prompt: str) -> PdfEditPlanOutput:
+        logger.debug("[pdf-edit selection] prompt:\n%s", prompt)
+        result = await self.agent.run(prompt)
+        logger.debug("[pdf-edit selection] output: %s", Pretty(result.output))
+        return result.output
+
+
 class PdfEditParameterSelector:
     def __init__(self, runtime: AppRuntime) -> None:
         self.runtime = runtime
@@ -34,7 +77,9 @@ class PdfEditParameterSelector:
             system_prompt=(
                 "Generate only the parameter object for the selected PDF operation. "
                 "Use reasonable defaults when the request does not specify optional details. "
-                "Only fill fields that belong to the selected operation's parameter model."
+                "Only fill fields that belong to the selected operation's parameter model. "
+                "When extracted page text is provided, use it to compute precise parameters "
+                "(e.g. exact page ranges that match a described pattern)."
             ),
             model_settings=runtime.smart_model_settings,
         )
@@ -48,14 +93,17 @@ class PdfEditParameterSelector:
     ) -> ParamToolModel:
         operation_id = operation_plan[operation_index]
         parameter_model = OPERATIONS[operation_id]
+        prompt = self._build_parameter_prompt(request, operation_plan, operation_index, generated_steps)
+        logger.debug("[pdf-edit params %s] prompt:\n%s", operation_id.name, prompt)
         parameter_result = await self.agent.run(
-            self._build_parameter_prompt(request, operation_plan, operation_index, generated_steps),
+            prompt,
             output_type=NativeOutput(parameter_model),
             instructions=(
                 f"Generate only the parameters for the PDF operation `{operation_id.name}`. "
                 "Do not include fields from any other operation."
             ),
         )
+        logger.debug("[pdf-edit params %s] output: %s", operation_id.name, Pretty(parameter_result.output))
         return parameter_result.output
 
     def _build_parameter_prompt(
@@ -76,12 +124,14 @@ class PdfEditParameterSelector:
             else "None"
         )
         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"Operation plan: {operation_list}\n"
             f"Selected operation index: {operation_index + 1} of {len(operation_plan)}\n"
             f"Selected operation: {operation_id.name}\n"
             f"Already generated steps:\n{generated_steps_text}\n"
+            f"Extracted page text:\n{format_page_text(request.page_text)}\n"
             "Return only the parameter object for the selected operation."
         )
 
@@ -91,32 +141,27 @@ class PdfEditAgent:
         self.runtime = runtime
         self.supported_operations = list(OPERATIONS)
         self.parameter_selector = PdfEditParameterSelector(runtime)
-        self.selection_agent = Agent(
-            model=runtime.smart_model,
-            output_type=NativeOutput(
-                [
-                    PdfEditPlanSelection,
-                    EditClarificationRequest,
-                    EditCannotDoResponse,
-                ]
-            ),
-            system_prompt=(
-                "Plan PDF edit requests. "
-                f"Supported operations are: {self._supported_operations_prompt()}. "
-                "Return an ordered list of one or more supported operations for the plan. "
-                "Do not produce operation parameters in this stage. "
-                "Return need_clarification when the request is genuinely ambiguous. "
-                "Return cannot_do when the request is outside the supported operations. "
-                "Return plan when a reasonable multi-step plan can be created. "
-                "Never return partial plans."
-            ),
-            model_settings=runtime.smart_model_settings,
-        )
 
-    async def handle(self, request: PdfEditRequest) -> PdfEditResponse:
-        selection = await self._select_plan(request)
+    @overload
+    async def handle(self, request: PdfEditRequest, allow_need_content: Literal[False]) -> PdfEditTerminalResponse: ...
+    @overload
+    async def handle(self, request: PdfEditRequest, allow_need_content: bool = True) -> PdfEditResponse: ...
+    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,
+            has_page_text(request.page_text),
+            allow_need_content,
+            request.user_message,
+        )
+        selection = await self._select_plan(request, allow_need_content=allow_need_content)
         if isinstance(selection, EditClarificationRequest | EditCannotDoResponse):
+            logger.info("[pdf-edit] selection -> %s: %s", selection.outcome, Pretty(selection))
             return selection
+        if isinstance(selection, NeedContentResponse):
+            logger.info("[pdf-edit] selection -> need_content: %s", selection.reason)
+            return self._fill_need_content_defaults(selection, request)
+        logger.info("[pdf-edit] plan: %s", [op.name for op in selection.operations])
         steps: list[ToolOperationStep] = []
         for operation_index, operation_id in enumerate(selection.operations):
             parameters = await self.parameter_selector.select(
@@ -140,9 +185,27 @@ class PdfEditAgent:
     async def _select_plan(
         self,
         request: PdfEditRequest,
-    ) -> PdfEditPlanSelection | EditClarificationRequest | EditCannotDoResponse:
-        selection_result = await self.selection_agent.run(self._build_selection_prompt(request))
-        return selection_result.output
+        allow_need_content: bool = True,
+    ) -> PdfEditPlanOutput:
+        can_request_content = allow_need_content and not has_page_text(request.page_text)
+        agent = self._build_selection_agent(allow_need_content=can_request_content)
+        return await agent.select(self._build_selection_prompt(request))
+
+    def _build_selection_agent(self, *, allow_need_content: bool) -> PdfEditSelectionAgent:
+        return PdfEditSelectionAgent(
+            self.runtime,
+            base_system_prompt=(
+                "Plan PDF edit requests. "
+                f"Supported operations are: {self._supported_operations_prompt()}. "
+                "Return an ordered list of one or more supported operations for the plan. "
+                "Do not produce operation parameters in this stage. "
+                "Return need_clarification when the request is genuinely ambiguous. "
+                "Return cannot_do when the request is outside the supported operations. "
+                "Return plan when a reasonable multi-step plan can be created. "
+                "Never return partial plans."
+            ),
+            allow_need_content=allow_need_content,
+        )
 
     def _build_selection_prompt(self, request: PdfEditRequest) -> str:
         file_names = ", ".join(request.file_names) if request.file_names else "No file names were provided."
@@ -151,8 +214,25 @@ class PdfEditAgent:
             f"User request: {request.user_message}\n"
             f"Files: {file_names}\n"
             f"Supported operations: {self._supported_operations_prompt()}\n"
-            "Plan an ordered list of supported PDF edit operations or return clarification/cannot_do."
+            f"Extracted page text:\n{format_page_text(request.page_text)}"
         )
 
     def _supported_operations_prompt(self) -> str:
         return ", ".join(f"{op.name} ({op.value})" for op in self.supported_operations)
+
+    def _fill_need_content_defaults(
+        self,
+        selection: NeedContentResponse,
+        request: PdfEditRequest,
+    ) -> NeedContentResponse:
+        files = selection.files or [
+            NeedContentFileRequest(file_name=file_name, content_types=[PdfContentType.PAGE_TEXT])
+            for file_name in request.file_names
+        ]
+        return NeedContentResponse(
+            resume_with=SupportedCapability.PDF_EDIT,
+            reason=selection.reason,
+            files=files,
+            max_pages=selection.max_pages or self.runtime.settings.max_pages,
+            max_characters=selection.max_characters or self.runtime.settings.max_characters,
+        )

engine/src/stirling/agents/pdf_questions.py

@@ -3,24 +3,22 @@ from __future__ import annotations
 from pydantic_ai import Agent
 from pydantic_ai.output import NativeOutput
 
+from stirling.agents._page_text import format_page_text, has_page_text
 from stirling.contracts import (
-    ExtractedFileText,
     NeedContentFileRequest,
+    NeedContentResponse,
     PdfContentType,
     PdfQuestionAnswerResponse,
-    PdfQuestionNeedContentResponse,
     PdfQuestionNotFoundResponse,
     PdfQuestionRequest,
     PdfQuestionResponse,
+    SupportedCapability,
     format_conversation_history,
 )
 from stirling.services import AppRuntime
 
 
 class PdfQuestionAgent:
-    DEFAULT_MAX_PAGES = 12
-    DEFAULT_MAX_CHARACTERS = 24_000
-
     def __init__(self, runtime: AppRuntime) -> None:
         self.runtime = runtime
         rag = runtime.rag_capability
@@ -44,8 +42,9 @@ class PdfQuestionAgent:
         )
 
     async def handle(self, request: PdfQuestionRequest) -> PdfQuestionResponse:
-        if not self._has_page_text(request.page_text):
-            return PdfQuestionNeedContentResponse(
+        if not has_page_text(request.page_text):
+            return NeedContentResponse(
+                resume_with=SupportedCapability.PDF_QUESTION,
                 reason="No extracted PDF page text was provided, so the question cannot be answered yet.",
                 files=[
                     NeedContentFileRequest(
@@ -54,8 +53,8 @@ class PdfQuestionAgent:
                     )
                     for file_name in request.file_names
                 ],
-                max_pages=self.DEFAULT_MAX_PAGES,
-                max_characters=self.DEFAULT_MAX_CHARACTERS,
+                max_pages=self.runtime.settings.max_pages,
+                max_characters=self.runtime.settings.max_characters,
             )
         return await self._run_answer_agent(request)
 
@@ -65,12 +64,7 @@ class PdfQuestionAgent:
 
     def _build_prompt(self, request: PdfQuestionRequest) -> str:
         file_names = ", ".join(request.file_names) if request.file_names else "Unknown files"
-        sections = [
-            f"[File: {file_text.file_name}, Page {selection.page_number or '?'}]\n{selection.text}"
-            for file_text in request.page_text
-            for selection in file_text.pages
-        ]
-        pages = "\n\n".join(sections)
+        pages = format_page_text(request.page_text, empty="")
         history = format_conversation_history(request.conversation_history)
         return (
             f"Conversation history:\n{history}\n"
@@ -78,6 +72,3 @@ class PdfQuestionAgent:
             f"Question: {request.question}\n"
             f"Extracted page text:\n{pages}"
         )
-
-    def _has_page_text(self, page_text: list[ExtractedFileText]) -> bool:
-        return any(selection.text.strip() for file_text in page_text for selection in file_text.pages)

engine/src/stirling/agents/user_spec.py

@@ -14,10 +14,9 @@ from stirling.contracts import (
     AgentRevisionWorkflowResponse,
     AiToolAgentStep,
     ConversationMessage,
-    EditCannotDoResponse,
-    EditClarificationRequest,
     EditPlanResponse,
     PdfEditRequest,
+    PdfEditTerminalResponse,
     format_conversation_history,
 )
 from stirling.models import ApiModel
@@ -102,7 +101,8 @@ class UserSpecAgent:
         self,
         user_message: str,
         conversation_history: list[ConversationMessage],
-    ) -> EditPlanResponse | EditClarificationRequest | EditCannotDoResponse:
+    ) -> PdfEditTerminalResponse:
         return await self.pdf_edit_agent.handle(
-            PdfEditRequest(user_message=user_message, conversation_history=conversation_history)
+            PdfEditRequest(user_message=user_message, conversation_history=conversation_history),
+            allow_need_content=False,
         )