Refactor code structure for improved readability and maintainability
This commit is contained in:
mi222eh
@@ -0,0 +1,219 @@
|
||||
"""Top-level coordinator: polls Gitea unread notifications, queues work, dispatches to agent."""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import datetime
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any, Optional
|
||||
|
||||
from core.queue import WorkQueue, WorkItem
|
||||
from core.dispatcher import AgentDispatcher
|
||||
from gitea.client import GiteaClient
|
||||
from gitea.models import IssueModel, PullRequestModel, RepositoryModel
|
||||
from gitea.tools.gitea_tools import GiteaTools
|
||||
from gitea.config import AGENT_MODEL_ID, AGENT_MAX_RETRIES
|
||||
from gitea.workspace import WorkspaceManager
|
||||
from core.factory import AgentFactory
|
||||
from core.notification_tools import NotificationTools
|
||||
from core.notification_agent import NotificationNoToolCalledError
|
||||
|
||||
logger: logging.Logger = logging.getLogger("agent-orchestrator")
|
||||
|
||||
|
||||
class AgentOrchestrator:
|
||||
"""Top-level coordinator: polls Gitea notifications, queues work, dispatches to agent."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
client: GiteaClient,
|
||||
tools: GiteaTools,
|
||||
model_name: str = AGENT_MODEL_ID,
|
||||
max_retries: int = AGENT_MAX_RETRIES,
|
||||
) -> None:
|
||||
self._client = client
|
||||
self._tools = tools
|
||||
self._model_name = model_name
|
||||
self._work_queue = WorkQueue()
|
||||
self._dispatcher = AgentDispatcher(client, tools, model_name, max_retries)
|
||||
self._notification_reader = AgentFactory.create_notification_reader_agent(model_name)
|
||||
self._max_retries = max_retries
|
||||
|
||||
|
||||
def _get_state_file_path(self) -> Path:
|
||||
"""Get the path to the persistent state file."""
|
||||
return Path(__file__).parent.parent / "agent_state.json"
|
||||
|
||||
def _read_last_checked(self) -> Optional[str]:
|
||||
"""Read the last checked timestamp from persistent storage."""
|
||||
state_file = self._get_state_file_path()
|
||||
if state_file.exists():
|
||||
try:
|
||||
with open(state_file, "r") as f:
|
||||
data = json.load(f)
|
||||
return data.get("last_checked")
|
||||
except Exception as e:
|
||||
logger.warning(f"Error reading agent_state.json: {e}")
|
||||
return None
|
||||
|
||||
def _write_last_checked(self, timestamp: str) -> None:
|
||||
"""Write the last checked timestamp to persistent storage."""
|
||||
state_file = self._get_state_file_path()
|
||||
try:
|
||||
with open(state_file, "w") as f:
|
||||
json.dump({"last_checked": timestamp}, f)
|
||||
except Exception as e:
|
||||
logger.error(f"Error writing agent_state.json: {e}")
|
||||
|
||||
async def poll_and_dispatch(self) -> None:
|
||||
"""Poll Gitea unread notifications, enqueue them, and dispatch to agent."""
|
||||
last_checked = self._read_last_checked()
|
||||
logger.info(f"Polling unread notifications since: {last_checked or 'beginning'}")
|
||||
|
||||
notifications = self._client.list_unread_notifications(since=last_checked)
|
||||
|
||||
if not notifications:
|
||||
logger.info("No new notifications found.")
|
||||
return
|
||||
|
||||
logger.info(f"Retrieved {len(notifications)} unread notifications.")
|
||||
|
||||
# Filter and enqueue tasks from notifications
|
||||
latest_timestamp = last_checked
|
||||
inspection_tools = [
|
||||
self._tools.get_issue,
|
||||
self._tools.get_pull_request,
|
||||
self._tools.get_issue_comments,
|
||||
self._tools.get_pull_request_comments,
|
||||
]
|
||||
|
||||
for n in notifications:
|
||||
notification_id = n.get("id")
|
||||
subject = n.get("subject") or {}
|
||||
subj_type = subject.get("type", "").lower()
|
||||
subj_url = subject.get("url", "")
|
||||
updated_at = n.get("updated_at")
|
||||
|
||||
# Track latest updated_at to advance checkpoint
|
||||
if updated_at and (not latest_timestamp or updated_at > latest_timestamp):
|
||||
latest_timestamp = updated_at
|
||||
|
||||
repo_info = n.get("repository") or {}
|
||||
repo_full_name = repo_info.get("full_name", "")
|
||||
if not repo_full_name or not subj_url:
|
||||
continue
|
||||
|
||||
owner, repo_name = repo_full_name.split("/")
|
||||
try:
|
||||
task_number = int(subj_url.rstrip("/").split("/")[-1])
|
||||
except (ValueError, IndexError):
|
||||
logger.warning(f"Could not parse task number from subject URL: {subj_url}")
|
||||
continue
|
||||
|
||||
# Run NotificationReaderAgent to pre-screen the notification
|
||||
mission = (
|
||||
f"Analyze Gitea notification ID {notification_id}.\n"
|
||||
f"Repository: {repo_full_name}\n"
|
||||
f"Subject Type: {subj_type}\n"
|
||||
f"Task Number: {task_number}\n"
|
||||
f"Please decide if we should process or skip this notification."
|
||||
)
|
||||
notification_tools = NotificationTools()
|
||||
attempt = 0
|
||||
attempt_limit = self._max_retries
|
||||
success = False
|
||||
while attempt < attempt_limit:
|
||||
attempt += 1
|
||||
try:
|
||||
await self._notification_reader.decide_notification(
|
||||
mission,
|
||||
inspection_tools,
|
||||
notification_tools
|
||||
)
|
||||
success = True
|
||||
break
|
||||
except NotificationNoToolCalledError as e:
|
||||
logger.error(f"Notification reader error on notification {notification_id} (attempt {attempt}/{attempt_limit}): {e}")
|
||||
|
||||
if not success or notification_tools.action == "SKIP":
|
||||
reason = notification_tools.arguments.get("reason", "Failed to call routing tool / default skip")
|
||||
logger.info(f"Skipping notification {notification_id} for {repo_full_name}#{task_number}. Reason: {reason}")
|
||||
if notification_id is not None:
|
||||
self._client.mark_notification_as_read(notification_id)
|
||||
logger.info(f"Marked skipped Gitea notification thread {notification_id} as read.")
|
||||
continue
|
||||
|
||||
# Route based on decided action
|
||||
if notification_tools.action == "PROCESS_ISSUE":
|
||||
try:
|
||||
issue = self._client.get_issue(owner, repo_name, task_number)
|
||||
if issue.repository is None:
|
||||
issue = issue.model_copy(update={"repository": RepositoryModel(**repo_info)})
|
||||
|
||||
item = WorkItem(
|
||||
repo_full_name=repo_full_name,
|
||||
task_type="issue",
|
||||
task_number=task_number,
|
||||
task_info=issue,
|
||||
notification_id=notification_id,
|
||||
priority=0
|
||||
)
|
||||
self._work_queue.enqueue(item)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to fetch issue #{task_number} for notification: {e}")
|
||||
elif notification_tools.action == "PROCESS_PR":
|
||||
try:
|
||||
pr = self._client.get_pull_request(owner, repo_name, task_number)
|
||||
if pr.repository is None:
|
||||
pr = pr.model_copy(update={"repository": RepositoryModel(**repo_info)})
|
||||
|
||||
item = WorkItem(
|
||||
repo_full_name=repo_full_name,
|
||||
task_type="pr",
|
||||
task_number=task_number,
|
||||
task_info=pr,
|
||||
notification_id=notification_id,
|
||||
priority=0
|
||||
)
|
||||
self._work_queue.enqueue(item)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to fetch PR #{task_number} for notification: {e}")
|
||||
|
||||
|
||||
# Process enqueued work
|
||||
if not self._work_queue.is_empty:
|
||||
await self._process_work()
|
||||
|
||||
# Update last checked timestamp checkpoint
|
||||
if latest_timestamp:
|
||||
self._write_last_checked(latest_timestamp)
|
||||
|
||||
async def _process_work(self) -> None:
|
||||
"""Process all queued work, repo by repo, then mark notifications as read."""
|
||||
while not self._work_queue.is_empty:
|
||||
repo = self._work_queue.get_next_repo()
|
||||
if not repo:
|
||||
break
|
||||
|
||||
work_items = self._work_queue.get_repo_work(repo)
|
||||
self._work_queue.remove_repo_work(repo)
|
||||
|
||||
# Ensure the workspace repository is cloned and sanitized
|
||||
workspace = WorkspaceManager()
|
||||
repo_path = workspace.get_repo_path(repo)
|
||||
if not repo_path.exists():
|
||||
workspace.clone_repo(repo)
|
||||
logger.info(f"Cloned {repo} to {repo_path}")
|
||||
else:
|
||||
workspace.sanitize_repo(repo, repo_path)
|
||||
logger.info(f"Sanitized existing repo at {repo_path}")
|
||||
|
||||
logger.info(f"Dispatching {len(work_items)} tasks for {repo}")
|
||||
results = await self._dispatcher.dispatch(repo, work_items)
|
||||
|
||||
for i, result in enumerate(results):
|
||||
item = work_items[i]
|
||||
logger.info(f"Completed {item.task_type} #{item.task_number}: {result[:200]}")
|
||||
if item.notification_id is not None:
|
||||
self._client.mark_notification_as_read(item.notification_id)
|
||||
logger.info(f"Marked Gitea notification thread {item.notification_id} as read.")
|
||||
Reference in New Issue
Block a user