fix exporter bugs

Author: yaozheng-fang

veadk/runner.py

@@ -27,7 +27,6 @@
 from veadk.agents.sequential_agent import SequentialAgent
 from veadk.evaluation import EvalSetRecorder
 from veadk.memory.short_term_memory import ShortTermMemory
-from veadk.tracing.base_tracer import UserMessagePlugin
 from veadk.types import MediaMessage
 from veadk.utils.logger import get_logger
 from veadk.utils.misc import read_png_to_bytes
@@ -73,17 +72,6 @@ def __init__(
         else:
             self.long_term_memory = None
 
-        # process plugins
-        try:
-            # try to detect tracer
-            _ = self.agent.tracers[0]
-            if not plugins:
-                plugins = [UserMessagePlugin(name="user_message_plugin")]
-            else:
-                plugins.append(UserMessagePlugin(name="user_message_plugin"))
-        except Exception:
-            logger.debug("Agent has no tracers, telemetry plugin not added.")
-
         self.runner = ADKRunner(
             app_name=self.app_name,
             agent=self.agent,
@@ -201,7 +189,7 @@ def _print_trace_id(self):
             return
 
         try:
-            trace_id = self.agent.tracers[0].get_trace_id()  # type: ignore
+            trace_id = self.agent.tracers[0].trace_id  # type: ignore
             logger.info(f"Trace id: {trace_id}")
         except Exception as e:
             logger.warning(f"Get tracer id failed as {e}")

veadk/tracing/base_tracer.py

@@ -13,69 +13,13 @@
 # limitations under the License.
 
 from abc import ABC, abstractmethod
-from typing import Optional
 
-from google.adk.agents.invocation_context import InvocationContext
-from google.adk.plugins.base_plugin import BasePlugin
-from google.genai import types
-from opentelemetry import trace
 
 from veadk.utils.logger import get_logger
 
 logger = get_logger(__name__)
 
 
-class UserMessagePlugin(BasePlugin):
-    def __init__(self, name: str):
-        super().__init__(name)
-
-    async def on_user_message_callback(
-        self,
-        *,
-        invocation_context: InvocationContext,
-        user_message: types.Content,
-    ) -> Optional[types.Content]:
-        """Callback executed when a user message is received before an invocation starts.
-
-        This callback helps logging and modifying the user message before the
-        runner starts the invocation.
-
-        Args:
-        invocation_context: The context for the entire invocation.
-        user_message: The message content input by user.
-
-        Returns:
-        An optional `types.Content` to be returned to the ADK. Returning a
-        value to replace the user message. Returning `None` to proceed
-        normally.
-        """
-        trace.get_tracer("gcp.vertex.agent")
-        span = trace.get_current_span()
-
-        logger.debug(f"User message plugin works, catch {span}")
-        span_name = getattr(span, "name", None)
-        if span_name and span_name.startswith("invocation"):
-            agent_name = invocation_context.agent.name
-            invoke_branch = (
-                invocation_context.branch if invocation_context.branch else agent_name
-            )
-            current_session = invocation_context.session
-
-            span.set_attribute("app.name", current_session.app_name)
-            span.set_attribute("user.id", current_session.user_id)
-            span.set_attribute("session.id", current_session.id)
-
-            span.set_attribute("agent.name", agent_name)
-            span.set_attribute("invoke.branch", invoke_branch)
-            span.set_attribute("gen_ai.system", "veadk")
-
-            logger.debug(
-                f"Add attributes to {span_name}: app_name={current_session.app_name}, user_id={current_session.user_id}, session_id={current_session.id}, agent_name={agent_name}, invoke_branch={invoke_branch}"
-            )
-
-        return None
-
-
 def replace_bytes_with_empty(data):
     """
     Recursively traverse the data structure and replace all bytes types with empty strings.

veadk/tracing/telemetry/attributes/attributes.py

@@ -1,22 +1,15 @@
 from veadk.tracing.telemetry.attributes.extractors.common_attributes_extractors import (
-    common_gen_ai_app_name,
-    common_gen_ai_session_id,
-    common_gen_ai_system,
-    common_gen_ai_system_version,
-    common_gen_ai_user_id,
+    COMMON_ATTRIBUTES,
 )
-from veadk.tracing.telemetry.attributes.extractors.llm_attributes_extrators import (
+from veadk.tracing.telemetry.attributes.extractors.llm_attributes_extractors import (
     LLM_ATTRIBUTES,
 )
+from veadk.tracing.telemetry.attributes.extractors.tool_attributes_extractors import (
+    TOOL_ATTRIBUTES,
+)
 
 ATTRIBUTES = {
-    "common": {
-        "gen_ai.system": common_gen_ai_system,
-        "gen_ai.system_version": common_gen_ai_system_version,
-        "gen_ai.app.name": common_gen_ai_app_name,
-        "gen_ai.user.id": common_gen_ai_user_id,
-        "gen_ai.session.id": common_gen_ai_session_id,
-    },
+    "common": COMMON_ATTRIBUTES,
     "llm": LLM_ATTRIBUTES,
-    "tool": ...,
+    "tool": TOOL_ATTRIBUTES,
 }

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

@@ -1,21 +1,33 @@
 from veadk.version import VERSION
 
 
-def common_gen_ai_system() -> str:
+def common_gen_ai_system(**kwargs) -> str:
     return "veadk"
 
 
-def common_gen_ai_system_version() -> str:
+def common_gen_ai_system_version(**kwargs) -> str:
     return VERSION
 
 
 def common_gen_ai_app_name(**kwargs) -> str:
-    return kwargs.get("app_name", "<unknown_app_name>")
+    app_name = kwargs.get("app_name")
+    return app_name or "<unknown_app_name>"
 
 
 def common_gen_ai_user_id(**kwargs) -> str:
-    return kwargs.get("user_id", "")
+    user_id = kwargs.get("user_id")
+    return user_id or "<unknown_user_id>"
 
 
 def common_gen_ai_session_id(**kwargs) -> str:
-    return kwargs.get("session_id", "<unknown_session_id>")
+    session_id = kwargs.get("session_id")
+    return session_id or "<unknown_session_id>"
+
+
+COMMON_ATTRIBUTES = {
+    "gen_ai.system": common_gen_ai_system,
+    "gen_ai.system.version": common_gen_ai_system_version,
+    "gen_ai.app.name": common_gen_ai_app_name,
+    "gen_ai.user.id": common_gen_ai_user_id,
+    "gen_ai.session.id": common_gen_ai_session_id,
+}

…s/extractors/llm_attributes_extrators.py …/extractors/llm_attributes_extractors.pyveadk/tracing/telemetry/attributes/extractors/llm_attributes_extrators.py renamed to veadk/tracing/telemetry/attributes/extractors/llm_attributes_extractors.py veadk/tracing/telemetry/attributes/extractors/llm_attributes_extrators.py renamed to veadk/tracing/telemetry/attributes/extractors/llm_attributes_extractors.py

@@ -0,0 +1,15 @@
+from typing import Any
+
+from attr import dataclass
+from google.adk.events import Event
+from google.adk.tools import BaseTool
+
+
+@dataclass
+class ToolAttributesParams:
+    tool: BaseTool
+    args: dict[str, Any]
+    function_response_event: Event
+
+
+TOOL_ATTRIBUTES = {}

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

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Any
+
 from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
 from opentelemetry.sdk.trace.export import BatchSpanProcessor
 from pydantic import BaseModel, Field
@@ -41,18 +43,23 @@ class CozeloopExporterConfig(BaseModel):
     )
 
 
+class _CozeloopExporter(OTLPSpanExporter):
+    pass
+
+
 class CozeloopExporter(BaseExporter):
     config: CozeloopExporterConfig = Field(default_factory=CozeloopExporterConfig)
 
-    def model_post_init(self) -> None:
+    def model_post_init(self, context: Any) -> None:
         headers = {
             "cozeloop-workspace-id": self.config.space_id,
+            "authorization": f"Bearer {self.config.token}",
         }
         self.headers |= headers
 
         self._exporter = OTLPSpanExporter(
             endpoint=self.config.endpoint,
-            headers=headers,
+            headers=self.headers,
             timeout=10,
         )
 

veadk/tracing/telemetry/exporters/cozeloop_exporter.py

@@ -23,46 +23,30 @@
 logger = get_logger(__name__)
 
 
+in_memory_exporter_instance = None
+
+
 # ======== Adapted from Google ADK ========
 class _InMemoryExporter(export.SpanExporter):
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__()
         self._spans = []
-        self.session_trace_dict = {}
         self.trace_id = ""
-        self.prompt_tokens = []
-        self.completion_tokens = []
+        self.span_dict = {}
+        self.session_trace_dict = {}
 
     @override
     def export(self, spans: Sequence[ReadableSpan]) -> export.SpanExportResult:
         for span in spans:
             if span.context:
-                trace_id = span.context.trace_id
-                self.trace_id = trace_id
+                self.trace_id = span.context.trace_id
+
+                span_id = span.context.span_id
+                self.span_dict[span_id] = span
             else:
                 logger.warning(
                     f"Span context is missing, failed to get `trace_id`. span: {span}"
                 )
-
-            if span.name == "call_llm":
-                attributes = dict(span.attributes or {})
-                prompt_token = attributes.get("gen_ai.usage.prompt_tokens", None)
-                completion_token = attributes.get(
-                    "gen_ai.usage.completion_tokens", None
-                )
-                if prompt_token:
-                    self.prompt_tokens.append(prompt_token)
-                if completion_token:
-                    self.completion_tokens.append(completion_token)
-
-            if span.name == "call_llm":
-                attributes = dict(span.attributes or {})
-                session_id = attributes.get("gcp.vertex.agent.session_id", None)
-                if session_id:
-                    if session_id not in self.session_trace_dict:
-                        self.session_trace_dict[session_id] = [trace_id]
-                    else:
-                        self.session_trace_dict[session_id] += [trace_id]
         self._spans.extend(spans)
         return export.SpanExportResult.SUCCESS
 

veadk/tracing/telemetry/exporters/inmemory_exporter.py

@@ -18,7 +18,7 @@
 import time
 from typing import Any
 
-from openinference.instrumentation.google_adk import GoogleADKInstrumentor
+# from openinference.instrumentation.google_adk import GoogleADKInstrumentor
 from opentelemetry import trace as trace_api
 from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
 from opentelemetry.sdk import trace as trace_sdk
@@ -79,7 +79,7 @@ def model_post_init(self, context: Any) -> None:
         # VeADK operates on global OpenTelemetry provider, return nothing
         self._init_global_tracer_provider()
 
-        GoogleADKInstrumentor().instrument()
+        # GoogleADKInstrumentor().instrument()
 
     def _init_global_tracer_provider(self) -> None:
         # set provider anyway, then get global provider
@@ -132,7 +132,7 @@ def trace_file_path(self) -> str:
     @property
     def trace_id(self) -> str:
         try:
-            trace_id = hex(int(self._inmemory_exporter._exporter.trace_id))[2:]
+            trace_id = hex(int(self._inmemory_exporter._exporter.trace_id))[2:]  # type: ignore
             return trace_id
         except Exception as e:
             logger.error(f"Failed to get trace_id from InMemoryExporter: {e}")
@@ -151,23 +151,11 @@ def dump(
             )
             return ""
 
-        prompt_tokens = self._inmemory_exporter._real_exporter.prompt_tokens
-        completion_tokens = self._inmemory_exporter._real_exporter.completion_tokens
-
-        # upload
-        for meter_uploader in self._meter_uploaders:
-            meter_uploader.record(
-                prompt_tokens=prompt_tokens, completion_tokens=completion_tokens
-            )
-        # clear tokens after dump
-        self._inmemory_exporter._real_exporter.completion_tokens = []
-        self._inmemory_exporter._real_exporter.prompt_tokens = []
-
         for processor in self._processors:
             time.sleep(0.05)  # give some time for the exporter to upload spans
             processor.force_flush()
 
-        spans = self._inmemory_exporter._real_exporter.get_finished_spans(
+        spans = self._inmemory_exporter._exporter.get_finished_spans(  # type: ignore
             session_id=session_id
         )
         if not spans:
@@ -186,14 +174,11 @@ def dump(
                 for s in spans
             ]
 
-        trace_id = hex(int(self._inmemory_exporter._real_exporter.trace_id))[2:]
-        self._trace_id = trace_id
-        file_path = f"{path}/{self.name}_{user_id}_{session_id}_{trace_id}.json"
+        file_path = f"{path}/{self.name}_{user_id}_{session_id}_{self.trace_id}.json"
         with open(file_path, "w") as f:
             json.dump(
                 data, f, indent=4, ensure_ascii=False
             )  # ensure_ascii=False to support Chinese characters
-
         self._trace_file_path = file_path
 
         for exporter in self.exporters:
@@ -203,6 +188,6 @@ def dump(
             f"OpenTelemetryTracer tracing done, trace id: {self._trace_id} (hex)"
         )
 
-        self._spans = spans
         logger.info(f"OpenTelemetryTracer dumps {len(spans)} spans to {file_path}")
+
         return file_path