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

@@ -0,0 +1,7 @@
+"""API surface for the Stirling AI service."""
+
+from .app import app
+
+__all__ = [
+    "app",
+]

engine/src/stirling/api/app.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+from contextlib import asynccontextmanager
+from typing import Annotated
+
+from fastapi import Depends, FastAPI
+
+from stirling.agents import ExecutionPlanningAgent, OrchestratorAgent, PdfEditAgent, PdfQuestionAgent, UserSpecAgent
+from stirling.api.routes import (
+    agent_draft_router,
+    execution_router,
+    orchestrator_router,
+    pdf_edit_router,
+    pdf_question_router,
+)
+from stirling.config import AppSettings, load_settings
+from stirling.contracts import HealthResponse
+from stirling.services import build_runtime
+
+
+def _load_startup_settings(fast_api: FastAPI) -> AppSettings:
+    override = fast_api.dependency_overrides.get(load_settings)
+    if override is not None:
+        return override()
+    return load_settings()
+
+
+@asynccontextmanager
+async def lifespan(fast_api: FastAPI):
+    # Load env vars on startup so we can immediately crash if required env vars aren't set
+    settings = _load_startup_settings(fast_api)
+    runtime = build_runtime(settings)
+    fast_api.state.settings = settings
+    fast_api.state.runtime = runtime
+    fast_api.state.orchestrator_agent = OrchestratorAgent(runtime)
+    fast_api.state.pdf_edit_agent = PdfEditAgent(runtime)
+    fast_api.state.pdf_question_agent = PdfQuestionAgent(runtime)
+    fast_api.state.user_spec_agent = UserSpecAgent(runtime)
+    fast_api.state.execution_planning_agent = ExecutionPlanningAgent(runtime)
+    yield
+
+
+app = FastAPI(title="Stirling AI Engine", lifespan=lifespan, version="0.1.0")
+app.include_router(orchestrator_router)
+app.include_router(pdf_edit_router)
+app.include_router(pdf_question_router)
+app.include_router(agent_draft_router)
+app.include_router(execution_router)
+
+
+@app.get("/health", response_model=HealthResponse)
+async def healthcheck(settings: Annotated[AppSettings, Depends(load_settings)]) -> HealthResponse:
+    return HealthResponse(
+        status="ok",
+        smart_model=settings.smart_model_name,
+        fast_model=settings.fast_model_name,
+    )

engine/src/stirling/api/dependencies.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+from fastapi import Request
+
+from stirling.agents import ExecutionPlanningAgent, OrchestratorAgent, PdfEditAgent, PdfQuestionAgent, UserSpecAgent
+from stirling.services import AppRuntime
+
+
+def get_runtime(request: Request) -> AppRuntime:
+    return request.app.state.runtime
+
+
+def get_orchestrator_agent(request: Request) -> OrchestratorAgent:
+    return request.app.state.orchestrator_agent
+
+
+def get_pdf_edit_agent(request: Request) -> PdfEditAgent:
+    return request.app.state.pdf_edit_agent
+
+
+def get_pdf_question_agent(request: Request) -> PdfQuestionAgent:
+    return request.app.state.pdf_question_agent
+
+
+def get_user_spec_agent(request: Request) -> UserSpecAgent:
+    return request.app.state.user_spec_agent
+
+
+def get_execution_planning_agent(request: Request) -> ExecutionPlanningAgent:
+    return request.app.state.execution_planning_agent

engine/src/stirling/api/routes/__init__.py

@@ -0,0 +1,13 @@
+from .agent_drafts import router as agent_draft_router
+from .execution import router as execution_router
+from .orchestrator import router as orchestrator_router
+from .pdf_edit import router as pdf_edit_router
+from .pdf_questions import router as pdf_question_router
+
+__all__ = [
+    "agent_draft_router",
+    "execution_router",
+    "orchestrator_router",
+    "pdf_edit_router",
+    "pdf_question_router",
+]

