fix tool common attr setting

Author: yaozheng-fang

veadk/tracing/base_tracer.py

@@ -19,25 +19,6 @@
 logger = get_logger(__name__)
 
 
-def replace_bytes_with_empty(data):
-    """
-    Recursively traverse the data structure and replace all bytes types with empty strings.
-    Supports handling any nested structure of lists and dictionaries.
-    """
-    if isinstance(data, dict):
-        # Handle dictionary: Recursively process each value
-        return {k: replace_bytes_with_empty(v) for k, v in data.items()}
-    elif isinstance(data, list):
-        # Handle list: Recursively process each element
-        return [replace_bytes_with_empty(item) for item in data]
-    elif isinstance(data, bytes):
-        # When encountering the bytes type, replace it with an empty string
-        return "<image data>"
-    else:
-        # Keep other types unchanged
-        return data
-
-
 class BaseTracer(ABC):
     def __init__(self, name: str):
         self.name = name

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

@@ -395,8 +395,8 @@ def llm_output_value(params: LLMAttributesParams) -> ExtractorResponse:
     # attributes
     "gen_ai.prompt": llm_gen_ai_prompt,
     "gen_ai.completion": llm_gen_ai_completion,
-    "input.value": llm_input_value,  # TLS required
-    "output.value": llm_output_value,  # TLS required
+    # "input.value": llm_input_value,
+    # "output.value": llm_output_value,
     # ===== usage =====
     "gen_ai.usage.input_tokens": llm_gen_ai_usage_input_tokens,
     "gen_ai.usage.output_tokens": llm_gen_ai_usage_output_tokens,

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

@@ -38,7 +38,7 @@ def tool_gen_ai_tool_message(params: ToolAttributesParams) -> ExtractorResponse:
     return ExtractorResponse(type="event", content=tool_input)
 
 
-def tool_cozeloop_input(params: ToolAttributesParams) -> ExtractorResponse:
+def tool_gen_ai_tool_input(params: ToolAttributesParams) -> ExtractorResponse:
     tool_input = {
         "name": params.tool.name,
         "description": params.tool.description,
@@ -53,7 +53,7 @@ def tool_gen_ai_tool_name(params: ToolAttributesParams) -> ExtractorResponse:
     return ExtractorResponse(content=params.tool.name or "<unknown_tool_name>")
 
 
-def tool_cozeloop_output(params: ToolAttributesParams) -> ExtractorResponse:
+def tool_gen_ai_tool_output(params: ToolAttributesParams) -> ExtractorResponse:
     function_response = params.function_response_event.get_function_responses()[
         0
     ].model_dump()
@@ -70,6 +70,8 @@ def tool_cozeloop_output(params: ToolAttributesParams) -> ExtractorResponse:
 TOOL_ATTRIBUTES = {
     "gen_ai.operation.name": tool_gen_ai_operation_name,
     "gen_ai.tool.name": tool_gen_ai_tool_name,  # TLS required
-    "cozeloop.input": tool_cozeloop_input,  # CozeLoop required
-    "cozeloop.output": tool_cozeloop_output,  # CozeLoop required
+    "gen_ai.tool.input": tool_gen_ai_tool_input,  # TLS required
+    "gen_ai.tool.output": tool_gen_ai_tool_output,  # TLS required
+    "cozeloop.input": tool_gen_ai_tool_input,  # CozeLoop required
+    "cozeloop.output": tool_gen_ai_tool_output,  # CozeLoop required
 }

veadk/tracing/telemetry/telemetry.py

@@ -51,10 +51,7 @@ def upload_metrics(
                     exporter.meter_uploader.record(llm_request, llm_response)
 
 
-def trace_send_data(): ...
-
-
-def set_common_attributes(
+def set_common_attributes_on_model_span(
     invocation_context: InvocationContext, current_span: _Span, **kwargs
 ) -> None:
     if current_span.context:
@@ -66,7 +63,7 @@ def set_common_attributes(
         return
 
     try:
-        spans = _INMEMORY_EXPORTER_INSTANCE.processor.spans  #  # type: ignore
+        spans = _INMEMORY_EXPORTER_INSTANCE.processor.spans  # type: ignore
 
         spans_in_current_trace = [
             span
@@ -76,24 +73,52 @@ def set_common_attributes(
 
         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)
+            if span.is_recording():
+                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}")
 
 
+def set_common_attributes_on_tool_span(current_span: _Span) -> None:
+    # find parent span (generally a llm span)
+    if not current_span.context:
+        logger.warning(
+            f"Get tool span's context failed. Skip setting common attributes for span {current_span.name}"
+        )
+        return
+
+    if not current_span.parent:
+        logger.warning(
+            f"Get tool span's parent failed. Skip setting common attributes for span {current_span.name}"
+        )
+        return
+
+    parent_span_id = current_span.parent.span_id
+    for span in _INMEMORY_EXPORTER_INSTANCE.processor.spans:  # type: ignore
+        if span.context.span_id == parent_span_id:
+            common_attributes = ATTRIBUTES.get("common", {})
+            for attr_name in common_attributes.keys():
+                current_span.set_attribute(attr_name, span.attributes[attr_name])
+
+
+def trace_send_data(): ...
+
+
 def trace_tool_call(
     tool: BaseTool,
     args: dict[str, Any],
     function_response_event: Event,
 ) -> None:
     span = trace.get_current_span()
 
+    set_common_attributes_on_tool_span(current_span=span)  # type: ignore
+
     tool_attributes_mapping = ATTRIBUTES.get("tool", {})
     params = ToolAttributesParams(tool, args, function_response_event)
 
@@ -112,7 +137,7 @@ def trace_call_llm(
 
     from veadk.agent import Agent
 
-    set_common_attributes(
+    set_common_attributes_on_model_span(
         invocation_context=invocation_context,
         current_span=span,  # type: ignore
         agent_name=invocation_context.agent.name,