fix: update telemetry

(cherry picked from commit 95d59b3f0eef649bebc022375293201d097e9417)

Author: feng-95

agentkit/apps/agent_server_app/middleware.py

@@ -19,8 +19,10 @@
 
 from agentkit.apps.agent_server_app.telemetry import telemetry
 
-
-trace_paths = ["/run_sse", "/run", "/invoke"]
+_EXCLUDED_HEADERS = {
+    "authorization",
+    "token"
+}
 
 
 class AgentkitTelemetryHTTPMiddleware:
@@ -36,19 +38,23 @@ async def __call__(self, scope, receive, send):
         path = scope.get("path", "")
         headers_list = scope.get("headers", [])
         headers = {k.decode("latin-1"): v.decode("latin-1") for k, v in headers_list}
-
         span = telemetry.tracer.start_span(name="agent_server_request")
         ctx = trace.set_span_in_context(span)
         context_api.attach(ctx)
+        headers = {
+            k: v for k, v in headers.items()
+            if k.lower() not in _EXCLUDED_HEADERS
+        }
 
+        # Currently unable to retrieve user_id and session_id from headers; keep logic for future use
         user_id = headers.get("user_id") or headers.get("x-user-id") or ""
         session_id = headers.get("session_id") or headers.get("x-session-id") or ""
-        headers_like = {"user_id": user_id, "session_id": session_id}
-
+        headers["user_id"] = user_id
+        headers["session_id"] = session_id
         telemetry.trace_agent_server(
             func_name=f"{method} {path}",
             span=span,
-            headers=headers_like,
+            headers=headers,
             text="",  # do not consume body in middleware
         )
 
@@ -58,7 +64,7 @@ async def send_wrapper(message):
                     more_body = message.get("more_body", False)
                     if not more_body:
                         telemetry.trace_agent_server_finish(
-                            func_result="", exception=None
+                            path=path, func_result="", exception=None
                         )
                 elif message.get("type") == "http.response.start":
                     # could record status code if needed
@@ -69,5 +75,5 @@ async def send_wrapper(message):
         try:
             await self.app(scope, receive, send_wrapper)
         except Exception as e:
-            telemetry.trace_agent_server_finish(func_result="", exception=e)
+            telemetry.trace_agent_server_finish(path=path,func_result="", exception=e)
             raise

agentkit/apps/agent_server_app/telemetry.py

@@ -23,6 +23,8 @@
 
 from agentkit.apps.utils import safe_serialize_to_json_string
 
+_INVOKE_PATH = ["/run_sse", "/run", "/invoke"]
+
 _GEN_AI_CLIENT_OPERATION_DURATION_BUCKETS = [
     0.01,
     0.02,
@@ -77,22 +79,26 @@ def trace_agent_server(
         if user_id:
             span.set_attribute(key="gen_ai.user.id", value=user_id)
 
-        span.set_attribute(
-            key="gen_ai.input", value=safe_serialize_to_json_string(text)
-        )
+        # Currently unable to retrieve input
+        # span.set_attribute(
+        #     key="gen_ai.input", value=safe_serialize_to_json_string(text)
+        # )
 
-        span.set_attribute(key="gen_ai.span.kind", value="workflow")
+        span.set_attribute(key="gen_ai.span.kind", value="agent_server")
         span.set_attribute(key="gen_ai.operation.name", value="invoke_agent")
         span.set_attribute(key="gen_ai.operation.type", value="agent_server")
 
     def trace_agent_server_finish(
         self,
+        path: str,
         func_result: str,
         exception: Optional[Exception],
     ) -> None:
         span = trace.get_current_span()
         if span and span.is_recording():
-            span.set_attribute(key="gen_ai.output", value=func_result)
+            # Currently unable to retrieve output
+            # span.set_attribute(key="gen_ai.output", value=func_result)
+
             attributes = {
                 "gen_ai_operation_name": "invoke_agent",
                 "gen_ai_operation_type": "agent_server",
@@ -101,7 +107,8 @@ def trace_agent_server_finish(
                 self.handle_exception(span, exception)
                 attributes["error_type"] = exception.__class__.__name__
 
-            if hasattr(span, "start_time") and self.latency_histogram:
+            # only record invoke request latency metrics
+            if hasattr(span, "start_time") and self.latency_histogram and path in _INVOKE_PATH:
                 duration = (time.time_ns() - span.start_time) / 1e9  # type: ignore
                 self.latency_histogram.record(duration, attributes)
             span.end()