Michael
71 lines
2.7 KiB
Python
71 lines
2.7 KiB
Python
import logging
|
|
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 BaseAgent:
|
|
"""Base AI agent implementing Pydantic AI interaction patterns."""
|
|
|
|
def __init__(self, model_name: str) -> None:
|
|
self.model_name: str = model_name
|
|
self.system_prompt: str = ""
|
|
self.pydantic_agent: Agent[Any, str] | None = None
|
|
|
|
async def initialize(self) -> None:
|
|
"""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.pydantic_agent is None:
|
|
await self.initialize()
|
|
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)})")
|
|
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[Callable[..., Any]]) -> str:
|
|
"""Run the agent with tool calling capability using Pydantic AI."""
|
|
if self.pydantic_agent is None:
|
|
await self.initialize()
|
|
|
|
try:
|
|
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}")
|
|
return f"Error in agent tool execution: {str(e)}"
|
|
|
|
|
|
class CavemanAgent(BaseAgent):
|
|
"""Caveman AI agent - minimal token usage variant."""
|
|
|
|
def __init__(self, model_name: str) -> None:
|
|
super().__init__(model_name)
|
|
self.system_prompt: str = CAVEMAN_PROMPT
|