Pdf comment agent (#6196)

Co-authored-by: James Brunton <[email protected]>

engine/src/stirling/contracts/__init__.py

@@ -8,10 +8,12 @@ from .agent_drafts import (
     AgentRevisionWorkflowResponse,
 )
 from .agent_specs import AgentSpec, AgentSpecStep, AiToolAgentStep
+from .comments import CommentSpec
 from .common import (
     ArtifactKind,
     ConversationMessage,
     ExtractedFileText,
+    MathAuditorToolReportArtifact,
     NeedContentFileRequest,
     NeedContentResponse,
     PdfContentType,
@@ -19,6 +21,7 @@ from .common import (
     StepKind,
     SupportedCapability,
     ToolOperationStep,
+    ToolReportArtifact,
     WorkflowOutcome,
     format_conversation_history,
 )
@@ -50,6 +53,13 @@ from .orchestrator import (
     UnsupportedCapabilityResponse,
     WorkflowArtifact,
 )
+from .pdf_comments import (
+    PdfCommentInstruction,
+    PdfCommentReport,
+    PdfCommentRequest,
+    PdfCommentResponse,
+    TextChunk,
+)
 from .pdf_edit import (
     EditCannotDoResponse,
     EditClarificationRequest,
@@ -92,6 +102,7 @@ __all__ = [
     "AiToolAgentStep",
     "ArtifactKind",
     "CannotContinueExecutionAction",
+    "CommentSpec",
     "CompletedExecutionAction",
     "ConversationMessage",
     "Discrepancy",
@@ -109,11 +120,16 @@ __all__ = [
     "FolioType",
     "format_conversation_history",
     "HealthResponse",
+    "MathAuditorToolReportArtifact",
     "NeedContentFileRequest",
     "NeedContentResponse",
     "NextExecutionAction",
     "OrchestratorRequest",
     "OrchestratorResponse",
+    "PdfCommentInstruction",
+    "PdfCommentReport",
+    "PdfCommentRequest",
+    "PdfCommentResponse",
     "PdfContentType",
     "PdfEditRequest",
     "PdfEditResponse",
@@ -136,8 +152,10 @@ __all__ = [
     "Severity",
     "StepKind",
     "SupportedCapability",
+    "TextChunk",
     "ToolCallExecutionAction",
     "ToolOperationStep",
+    "ToolReportArtifact",
     "UnsupportedCapabilityResponse",
     "Verdict",
     "WorkflowArtifact",

engine/src/stirling/contracts/comments.py

@@ -0,0 +1,36 @@
+"""Structured sticky-note comment specs for the ``add-comments`` tool.
+
+The ``/api/v1/misc/add-comments`` tool takes a JSON string of comment specs
+(see :class:`stirling.models.tool_models.AddCommentsParams`). This module
+defines the typed Python shape we serialise into that string so callers
+don't have to hand-roll dictionaries.
+"""
+
+from __future__ import annotations
+
+from pydantic import Field
+
+from stirling.models import ApiModel
+
+
+class CommentSpec(ApiModel):
+    """Sticky-note spec serialised into the ``comments`` JSON string sent to
+    ``/api/v1/misc/add-comments``. The backend's tool contract takes the JSON
+    string form, not this type; this is the engine-side structured representation.
+    """
+
+    page_index: int = Field(description="0-indexed page number.")
+    x: float = Field(description="Bottom-left x coord of the icon (PDF user-space).")
+    y: float = Field(description="Bottom-left y coord of the icon (PDF user-space).")
+    width: float = Field(description="Width of the icon in user-space units.")
+    height: float = Field(description="Height of the icon in user-space units.")
+    text: str = Field(description="Comment body shown in the popup.")
+    author: str | None = Field(default=None)
+    subject: str | None = Field(default=None)
+    anchor_text: str | None = Field(
+        default=None,
+        description=(
+            "Optional text snippet to locate on the page; when set, the server anchors"
+            " the icon at the first matching line and ignores the x/y coords."
+        ),
+    )

engine/src/stirling/contracts/common.py

@@ -5,6 +5,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.agent_tool_models import AGENT_OPERATIONS, AgentToolId, AnyParamModel, AnyToolId
 
@@ -67,6 +68,7 @@ class ArtifactKind(StrEnum):
     """
 
     EXTRACTED_TEXT = "extracted_text"
+    TOOL_REPORT = "tool_report"
 
 
 class StepKind(StrEnum):
@@ -80,6 +82,7 @@ class SupportedCapability(StrEnum):
     ORCHESTRATE = "orchestrate"
     PDF_EDIT = "pdf_edit"
     PDF_QUESTION = "pdf_question"
+    PDF_REVIEW = "pdf_review"
     AGENT_DRAFT = "agent_draft"
     AGENT_REVISE = "agent_revise"
     AGENT_NEXT_ACTION = "agent_next_action"
@@ -122,6 +125,27 @@ class NeedContentResponse(ApiModel):
     max_characters: int
 
 
+class MathAuditorToolReportArtifact(ApiModel):
+    """Structured Verdict produced by the math-auditor on a previous orchestrator turn.
+
+    New specialists that the orchestrator needs to digest on a resume turn
+    should add a sibling artifact type here and lift this into a discriminated
+    union keyed on ``source_tool``.
+
+    Java counterpart: {@code PdfContentExtractor.ToolReportArtifact}.
+    """
+
+    kind: Literal[ArtifactKind.TOOL_REPORT] = ArtifactKind.TOOL_REPORT
+    source_tool: Literal[AgentToolId.MATH_AUDITOR_AGENT] = AgentToolId.MATH_AUDITOR_AGENT
+    report: Verdict
+
+
+# Type alias kept around so callers don't have to know there's only one variant
+# today; lifts into a discriminated union when a second consumer-side report
+# appears.
+ToolReportArtifact = MathAuditorToolReportArtifact
+
+
 class ToolOperationStep(ApiModel):
     kind: Literal[StepKind.TOOL] = StepKind.TOOL
     tool: AnyToolId

engine/src/stirling/contracts/orchestrator.py

@@ -13,6 +13,7 @@ from .common import (
     ExtractedFileText,
     NeedContentResponse,
     SupportedCapability,
+    ToolReportArtifact,
     WorkflowOutcome,
 )
 from .execution import NextExecutionAction
@@ -25,7 +26,7 @@ class ExtractedTextArtifact(ApiModel):
     files: list[ExtractedFileText] = Field(default_factory=list)
 
 
-WorkflowArtifact = Annotated[ExtractedTextArtifact, Field(discriminator="kind")]
+WorkflowArtifact = Annotated[ExtractedTextArtifact | ToolReportArtifact, Field(discriminator="kind")]
 
 
 class OrchestratorRequest(ApiModel):

engine/src/stirling/contracts/pdf_comments.py

@@ -0,0 +1,150 @@
+"""
+PDF Comment Agent — shared models for the Java-Python protocol.
+
+The Java backend extracts positioned text chunks from a PDF and sends them
+along with a user prompt to the Python engine. Python selects the chunks
+that warrant a comment and returns an instruction list; Java then applies
+the actual PDF sticky-note annotations.
+
+Python never touches the PDF bytes. It only sees pre-extracted text with
+stable ids and must echo those ids back so Java can resolve each comment
+to its anchor.
+"""
+
+from __future__ import annotations
+
+from pydantic import Field
+
+from stirling.models import ApiModel
+
+# Bounds shared between the on-wire contract (enforced by pydantic) and any
+# Python-side defence-in-depth validation. Java enforces its own caps before
+# sending, but a malicious or buggy direct caller could otherwise ship an
+# unbounded payload.
+MAX_USER_MESSAGE_LENGTH = 4_000
+MAX_CHUNK_TEXT_LENGTH = 1_000
+MAX_COMMENT_TEXT_LENGTH = 2_000
+MAX_CHUNKS_PER_REQUEST = 2_500  # a hair above Java's 2000 cap — soft ceiling
+
+
+class TextChunk(ApiModel):
+    """One positioned text chunk extracted from a PDF page by Java.
+
+    The ``id`` is the stable handle used to anchor a comment to this chunk;
+    Python must echo it back verbatim on any comment that targets this chunk.
+    The bounding box is in PDF user-space (origin = bottom-left of the page).
+    """
+
+    id: str = Field(
+        min_length=1,
+        max_length=64,
+        description="Stable id, typically 'p{page}-c{chunk}'. Must be echoed unchanged on returned comments.",
+    )
+    page: int = Field(ge=0, description="0-indexed page number this chunk lives on.")
+    x: float = Field(description="PDF user-space x of the chunk's bounding box (bottom-left origin).")
+    y: float = Field(description="PDF user-space y of the chunk's bounding box (bottom-left origin).")
+    width: float = Field(ge=0, description="Width of the chunk's bounding box, in PDF user-space units.")
+    height: float = Field(ge=0, description="Height of the chunk's bounding box, in PDF user-space units.")
+    text: str = Field(
+        min_length=1,
+        max_length=MAX_CHUNK_TEXT_LENGTH,
+        description="The extracted text for this chunk. Typically one line.",
+    )
+
+
+class PdfCommentRequest(ApiModel):
+    """Request body Java sends to POST /api/v1/ai/pdf-comment-agent/generate.
+
+    Carries the user's natural-language instruction plus the list of text
+    chunks Java was able to extract from the PDF.
+    """
+
+    session_id: str = Field(
+        min_length=1,
+        max_length=128,
+        description="Opaque handle Java uses to correlate the request with its in-flight PDF job.",
+    )
+    user_message: str = Field(
+        min_length=1,
+        max_length=MAX_USER_MESSAGE_LENGTH,
+        description="The end-user prompt describing what the AI should comment on.",
+    )
+    chunks: list[TextChunk] = Field(
+        default_factory=list,
+        max_length=MAX_CHUNKS_PER_REQUEST,
+        description="All positioned text chunks Java extracted from the PDF; may be empty if the PDF has no text.",
+    )
+
+
+class PdfCommentInstruction(ApiModel):
+    """One review comment the agent wants Java to apply to the PDF.
+
+    ``chunk_id`` MUST match the id of a chunk that appeared in the request;
+    Java uses it to resolve the bounding box and anchor the sticky-note
+    annotation. Comments referencing an unknown id are dropped.
+    """
+
+    chunk_id: str = Field(
+        min_length=1,
+        max_length=64,
+        description="Id of the input chunk this comment anchors to. Must match an input chunk.id.",
+    )
+    comment_text: str = Field(
+        min_length=1,
+        max_length=MAX_COMMENT_TEXT_LENGTH,
+        description="The comment body shown in the sticky-note popup. One or two sentences.",
+    )
+    author: str | None = Field(
+        default=None,
+        max_length=128,
+        description="Optional author label; Java falls back to a default when absent.",
+    )
+    subject: str | None = Field(
+        default=None,
+        max_length=256,
+        description="Optional short subject/title for the comment popup; Java falls back to a default when absent.",
+    )
+
+
+class PdfCommentResponse(ApiModel):
+    """Response body the agent returns for POST /api/v1/ai/pdf-comment-agent/generate.
+
+    ``session_id`` is echoed from the request so Java can match the reply to
+    its pending job. ``comments`` is the (possibly filtered) list of review
+    instructions Java should apply as PDF Text annotations.
+    """
+
+    session_id: str = Field(
+        min_length=1,
+        max_length=128,
+        description="Echoed from the request so Java can match the reply to its pending job.",
+    )
+    comments: list[PdfCommentInstruction] = Field(
+        default_factory=list,
+        description="Review comments to apply. Each chunk_id is guaranteed to match an input chunk.",
+    )
+    rationale: str = Field(
+        max_length=1_000,
+        description="One-sentence summary describing the agent's overall approach for traceability/logging.",
+    )
+
+
+class PdfCommentReport(ApiModel):
+    """Structured report surfaced by the pdf-comment-agent tool alongside the
+    annotated PDF body. Mirrors the JSON shape the controller builds in
+    ``PdfCommentAgentController.buildReportHeader``.
+
+    Lands as the top-level ``AiWorkflowResponse.report`` on the COMPLETED
+    outcome (the pdf-comment-agent flow terminates without ``resume_with``,
+    so this never re-enters the orchestrator as a resume artifact).
+    """
+
+    annotations_applied: int = Field(
+        ge=0, description="Number of sticky-note annotations actually written into the PDF."
+    )
+    instructions_received: int = Field(
+        ge=0, description="Number of comment instructions the engine produced before filtering."
+    )
+    rationale: str | None = Field(
+        default=None, description="One-sentence summary the engine emitted alongside the comments."
+    )

engine/src/stirling/contracts/pdf_edit.py

@@ -6,7 +6,14 @@ from pydantic import Field
 
 from stirling.models import ApiModel
 
-from .common import ConversationMessage, ExtractedFileText, NeedContentResponse, ToolOperationStep, WorkflowOutcome
+from .common import (
+    ConversationMessage,
+    ExtractedFileText,
+    NeedContentResponse,
+    SupportedCapability,
+    ToolOperationStep,
+    WorkflowOutcome,
+)
 
 
 class PdfEditRequest(ApiModel):
@@ -21,6 +28,15 @@ class EditPlanResponse(ApiModel):
     summary: str
     rationale: str | None = None
     steps: list[ToolOperationStep]
+    resume_with: SupportedCapability | None = Field(
+        default=None,
+        description=(
+            "Optional: if set, Java runs the plan steps then re-invokes the orchestrator with"
+            " the captured tool reports attached as ToolReportArtifacts and"
+            " resume_with set to this capability. Used by meta-agents that need to digest a"
+            " specialist's output (e.g. pdf_review consulting math-auditor)."
+        ),
+    )
 
 
 class EditClarificationRequest(ApiModel):

engine/src/stirling/contracts/pdf_questions.py

@@ -12,6 +12,7 @@ from .common import (
     NeedContentResponse,
     WorkflowOutcome,
 )
+from .pdf_edit import EditPlanResponse
 
 
 class PdfQuestionRequest(ApiModel):
@@ -25,6 +26,15 @@ 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):