Redesign Python AI engine (#5991)

# Description of Changes
Redesign the Python AI engine to be properly agentic and make use of
`pydantic-ai` instead of `langchain` for correctness and ergonomics.
This should be a good foundation for us to build our AI engine on going
forwards.

engine/src/stirling/contracts/__init__.py

@@ -0,0 +1,74 @@
+from .agent_drafts import (
+    AgentDraft,
+    AgentDraftRequest,
+    AgentDraftResponse,
+    AgentDraftWorkflowResponse,
+    AgentRevisionRequest,
+    AgentRevisionResponse,
+    AgentRevisionWorkflowResponse,
+)
+from .agent_specs import AgentSpec, AgentSpecStep, AiToolAgentStep
+from .common import ConversationMessage, PdfTextSelection, ToolOperationStep
+from .execution import (
+    AgentExecutionRequest,
+    CannotContinueExecutionAction,
+    CompletedExecutionAction,
+    ExecutionContext,
+    ExecutionStepResult,
+    NextExecutionAction,
+    ToolCallExecutionAction,
+)
+from .health import HealthResponse
+from .orchestrator import OrchestratorRequest, OrchestratorResponse, SupportedCapability, UnsupportedCapabilityResponse
+from .pdf_edit import (
+    EditCannotDoResponse,
+    EditClarificationRequest,
+    EditPlanResponse,
+    PdfEditRequest,
+    PdfEditResponse,
+)
+from .pdf_questions import (
+    PdfQuestionAnswerResponse,
+    PdfQuestionNeedTextResponse,
+    PdfQuestionNotFoundResponse,
+    PdfQuestionRequest,
+    PdfQuestionResponse,
+)
+
+__all__ = [
+    "AgentDraft",
+    "AgentDraftRequest",
+    "AgentDraftResponse",
+    "AgentDraftWorkflowResponse",
+    "AgentExecutionRequest",
+    "AgentRevisionRequest",
+    "AgentRevisionResponse",
+    "AgentRevisionWorkflowResponse",
+    "AgentSpec",
+    "AgentSpecStep",
+    "AiToolAgentStep",
+    "CannotContinueExecutionAction",
+    "ConversationMessage",
+    "CompletedExecutionAction",
+    "EditCannotDoResponse",
+    "EditClarificationRequest",
+    "EditPlanResponse",
+    "ExecutionContext",
+    "ExecutionStepResult",
+    "HealthResponse",
+    "NextExecutionAction",
+    "OrchestratorRequest",
+    "OrchestratorResponse",
+    "PdfEditRequest",
+    "PdfEditResponse",
+    "PdfQuestionAnswerResponse",
+    "PdfQuestionNotFoundResponse",
+    "PdfQuestionNeedTextResponse",
+    "PdfQuestionRequest",
+    "PdfQuestionResponse",
+    "PdfTextSelection",
+    "SupportedCapability",
+    "ToolOperationStep",
+    "ToolCallExecutionAction",
+    "UnsupportedCapabilityResponse",
+]

engine/src/stirling/contracts/agent_drafts.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+from typing import Annotated, Literal
+
+from pydantic import Field
+
+from stirling.models import ApiModel
+
+from .agent_specs import AgentSpecStep
+from .common import ConversationMessage
+from .pdf_edit import EditCannotDoResponse, EditClarificationRequest
+
+
+class AgentDraftStep(ApiModel):
+    kind: Literal["tool", "ai_tool"]
+    title: str
+    description: str
+
+
+class AgentDraft(ApiModel):
+    name: str
+    description: str
+    objective: str
+    steps: list[AgentSpecStep] = Field(default_factory=list)
+
+
+class AgentDraftRequest(ApiModel):
+    user_message: str
+    conversation_history: list[ConversationMessage] = Field(default_factory=list)
+
+
+class AgentDraftResponse(ApiModel):
+    outcome: Literal["draft"] = "draft"
+    draft: AgentDraft
+
+
+class AgentRevisionRequest(ApiModel):
+    user_message: str
+    conversation_history: list[ConversationMessage] = Field(default_factory=list)
+    current_draft: AgentDraft
+
+
+class AgentRevisionResponse(ApiModel):
+    outcome: Literal["draft"] = "draft"
+    draft: AgentDraft
+
+
+AgentDraftWorkflowResponse = Annotated[
+    AgentDraftResponse | EditClarificationRequest | EditCannotDoResponse,
+    Field(discriminator="outcome"),
+]
+
+
+AgentRevisionWorkflowResponse = Annotated[
+    AgentRevisionResponse | EditClarificationRequest | EditCannotDoResponse,
+    Field(discriminator="outcome"),
+]

engine/src/stirling/contracts/agent_specs.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+from typing import Annotated, Literal
+
+from pydantic import Field
+
+from stirling.models import ApiModel, OperationId
+
+from .common import ToolOperationStep
+
+
+class AiToolAgentStep(ApiModel):
+    kind: Literal["ai_tool"] = "ai_tool"
+    title: str
+    description: str
+    tool: OperationId
+    instruction: str
+
+
+AgentSpecStep = Annotated[ToolOperationStep | AiToolAgentStep, Field(discriminator="kind")]
+
+
+class AgentSpec(ApiModel):
+    name: str
+    description: str
+    objective: str
+    steps: list[AgentSpecStep] = Field(default_factory=list)

engine/src/stirling/contracts/common.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+from typing import Literal
+
+from pydantic import model_validator
+
+from stirling.models import OPERATIONS, ApiModel, OperationId, ParamToolModel
+
+
+class ConversationMessage(ApiModel):
+    role: str
+    content: str
+
+
+class PdfTextSelection(ApiModel):
+    page_number: int | None = None
+    text: str
+
+
+class ToolOperationStep(ApiModel):
+    kind: Literal["tool"] = "tool"
+    tool: OperationId
+    parameters: ParamToolModel
+
+    @model_validator(mode="after")
+    def validate_tool_parameter_pairing(self) -> ToolOperationStep:
+        expected_type = OPERATIONS[self.tool]
+        if not isinstance(self.parameters, expected_type):
+            actual_type = type(self.parameters).__name__
+            expected_type_name = expected_type.__name__
+            raise ValueError(f"Parameters for tool {self.tool.value} must be {expected_type_name}, got {actual_type}.")
+        return self

engine/src/stirling/contracts/execution.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+from typing import Annotated, Any, Literal
+
+from pydantic import Field
+
+from stirling.models import ApiModel, OperationId, ParamToolModel
+
+from .agent_specs import AgentSpec
+
+
+class ExecutionStepResult(ApiModel):
+    step_index: int
+    tool: OperationId | None = None
+    success: bool
+    output_summary: str | None = None
+    output_data: dict[str, Any] = Field(default_factory=dict)
+
+
+class ExecutionContext(ApiModel):
+    trigger_type: str | None = None
+    input_files: list[str] = Field(default_factory=list)
+    metadata: dict[str, Any] = Field(default_factory=dict)
+
+
+class AgentExecutionRequest(ApiModel):
+    agent_spec: AgentSpec
+    current_step_index: int
+    execution_context: ExecutionContext
+    previous_step_results: list[ExecutionStepResult] = Field(default_factory=list)
+
+
+class ToolCallExecutionAction(ApiModel):
+    outcome: Literal["tool_call"] = "tool_call"
+    tool: OperationId
+    parameters: ParamToolModel
+    rationale: str | None = None
+
+
+class CompletedExecutionAction(ApiModel):
+    outcome: Literal["completed"] = "completed"
+    summary: str
+
+
+class CannotContinueExecutionAction(ApiModel):
+    outcome: Literal["cannot_continue"] = "cannot_continue"
+    reason: str
+
+
+NextExecutionAction = Annotated[
+    ToolCallExecutionAction | CompletedExecutionAction | CannotContinueExecutionAction,
+    Field(discriminator="outcome"),
+]

engine/src/stirling/contracts/health.py

@@ -0,0 +1,9 @@
+from __future__ import annotations
+
+from stirling.models import ApiModel
+
+
+class HealthResponse(ApiModel):
+    status: str
+    smart_model: str
+    fast_model: str

engine/src/stirling/contracts/orchestrator.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+from enum import StrEnum
+from typing import Annotated, Literal
+
+from pydantic import Field
+
+from stirling.models import ApiModel
+
+from .agent_drafts import AgentDraftResponse
+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 OrchestratorRequest(ApiModel):
+    user_message: str
+    conversation_id: str | None = None
+
+
+class UnsupportedCapabilityResponse(ApiModel):
+    outcome: Literal["unsupported_capability"] = "unsupported_capability"
+    capability: str
+    message: str
+
+
+OrchestratorResponse = Annotated[
+    PdfEditResponse | PdfQuestionResponse | AgentDraftResponse | NextExecutionAction | UnsupportedCapabilityResponse,
+    Field(discriminator="outcome"),
+]

engine/src/stirling/contracts/pdf_edit.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+from typing import Annotated, Literal
+
+from pydantic import Field
+
+from stirling.models import ApiModel
+
+from .common import ToolOperationStep
+
+
+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"
+    summary: str
+    rationale: str | None = None
+    steps: list[ToolOperationStep]
+
+
+class EditClarificationRequest(ApiModel):
+    outcome: Literal["need_clarification"] = "need_clarification"
+    question: str
+    reason: str
+
+
+class EditCannotDoResponse(ApiModel):
+    outcome: Literal["cannot_do"] = "cannot_do"
+    reason: str
+
+
+PdfEditResponse = Annotated[
+    EditPlanResponse | EditClarificationRequest | EditCannotDoResponse,
+    Field(discriminator="outcome"),
+]

engine/src/stirling/contracts/pdf_questions.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+from typing import Annotated, Literal
+
+from pydantic import Field
+
+from stirling.models import ApiModel
+
+
+class PdfQuestionRequest(ApiModel):
+    question: str
+    conversation_id: str | None = None
+    extracted_text: str = ""
+    file_name: str | None = None
+
+
+class PdfQuestionAnswerResponse(ApiModel):
+    outcome: Literal["answer"] = "answer"
+    answer: str
+    evidence: list[str] = Field(default_factory=list)
+
+
+class PdfQuestionNeedTextResponse(ApiModel):
+    outcome: Literal["need_text"] = "need_text"
+    reason: str
+
+
+class PdfQuestionNotFoundResponse(ApiModel):
+    outcome: Literal["not_found"] = "not_found"
+    reason: str
+
+
+PdfQuestionResponse = Annotated[
+    PdfQuestionAnswerResponse | PdfQuestionNeedTextResponse | PdfQuestionNotFoundResponse,
+    Field(discriminator="outcome"),
+]