veadk-python 0.2.2__py3-none-any.whl → 0.2.5__py3-none-any.whl

veadk/tracing/telemetry/opentelemetry_tracer.py

@@ -18,96 +18,74 @@ import json
 import time
 from typing import Any
 
-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
 from opentelemetry.sdk.resources import Resource
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import BatchSpanProcessor, SimpleSpanProcessor
-from pydantic import BaseModel, ConfigDict, Field
+from pydantic import BaseModel, ConfigDict, Field, field_validator
 from typing_extensions import override
 
 from veadk.tracing.base_tracer import BaseTracer
 from veadk.tracing.telemetry.exporters.apmplus_exporter import APMPlusExporter
 from veadk.tracing.telemetry.exporters.base_exporter import BaseExporter
 from veadk.tracing.telemetry.exporters.inmemory_exporter import InMemoryExporter
-from veadk.tracing.telemetry.metrics.opentelemetry_metrics import MeterUploader
+from veadk.tracing.telemetry.exporters.inmemory_exporter import (
+    _INMEMORY_EXPORTER_INSTANCE,
+)
 from veadk.utils.logger import get_logger
+from veadk.utils.patches import patch_google_adk_telemetry
 
 logger = get_logger(__name__)
 
-DEFAULT_VEADK_TRACER_NAME = "veadk_global_tracer"
 
-
-def update_resource_attributions(provider: TracerProvider, resource_attributes: dict):
+def _update_resource_attributions(
+    provider: TracerProvider, resource_attributes: dict
+) -> None:
     provider._resource = provider._resource.merge(Resource.create(resource_attributes))
 
 
 class OpentelemetryTracer(BaseModel, BaseTracer):
     model_config = ConfigDict(arbitrary_types_allowed=True)
-    exporters: list[BaseExporter] = Field(
-        default=[],
-        description="The exporters to export spans.",
-    )
+
     name: str = Field(
-        DEFAULT_VEADK_TRACER_NAME, description="The identifier of tracer."
+        default="veadk_opentelemetry_tracer", description="The identifier of tracer."
     )
 
