mirror of
https://github.com/arsvendg/Stirling-PDF.git
synced 2026-07-16 03:20:46 +02:00
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.
This commit is contained in:
James Brunton
@@ -0,0 +1,17 @@
|
||||
"""Agent modules for Stirling AI reasoning flows."""
|
||||
|
||||
from .execution import ExecutionPlanningAgent
|
||||
from .orchestrator import OrchestratorAgent
|
||||
from .pdf_edit import PdfEditAgent, PdfEditParameterSelector, PdfEditPlanSelection
|
||||
from .pdf_questions import PdfQuestionAgent
|
||||
from .user_spec import UserSpecAgent
|
||||
|
||||
__all__ = [
|
||||
"ExecutionPlanningAgent",
|
||||
"OrchestratorAgent",
|
||||
"PdfEditAgent",
|
||||
"PdfEditParameterSelector",
|
||||
"PdfEditPlanSelection",
|
||||
"PdfQuestionAgent",
|
||||
"UserSpecAgent",
|
||||
]
|
||||
@@ -0,0 +1,14 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from stirling.contracts import AgentExecutionRequest, CannotContinueExecutionAction, NextExecutionAction
|
||||
from stirling.services import AppRuntime
|
||||
|
||||
|
||||
class ExecutionPlanningAgent:
|
||||
def __init__(self, runtime: AppRuntime) -> None:
|
||||
self.runtime = runtime
|
||||
|
||||
async def next_action(self, request: AgentExecutionRequest) -> NextExecutionAction:
|
||||
return CannotContinueExecutionAction(
|
||||
reason=f"Execution planning is not implemented yet for step {request.current_step_index}."
|
||||
)
|
||||
@@ -0,0 +1,100 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai.output import ToolOutput
|
||||
from pydantic_ai.tools import RunContext
|
||||
|
||||
from stirling.agents.pdf_edit import PdfEditAgent
|
||||
from stirling.agents.pdf_questions import PdfQuestionAgent
|
||||
from stirling.agents.user_spec import UserSpecAgent
|
||||
from stirling.contracts import (
|
||||
AgentDraftRequest,
|
||||
AgentDraftWorkflowResponse,
|
||||
OrchestratorRequest,
|
||||
OrchestratorResponse,
|
||||
PdfEditRequest,
|
||||
PdfEditResponse,
|
||||
PdfQuestionRequest,
|
||||
PdfQuestionResponse,
|
||||
UnsupportedCapabilityResponse,
|
||||
)
|
||||
from stirling.services import AppRuntime
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class OrchestratorDeps:
|
||||
runtime: AppRuntime
|
||||
request: OrchestratorRequest
|
||||
|
||||
|
||||
class OrchestratorAgent:
|
||||
def __init__(self, runtime: AppRuntime) -> None:
|
||||
self.runtime = runtime
|
||||
self.agent = Agent(
|
||||
model=runtime.fast_model,
|
||||
output_type=[
|
||||
ToolOutput(
|
||||
self.delegate_pdf_edit,
|
||||
name="delegate_pdf_edit",
|
||||
description="Delegate requests for PDF modifications and return the PDF edit result.",
|
||||
),
|
||||
ToolOutput(
|
||||
self.delegate_pdf_question,
|
||||
name="delegate_pdf_question",
|
||||
description="Delegate questions about PDF contents and return the PDF question result.",
|
||||
),
|
||||
ToolOutput(
|
||||
self.delegate_user_spec,
|
||||
name="delegate_user_spec",
|
||||
description="Delegate requests to create or revise a user agent spec and return the draft result.",
|
||||
),
|
||||
ToolOutput(
|
||||
self.unsupported_capability,
|
||||
name="unsupported_capability",
|
||||
description="Return this when none of the delegate outputs fit the request.",
|
||||
),
|
||||
],
|
||||
deps_type=OrchestratorDeps,
|
||||
system_prompt=(
|
||||
"You are the top-level orchestrator. "
|
||||
"Choose exactly one output function that best handles the request. "
|
||||
"Use delegate_pdf_edit for requested PDF modifications. "
|
||||
"Use delegate_pdf_question for questions about the contents of a PDF. "
|
||||
"Use delegate_user_spec for requests to create or define an agent spec. "
|
||||
"Use unsupported_capability only when none of the other outputs fit."
|
||||
),
|
||||
model_settings=runtime.fast_model_settings,
|
||||
)
|
||||
|
||||
async def handle(self, request: OrchestratorRequest) -> OrchestratorResponse:
|
||||
result = await self.agent.run(
|
||||
request.user_message,
|
||||
deps=OrchestratorDeps(runtime=self.runtime, request=request),
|
||||
)
|
||||
return result.output
|
||||
|
||||
async def delegate_pdf_edit(self, ctx: RunContext[OrchestratorDeps]) -> PdfEditResponse:
|
||||
request = ctx.deps.request
|
||||
return await PdfEditAgent(ctx.deps.runtime).handle(
|
||||
PdfEditRequest(user_message=request.user_message, conversation_id=request.conversation_id)
|
||||
)
|
||||
|
||||
async def delegate_pdf_question(self, ctx: RunContext[OrchestratorDeps]) -> PdfQuestionResponse:
|
||||
request = ctx.deps.request
|
||||
return await PdfQuestionAgent(ctx.deps.runtime).handle(
|
||||
PdfQuestionRequest(question=request.user_message, conversation_id=request.conversation_id)
|
||||
)
|
||||
|
||||
async def delegate_user_spec(self, ctx: RunContext[OrchestratorDeps]) -> AgentDraftWorkflowResponse:
|
||||
request = ctx.deps.request
|
||||
return await UserSpecAgent(ctx.deps.runtime).draft(AgentDraftRequest(user_message=request.user_message))
|
||||
|
||||
async def unsupported_capability(
|
||||
self,
|
||||
ctx: RunContext[OrchestratorDeps],
|
||||
capability: str,
|
||||
message: str,
|
||||
) -> UnsupportedCapabilityResponse:
|
||||
return UnsupportedCapabilityResponse(capability=capability, message=message)
|
||||
@@ -0,0 +1,156 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Literal
|
||||
|
||||
from pydantic import Field
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai.output import NativeOutput
|
||||
|
||||
from stirling.contracts import (
|
||||
EditCannotDoResponse,
|
||||
EditClarificationRequest,
|
||||
EditPlanResponse,
|
||||
PdfEditRequest,
|
||||
PdfEditResponse,
|
||||
ToolOperationStep,
|
||||
)
|
||||
from stirling.models import OPERATIONS, ApiModel, OperationId, ParamToolModel
|
||||
from stirling.services import AppRuntime
|
||||
|
||||
|
||||
class PdfEditPlanSelection(ApiModel):
|
||||
outcome: Literal["plan"] = "plan"
|
||||
operations: list[OperationId] = Field(min_length=1)
|
||||
summary: str
|
||||
rationale: str | None = None
|
||||
|
||||
|
||||
class PdfEditParameterSelector:
|
||||
def __init__(self, runtime: AppRuntime) -> None:
|
||||
self.runtime = runtime
|
||||
self.agent = Agent(
|
||||
model=runtime.smart_model,
|
||||
system_prompt=(
|
||||
"Generate only the parameter object for the selected PDF operation. "
|
||||
"Use reasonable defaults when the request does not specify optional details. "
|
||||
"Only fill fields that belong to the selected operation's parameter model."
|
||||
),
|
||||
model_settings=runtime.smart_model_settings,
|
||||
)
|
||||
|
||||
async def select(
|
||||
self,
|
||||
request: PdfEditRequest,
|
||||
operation_plan: list[OperationId],
|
||||
operation_index: int,
|
||||
generated_steps: list[ToolOperationStep],
|
||||
) -> ParamToolModel:
|
||||
operation_id = operation_plan[operation_index]
|
||||
parameter_model = OPERATIONS[operation_id]
|
||||
parameter_result = await self.agent.run(
|
||||
self._build_parameter_prompt(request, operation_plan, operation_index, generated_steps),
|
||||
output_type=NativeOutput(parameter_model),
|
||||
instructions=(
|
||||
f"Generate only the parameters for the PDF operation `{operation_id.value}`. "
|
||||
"Do not include fields from any other operation."
|
||||
),
|
||||
)
|
||||
return parameter_result.output
|
||||
|
||||
def _build_parameter_prompt(
|
||||
self,
|
||||
request: PdfEditRequest,
|
||||
operation_plan: list[OperationId],
|
||||
operation_index: int,
|
||||
generated_steps: list[ToolOperationStep],
|
||||
) -> str:
|
||||
operation_id = operation_plan[operation_index]
|
||||
operation_list = ", ".join(operation.value for operation in operation_plan)
|
||||
file_names = ", ".join(request.file_names) if request.file_names else "No file names were provided."
|
||||
generated_steps_text = (
|
||||
"\n".join(
|
||||
f"- Step {step_index + 1}: {step.model_dump_json()}" for step_index, step in enumerate(generated_steps)
|
||||
)
|
||||
if generated_steps
|
||||
else "None"
|
||||
)
|
||||
return (
|
||||
f"User request: {request.user_message}\n"
|
||||
f"Files: {file_names}\n"
|
||||
f"Operation plan: {operation_list}\n"
|
||||
f"Selected operation index: {operation_index + 1} of {len(operation_plan)}\n"
|
||||
f"Selected operation: {operation_id.value}\n"
|
||||
f"Already generated steps:\n{generated_steps_text}\n"
|
||||
"Return only the parameter object for the selected operation."
|
||||
)
|
||||
|
||||
|
||||
class PdfEditAgent:
|
||||
def __init__(self, runtime: AppRuntime) -> None:
|
||||
self.runtime = runtime
|
||||
self.supported_operations = list(OPERATIONS)
|
||||
self.parameter_selector = PdfEditParameterSelector(runtime)
|
||||
self.selection_agent = Agent(
|
||||
model=runtime.smart_model,
|
||||
output_type=NativeOutput(
|
||||
[
|
||||
PdfEditPlanSelection,
|
||||
EditClarificationRequest,
|
||||
EditCannotDoResponse,
|
||||
]
|
||||
),
|
||||
system_prompt=(
|
||||
"Plan PDF edit requests. "
|
||||
f"Supported operations are: {self._supported_operations_prompt()}. "
|
||||
"Return an ordered list of one or more supported operations for the plan. "
|
||||
"Do not produce operation parameters in this stage. "
|
||||
"Return need_clarification when the request is genuinely ambiguous. "
|
||||
"Return cannot_do when the request is outside the supported operations. "
|
||||
"Return plan when a reasonable multi-step plan can be created. "
|
||||
"Never return partial plans."
|
||||
),
|
||||
model_settings=runtime.smart_model_settings,
|
||||
)
|
||||
|
||||
async def handle(self, request: PdfEditRequest) -> PdfEditResponse:
|
||||
selection = await self._select_plan(request)
|
||||
if isinstance(selection, EditClarificationRequest | EditCannotDoResponse):
|
||||
return selection
|
||||
steps: list[ToolOperationStep] = []
|
||||
for operation_index, operation_id in enumerate(selection.operations):
|
||||
parameters = await self.parameter_selector.select(
|
||||
request,
|
||||
selection.operations,
|
||||
operation_index,
|
||||
steps,
|
||||
)
|
||||
steps.append(
|
||||
ToolOperationStep(
|
||||
tool=operation_id,
|
||||
parameters=parameters,
|
||||
)
|
||||
)
|
||||
return EditPlanResponse(
|
||||
summary=selection.summary,
|
||||
rationale=selection.rationale,
|
||||
steps=steps,
|
||||
)
|
||||
|
||||
async def _select_plan(
|
||||
self,
|
||||
request: PdfEditRequest,
|
||||
) -> PdfEditPlanSelection | EditClarificationRequest | EditCannotDoResponse:
|
||||
selection_result = await self.selection_agent.run(self._build_selection_prompt(request))
|
||||
return selection_result.output
|
||||
|
||||
def _build_selection_prompt(self, request: PdfEditRequest) -> str:
|
||||
file_names = ", ".join(request.file_names) if request.file_names else "No file names were provided."
|
||||
return (
|
||||
f"User request: {request.user_message}\n"
|
||||
f"Files: {file_names}\n"
|
||||
f"Supported operations: {self._supported_operations_prompt()}\n"
|
||||
"Plan an ordered list of supported PDF edit operations or return clarification/cannot_do."
|
||||
)
|
||||
|
||||
def _supported_operations_prompt(self) -> str:
|
||||
return ", ".join(operation_id.value for operation_id in self.supported_operations)
|
||||
@@ -0,0 +1,49 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai.output import NativeOutput
|
||||
|
||||
from stirling.contracts import (
|
||||
PdfQuestionAnswerResponse,
|
||||
PdfQuestionNeedTextResponse,
|
||||
PdfQuestionNotFoundResponse,
|
||||
PdfQuestionRequest,
|
||||
PdfQuestionResponse,
|
||||
)
|
||||
from stirling.services import AppRuntime
|
||||
|
||||
|
||||
class PdfQuestionAgent:
|
||||
def __init__(self, runtime: AppRuntime) -> None:
|
||||
self.runtime = runtime
|
||||
self.agent = Agent(
|
||||
model=runtime.smart_model,
|
||||
output_type=NativeOutput(
|
||||
[
|
||||
PdfQuestionAnswerResponse,
|
||||
PdfQuestionNotFoundResponse,
|
||||
]
|
||||
),
|
||||
system_prompt=(
|
||||
"Answer questions about a PDF using only the extracted text provided in the prompt. "
|
||||
"Do not guess or use outside knowledge. "
|
||||
"If the answer is not supported by the provided text, return not_found. "
|
||||
"When answering, include a short list of evidence snippets copied from the provided text."
|
||||
),
|
||||
model_settings=runtime.smart_model_settings,
|
||||
)
|
||||
|
||||
async def handle(self, request: PdfQuestionRequest) -> PdfQuestionResponse:
|
||||
if not request.extracted_text.strip():
|
||||
return PdfQuestionNeedTextResponse(
|
||||
reason="No extracted PDF text was provided, so the question cannot be answered yet."
|
||||
)
|
||||
return await self._run_answer_agent(request)
|
||||
|
||||
async def _run_answer_agent(self, request: PdfQuestionRequest) -> PdfQuestionResponse:
|
||||
result = await self.agent.run(self._build_prompt(request))
|
||||
return result.output
|
||||
|
||||
def _build_prompt(self, request: PdfQuestionRequest) -> str:
|
||||
file_name = request.file_name or "Unknown file"
|
||||
return f"File: {file_name}\nQuestion: {request.question}\nExtracted text:\n{request.extracted_text}"
|
||||
@@ -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))
|
||||
Reference in New Issue
Block a user