166 lines
5.1 KiB
Python
166 lines
5.1 KiB
Python
"""Handoff manager for task coordination between agents."""
|
|
import json
|
|
import time
|
|
import uuid
|
|
import logging
|
|
from typing import Dict, Optional
|
|
from pydantic import BaseModel
|
|
from .redis_client import redis_client
|
|
from .agent_registry import agent_registry, AgentStatus
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class HandoffRequest(BaseModel):
|
|
"""Handoff request model."""
|
|
type: str = "handoff"
|
|
source_agent_id: str
|
|
target_agent_id: str
|
|
task_context: Dict
|
|
|
|
|
|
class HandoffRecord(BaseModel):
|
|
"""Handoff history record."""
|
|
handoff_id: str
|
|
source_agent_id: str
|
|
target_agent_id: str
|
|
timestamp: float
|
|
task_context: Dict
|
|
status: str # pending, completed, failed
|
|
|
|
|
|
class HandoffManager:
|
|
"""Manages task handoffs between agents."""
|
|
|
|
HANDOFF_KEY_PREFIX = "handoff:"
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
async def initiate_handoff(
|
|
self, source_agent_id: str, target_agent_id: str, task_context: Dict
|
|
) -> Optional[str]:
|
|
"""Initiate a handoff from source to target agent."""
|
|
|
|
# Validate source agent exists and is busy
|
|
source_agent = await agent_registry.get_agent(source_agent_id)
|
|
if not source_agent:
|
|
logger.error(f"Source agent {source_agent_id} not found")
|
|
return None
|
|
|
|
# Validate target agent exists and is idle
|
|
target_agent = await agent_registry.get_agent(target_agent_id)
|
|
if not target_agent:
|
|
logger.error(f"Target agent {target_agent_id} not found")
|
|
return None
|
|
|
|
if target_agent.status != AgentStatus.IDLE:
|
|
logger.error(
|
|
f"Target agent {target_agent_id} is not idle (status: {target_agent.status})"
|
|
)
|
|
return None
|
|
|
|
# Create handoff record
|
|
handoff_id = str(uuid.uuid4())
|
|
record = HandoffRecord(
|
|
handoff_id=handoff_id,
|
|
source_agent_id=source_agent_id,
|
|
target_agent_id=target_agent_id,
|
|
timestamp=time.time(),
|
|
task_context=task_context,
|
|
status="pending",
|
|
)
|
|
|
|
# Store in Redis
|
|
key = f"{self.HANDOFF_KEY_PREFIX}{handoff_id}"
|
|
await redis_client.set(key, record.model_dump_json())
|
|
|
|
# Update agent statuses
|
|
await agent_registry.update_status(source_agent_id, AgentStatus.IDLE)
|
|
await agent_registry.update_status(
|
|
target_agent_id, AgentStatus.HANDOFF_PENDING, task_context.get("task_id")
|
|
)
|
|
|
|
logger.info(
|
|
f"Handoff {handoff_id} initiated: {source_agent_id} -> {target_agent_id}"
|
|
)
|
|
return handoff_id
|
|
|
|
async def complete_handoff(self, handoff_id: str) -> bool:
|
|
"""Mark handoff as completed."""
|
|
key = f"{self.HANDOFF_KEY_PREFIX}{handoff_id}"
|
|
data = await redis_client.get(key)
|
|
|
|
if not data:
|
|
logger.error(f"Handoff {handoff_id} not found")
|
|
return False
|
|
|
|
record = HandoffRecord.model_validate_json(data)
|
|
record.status = "completed"
|
|
|
|
await redis_client.set(key, record.model_dump_json())
|
|
|
|
# Update target agent to busy
|
|
await agent_registry.update_status(
|
|
record.target_agent_id, AgentStatus.BUSY, record.task_context.get("task_id")
|
|
)
|
|
|
|
logger.info(f"Handoff {handoff_id} completed")
|
|
return True
|
|
|
|
async def fail_handoff(self, handoff_id: str, reason: str) -> bool:
|
|
"""Mark handoff as failed."""
|
|
key = f"{self.HANDOFF_KEY_PREFIX}{handoff_id}"
|
|
data = await redis_client.get(key)
|
|
|
|
if not data:
|
|
logger.error(f"Handoff {handoff_id} not found")
|
|
return False
|
|
|
|
record = HandoffRecord.model_validate_json(data)
|
|
record.status = "failed"
|
|
|
|
await redis_client.set(key, record.model_dump_json())
|
|
|
|
# Revert target agent to idle
|
|
await agent_registry.update_status(record.target_agent_id, AgentStatus.IDLE)
|
|
|
|
logger.warning(f"Handoff {handoff_id} failed: {reason}")
|
|
return True
|
|
|
|
async def get_handoff(self, handoff_id: str) -> Optional[HandoffRecord]:
|
|
"""Get handoff record by ID."""
|
|
key = f"{self.HANDOFF_KEY_PREFIX}{handoff_id}"
|
|
data = await redis_client.get(key)
|
|
|
|
if not data:
|
|
return None
|
|
|
|
return HandoffRecord.model_validate_json(data)
|
|
|
|
async def get_handoff_history(
|
|
self, agent_id: Optional[str] = None, limit: int = 100
|
|
) -> list[HandoffRecord]:
|
|
"""Get handoff history, optionally filtered by agent."""
|
|
pattern = f"{self.HANDOFF_KEY_PREFIX}*"
|
|
keys = await redis_client.keys(pattern)
|
|
|
|
records = []
|
|
for key in keys[:limit]:
|
|
data = await redis_client.get(key)
|
|
if data:
|
|
record = HandoffRecord.model_validate_json(data)
|
|
if agent_id is None or (
|
|
record.source_agent_id == agent_id
|
|
or record.target_agent_id == agent_id
|
|
):
|
|
records.append(record)
|
|
|
|
# Sort by timestamp descending
|
|
records.sort(key=lambda r: r.timestamp, reverse=True)
|
|
return records[:limit]
|
|
|
|
|
|
# Global handoff manager instance
|
|
handoff_manager = HandoffManager()
|