Add Java orchestrator to connect to the AI engine (#6003)

# Description of Changes
Add Java orchestration layer which can connect and go back and forth
with the AI engine to get results for the user. It's expected that the
AI engine will not be publicly available and this Java layer will always
be in front of it, to manage sessions and auth etc.

engine/src/stirling/contracts/__init__.py

@@ -8,7 +8,17 @@ from .agent_drafts import (
     AgentRevisionWorkflowResponse,
 )
 from .agent_specs import AgentSpec, AgentSpecStep, AiToolAgentStep
-from .common import ConversationMessage, PdfTextSelection, ToolOperationStep
+from .common import (
+    ArtifactKind,
+    ConversationMessage,
+    ExtractedFileText,
+    PdfContentType,
+    PdfTextSelection,
+    StepKind,
+    SupportedCapability,
+    ToolOperationStep,
+    WorkflowOutcome,
+)
 from .execution import (
     AgentExecutionRequest,
     CannotContinueExecutionAction,
@@ -19,7 +29,13 @@ from .execution import (
     ToolCallExecutionAction,
 )
 from .health import HealthResponse
-from .orchestrator import OrchestratorRequest, OrchestratorResponse, SupportedCapability, UnsupportedCapabilityResponse
+from .orchestrator import (
+    ExtractedTextArtifact,
+    OrchestratorRequest,
+    OrchestratorResponse,
+    UnsupportedCapabilityResponse,
+    WorkflowArtifact,
+)
 from .pdf_edit import (
     EditCannotDoResponse,
     EditClarificationRequest,
@@ -28,14 +44,16 @@ from .pdf_edit import (
     PdfEditResponse,
 )
 from .pdf_questions import (
+    NeedContentFileRequest,
     PdfQuestionAnswerResponse,
-    PdfQuestionNeedTextResponse,
+    PdfQuestionNeedContentResponse,
     PdfQuestionNotFoundResponse,
     PdfQuestionRequest,
     PdfQuestionResponse,
 )
 
 __all__ = [
+    "ArtifactKind",
     "AgentDraft",
     "AgentDraftRequest",
     "AgentDraftResponse",
@@ -49,6 +67,7 @@ __all__ = [
     "AiToolAgentStep",
     "CannotContinueExecutionAction",
     "ConversationMessage",
+    "ExtractedFileText",
     "CompletedExecutionAction",
     "EditCannotDoResponse",
     "EditClarificationRequest",
@@ -56,19 +75,25 @@ __all__ = [
     "ExecutionContext",
     "ExecutionStepResult",
     "HealthResponse",
+    "NeedContentFileRequest",
     "NextExecutionAction",
+    "ExtractedTextArtifact",
     "OrchestratorRequest",
     "OrchestratorResponse",
     "PdfEditRequest",
     "PdfEditResponse",
     "PdfQuestionAnswerResponse",
     "PdfQuestionNotFoundResponse",
-    "PdfQuestionNeedTextResponse",
+    "PdfContentType",
+    "PdfQuestionNeedContentResponse",
     "PdfQuestionRequest",
     "PdfQuestionResponse",
     "PdfTextSelection",
+    "StepKind",
     "SupportedCapability",
     "ToolOperationStep",
     "ToolCallExecutionAction",
+    "WorkflowOutcome",
     "UnsupportedCapabilityResponse",
+    "WorkflowArtifact",
 ]

engine/src/stirling/contracts/agent_drafts.py

@@ -7,12 +7,12 @@ from pydantic import Field
 from stirling.models import ApiModel
 
 from .agent_specs import AgentSpecStep
-from .common import ConversationMessage
+from .common import ConversationMessage, StepKind, WorkflowOutcome
 from .pdf_edit import EditCannotDoResponse, EditClarificationRequest
 
 
 class AgentDraftStep(ApiModel):
-    kind: Literal["tool", "ai_tool"]
+    kind: Literal[StepKind.TOOL, StepKind.AI_TOOL]
     title: str
     description: str
 
@@ -30,7 +30,7 @@ class AgentDraftRequest(ApiModel):
 
 
 class AgentDraftResponse(ApiModel):
-    outcome: Literal["draft"] = "draft"
+    outcome: Literal[WorkflowOutcome.DRAFT] = WorkflowOutcome.DRAFT
     draft: AgentDraft
 
 
@@ -41,7 +41,7 @@ class AgentRevisionRequest(ApiModel):
 
 
 class AgentRevisionResponse(ApiModel):
-    outcome: Literal["draft"] = "draft"
+    outcome: Literal[WorkflowOutcome.DRAFT] = WorkflowOutcome.DRAFT
     draft: AgentDraft
 
 

engine/src/stirling/contracts/agent_specs.py

@@ -6,11 +6,11 @@ from pydantic import Field
 
 from stirling.models import ApiModel, OperationId
 
-from .common import ToolOperationStep
+from .common import StepKind, ToolOperationStep
 
 
 class AiToolAgentStep(ApiModel):
-    kind: Literal["ai_tool"] = "ai_tool"
+    kind: Literal[StepKind.AI_TOOL] = StepKind.AI_TOOL
     title: str
     description: str
     tool: OperationId

engine/src/stirling/contracts/common.py

@@ -1,12 +1,89 @@
 from __future__ import annotations
 
+from enum import StrEnum
 from typing import Literal
 
-from pydantic import model_validator
+from pydantic import Field, model_validator
 
 from stirling.models import OPERATIONS, ApiModel, OperationId, ParamToolModel
 
 
+class PdfContentType(StrEnum):
+    """Types of content that can be extracted from a PDF and sent to the AI.
+
+    Java counterpart: AiPdfContentType.java - values must stay in sync.
+    """
+
+    # Document-level structured data
+    PAGE_LAYOUT = "page_layout"
+    DOCUMENT_METADATA = "document_metadata"
+    ENCRYPTION_INFO = "encryption_info"
+    BOOKMARKS = "bookmarks"
+    LAYERS = "layers"
+    EMBEDDED_FILES = "embedded_files"
+    JAVASCRIPT = "javascript"
+    LINKS = "links"
+    IMAGE_INFO = "image_info"
+    FONTS = "fonts"
+
+    # Text and content
+    PAGE_TEXT = "page_text"
+    FULL_TEXT = "full_text"
+    FORM_FIELDS = "form_fields"
+    ANNOTATIONS = "annotations"
+    SIGNATURES = "signatures"
+    STRUCTURE_TREE = "structure_tree"
+    XMP_METADATA = "xmp_metadata"
+
+    # Heavy content
+    COMPLIANCE = "compliance"
+    IMAGES = "images"
+
+
+class WorkflowOutcome(StrEnum):
+    """Discriminator values for all workflow response unions (outcome field).
+
+    Java counterpart: AiWorkflowOutcome.java - values must stay in sync.
+    """
+
+    ANSWER = "answer"
+    NEED_CONTENT = "need_content"
+    NOT_FOUND = "not_found"
+    PLAN = "plan"
+    NEED_CLARIFICATION = "need_clarification"
+    CANNOT_DO = "cannot_do"
+    DRAFT = "draft"
+    TOOL_CALL = "tool_call"
+    COMPLETED = "completed"
+    CANNOT_CONTINUE = "cannot_continue"
+    UNSUPPORTED_CAPABILITY = "unsupported_capability"
+
+
+class ArtifactKind(StrEnum):
+    """Discriminator values for WorkflowArtifact unions (kind field).
+
+    Java counterpart: PdfContentExtractor.ArtifactKind - values must stay in sync.
+    """
+
+    EXTRACTED_TEXT = "extracted_text"
+
+
+class StepKind(StrEnum):
+    """Discriminator values for AgentSpecStep unions (kind field)."""
+
+    TOOL = "tool"
+    AI_TOOL = "ai_tool"
+
+
+class SupportedCapability(StrEnum):
+    ORCHESTRATE = "orchestrate"
+    PDF_EDIT = "pdf_edit"
+    PDF_QUESTION = "pdf_question"
+    AGENT_DRAFT = "agent_draft"
+    AGENT_REVISE = "agent_revise"
+    AGENT_NEXT_ACTION = "agent_next_action"
+
+
 class ConversationMessage(ApiModel):
     role: str
     content: str
@@ -17,8 +94,13 @@ class PdfTextSelection(ApiModel):
     text: str
 
 
+class ExtractedFileText(ApiModel):
+    file_name: str
+    pages: list[PdfTextSelection] = Field(default_factory=list)
+
+
 class ToolOperationStep(ApiModel):
-    kind: Literal["tool"] = "tool"
+    kind: Literal[StepKind.TOOL] = StepKind.TOOL
     tool: OperationId
     parameters: ParamToolModel
 

engine/src/stirling/contracts/execution.py

@@ -7,6 +7,7 @@ from pydantic import Field
 from stirling.models import ApiModel, OperationId, ParamToolModel
 
 from .agent_specs import AgentSpec
+from .common import WorkflowOutcome
 
 
 class ExecutionStepResult(ApiModel):
@@ -31,19 +32,19 @@ class AgentExecutionRequest(ApiModel):
 
 
 class ToolCallExecutionAction(ApiModel):
-    outcome: Literal["tool_call"] = "tool_call"
+    outcome: Literal[WorkflowOutcome.TOOL_CALL] = WorkflowOutcome.TOOL_CALL
     tool: OperationId
     parameters: ParamToolModel
     rationale: str | None = None
 
 
 class CompletedExecutionAction(ApiModel):
-    outcome: Literal["completed"] = "completed"
+    outcome: Literal[WorkflowOutcome.COMPLETED] = WorkflowOutcome.COMPLETED
     summary: str
 
 
 class CannotContinueExecutionAction(ApiModel):
-    outcome: Literal["cannot_continue"] = "cannot_continue"
+    outcome: Literal[WorkflowOutcome.CANNOT_CONTINUE] = WorkflowOutcome.CANNOT_CONTINUE
     reason: str
 
 

engine/src/stirling/contracts/orchestrator.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-from enum import StrEnum
 from typing import Annotated, Literal
 
 from pydantic import Field
@@ -8,27 +7,29 @@ from pydantic import Field
 from stirling.models import ApiModel
 
 from .agent_drafts import AgentDraftResponse
+from .common import ArtifactKind, ExtractedFileText, SupportedCapability, WorkflowOutcome
 from .execution import NextExecutionAction
 from .pdf_edit import PdfEditResponse
 from .pdf_questions import PdfQuestionResponse
 
 
-class SupportedCapability(StrEnum):
-    ORCHESTRATE = "orchestrate"
-    PDF_EDIT = "pdf_edit"
-    PDF_QUESTION = "pdf_question"
-    AGENT_DRAFT = "agent_draft"
-    AGENT_REVISE = "agent_revise"
-    AGENT_NEXT_ACTION = "agent_next_action"
+class ExtractedTextArtifact(ApiModel):
+    kind: Literal[ArtifactKind.EXTRACTED_TEXT] = ArtifactKind.EXTRACTED_TEXT
+    files: list[ExtractedFileText] = Field(default_factory=list)
+
+
+WorkflowArtifact = Annotated[ExtractedTextArtifact, Field(discriminator="kind")]
 
 
 class OrchestratorRequest(ApiModel):
     user_message: str
-    conversation_id: str | None = None
+    file_names: list[str]
+    artifacts: list[WorkflowArtifact] = Field(default_factory=list)
+    resume_with: SupportedCapability | None = None
 
 
 class UnsupportedCapabilityResponse(ApiModel):
-    outcome: Literal["unsupported_capability"] = "unsupported_capability"
+    outcome: Literal[WorkflowOutcome.UNSUPPORTED_CAPABILITY] = WorkflowOutcome.UNSUPPORTED_CAPABILITY
     capability: str
     message: str
 

engine/src/stirling/contracts/pdf_edit.py

@@ -6,30 +6,29 @@ from pydantic import Field
 
 from stirling.models import ApiModel
 
-from .common import ToolOperationStep
+from .common import ToolOperationStep, WorkflowOutcome
 
 
 class PdfEditRequest(ApiModel):
     user_message: str
-    conversation_id: str | None = None
     file_names: list[str] = Field(default_factory=list)
 
 
 class EditPlanResponse(ApiModel):
-    outcome: Literal["plan"] = "plan"
+    outcome: Literal[WorkflowOutcome.PLAN] = WorkflowOutcome.PLAN
     summary: str
     rationale: str | None = None
     steps: list[ToolOperationStep]
 
 
 class EditClarificationRequest(ApiModel):
-    outcome: Literal["need_clarification"] = "need_clarification"
+    outcome: Literal[WorkflowOutcome.NEED_CLARIFICATION] = WorkflowOutcome.NEED_CLARIFICATION
     question: str
     reason: str
 
 
 class EditCannotDoResponse(ApiModel):
-    outcome: Literal["cannot_do"] = "cannot_do"
+    outcome: Literal[WorkflowOutcome.CANNOT_DO] = WorkflowOutcome.CANNOT_DO
     reason: str
 
 

engine/src/stirling/contracts/pdf_questions.py

@@ -6,31 +6,42 @@ from pydantic import Field
 
 from stirling.models import ApiModel
 
+from .common import ExtractedFileText, PdfContentType, SupportedCapability, WorkflowOutcome
+
 
 class PdfQuestionRequest(ApiModel):
     question: str
-    conversation_id: str | None = None
-    extracted_text: str = ""
-    file_name: str | None = None
+    page_text: list[ExtractedFileText] = Field(default_factory=list)
+    file_names: list[str]
 
 
 class PdfQuestionAnswerResponse(ApiModel):
-    outcome: Literal["answer"] = "answer"
+    outcome: Literal[WorkflowOutcome.ANSWER] = WorkflowOutcome.ANSWER
     answer: str
-    evidence: list[str] = Field(default_factory=list)
+    evidence: list[ExtractedFileText] = Field(default_factory=list)
 
 
-class PdfQuestionNeedTextResponse(ApiModel):
-    outcome: Literal["need_text"] = "need_text"
+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["not_found"] = "not_found"
+    outcome: Literal[WorkflowOutcome.NOT_FOUND] = WorkflowOutcome.NOT_FOUND
     reason: str
 
 
 PdfQuestionResponse = Annotated[
-    PdfQuestionAnswerResponse | PdfQuestionNeedTextResponse | PdfQuestionNotFoundResponse,
+    PdfQuestionAnswerResponse | PdfQuestionNeedContentResponse | PdfQuestionNotFoundResponse,
     Field(discriminator="outcome"),
 ]