agents/protocol.py

"""
Agent Communication Protocol — message types, queues, and metrics for ArcisVLM agents.

Provides:
- TaskMessage / ResultMessage / HeartbeatMessage for inter-agent communication
- JSON serialization round-tripping
- MessageQueue interface with an in-memory priority-queue implementation
- AgentMetrics for tracking latency, throughput, queue depth, and expert activation
"""

from __future__ import annotations

import json
import time
import heapq
import threading
import logging
from abc import ABC, abstractmethod
from dataclasses import dataclass, field, asdict
from typing import Any

from agents.base import Task, Result, PRIORITY_ORDER

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Message types
# ---------------------------------------------------------------------------

@dataclass
class TaskMessage:
    """Wraps a Task for routing between agents."""

    task: Task
    sender_id: str
    target_agent_type: str  # expert type the task should be routed to
    timestamp: float = field(default_factory=time.time)
    message_id: str = ""

    def __post_init__(self) -> None:
        if not self.message_id:
            self.message_id = f"tmsg-{int(self.timestamp * 1000)}-{id(self) % 0xFFFF:04x}"

    def to_json(self) -> str:
        return json.dumps({
            "type": "task",
            "message_id": self.message_id,
            "sender_id": self.sender_id,
            "target_agent_type": self.target_agent_type,
            "timestamp": self.timestamp,
            "task": {
                "task_id": self.task.task_id,
                "type": self.task.type,
                "payload": self.task.payload,
                "priority": self.task.priority,
                "deadline_ms": self.task.deadline_ms,
                "source": self.task.source,
                "expert_hint": self.task.expert_hint,
                "created_at": self.task.created_at,
            },
        })

    @classmethod
    def from_json(cls, raw: str) -> TaskMessage:
        d = json.loads(raw)
        task_data = d["task"]
        task = Task(
            type=task_data["type"],
            payload=task_data["payload"],
            priority=task_data["priority"],
            deadline_ms=task_data["deadline_ms"],
            source=task_data["source"],
            expert_hint=task_data["expert_hint"],
            task_id=task_data["task_id"],
            created_at=task_data["created_at"],
        )
        return cls(
            task=task,
            sender_id=d["sender_id"],
            target_agent_type=d["target_agent_type"],
            timestamp=d["timestamp"],
            message_id=d["message_id"],
        )


@dataclass
class ResultMessage:
    """Wraps a Result returned from a child agent back to the Mother orchestrator."""

    result: Result
    sender_id: str
    timestamp: float = field(default_factory=time.time)
    message_id: str = ""

    def __post_init__(self) -> None:
        if not self.message_id:
            self.message_id = f"rmsg-{int(self.timestamp * 1000)}-{id(self) % 0xFFFF:04x}"

    def to_json(self) -> str:
        return json.dumps({
            "type": "result",
            "message_id": self.message_id,
            "sender_id": self.sender_id,
            "timestamp": self.timestamp,
            "result": {
                "answer": self.result.answer,
                "confidence": self.result.confidence,
                "processing_time_ms": self.result.processing_time_ms,
                "expert_used": self.result.expert_used,
                "metadata": self.result.metadata,
                "task_id": self.result.task_id,
                "agent_id": self.result.agent_id,
                "error": self.result.error,
            },
        })

    @classmethod
    def from_json(cls, raw: str) -> ResultMessage:
        d = json.loads(raw)
        rd = d["result"]
        result = Result(
            answer=rd["answer"],
            confidence=rd["confidence"],
            processing_time_ms=rd["processing_time_ms"],
            expert_used=rd["expert_used"],
            metadata=rd["metadata"],
            task_id=rd["task_id"],
            agent_id=rd["agent_id"],
            error=rd.get("error"),
        )
        return cls(
            result=result,
            sender_id=d["sender_id"],
            timestamp=d["timestamp"],
            message_id=d["message_id"],
        )


@dataclass
class HeartbeatMessage:
    """Periodic health signal from an agent."""

    agent_id: str
    agent_type: str
    status: str  # "idle" | "processing" | "error"
    tasks_processed: int = 0
    avg_latency_ms: float = 0.0
    timestamp: float = field(default_factory=time.time)
    message_id: str = ""

    def __post_init__(self) -> None:
        if not self.message_id:
            self.message_id = f"hb-{int(self.timestamp * 1000)}-{id(self) % 0xFFFF:04x}"

    def to_json(self) -> str:
        return json.dumps({
            "type": "heartbeat",
            "message_id": self.message_id,
            "agent_id": self.agent_id,
            "agent_type": self.agent_type,
            "status": self.status,
            "tasks_processed": self.tasks_processed,
            "avg_latency_ms": self.avg_latency_ms,
            "timestamp": self.timestamp,
        })

    @classmethod
    def from_json(cls, raw: str) -> HeartbeatMessage:
        d = json.loads(raw)
        return cls(
            agent_id=d["agent_id"],
            agent_type=d["agent_type"],
            status=d["status"],
            tasks_processed=d["tasks_processed"],
            avg_latency_ms=d["avg_latency_ms"],
            timestamp=d["timestamp"],
            message_id=d["message_id"],
        )


# ---------------------------------------------------------------------------
# Generic message deserialization
# ---------------------------------------------------------------------------

def deserialize_message(raw: str) -> TaskMessage | ResultMessage | HeartbeatMessage:
    """Deserialize a JSON string into the appropriate message type."""
    d = json.loads(raw)
    msg_type = d.get("type")
    if msg_type == "task":
        return TaskMessage.from_json(raw)
    elif msg_type == "result":
        return ResultMessage.from_json(raw)
    elif msg_type == "heartbeat":
        return HeartbeatMessage.from_json(raw)
    else:
        raise ValueError(f"Unknown message type: {msg_type}")


