From be0d8dc24b0f71b3a10887a461987d2ca8eac81c Mon Sep 17 00:00:00 2001 From: Michael Date: Sun, 2 Aug 2026 17:04:44 +0200 Subject: [PATCH] feat: migrate agent core to pydantic-ai framework --- core/agent.py | 108 ++++++++++------------------------- core/coding_agent.py | 6 +- core/coordinator_agent.py | 6 +- core/notification_agent.py | 6 +- core/planning_agent.py | 4 +- core/schemas.py | 36 ++++++++++++ tests/test_best_practices.py | 2 +- tests/test_dispatcher.py | 8 ++- tests/test_main.py | 10 +--- tests/test_orchestrator.py | 40 ++++++------- 10 files changed, 109 insertions(+), 117 deletions(-) create mode 100644 core/schemas.py diff --git a/core/agent.py b/core/agent.py index f5a0c74..c8db501 100644 --- a/core/agent.py +++ b/core/agent.py @@ -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 - diff --git a/core/coding_agent.py b/core/coding_agent.py index f392336..8a4c13c 100644 --- a/core/coding_agent.py +++ b/core/coding_agent.py @@ -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 - diff --git a/core/coordinator_agent.py b/core/coordinator_agent.py index 7ee8204..4ced25a 100644 --- a/core/coordinator_agent.py +++ b/core/coordinator_agent.py @@ -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, diff --git a/core/notification_agent.py b/core/notification_agent.py index e7fb840..8f44d53 100644 --- a/core/notification_agent.py +++ b/core/notification_agent.py @@ -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, diff --git a/core/planning_agent.py b/core/planning_agent.py index 3a30eb7..43fc3f4 100644 --- a/core/planning_agent.py +++ b/core/planning_agent.py @@ -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 - diff --git a/core/schemas.py b/core/schemas.py new file mode 100644 index 0000000..f88570c --- /dev/null +++ b/core/schemas.py @@ -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") diff --git a/tests/test_best_practices.py b/tests/test_best_practices.py index d6326a5..fb6d9e6 100644 --- a/tests/test_best_practices.py +++ b/tests/test_best_practices.py @@ -107,7 +107,7 @@ async def test_dispatch_planning_and_coding_phases(mock_planning_class: MagicMoc mock_client.issues.get_issue_comments.return_value = [] from gitea.models import UserModel - mock_client.get_authenticated_user.return_value = UserModel(login="unknown-ai") + mock_client.repos.get_authenticated_user.return_value = UserModel(login="unknown-ai") mock_pr = PullRequestModel( number=42, diff --git a/tests/test_dispatcher.py b/tests/test_dispatcher.py index f8d5deb..1b5f9f4 100644 --- a/tests/test_dispatcher.py +++ b/tests/test_dispatcher.py @@ -275,7 +275,13 @@ def _make_comment(login: str, body: str) -> CommentModel: def _make_dispatcher_for_reply_tests() -> AgentDispatcher: mock_client = MagicMock() mock_client.repos.get_authenticated_user.return_value = UserModel(login="unknown-ai") - return AgentDispatcher(client=mock_client, tools=MagicMock()) + return AgentDispatcher( + client=mock_client, + issue_tools=MagicMock(), + pr_tools=MagicMock(), + file_tools=MagicMock(), + git_tools=MagicMock(), + ) def test_is_awaiting_reply_no_comments() -> None: diff --git a/tests/test_main.py b/tests/test_main.py index b9e144c..2fc8d8a 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -8,11 +8,9 @@ pytestmark = pytest.mark.anyio @patch("main.load_dotenv") @patch("main.os.chdir") @patch("main.GiteaClient") -@patch("main.GiteaTools") @patch("main.AgentOrchestrator") async def test_main_startup_success( mock_orchestrator_class: MagicMock, - mock_tools_class: MagicMock, mock_client_class: MagicMock, mock_chdir: MagicMock, mock_load_dotenv: MagicMock @@ -21,7 +19,7 @@ async def test_main_startup_success( mock_client_class.return_value = mock_client mock_user = MagicMock() mock_user.login = "agent-test" - mock_client.get_authenticated_user.return_value = mock_user + mock_client.repos.get_authenticated_user.return_value = mock_user mock_orchestrator = MagicMock() mock_orchestrator.poll_and_dispatch = AsyncMock(side_effect=KeyboardInterrupt()) @@ -30,18 +28,16 @@ async def test_main_startup_success( # Run main; it should exit gracefully on KeyboardInterrupt await main() - mock_client.get_authenticated_user.assert_called_once() + mock_client.repos.get_authenticated_user.assert_called_once() mock_orchestrator.poll_and_dispatch.assert_called_once() @patch("main.load_dotenv") @patch("main.os.chdir") @patch("main.GiteaClient") -@patch("main.GiteaTools") @patch("main.AgentOrchestrator") async def test_main_startup_fails_no_authenticated_user( mock_orchestrator_class: MagicMock, - mock_tools_class: MagicMock, mock_client_class: MagicMock, mock_chdir: MagicMock, mock_load_dotenv: MagicMock @@ -49,7 +45,7 @@ async def test_main_startup_fails_no_authenticated_user( mock_client = MagicMock() mock_client_class.return_value = mock_client # Simulate no user returned - mock_client.get_authenticated_user.return_value = None + mock_client.repos.get_authenticated_user.return_value = None with pytest.raises(SystemExit) as exc_info: await main() diff --git a/tests/test_orchestrator.py b/tests/test_orchestrator.py index 996c4ca..64f3f47 100644 --- a/tests/test_orchestrator.py +++ b/tests/test_orchestrator.py @@ -4,7 +4,6 @@ import pytest from unittest.mock import MagicMock, AsyncMock, patch from core.orchestrator import AgentOrchestrator -from gitea.client import GiteaClient from gitea.models import IssueModel, PullRequestModel, RepositoryModel pytestmark = pytest.mark.anyio @@ -29,16 +28,17 @@ async def test_poll_and_dispatch_no_notifications( temp_state_file: Path ) -> None: mock_get_path.return_value = temp_state_file - mock_client = MagicMock(spec=GiteaClient) - mock_tools = MagicMock() + mock_client = MagicMock() # Return no notifications - mock_client.list_unread_notifications.return_value = [] + mock_client.notifications.list_unread_notifications.return_value = [] - orchestrator = AgentOrchestrator(mock_client, mock_tools) + orchestrator = AgentOrchestrator( + mock_client, MagicMock(), MagicMock(), MagicMock(), MagicMock() + ) await orchestrator.poll_and_dispatch() - mock_client.list_unread_notifications.assert_called_once_with(since=None) + mock_client.notifications.list_unread_notifications.assert_called_once_with(since=None) assert not temp_state_file.exists() @@ -66,8 +66,7 @@ async def test_poll_and_dispatch_with_notifications( return "Decided" mock_reader.decide_notification = AsyncMock(side_effect=mock_decide_notification) mock_notification_reader_class.return_value = mock_reader - mock_client = MagicMock(spec=GiteaClient) - mock_tools = MagicMock() + mock_client = MagicMock() # Set up mock Gitea notifications notifications = [ @@ -98,13 +97,14 @@ async def test_poll_and_dispatch_with_notifications( } } ] - mock_client.list_unread_notifications.return_value = notifications + mock_client.notifications.list_unread_notifications.return_value = notifications - # Mock issue and PR get methods + # Mock issue and PR get methods on client issue_model = IssueModel(number=42, title="Bug issue", repository=RepositoryModel(name="repo1", full_name="meeks/repo1")) pr_model = PullRequestModel(number=10, title="Fix PR", repository=RepositoryModel(name="repo1", full_name="meeks/repo1")) - mock_client.get_issue.return_value = issue_model - mock_client.get_pull_request.return_value = pr_model + + mock_client.issues.get_issue.return_value = issue_model + mock_client.prs.get_pull_request.return_value = pr_model # Mock dispatcher and workspace path mock_dispatcher_instance = MagicMock() @@ -116,15 +116,13 @@ async def test_poll_and_dispatch_with_notifications( mock_workspace_class.return_value = mock_workspace_instance # Create orchestrator and poll - orchestrator = AgentOrchestrator(mock_client, mock_tools) + orchestrator = AgentOrchestrator( + mock_client, MagicMock(), MagicMock(), MagicMock(), MagicMock() + ) await orchestrator.poll_and_dispatch() # Assert notifications were checked with None (first execution) - mock_client.list_unread_notifications.assert_called_once_with(since=None) - - # Assert issue and PR details were fetched - mock_client.get_issue.assert_called_once_with("meeks", "repo1", 42) - mock_client.get_pull_request.assert_called_once_with("meeks", "repo1", 10) + mock_client.notifications.list_unread_notifications.assert_called_once_with(since=None) # Assert work was processed by dispatcher mock_dispatcher_instance.dispatch.assert_called_once() @@ -136,9 +134,9 @@ async def test_poll_and_dispatch_with_notifications( assert work_items[1].notification_id == 102 # Assert notifications were marked as read - mock_client.mark_notification_as_read.assert_any_call(101) - mock_client.mark_notification_as_read.assert_any_call(102) - assert mock_client.mark_notification_as_read.call_count == 2 + mock_client.notifications.mark_notification_as_read.assert_any_call(101) + mock_client.notifications.mark_notification_as_read.assert_any_call(102) + assert mock_client.notifications.mark_notification_as_read.call_count == 2 # Assert checkpoint date was persisted assert temp_state_file.exists()