update attributes

Author: yaozheng-fang

veadk/tools/builtin_tools/lark.py

@@ -46,7 +46,7 @@ def check_env():
         ) from e
 
 
-# check_env()
+check_env()
 
 lark_tools = MCPToolset(
     connection_params=StdioServerParameters(

veadk/tracing/telemetry/attributes/extractors/common_attributes_extractors.py

@@ -31,14 +31,20 @@ def common_gen_ai_session_id(**kwargs) -> str:
     return session_id or "<unknown_session_id>"
 
 
+def common_cozeloop_report_source(**kwargs) -> str:
+    return "veadk"
+
+
 COMMON_ATTRIBUTES = {
     "gen_ai.system": common_gen_ai_system,
     "gen_ai.system.version": common_gen_ai_system_version,
     "gen_ai.agent.name": common_gen_ai_agent_name,
     "gen_ai.app.name": common_gen_ai_app_name,  # APMPlus required
     "gen_ai.user.id": common_gen_ai_user_id,  # APMPlus required
     "gen_ai.session.id": common_gen_ai_session_id,  # APMPlus required
-    "app.name": common_gen_ai_app_name,  # CozeLoop required
+    "agent_name": common_gen_ai_agent_name,  # CozeLoop required
+    "app_name": common_gen_ai_app_name,  # CozeLoop required
     "user.id": common_gen_ai_user_id,  # CozeLoop required
     "session.id": common_gen_ai_session_id,  # CozeLoop required
+    "cozeloop.report.source": common_cozeloop_report_source,  # CozeLoop required
 }

veadk/tracing/telemetry/attributes/extractors/llm_attributes_extractors.py

@@ -6,30 +6,12 @@
 )
 
 
-def flatten(d, parent_key="", sep="."):
-    items = []
-    if isinstance(d, dict):
-        for k, v in d.items():
-            new_key = f"{parent_key}{sep}{k}" if parent_key else k
-            items.extend(flatten(v, new_key, sep=sep).items())
-    elif isinstance(d, list):
-        if not d:  # empty list
-            items.append((parent_key, ""))
-        else:
-            for i, v in enumerate(d):
-                new_key = f"{parent_key}{sep}{i}" if parent_key else str(i)
-                items.extend(flatten(v, new_key, sep=sep).items())
-    else:
-        items.append((parent_key, d))
-    return dict(items)
-
-
 def llm_gen_ai_request_model(params: LLMAttributesParams) -> ExtractorResponse:
     return ExtractorResponse(content=params.llm_request.model or "<unknown_model_name>")
 
 
 def llm_gen_ai_request_type(params: LLMAttributesParams) -> ExtractorResponse:
-    return ExtractorResponse(content=type or "<unknown_type>")
+    return ExtractorResponse(content="chat" or "<unknown_type>")
 
 
 def llm_gen_ai_response_model(params: LLMAttributesParams) -> ExtractorResponse:

veadk/tracing/telemetry/attributes/extractors/tool_attributes_extractors.py

@@ -1,3 +1,5 @@
+import json
+
 from veadk.tracing.telemetry.attributes.extractors.types import (
     ExtractorResponse,
     ToolAttributesParams,
@@ -8,6 +10,32 @@ def tool_gen_ai_operation_name(params: ToolAttributesParams) -> ExtractorRespons
     return ExtractorResponse(content="execute_tool")
 
 
+def tool_gen_ai_tool_message(params: ToolAttributesParams) -> ExtractorResponse:
+    tool_input = {
+        "id": "123",
+        "role": "tool",
+        "content": json.dumps(
+            {
+                "name": params.tool.name,
+                "description": params.tool.description,
+                "parameters": params.args,
+            }
+        ),
+    }
+    return ExtractorResponse(type="event", content=tool_input)
+
+
+# def tool_gen_ai_tool_message(params: ToolAttributesParams) -> ExtractorResponse:
+#     # tool_output = {
+#     #     "id": params.function_response_event.,
+#     #     "name": ...,
+#     #     "response": ...,
+#     # }
+#     print(params.function_response_event)
+#     # return ExtractorResponse(content=json.dumps(tool_output) or "<unknown_tool_output>")
+
+
 TOOL_ATTRIBUTES = {
     "gen_ai.operation.name": tool_gen_ai_operation_name,
+    "gen_ai.tool.message": tool_gen_ai_tool_message,
 }

veadk/tracing/telemetry/attributes/extractors/types.py

@@ -12,13 +12,13 @@
 
 @dataclass
 class ExtractorResponse:
-    content: Any
+    content: list | dict | None | str | int | float
 
     type: Literal["attribute", "event"] = "attribute"
     """Type of extractor response.
     
-    `raw`: span[...] = raw
-    `event`: span.add_event(...) = ...
+    `attribute`: span.add_attribute(attr_name, attr_value)
+    `event`: span.add_event(...)
     """
 
     @staticmethod
@@ -33,14 +33,16 @@ def update_span(
                         for k, v in _res.items():
                             span.set_attribute(k, v)
             else:
-                span.set_attribute(attr_name, res)
+                # set anyway
+                span.set_attribute(attr_name, res)  # type: ignore
         elif response.type == "event":
             if isinstance(response.content, dict):
                 span.add_event(attr_name, response.content)
-            else:  # list
+            elif isinstance(response.content, list):
                 for event in response.content:
                     span.add_event(attr_name, event)
         else:
+            # Unsupported response type, discard it.
             pass
 
 

veadk/tracing/telemetry/exporters/inmemory_exporter.py

@@ -29,9 +29,6 @@
 logger = get_logger(__name__)
 
 
-inmemory_span_processor = None
-
-
 # ======== Adapted from Google ADK ========
 class _InMemoryExporter(export.SpanExporter):
     def __init__(self) -> None:
@@ -109,6 +106,3 @@ def __init__(self, name: str = "inmemory_exporter") -> None:
 
         self._exporter = _InMemoryExporter()
         self.processor = _InMemorySpanProcessor(self._exporter)
-
-        global inmemory_span_processor
-        inmemory_span_processor = self.processor

veadk/tracing/telemetry/telemetry.py

@@ -14,7 +14,6 @@
     LLMAttributesParams,
     ToolAttributesParams,
 )
-from veadk.tracing.telemetry.exporters import inmemory_exporter
 from veadk.utils.logger import get_logger
 
 logger = get_logger(__name__)
@@ -23,27 +22,47 @@
 def trace_send_data(): ...
 
 
-def set_common_attributes(current_span: _Span, **kwargs) -> None:
+def set_common_attributes(
+    invocation_context: InvocationContext, current_span: _Span, **kwargs
+) -> None:
+    from veadk.agent import Agent
+
     if current_span.context:
         current_span_id = current_span.context.trace_id
-
-    spans = inmemory_exporter.inmemory_span_processor.spans
-
-    spans_in_current_trace = [
-        span
-        for span in spans
-        if span.context and span.context.trace_id == current_span_id
-    ]
-
-    common_attributes = ATTRIBUTES.get("common", {})
-    for span in spans_in_current_trace:
-        if span.name.startswith("invocation"):
-            span.set_attribute("gen_ai.operation.name", "chain")
-        elif span.name.startswith("agent_run"):
-            span.set_attribute("gen_ai.operation.name", "agent")
-        for attr_name, attr_extractor in common_attributes.items():
-            value = attr_extractor(**kwargs)
-            span.set_attribute(attr_name, value)
+    else:
+        logger.warning(
+            "Current span context is missing, failed to get `trace_id` to set common attributes."
+        )
+        return
+
+    if isinstance(invocation_context.agent, Agent):
+        try:
+            from veadk.tracing.telemetry.opentelemetry_tracer import OpentelemetryTracer
+
+            tracer: OpentelemetryTracer = invocation_context.agent.tracers[0]  # type: ignore
+            spans = tracer._inmemory_exporter.processor.spans  #  # type: ignore
+
+            spans_in_current_trace = [
+                span
+                for span in spans
+                if span.context and span.context.trace_id == current_span_id
+            ]
+
+            common_attributes = ATTRIBUTES.get("common", {})
+            for span in spans_in_current_trace:
+                if span.name.startswith("invocation"):
+                    span.set_attribute("gen_ai.operation.name", "chain")
+                elif span.name.startswith("agent_run"):
+                    span.set_attribute("gen_ai.operation.name", "agent")
+                for attr_name, attr_extractor in common_attributes.items():
+                    value = attr_extractor(**kwargs)
+                    span.set_attribute(attr_name, value)
+        except Exception as e:
+            logger.error(f"Failed to set common attributes for spans: {e}")
+    else:
+        logger.warning(
+            "Failed to set common attributes for spans as your agent is not VeADK Agent. Skip this."
+        )
 
 
 def trace_tool_call(
@@ -72,6 +91,7 @@ def trace_call_llm(
     from veadk.agent import Agent
 
     set_common_attributes(
+        invocation_context=invocation_context,
         current_span=span,  # type: ignore
         agent_name=invocation_context.agent.name,
         user_id=invocation_context.user_id,