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/contracts/__init__.py

@@ -12,6 +12,8 @@ from .common import (
     ArtifactKind,
     ConversationMessage,
     ExtractedFileText,
+    NeedContentFileRequest,
+    NeedContentResponse,
     PdfContentType,
     PdfTextSelection,
     StepKind,
@@ -54,14 +56,14 @@ from .pdf_edit import (
     EditPlanResponse,
     PdfEditRequest,
     PdfEditResponse,
+    PdfEditTerminalResponse,
 )
 from .pdf_questions import (
-    NeedContentFileRequest,
     PdfQuestionAnswerResponse,
-    PdfQuestionNeedContentResponse,
     PdfQuestionNotFoundResponse,
     PdfQuestionRequest,
     PdfQuestionResponse,
+    PdfQuestionTerminalResponse,
 )
 from .rag import (
     MAX_INDEX_TEXT_LENGTH,
@@ -108,17 +110,19 @@ __all__ = [
     "format_conversation_history",
     "HealthResponse",
     "NeedContentFileRequest",
+    "NeedContentResponse",
     "NextExecutionAction",
     "OrchestratorRequest",
     "OrchestratorResponse",
     "PdfContentType",
     "PdfEditRequest",
     "PdfEditResponse",
+    "PdfEditTerminalResponse",
     "PdfQuestionAnswerResponse",
-    "PdfQuestionNeedContentResponse",
     "PdfQuestionNotFoundResponse",
     "PdfQuestionRequest",
     "PdfQuestionResponse",
+    "PdfQuestionTerminalResponse",
     "PdfTextSelection",
     "RagCollectionsResponse",
     "RagDeleteCollectionResponse",

engine/src/stirling/contracts/common.py

@@ -107,6 +107,21 @@ class ExtractedFileText(ApiModel):
     pages: list[PdfTextSelection] = Field(default_factory=list)
 
 
+class NeedContentFileRequest(ApiModel):
+    file_name: str
+    page_numbers: list[int] = Field(default_factory=list)
+    content_types: list[PdfContentType]
+
+
+class NeedContentResponse(ApiModel):
+    outcome: Literal[WorkflowOutcome.NEED_CONTENT] = WorkflowOutcome.NEED_CONTENT
+    resume_with: SupportedCapability
+    reason: str
+    files: list[NeedContentFileRequest] = Field(default_factory=list)
+    max_pages: int
+    max_characters: int
+
+
 class ToolOperationStep(ApiModel):
     kind: Literal[StepKind.TOOL] = StepKind.TOOL
     tool: AnyToolId

engine/src/stirling/contracts/orchestrator.py

@@ -11,12 +11,13 @@ from .common import (
     ArtifactKind,
     ConversationMessage,
     ExtractedFileText,
+    NeedContentResponse,
     SupportedCapability,
     WorkflowOutcome,
 )
 from .execution import NextExecutionAction
-from .pdf_edit import PdfEditResponse
-from .pdf_questions import PdfQuestionResponse
+from .pdf_edit import PdfEditTerminalResponse
+from .pdf_questions import PdfQuestionTerminalResponse
 
 
 class ExtractedTextArtifact(ApiModel):
@@ -41,7 +42,12 @@ class UnsupportedCapabilityResponse(ApiModel):
     message: str
 
 
-OrchestratorResponse = Annotated[
-    PdfEditResponse | PdfQuestionResponse | AgentDraftResponse | NextExecutionAction | UnsupportedCapabilityResponse,
+type OrchestratorResponse = Annotated[
+    PdfEditTerminalResponse
+    | PdfQuestionTerminalResponse
+    | NeedContentResponse
+    | AgentDraftResponse
+    | NextExecutionAction
+    | UnsupportedCapabilityResponse,
     Field(discriminator="outcome"),
 ]

engine/src/stirling/contracts/pdf_edit.py

@@ -6,13 +6,14 @@ from pydantic import Field
 
 from stirling.models import ApiModel
 
-from .common import ConversationMessage, ToolOperationStep, WorkflowOutcome
+from .common import ConversationMessage, ExtractedFileText, NeedContentResponse, ToolOperationStep, WorkflowOutcome
 
 
 class PdfEditRequest(ApiModel):
     user_message: str
     file_names: list[str] = Field(default_factory=list)
     conversation_history: list[ConversationMessage] = Field(default_factory=list)
+    page_text: list[ExtractedFileText] = Field(default_factory=list)
 
 
 class EditPlanResponse(ApiModel):
@@ -33,7 +34,8 @@ class EditCannotDoResponse(ApiModel):
     reason: str
 
 
-PdfEditResponse = Annotated[
-    EditPlanResponse | EditClarificationRequest | EditCannotDoResponse,
+type PdfEditTerminalResponse = EditPlanResponse | EditClarificationRequest | EditCannotDoResponse
+type PdfEditResponse = Annotated[
+    PdfEditTerminalResponse | NeedContentResponse,
     Field(discriminator="outcome"),
 ]

engine/src/stirling/contracts/pdf_questions.py

@@ -9,8 +9,7 @@ from stirling.models import ApiModel
 from .common import (
     ConversationMessage,
     ExtractedFileText,
-    PdfContentType,
-    SupportedCapability,
+    NeedContentResponse,
     WorkflowOutcome,
 )
 
@@ -28,27 +27,13 @@ class PdfQuestionAnswerResponse(ApiModel):
     evidence: list[ExtractedFileText] = Field(default_factory=list)
 
 
-class NeedContentFileRequest(ApiModel):
-    file_name: str
-    page_numbers: list[int] = Field(default_factory=list)
-    content_types: list[PdfContentType]
-
-
-class PdfQuestionNeedContentResponse(ApiModel):
-    outcome: Literal[WorkflowOutcome.NEED_CONTENT] = WorkflowOutcome.NEED_CONTENT
-    resume_with: SupportedCapability = SupportedCapability.PDF_QUESTION
-    reason: str
-    files: list[NeedContentFileRequest] = Field(default_factory=list)
-    max_pages: int
-    max_characters: int
-
-
 class PdfQuestionNotFoundResponse(ApiModel):
     outcome: Literal[WorkflowOutcome.NOT_FOUND] = WorkflowOutcome.NOT_FOUND
     reason: str
 
 
-PdfQuestionResponse = Annotated[
-    PdfQuestionAnswerResponse | PdfQuestionNeedContentResponse | PdfQuestionNotFoundResponse,
+type PdfQuestionTerminalResponse = PdfQuestionAnswerResponse | PdfQuestionNotFoundResponse
+type PdfQuestionResponse = Annotated[
+    PdfQuestionTerminalResponse | NeedContentResponse,
     Field(discriminator="outcome"),
 ]