101 lines
3.1 KiB
Python
101 lines
3.1 KiB
Python
"""Structured metrics for agent tool calls."""
|
|
|
|
import asyncio
|
|
import logging
|
|
import time
|
|
from typing import Any, Callable
|
|
|
|
from langchain.agents.middleware import AgentMiddleware
|
|
from langchain.tools.tool_node import ToolCallRequest
|
|
|
|
from agent.agent_config import AgentConfig
|
|
from utils.structured_log import emit_question_metric
|
|
|
|
logger = logging.getLogger("app")
|
|
|
|
|
|
class ToolMetricsMiddleware(AgentMiddleware):
|
|
"""Emit structured timing metrics for every tool call."""
|
|
|
|
def __init__(self, config: AgentConfig):
|
|
self.config = config
|
|
|
|
def _emit_tool_metric(
|
|
self,
|
|
request: ToolCallRequest,
|
|
*,
|
|
started_at: float,
|
|
status: str,
|
|
error_type: str | None = None,
|
|
) -> None:
|
|
tool_call = request.tool_call or {}
|
|
tool_name = tool_call.get("name") or "unknown_tool"
|
|
tool_call_id = tool_call.get("id")
|
|
duration_ms = max(int((time.monotonic() - started_at) * 1000), 0)
|
|
|
|
try:
|
|
emit_question_metric(
|
|
stage="catalog_agent.tool_call",
|
|
status=status,
|
|
duration_ms=duration_ms,
|
|
trace_id=self.config.trace_id,
|
|
ai_id=self.config.bot_id,
|
|
session_id=self.config.session_id,
|
|
robot_type="agent",
|
|
model=self.config.model_name,
|
|
stream=self.config.stream,
|
|
error_type=error_type,
|
|
extra={
|
|
"bot_id": self.config.bot_id,
|
|
"tool_name": tool_name,
|
|
"tool_call_id": tool_call_id,
|
|
"tool_response": self.config.tool_response,
|
|
"enable_thinking": self.config.enable_thinking,
|
|
},
|
|
)
|
|
except Exception:
|
|
logger.exception("Failed to emit tool metric for tool_name=%s", tool_name)
|
|
|
|
def wrap_tool_call(
|
|
self,
|
|
request: ToolCallRequest,
|
|
handler: Callable[[ToolCallRequest], Any],
|
|
) -> Any:
|
|
started_at = time.monotonic()
|
|
try:
|
|
result = handler(request)
|
|
except Exception as exc:
|
|
self._emit_tool_metric(
|
|
request,
|
|
started_at=started_at,
|
|
status="error",
|
|
error_type=type(exc).__name__,
|
|
)
|
|
raise
|
|
|
|
self._emit_tool_metric(request, started_at=started_at, status="success")
|
|
return result
|
|
|
|
async def awrap_tool_call(
|
|
self,
|
|
request: ToolCallRequest,
|
|
handler: Callable[[ToolCallRequest], Any],
|
|
) -> Any:
|
|
started_at = time.monotonic()
|
|
try:
|
|
result = await handler(request)
|
|
except asyncio.CancelledError:
|
|
self._emit_tool_metric(request, started_at=started_at, status="cancel")
|
|
raise
|
|
except Exception as exc:
|
|
self._emit_tool_metric(
|
|
request,
|
|
started_at=started_at,
|
|
status="error",
|
|
error_type=type(exc).__name__,
|
|
)
|
|
raise
|
|
|
|
self._emit_tool_metric(request, started_at=started_at, status="success")
|
|
return result
|