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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user