Add tool call metrics middleware
This commit is contained in:
parent
63f1f836c4
commit
ecc2687f7b
@ -23,6 +23,7 @@ from utils.fastapi_utils import detect_provider, sanitize_model_kwargs
|
||||
from .guideline_middleware import GuidelineMiddleware
|
||||
from .tool_output_length_middleware import ToolOutputLengthMiddleware
|
||||
from .tool_use_cleanup_middleware import ToolUseCleanupMiddleware
|
||||
from .tool_metrics_middleware import ToolMetricsMiddleware
|
||||
from .filepath_fix_middleware import FilePathFixMiddleware
|
||||
from .mcp_trace_meta import patch_mcp_client_session_trace_meta
|
||||
from utils.settings import (
|
||||
@ -256,6 +257,7 @@ async def init_agent(config: AgentConfig):
|
||||
# Build the middleware list
|
||||
middleware = []
|
||||
middleware.append(EmptyResponseRetryMiddleware())
|
||||
middleware.append(ToolMetricsMiddleware(config))
|
||||
middleware.append(ToolUseCleanupMiddleware())
|
||||
# tool_output_middleware = ToolOutputLengthMiddleware(
|
||||
# max_length=(getattr(config.generate_cfg, 'tool_output_max_length', None) if config.generate_cfg else None) or TOOL_OUTPUT_MAX_LENGTH,
|
||||
|
||||
100
agent/tool_metrics_middleware.py
Normal file
100
agent/tool_metrics_middleware.py
Normal file
@ -0,0 +1,100 @@
|
||||
"""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
|
||||
Loading…
Reference in New Issue
Block a user