feat: migrate agent core to pydantic-ai framework

core/agent.py

@@ -1,108 +1,61 @@
-import asyncio
 import logging
-import lmstudio as lms
+from typing import Any, Callable
-from typing import Callable
+from pydantic_ai import Agent
-from .prompt import CAVEMAN_PROMPT
+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."""
+        """Initialize the Pydantic AI agent instance."""
-        logger.info(f"Initializing agent with model: {self.model_name}")
+        logger.info(f"Initializing Pydantic AI agent with model: {self.model_name}")
-        self.model = lms.llm(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:
+        if self.pydantic_agent is None:
-            raise RuntimeError("Model initialization failed: model is None")
+            raise RuntimeError("Model initialization failed: pydantic_agent is None")
-
-        messages: list[dict[str, str]] = [
-            {"role": "system", "content": self.system_prompt},
-            {"role": "user", "content": user_input},
-        ]
 
         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:
+    async def run_with_tools(self, user_input: str, tools: list[Callable[..., Any]]) -> str:
-        """Run the agent with tool calling capability."""
+        """Run the agent with tool calling capability using Pydantic AI."""
-        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")
 
         try:
-            capture = _ActResponseCapture()
+            logger.info(f"Calling Pydantic AI agent with {len(tools)} tools...")
-            logger.info(f"Calling LMStudio act() on agent with {len(tools)} tools...")
+            model_str: str = self.model_name if ":" in self.model_name else f"openai:{self.model_name}"
-            result: lms.ActResult = self.model.act(user_input, tools=tools, on_message=capture)
+            agent: Agent[Any, str] = Agent(
-            logger.info(f"act() on agent returned: {result}")
+                model_str,
-            response: str = capture.full_response
+                system_prompt=self.system_prompt,
-            if not response or response == "No response captured.":
+                tools=tools,
-                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."
+            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
-

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
-

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,

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,

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
-

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")

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,

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:

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()

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_client = MagicMock()
-    mock_tools = 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_client = MagicMock()
-    mock_tools = 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)
+    mock_client.notifications.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)
 
     # 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.notifications.mark_notification_as_read.assert_any_call(101)
-    mock_client.mark_notification_as_read.assert_any_call(102)
+    mock_client.notifications.mark_notification_as_read.assert_any_call(102)
-    assert mock_client.mark_notification_as_read.call_count == 2
+    assert mock_client.notifications.mark_notification_as_read.call_count == 2
 
     # Assert checkpoint date was persisted
     assert temp_state_file.exists()