diff --git a/agent/deep_assistant.py b/agent/deep_assistant.py index 957b0d0..ec11513 100644 --- a/agent/deep_assistant.py +++ b/agent/deep_assistant.py @@ -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, diff --git a/agent/tool_metrics_middleware.py b/agent/tool_metrics_middleware.py new file mode 100644 index 0000000..3f96b84 --- /dev/null +++ b/agent/tool_metrics_middleware.py @@ -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