# ---------------------------------------------------------------------------
# Message Queue — abstract interface + in-memory implementation
# ---------------------------------------------------------------------------

class MessageQueue(ABC):
    """Abstract message queue that routes TaskMessages between agents."""

    @abstractmethod
    def put(self, message: TaskMessage) -> None:
        """Enqueue a task message."""
        ...

    @abstractmethod
    def get(self, agent_type: str, timeout_s: float | None = None) -> TaskMessage | None:
        """
        Dequeue the highest-priority TaskMessage for the given agent type.

        Returns None if no message is available within the timeout.
        """
        ...

    @abstractmethod
    def depth(self, agent_type: str | None = None) -> int:
        """
        Return number of pending messages.
        If agent_type is specified, return depth for that type only.
        """
        ...


class InMemoryQueue(MessageQueue):
    """
    Thread-safe, priority-ordered in-memory message queue.

    Messages are stored per agent_type in min-heaps keyed by (priority_rank, timestamp).
    """

    def __init__(self) -> None:
        self._lock = threading.Lock()
        self._condition = threading.Condition(self._lock)
        # agent_type -> list of (priority_rank, timestamp, TaskMessage)
        self._queues: dict[str, list[tuple[int, float, TaskMessage]]] = {}
        self._total_enqueued = 0
        self._total_dequeued = 0

    def put(self, message: TaskMessage) -> None:
        with self._condition:
            agent_type = message.target_agent_type
            if agent_type not in self._queues:
                self._queues[agent_type] = []
            priority_rank = PRIORITY_ORDER.get(message.task.priority, 1)
            heapq.heappush(
                self._queues[agent_type],
                (priority_rank, message.timestamp, message),
            )
            self._total_enqueued += 1
            self._condition.notify_all()

    def get(self, agent_type: str, timeout_s: float | None = None) -> TaskMessage | None:
        deadline = None if timeout_s is None else time.monotonic() + timeout_s
        with self._condition:
            while True:
                q = self._queues.get(agent_type, [])
                if q:
                    _, _, message = heapq.heappop(q)
                    self._total_dequeued += 1
                    return message
                if timeout_s is not None:
                    remaining = deadline - time.monotonic()
                    if remaining <= 0:
                        return None
                    self._condition.wait(timeout=remaining)
                else:
                    return None

    def depth(self, agent_type: str | None = None) -> int:
        with self._lock:
            if agent_type is not None:
                return len(self._queues.get(agent_type, []))
            return sum(len(q) for q in self._queues.values())

    @property
    def total_enqueued(self) -> int:
        return self._total_enqueued

    @property
    def total_dequeued(self) -> int:
        return self._total_dequeued


# ---------------------------------------------------------------------------
# Metrics tracker
# ---------------------------------------------------------------------------

class AgentMetrics:
    """
    Collects operational metrics across the agent system.

    Thread-safe.  Intended to be a singleton owned by the Mother orchestrator.
    """

    def __init__(self) -> None:
        self._lock = threading.Lock()
        # Per-agent latency tracking: agent_id -> list of latencies in ms
        self._latencies: dict[str, list[float]] = {}
        # Per-expert activation counts
        self._expert_activations: dict[str, int] = {}
        # Task counters
        self._tasks_submitted = 0
        self._tasks_completed = 0
        self._tasks_failed = 0

    # -- Recording -----------------------------------------------------------

    def record_task_submitted(self) -> None:
        with self._lock:
            self._tasks_submitted += 1

    def record_task_completed(self, agent_id: str, expert_used: str, latency_ms: float) -> None:
        with self._lock:
            self._tasks_completed += 1
            self._latencies.setdefault(agent_id, []).append(latency_ms)
            self._expert_activations[expert_used] = self._expert_activations.get(expert_used, 0) + 1

    def record_task_failed(self) -> None:
        with self._lock:
            self._tasks_failed += 1

    # -- Queries -------------------------------------------------------------

    def avg_latency_ms(self, agent_id: str | None = None) -> float:
        """Average latency.  If agent_id is None, compute across all agents."""
        with self._lock:
            if agent_id is not None:
                lats = self._latencies.get(agent_id, [])
            else:
                lats = [l for vals in self._latencies.values() for l in vals]
            return sum(lats) / len(lats) if lats else 0.0

    def p95_latency_ms(self, agent_id: str | None = None) -> float:
        """P95 latency."""
        with self._lock:
            if agent_id is not None:
                lats = sorted(self._latencies.get(agent_id, []))
            else:
                lats = sorted(l for vals in self._latencies.values() for l in vals)
            if not lats:
                return 0.0
            idx = int(0.95 * len(lats))
            return lats[min(idx, len(lats) - 1)]

    def expert_activation_counts(self) -> dict[str, int]:
        with self._lock:
            return dict(self._expert_activations)

    def queue_depth(self, queue: MessageQueue, agent_type: str | None = None) -> int:
        return queue.depth(agent_type)

    def summary(self, queue: MessageQueue | None = None) -> dict[str, Any]:
        """Return a snapshot dict of all metrics."""
        with self._lock:
            result: dict[str, Any] = {
                "tasks_submitted": self._tasks_submitted,
                "tasks_completed": self._tasks_completed,
                "tasks_failed": self._tasks_failed,
                "avg_latency_ms": self.avg_latency_ms(),
                "p95_latency_ms": self.p95_latency_ms(),
                "expert_activations": dict(self._expert_activations),
            }
        if queue is not None:
            result["queue_depth"] = queue.depth()
        return result