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.
2026-07-16 03:20:46 +02:00 · 2026-03-26 10:35:47 +00:00
parent 9500acd69f
commit e10c5f6283
211 changed files with 3891 additions and 27744 deletions
@@ -0,0 +1,108 @@
+from __future__ import annotations
+
+from pydantic_ai import Agent
+from pydantic_ai.output import NativeOutput
+
+from stirling.agents.pdf_edit import PdfEditAgent
+from stirling.contracts import (
+    AgentDraft,
+    AgentDraftRequest,
+    AgentDraftResponse,
+    AgentDraftWorkflowResponse,
+    AgentRevisionRequest,
+    AgentRevisionResponse,
+    AgentRevisionWorkflowResponse,
+    AiToolAgentStep,
+    ConversationMessage,
+    EditCannotDoResponse,
+    EditClarificationRequest,
+    EditPlanResponse,
+    PdfEditRequest,
+)
+from stirling.models import ApiModel
+from stirling.services import AppRuntime
+
+
+class UserSpecMetadata(ApiModel):
+    name: str
+    description: str
+    objective: str
+
+
+class UserSpecAgent:
+    def __init__(self, runtime: AppRuntime) -> None:
+        self.runtime = runtime
+        self.pdf_edit_agent = PdfEditAgent(runtime)
+        self.agent = Agent(
+            model=runtime.smart_model,
+            output_type=NativeOutput(UserSpecMetadata),
+            system_prompt=(
+                "Create or revise a saved agent draft from the provided request and edit plan. "
+                "Return a concise name, description, and objective. "
+                "Keep the workflow grounded and practical."
+            ),
+            model_settings=runtime.smart_model_settings,
+        )
+
+    async def draft(self, request: AgentDraftRequest) -> AgentDraftWorkflowResponse:
+        edit_plan = await self._build_edit_plan(request.user_message)
+        if not isinstance(edit_plan, EditPlanResponse):
+            return edit_plan
+        return AgentDraftResponse(draft=await self._run_draft_agent(request, edit_plan))
+
+    async def revise(self, request: AgentRevisionRequest) -> AgentRevisionWorkflowResponse:
+        edit_plan = await self._build_edit_plan(
+            f"Current objective: {request.current_draft.objective}\nRevision request: {request.user_message}"
+        )
+        if not isinstance(edit_plan, EditPlanResponse):
+            return edit_plan
+        return AgentRevisionResponse(draft=await self._run_revision_agent(request, edit_plan))
+
+    async def _run_draft_agent(self, request: AgentDraftRequest, edit_plan: EditPlanResponse) -> AgentDraft:
+        metadata = (await self.agent.run(self._build_draft_prompt(request, edit_plan))).output
+        return AgentDraft(
+            name=metadata.name,
+            description=metadata.description,
+            objective=metadata.objective,
+            steps=[*edit_plan.steps],
+        )
+
+    async def _run_revision_agent(self, request: AgentRevisionRequest, edit_plan: EditPlanResponse) -> AgentDraft:
+        metadata = (await self.agent.run(self._build_revision_prompt(request, edit_plan))).output
+        preserved_ai_steps = [step for step in request.current_draft.steps if isinstance(step, AiToolAgentStep)]
+        return AgentDraft(
+            name=metadata.name,
+            description=metadata.description,
+            objective=metadata.objective,
+            steps=[*edit_plan.steps, *preserved_ai_steps],
+        )
+
+    def _build_draft_prompt(self, request: AgentDraftRequest, edit_plan: EditPlanResponse) -> str:
+        return (
+            f"User request:\n{request.user_message}\n\n"
+            f"Conversation history:\n{self._format_conversation_history(request.conversation_history)}\n\n"
+            f"Edit plan summary:\n{edit_plan.summary}\n\n"
+            f"Edit plan rationale:\n{edit_plan.rationale or 'None'}\n\n"
+            f"Edit plan steps:\n{edit_plan.model_dump_json(indent=2)}"
+        )
+
+    def _build_revision_prompt(self, request: AgentRevisionRequest, edit_plan: EditPlanResponse) -> str:
+        return (
+            f"Revision request:\n{request.user_message}\n\n"
+            f"Conversation history:\n{self._format_conversation_history(request.conversation_history)}\n\n"
+            f"Current draft:\n{request.current_draft.model_dump_json(indent=2)}\n\n"
+            f"Edit plan summary:\n{edit_plan.summary}\n\n"
+            f"Edit plan rationale:\n{edit_plan.rationale or 'None'}\n\n"
+            f"Edit plan steps:\n{edit_plan.model_dump_json(indent=2)}"
+        )
+
+    def _format_conversation_history(self, conversation_history: list[ConversationMessage]) -> str:
+        if not conversation_history:
+            return "None"
+        return "\n".join(f"- {message.role}: {message.content}" for message in conversation_history)
+
+    async def _build_edit_plan(
+        self,
+        user_message: str,
+    ) -> EditPlanResponse | EditClarificationRequest | EditCannotDoResponse:
+        return await self.pdf_edit_agent.handle(PdfEditRequest(user_message=user_message))