Michael
49 lines
1.8 KiB
Python
49 lines
1.8 KiB
Python
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")
|
|
|
|
|
|
class NotificationNoToolCalledError(Exception):
|
|
"""Raised when the Notification Reader Agent completes execution without calling any routing tool."""
|
|
pass
|
|
|
|
|
|
class NotificationReaderAgent(BaseAgent):
|
|
"""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,
|
|
inspection_tools: list[Callable[..., Any]],
|
|
notification_tools: NotificationTools,
|
|
) -> str:
|
|
"""Run the Notification Reader Agent and ensure a decision tool is called."""
|
|
decision_tools: list[Callable[..., Any]] = [
|
|
notification_tools.process_issue,
|
|
notification_tools.process_pr,
|
|
notification_tools.skip_notification,
|
|
]
|
|
combined_tools: list[Callable[..., Any]] = inspection_tools + decision_tools
|
|
|
|
logger.info("Running NotificationReaderAgent to decide action...")
|
|
response_text: str = await self.run_with_tools(mission, combined_tools)
|
|
|
|
if not notification_tools.tool_called:
|
|
logger.warning("NotificationReaderAgent did not call any tools!")
|
|
raise NotificationNoToolCalledError(
|
|
"NotificationReaderAgent failed to call a routing tool during execution."
|
|
)
|
|
|
|
return response_text
|