Merge pull request #42 from sparticleinc/codex/tool-metrics-for-agent-first-char-staging

Add tool call metrics middleware to staging

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,

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