James BruntonandGitHub
3e94157137
# Description of Changes Adds the ability for the Edit agent to request the content of the document before it decides which parameters it needs. This makes it able to process requests like `Split the document after the page containing the "My Section" section`, allowing for document context-based requests for all[^1] tools. I had to make a few changes elsewhere to make this work, including: - Moving the requesting of content out of the Question Agent and into a common location - Added specific API docs for the Split param because the generic ones were not specific enough for the AI to be able to reliably perform the correct operation - Fixed an issue in the tool models generator which caused the Redact params to only be half-generated (causing Pydantic to crash when the AI tried to run Redact) - Added missing logging to a bunch of tools and hooked it up properly so it'll print to stderr - Made the limits for the max pages/chars to extract from PDFs configurable via env var [^1]: Many of the tools can't actually do anything useful with the context at this stage, but will just need the tool API to be extended with new features like page-specific operations to be automatically able to do smart operations without needing to change the Edit agent itself.
109 lines
4.4 KiB
Python
109 lines
4.4 KiB
Python
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,
|
|
EditPlanResponse,
|
|
PdfEditRequest,
|
|
PdfEditTerminalResponse,
|
|
format_conversation_history,
|
|
)
|
|
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, request.conversation_history)
|
|
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}",
|
|
request.conversation_history,
|
|
)
|
|
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{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{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)}"
|
|
)
|
|
|
|
async def _build_edit_plan(
|
|
self,
|
|
user_message: str,
|
|
conversation_history: list[ConversationMessage],
|
|
) -> PdfEditTerminalResponse:
|
|
return await self.pdf_edit_agent.handle(
|
|
PdfEditRequest(user_message=user_message, conversation_history=conversation_history),
|
|
allow_need_content=False,
|
|
)
|