engine/src/stirling/api/routes/agent_drafts.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+from typing import Annotated
+
+from fastapi import APIRouter, Depends
+
+from stirling.agents import UserSpecAgent
+from stirling.api.dependencies import get_user_spec_agent
+from stirling.contracts import (
+    AgentDraftRequest,
+    AgentDraftWorkflowResponse,
+    AgentRevisionRequest,
+    AgentRevisionWorkflowResponse,
+)
+
+router = APIRouter(prefix="/api/v1/agents", tags=["agents"])
+
+
+@router.post("/draft", response_model=AgentDraftWorkflowResponse)
+async def draft_agent(
+    request: AgentDraftRequest,
+    agent: Annotated[UserSpecAgent, Depends(get_user_spec_agent)],
+) -> AgentDraftWorkflowResponse:
+    return await agent.draft(request)
+
+
+@router.post("/revise", response_model=AgentRevisionWorkflowResponse)
+async def revise_agent(
+    request: AgentRevisionRequest,
+    agent: Annotated[UserSpecAgent, Depends(get_user_spec_agent)],
+) -> AgentRevisionWorkflowResponse:
+    return await agent.revise(request)

engine/src/stirling/api/routes/execution.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+from typing import Annotated
+
+from fastapi import APIRouter, Depends
+
+from stirling.agents import ExecutionPlanningAgent
+from stirling.api.dependencies import get_execution_planning_agent
+from stirling.contracts import AgentExecutionRequest, NextExecutionAction
+
+router = APIRouter(prefix="/api/v1/agents", tags=["agents"])
+
+
+@router.post("/next-action", response_model=NextExecutionAction)
+async def next_action(
+    request: AgentExecutionRequest,
+    agent: Annotated[ExecutionPlanningAgent, Depends(get_execution_planning_agent)],
+) -> NextExecutionAction:
+    return await agent.next_action(request)

engine/src/stirling/api/routes/orchestrator.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+from typing import Annotated
+
+from fastapi import APIRouter, Depends
+
+from stirling.agents import OrchestratorAgent
+from stirling.api.dependencies import get_orchestrator_agent
+from stirling.contracts import OrchestratorRequest, OrchestratorResponse
+
+router = APIRouter(prefix="/api/v1/orchestrator", tags=["orchestrator"])
+
+
+@router.post("", response_model=OrchestratorResponse)
+async def orchestrate(
+    request: OrchestratorRequest,
+    agent: Annotated[OrchestratorAgent, Depends(get_orchestrator_agent)],
+) -> OrchestratorResponse:
+    return await agent.handle(request)

engine/src/stirling/api/routes/pdf_edit.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+from typing import Annotated
+
+from fastapi import APIRouter, Depends
+
+from stirling.agents import PdfEditAgent
+from stirling.api.dependencies import get_pdf_edit_agent
+from stirling.contracts import PdfEditRequest, PdfEditResponse
+
+router = APIRouter(prefix="/api/v1/pdf/edit", tags=["pdf-edit"])
+
+
+@router.post("", response_model=PdfEditResponse)
+async def pdf_edit(
+    request: PdfEditRequest,
+    agent: Annotated[PdfEditAgent, Depends(get_pdf_edit_agent)],
+) -> PdfEditResponse:
+    return await agent.handle(request)

engine/src/stirling/api/routes/pdf_questions.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+from typing import Annotated
+
+from fastapi import APIRouter, Depends
+
+from stirling.agents import PdfQuestionAgent
+from stirling.api.dependencies import get_pdf_question_agent
+from stirling.contracts import PdfQuestionRequest, PdfQuestionResponse
+
+router = APIRouter(prefix="/api/v1/pdf/questions", tags=["pdf-questions"])
+
+
+@router.post("", response_model=PdfQuestionResponse)
+async def pdf_questions(
+    request: PdfQuestionRequest,
+    agent: Annotated[PdfQuestionAgent, Depends(get_pdf_question_agent)],
+) -> PdfQuestionResponse:
+    return await agent.handle(request)