-    app_name: str = Field(
-        "veadk_app",
-        description="The identifier of app.",
+    exporters: list[BaseExporter] = Field(
+        default_factory=list,
+        description="The exporters to export spans.",
     )
 
-    def model_post_init(self, context: Any, /) -> None:
-        self._processors = []
-        self._inmemory_exporter: InMemoryExporter | None = None
+    @field_validator("exporters")
+    @classmethod
+    def forbid_inmemory_exporter(cls, v: list[BaseExporter]) -> list[BaseExporter]:
+        for e in v:
+            if isinstance(e, InMemoryExporter):
+                raise ValueError("InMemoryExporter is not allowed in exporters list")
+        return v
 
-        # InMemoryExporter is a default exporter for exporting local tracing file
-        for exporter in self.exporters:
-            if isinstance(exporter, InMemoryExporter):
-                self._inmemory_exporter = exporter
+    def model_post_init(self, context: Any) -> None:
+        patch_google_adk_telemetry()
+        self._init_global_tracer_provider()
 
-        if self._inmemory_exporter is None:
-            self._inmemory_exporter = InMemoryExporter()
-            self.exporters.append(self._inmemory_exporter)
-        # ========================================================================
-
-        # Process meter-related attributes
-        self._meter_contexts = []
-        self._meter_uploaders = []
-        for exporter in self.exporters:
-            meter_context = exporter.get_meter_context()
-            if meter_context is not None:
-                self._meter_contexts.append(meter_context)
+    #  GoogleADKInstrumentor().instrument()
 
-        for meter_context in self._meter_contexts:
-            meter_uploader = MeterUploader(meter_context)
-            self._meter_uploaders.append(meter_uploader)
-        # ================================
-
-        # init tracer provider
-        # VeADK operates on global OpenTelemetry provider, hence return nothing
-        self._init_tracer_provider()
-
-        # just for debug
-        self._trace_file_path = ""
-
-        GoogleADKInstrumentor().instrument()
+    def _init_global_tracer_provider(self) -> None:
+        self._processors = []
 
-    def _init_tracer_provider(self) -> None:
-        # set provider anyway
-        # finally, get global provider
-        tracer_provider = trace_sdk.TracerProvider()
-        trace_api.set_tracer_provider(tracer_provider)
+        # set provider anyway, then get global provider
+        trace_api.set_tracer_provider(trace_sdk.TracerProvider())
         global_tracer_provider: TracerProvider = trace_api.get_tracer_provider()  # type: ignore
 
-        have_apmplus_exporter = False
-        for processor in global_tracer_provider._active_span_processor._span_processors:
-            if isinstance(processor, (BatchSpanProcessor, SimpleSpanProcessor)):
-                if isinstance(processor.span_exporter, OTLPSpanExporter):
-                    if "apmplus" in processor.span_exporter._endpoint:
-                        have_apmplus_exporter = True
+        span_processors = global_tracer_provider._active_span_processor._span_processors
+        have_apmplus_exporter = any(
+            isinstance(p, (BatchSpanProcessor, SimpleSpanProcessor))
+            and isinstance(p.span_exporter, OTLPSpanExporter)
+            and "apmplus" in p.span_exporter._endpoint
+            for p in span_processors
+        )
 
         if have_apmplus_exporter:
             self.exporters = [
@@ -115,54 +93,83 @@ class OpentelemetryTracer(BaseModel, BaseTracer):
             ]
 
         for exporter in self.exporters:
-            processor, resource_attributes = exporter.get_processor()
-            if resource_attributes is not None:
-                update_resource_attributions(
+            processor = exporter.processor
+            resource_attributes = exporter.resource_attributes
+
+            if resource_attributes:
+                _update_resource_attributions(
                     global_tracer_provider, resource_attributes
                 )
-            global_tracer_provider.add_span_processor(processor)
-            logger.debug(
-                f"Add exporter `{exporter.__class__.__name__}` to OpentelemetryTracer."
+
+            if processor:
+                global_tracer_provider.add_span_processor(processor)
+                self._processors.append(processor)
+
+                logger.debug(
+                    f"Add span processor for exporter `{exporter.__class__.__name__}` to OpentelemetryTracer."
+                )
+            else:
+                logger.error(
+                    f"Add span processor for exporter `{exporter.__class__.__name__}` to OpentelemetryTracer failed."
+                )
+
+        self._inmemory_exporter = _INMEMORY_EXPORTER_INSTANCE
+        if self._inmemory_exporter.processor:
+            # make sure the in memory exporter processor is added at index 0
+            # because we use this to record all spans
+            global_tracer_provider._active_span_processor._span_processors = (
+                self._inmemory_exporter.processor,
+            ) + global_tracer_provider._active_span_processor._span_processors
+
+            self._processors.append(self._inmemory_exporter.processor)
+        else:
+            logger.warning(
+                "InMemoryExporter processor is not initialized, cannot add to OpentelemetryTracer."
             )
-            self._processors.append(processor)
-        logger.debug(f"Init OpentelemetryTracer with {len(self.exporters)} exporters.")
+
+        logger.info(f"Init OpentelemetryTracer with {len(self._processors)} exporters.")
+
+    @property
+    def trace_file_path(self) -> str:
+        return self._trace_file_path
+
+    @property
+    def trace_id(self) -> str:
+        try:
+            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}")
+            return self._trace_id
+
+    def force_export(self) -> None:
+        """Force to export spans in all processors."""
+        for processor in self._processors:
+            time.sleep(0.05)
+            processor.force_flush()
 
     @override
     def dump(
         self,
-        user_id: str,
-        session_id: str,
+        user_id: str = "unknown_user_id",
+        session_id: str = "unknown_session_id",
         path: str = "/tmp",
     ) -> str:
+        def _build_trace_file_path(path: str, user_id: str, session_id: str) -> str:
+            return f"{path}/{self.name}_{user_id}_{session_id}_{self.trace_id}.json"
+
         if not self._inmemory_exporter:
             logger.warning(
                 "InMemoryExporter is not initialized. Please check your tracer exporters."
             )
             return ""
+        self.force_export()
 
-        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:
-            data = []
-        else:
-            data = [
+        data = (
+            [
                 {
                     "name": s.name,
                     "span_id": s.context.span_id,
@@ -174,22 +181,18 @@ class OpentelemetryTracer(BaseModel, BaseTracer):
                 }
                 for s in spans
             ]
+            if spans
+            else []
+        )
 
-        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"
-        with open(file_path, "w") as f:
-            json.dump(data, f, indent=4)
-
-        self._trace_file_path = file_path
+        self._trace_file_path = _build_trace_file_path(path, user_id, session_id)
+        with open(self._trace_file_path, "w") as f:
+            json.dump(
+                data, f, indent=4, ensure_ascii=False
+            )  # ensure_ascii=False to support Chinese characters
 
-        for exporter in self.exporters:
-            if not isinstance(exporter, InMemoryExporter):
-                exporter.export()
         logger.info(
-            f"OpenTelemetryTracer tracing done, trace id: {self._trace_id} (hex)"
+            f"OpenTelemetryTracer dumps {len(spans)} spans to {self._trace_file_path}. Trace id: {self.trace_id} (hex)"
         )
 
-        self._spans = spans
-        logger.info(f"OpenTelemetryTracer dumps {len(spans)} spans to {file_path}")
-        return file_path
+        return self._trace_file_path

veadk/tracing/telemetry/telemetry.py

@@ -0,0 +1,238 @@
+# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any
+
+from google.adk.agents.invocation_context import InvocationContext
+from google.adk.events import Event
+from google.adk.models.llm_request import LlmRequest
+from google.adk.models.llm_response import LlmResponse
+from google.adk.tools import BaseTool
+from opentelemetry import trace
+from opentelemetry.sdk.trace import _Span
+
+from veadk.tracing.telemetry.attributes.attributes import ATTRIBUTES
+from veadk.tracing.telemetry.attributes.extractors.types import (
+    ExtractorResponse,
+    LLMAttributesParams,
+    ToolAttributesParams,
+)
+from veadk.tracing.telemetry.exporters.inmemory_exporter import (
+    _INMEMORY_EXPORTER_INSTANCE,
+)
+from veadk.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+def upload_metrics(
+    invocation_context: InvocationContext,
+    llm_request: LlmRequest,
+    llm_response: LlmResponse,
+) -> None:
+    from veadk.agent import Agent
+
+    if isinstance(invocation_context.agent, Agent):
+        tracers = invocation_context.agent.tracers
+        for tracer in tracers:
+            for exporter in getattr(tracer, "exporters", []):
+                if getattr(exporter, "meter_uploader", None):
+                    exporter.meter_uploader.record(llm_request, llm_response)
+
+
+def _set_agent_input_attribute(
+    span: _Span, invocation_context: InvocationContext
+) -> None:
+    # We only save the original user input as the agent input
+    # hence once the `agent.input` has been set, we don't overwrite it
+    event_names = [event.name for event in span.events]
+    if "gen_ai.user.message" in event_names:
+        return
+
+    # input = {
+    #     "agent_name": invocation_context.agent.name,
+    #     "app_name": invocation_context.session.app_name,
+    #     "user_id": invocation_context.user_id,
+    #     "session_id": invocation_context.session.id,
+    #     "input": invocation_context.user_content.model_dump(exclude_none=True)
+    #     if invocation_context.user_content
+    #     else None,
+    # }
+
+    user_content = invocation_context.user_content
+    if user_content and user_content.parts:
+        span.add_event(
+            "gen_ai.user.message",
+            {
+                "agent_name": invocation_context.agent.name,
+                "app_name": invocation_context.session.app_name,
+                "user_id": invocation_context.user_id,
+                "session_id": invocation_context.session.id,
+            },
+        )
+        for idx, part in enumerate(user_content.parts):
+            if part.text:
+                span.add_event(
+                    "gen_ai.user.message",
+                    {f"parts.{idx}.type": "text", f"parts.{idx}.content": part.text},
+                )
+            if part.inline_data:
+                span.add_event(
+                    "gen_ai.user.message",
+                    {
+                        f"parts.{idx}.type": "image_url",
+                        f"parts.{idx}.image_url.name": part.inline_data.display_name.split(
+                            "/"
+                        )[-1],
+                        f"parts.{idx}.image_url.url": part.inline_data.display_name,
+                    },
+                )
+
+
+def _set_agent_output_attribute(span: _Span, llm_response: LlmResponse) -> None:
+    content = llm_response.content
+    if content and content.parts:
+        for idx, part in enumerate(content.parts):
+            if part.text:
+                span.add_event(
+                    "gen_ai.choice",
+                    {
+                        f"message.parts.{idx}.type": "text",
+                        f"message.parts.{idx}.text": part.text,
+                    },
+                )
+
+
+def set_common_attributes_on_model_span(
+    invocation_context: InvocationContext,
+    llm_response: LlmResponse,
+    current_span: _Span,
+    **kwargs,
+) -> None:
+    if current_span.context:
+        current_span_id = current_span.context.trace_id
+    else:
+        logger.warning(
+            "Current span context is missing, failed to get `trace_id` to set common attributes."
+        )
+        return
+
+    try:
+        spans = _INMEMORY_EXPORTER_INSTANCE.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.is_recording():
+                if span.name.startswith("invocation"):
+                    span.set_attribute("gen_ai.operation.name", "chain")
+                    _set_agent_input_attribute(span, invocation_context)
+                    _set_agent_output_attribute(span, llm_response)
+                elif span.name.startswith("agent_run"):
+                    span.set_attribute("gen_ai.operation.name", "agent")
+                    _set_agent_input_attribute(span, invocation_context)
+                    _set_agent_output_attribute(span, llm_response)
+                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)
+
+    for attr_name, attr_extractor in tool_attributes_mapping.items():
+        response: ExtractorResponse = attr_extractor(params)
+        ExtractorResponse.update_span(span, attr_name, response)
+
+
+def trace_call_llm(
+    invocation_context: InvocationContext,
+    event_id: str,
+    llm_request: LlmRequest,
+    llm_response: LlmResponse,
+) -> None:
+    span = trace.get_current_span()
+
+    from veadk.agent import Agent
+
+    set_common_attributes_on_model_span(
+        invocation_context=invocation_context,
+        llm_response=llm_response,
+        current_span=span,  # type: ignore
+        agent_name=invocation_context.agent.name,
+        user_id=invocation_context.user_id,
+        app_name=invocation_context.app_name,
+        session_id=invocation_context.session.id,
+        model_provider=invocation_context.agent.model_provider
+        if isinstance(invocation_context.agent, Agent)
+        else "",
+        model_name=invocation_context.agent.model_name
+        if isinstance(invocation_context.agent, Agent)
+        else "",
+    )
+
+    llm_attributes_mapping = ATTRIBUTES.get("llm", {})
+    params = LLMAttributesParams(
+        invocation_context=invocation_context,
+        event_id=event_id,
+        llm_request=llm_request,
+        llm_response=llm_response,
+    )
+
+    for attr_name, attr_extractor in llm_attributes_mapping.items():
+        response: ExtractorResponse = attr_extractor(params)
+        ExtractorResponse.update_span(span, attr_name, response)
+
+    upload_metrics(invocation_context, llm_request, llm_response)

veadk/types.py

@@ -15,6 +15,9 @@
 from pydantic import BaseModel, Field
 
 from veadk.agent import Agent
+from veadk.agents.loop_agent import LoopAgent
+from veadk.agents.parallel_agent import ParallelAgent
+from veadk.agents.sequential_agent import SequentialAgent
 from veadk.memory.short_term_memory import ShortTermMemory
 
 
@@ -35,7 +38,9 @@ class AgentRunConfig(BaseModel):
         default="veadk_vefaas_app", description="The name of the application"
     )
 
-    agent: Agent = Field(..., description="The root agent instance")
+    agent: Agent | SequentialAgent | ParallelAgent | LoopAgent = Field(
+        ..., description="The root agent instance"
+    )
 
     short_term_memory: ShortTermMemory = Field(
         default_factory=ShortTermMemory, description="The short-term memory instance"

veadk/utils/misc.py

@@ -13,9 +13,11 @@
 # limitations under the License.
 
 import importlib.util
+import json
 import sys
 import time
 import types
+from typing import Any, Dict, List, MutableMapping, Tuple
 
 import requests
 
@@ -27,7 +29,7 @@ def read_file(file_path):
     return data
 
 
-def formatted_timestamp():
+def formatted_timestamp() -> str:
     # YYYYMMDDHHMMSS
     return time.strftime("%Y%m%d%H%M%S", time.localtime())
 
@@ -60,3 +62,41 @@ def load_module_from_file(module_name: str, file_path: str) -> types.ModuleType:
             )
     else:
         raise ImportError(f"Could not load module {module_name} from {file_path}")
+
+
+def flatten_dict(
+    d: MutableMapping[str, Any], parent_key: str = "", sep: str = "_"
+) -> Dict[str, Any]:
+    """Flatten a nested dictionary.
+
+    Input:
+        {"a": {"b": 1}}
+    Output:
+        {"a_b": 1}
+    """
+    items: List[Tuple[str, Any]] = []
+    for k, v in d.items():
+        new_key = f"{parent_key}{sep}{k}" if parent_key else k
+        if isinstance(v, MutableMapping):
+            items.extend(flatten_dict(v, new_key, sep=sep).items())
+        else:
+            items.append((new_key, v))
+    return dict(items)
+
+
+def safe_json_serialize(obj) -> str:
+    """Convert any Python object to a JSON-serializable type or string.
+
+    Args:
+      obj: The object to serialize.
+
+    Returns:
+      The JSON-serialized object string or <non-serializable> if the object cannot be serialized.
+    """
+
+    try:
+        return json.dumps(
+            obj, ensure_ascii=False, default=lambda o: "<not serializable>"
+        )
+    except (TypeError, OverflowError):
+        return "<not serializable>"

veadk/utils/patches.py

@@ -13,7 +13,14 @@
 # limitations under the License.
 
 import asyncio
+import sys
+from typing import Callable
 
+from veadk.tracing.telemetry.telemetry import (
+    trace_call_llm,
+    trace_send_data,
+    trace_tool_call,
+)
 from veadk.utils.logger import get_logger
 
 logger = get_logger(__name__)
@@ -53,3 +60,21 @@ def patch_asyncio():
             raise
 
     CancelScope.__exit__ = patched_cancel_scope_exit
+
+
+def patch_google_adk_telemetry() -> None:
+    trace_functions = {
+        "trace_tool_call": trace_tool_call,
+        "trace_call_llm": trace_call_llm,
+        "trace_send_data": trace_send_data,
+    }
+
+    for mod_name, mod in sys.modules.items():
+        if mod_name.startswith("google.adk"):
+            for var_name in dir(mod):
+                var = getattr(mod, var_name, None)
+                if var_name in trace_functions.keys() and isinstance(var, Callable):
+                    setattr(mod, var_name, trace_functions[var_name])
+                    logger.debug(
+                        f"Patch {mod_name} {var_name} with {trace_functions[var_name]}"
+                    )

veadk/version.py

@@ -12,4 +12,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-VERSION = "0.2.2"
+VERSION = "0.2.5"