feat: migrate agent core to pydantic-ai framework
This commit is contained in:
Michael
+30
-78
@@ -1,108 +1,61 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import lmstudio as lms
|
||||
from typing import Callable
|
||||
from .prompt import CAVEMAN_PROMPT
|
||||
from typing import Any, Callable
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai.exceptions import ModelRetry
|
||||
from core.prompt import CAVEMAN_PROMPT
|
||||
|
||||
logger: logging.Logger = logging.getLogger("agent-base")
|
||||
|
||||
|
||||
class _ActResponseCapture:
|
||||
"""Captures the AI response from LMStudio act() callback."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.responses: list[str] = []
|
||||
|
||||
def __call__(self, message: object) -> None:
|
||||
content: str = ""
|
||||
if hasattr(message, 'content'):
|
||||
content = message.content
|
||||
elif hasattr(message, 'text'):
|
||||
content = message.text
|
||||
elif hasattr(message, 'response'):
|
||||
content = message.response
|
||||
elif hasattr(message, 'message'):
|
||||
content = message.message
|
||||
else:
|
||||
return
|
||||
|
||||
if isinstance(content, list):
|
||||
parts: list[str] = []
|
||||
for item in content:
|
||||
if isinstance(item, dict):
|
||||
text: str = item.get('text', '')
|
||||
if isinstance(text, list):
|
||||
parts.extend([str(t) for t in text])
|
||||
else:
|
||||
parts.append(str(text))
|
||||
elif isinstance(item, str):
|
||||
parts.append(item)
|
||||
elif hasattr(item, 'text'):
|
||||
parts.append(str(item.text))
|
||||
elif hasattr(item, 'content'):
|
||||
parts.append(str(item.content))
|
||||
content = ''.join(parts)
|
||||
elif not isinstance(content, str):
|
||||
content = str(content)
|
||||
|
||||
if content.strip():
|
||||
self.responses.append(content.strip())
|
||||
|
||||
@property
|
||||
def full_response(self) -> str:
|
||||
return '\n'.join(self.responses) if self.responses else "No response captured."
|
||||
|
||||
|
||||
class BaseAgent:
|
||||
"""Base AI agent implementing common LMStudio interaction patterns."""
|
||||
"""Base AI agent implementing Pydantic AI interaction patterns."""
|
||||
|
||||
def __init__(self, model_name: str) -> None:
|
||||
self.model_name: str = model_name
|
||||
self.model: object | None = None
|
||||
self.system_prompt: str = ""
|
||||
self.pydantic_agent: Agent[Any, str] | None = None
|
||||
|
||||
async def initialize(self) -> None:
|
||||
"""Initialize the LM Studio model."""
|
||||
logger.info(f"Initializing agent with model: {self.model_name}")
|
||||
self.model = lms.llm(self.model_name)
|
||||
"""Initialize the Pydantic AI agent instance."""
|
||||
logger.info(f"Initializing Pydantic AI agent with model: {self.model_name}")
|
||||
model_str: str = self.model_name if ":" in self.model_name else f"openai:{self.model_name}"
|
||||
self.pydantic_agent = Agent(
|
||||
model_str,
|
||||
system_prompt=self.system_prompt,
|
||||
)
|
||||
|
||||
async def run(self, user_input: str) -> str:
|
||||
"""Run a single interaction with the agent."""
|
||||
if self.model is None:
|
||||
if self.pydantic_agent is None:
|
||||
await self.initialize()
|
||||
if self.model is None:
|
||||
raise RuntimeError("Model initialization failed: model is None")
|
||||
|
||||
messages: list[dict[str, str]] = [
|
||||
{"role": "system", "content": self.system_prompt},
|
||||
{"role": "user", "content": user_input},
|
||||
]
|
||||
if self.pydantic_agent is None:
|
||||
raise RuntimeError("Model initialization failed: pydantic_agent is None")
|
||||
|
||||
try:
|
||||
logger.info(f"Running agent interactively (input length: {len(user_input)})")
|
||||
response = await self.model.respond(user_input, messages=messages)
|
||||
result = await self.pydantic_agent.run(user_input)
|
||||
response: str = str(result.data)
|
||||
logger.info(f"Agent responded successfully (response length: {len(response)})")
|
||||
return response
|
||||
except Exception as e:
|
||||
logger.error(f"Agent execution error: {e}")
|
||||
return f"Error in agent execution: {str(e)}"
|
||||
|
||||
async def run_with_tools(self, user_input: str, tools: list[object]) -> str:
|
||||
"""Run the agent with tool calling capability."""
|
||||
if self.model is None:
|
||||
async def run_with_tools(self, user_input: str, tools: list[Callable[..., Any]]) -> str:
|
||||
"""Run the agent with tool calling capability using Pydantic AI."""
|
||||
if self.pydantic_agent is None:
|
||||
await self.initialize()
|
||||
if self.model is None:
|
||||
raise RuntimeError("Model initialization failed: model is None")
|
||||
|
||||
try:
|
||||
capture = _ActResponseCapture()
|
||||
logger.info(f"Calling LMStudio act() on agent with {len(tools)} tools...")
|
||||
result: lms.ActResult = self.model.act(user_input, tools=tools, on_message=capture)
|
||||
logger.info(f"act() on agent returned: {result}")
|
||||
response: str = capture.full_response
|
||||
if not response or response == "No response captured.":
|
||||
logger.warning(f"Act completed with {result.rounds} rounds but no response was captured.")
|
||||
return f"Act completed with {result.rounds} rounds but no response captured."
|
||||
logger.info(f"Calling Pydantic AI agent with {len(tools)} tools...")
|
||||
model_str: str = self.model_name if ":" in self.model_name else f"openai:{self.model_name}"
|
||||
agent: Agent[Any, str] = Agent(
|
||||
model_str,
|
||||
system_prompt=self.system_prompt,
|
||||
tools=tools,
|
||||
)
|
||||
result = await agent.run(user_input)
|
||||
response: str = str(result.data)
|
||||
return response
|
||||
except Exception as e:
|
||||
logger.error(f"Agent tool execution error: {e}")
|
||||
@@ -115,4 +68,3 @@ class CavemanAgent(BaseAgent):
|
||||
def __init__(self, model_name: str) -> None:
|
||||
super().__init__(model_name)
|
||||
self.system_prompt: str = CAVEMAN_PROMPT
|
||||
|
||||
|
||||
@@ -1,14 +1,14 @@
|
||||
import logging
|
||||
from core.agent import BaseAgent
|
||||
from .coding_prompt import CODING_AGENT_SYSTEM_PROMPT
|
||||
from core.coding_prompt import CODING_AGENT_SYSTEM_PROMPT
|
||||
from core.schemas import CodingTaskResult
|
||||
|
||||
logger: logging.Logger = logging.getLogger("agent-coding")
|
||||
|
||||
|
||||
class CodingAgent(BaseAgent):
|
||||
"""AI agent that interacts with LMStudio models and tools for coding tasks."""
|
||||
"""AI agent that interacts with Pydantic AI models and tools for coding tasks."""
|
||||
|
||||
def __init__(self, model_name: str) -> None:
|
||||
super().__init__(model_name)
|
||||
self.system_prompt: str = CODING_AGENT_SYSTEM_PROMPT
|
||||
|
||||
|
||||
@@ -1,8 +1,11 @@
|
||||
import logging
|
||||
from typing import Any, Callable
|
||||
from pydantic_ai import Agent, RunContext
|
||||
from pydantic_ai.exceptions import ModelRetry
|
||||
from core.agent import BaseAgent
|
||||
from core.prompts import COORDINATOR_SYSTEM_PROMPT
|
||||
from core.coordinator_tools import CoordinatorTools
|
||||
from core.schemas import CoordinatorDecision
|
||||
|
||||
logger: logging.Logger = logging.getLogger("agent-coordinator")
|
||||
|
||||
@@ -13,13 +16,12 @@ class CoordinatorNoToolCalledError(Exception):
|
||||
|
||||
|
||||
class CoordinatorAgent(BaseAgent):
|
||||
"""AI agent that coordinates Gitea issues and decides the next action."""
|
||||
"""AI agent that coordinates Gitea issues and decides the next action using Pydantic AI."""
|
||||
|
||||
def __init__(self, model_name: str) -> None:
|
||||
super().__init__(model_name)
|
||||
self.system_prompt: str = COORDINATOR_SYSTEM_PROMPT
|
||||
|
||||
|
||||
async def decide_action(
|
||||
self,
|
||||
mission: str,
|
||||
|
||||
@@ -1,8 +1,11 @@
|
||||
import logging
|
||||
from typing import Any, Callable
|
||||
from pydantic_ai import Agent, RunContext
|
||||
from pydantic_ai.exceptions import ModelRetry
|
||||
from core.agent import BaseAgent
|
||||
from core.prompts import NOTIFICATION_READER_SYSTEM_PROMPT
|
||||
from core.notification_tools import NotificationTools
|
||||
from core.schemas import NotificationDecision
|
||||
|
||||
logger: logging.Logger = logging.getLogger("agent-notification-reader")
|
||||
|
||||
@@ -13,13 +16,12 @@ class NotificationNoToolCalledError(Exception):
|
||||
|
||||
|
||||
class NotificationReaderAgent(BaseAgent):
|
||||
"""AI agent that reviews Gitea notifications and decides how to route them."""
|
||||
"""AI agent that reviews Gitea notifications and decides how to route them using Pydantic AI."""
|
||||
|
||||
def __init__(self, model_name: str) -> None:
|
||||
super().__init__(model_name)
|
||||
self.system_prompt: str = NOTIFICATION_READER_SYSTEM_PROMPT
|
||||
|
||||
|
||||
async def decide_notification(
|
||||
self,
|
||||
mission: str,
|
||||
|
||||
@@ -1,14 +1,14 @@
|
||||
import logging
|
||||
from core.agent import BaseAgent
|
||||
from core.prompts import PLANNING_AGENT_SYSTEM_PROMPT
|
||||
from core.schemas import ExecutionPlan
|
||||
|
||||
logger: logging.Logger = logging.getLogger("agent-planning")
|
||||
|
||||
|
||||
class PlanningAgent(BaseAgent):
|
||||
"""AI agent that analyzes a PR/issue and builds an implementation plan."""
|
||||
"""AI agent that analyzes a PR/issue and builds an implementation plan using Pydantic AI."""
|
||||
|
||||
def __init__(self, model_name: str) -> None:
|
||||
super().__init__(model_name)
|
||||
self.system_prompt: str = PLANNING_AGENT_SYSTEM_PROMPT
|
||||
|
||||
|
||||
@@ -0,0 +1,36 @@
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import Optional
|
||||
|
||||
|
||||
class NotificationDecision(BaseModel):
|
||||
"""Structured decision returned by NotificationReaderAgent."""
|
||||
action: str = Field(..., description="Action to take: 'PROCESS_ISSUE', 'PROCESS_PR', or 'SKIP'")
|
||||
owner: str = Field(default="", description="Repository owner")
|
||||
repo: str = Field(default="", description="Repository name")
|
||||
number: int = Field(default=0, description="Issue or PR number")
|
||||
reason: str = Field(..., description="Reason for the routing decision")
|
||||
|
||||
|
||||
class CoordinatorDecision(BaseModel):
|
||||
"""Structured decision returned by CoordinatorAgent."""
|
||||
action: str = Field(..., description="Action to take: 'PROPOSE_PLAN', 'EXECUTE_PLAN', 'ANSWER_QUESTION', 'CLOSE_ISSUE', 'NO_ACTION'")
|
||||
issue_number: int = Field(default=0, description="Issue number")
|
||||
plan: Optional[str] = Field(default=None, description="Proposed or approved implementation plan")
|
||||
answer: Optional[str] = Field(default=None, description="Answer to question")
|
||||
comment: Optional[str] = Field(default=None, description="Closing comment")
|
||||
|
||||
|
||||
class ExecutionPlan(BaseModel):
|
||||
"""Structured implementation plan generated by PlanningAgent."""
|
||||
issue_number: int = Field(..., description="Target issue number")
|
||||
title: str = Field(..., description="Plan title")
|
||||
steps: list[str] = Field(default_factory=list, description="Step-by-step implementation tasks")
|
||||
summary: str = Field(default="", description="Summary of proposed changes")
|
||||
|
||||
|
||||
class CodingTaskResult(BaseModel):
|
||||
"""Structured execution result returned by CodingAgent."""
|
||||
status: str = Field(..., description="Execution status: 'SUCCESS', 'FAILED', or 'PARTIAL'")
|
||||
summary: str = Field(default="", description="Summary of completed coding work")
|
||||
modified_files: list[str] = Field(default_factory=list, description="List of modified or created files")
|
||||
error_message: Optional[str] = Field(default=None, description="Error message if execution failed")
|
||||
Reference in New Issue
Block a user