Tool_Agent/src/tool_agent/router/middleware/metrics.py

"""调用计量中间件"""

from __future__ import annotations

import time
from dataclasses import dataclass, field


@dataclass
class ToolMetrics:
    call_count: int = 0
    total_latency_ms: float = 0.0
    error_count: int = 0

    @property
    def avg_latency_ms(self) -> float:
        return self.total_latency_ms / self.call_count if self.call_count else 0.0

    @property
    def error_rate(self) -> float:
        return self.error_count / self.call_count if self.call_count else 0.0


class MetricsCollector:
    """记录每个工具的调用频次、耗时、错误率"""

    def __init__(self) -> None:
        self._metrics: dict[str, ToolMetrics] = {}

    def record(self, tool_id: str, latency_ms: float, error: bool = False) -> None:
        if tool_id not in self._metrics:
            self._metrics[tool_id] = ToolMetrics()
        m = self._metrics[tool_id]
        m.call_count += 1
        m.total_latency_ms += latency_ms
        if error:
            m.error_count += 1

    def get(self, tool_id: str) -> ToolMetrics | None:
        return self._metrics.get(tool_id)

    def all(self) -> dict[str, ToolMetrics]:
        return dict(self._metrics)