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

veadk/tracing/telemetry/exporters/apmplus_exporter.py

@@ -12,10 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Any
+
+from google.adk.models.llm_request import LlmRequest
+from google.adk.models.llm_response import LlmResponse
 from opentelemetry import metrics
+from opentelemetry import metrics as metrics_api
 from opentelemetry.exporter.otlp.proto.grpc.metric_exporter import OTLPMetricExporter
 from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
-from opentelemetry.sdk.metrics import MeterProvider
+from opentelemetry.metrics._internal import Meter
+from opentelemetry.sdk import metrics as metrics_sdk
 from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader
 from opentelemetry.sdk.resources import Resource
 from opentelemetry.sdk.trace.export import BatchSpanProcessor
@@ -24,12 +30,76 @@ from typing_extensions import override
 
 from veadk.config import getenv
 from veadk.tracing.telemetry.exporters.base_exporter import BaseExporter
-from veadk.tracing.telemetry.metrics.opentelemetry_metrics import MeterContext
 from veadk.utils.logger import get_logger
 
 logger = get_logger(__name__)
 
 
+class MeterUploader:
+    def __init__(
+        self, name: str, endpoint: str, headers: dict, resource_attributes: dict
+    ) -> None:
+        # global_metrics_provider -> global_tracer_provider
+        # exporter -> exporter
+        # metric_reader -> processor
+        global_metrics_provider = metrics_api.get_meter_provider()
+
+        # 1. init resource
+        if hasattr(global_metrics_provider, "_sdk_config"):
+            global_resource = global_metrics_provider._sdk_config.resource  # type: ignore
+        else:
+            global_resource = Resource.create()
+
+        resource = global_resource.merge(Resource.create(resource_attributes))
+
+        # 2. init exporter and reader
+        exporter = OTLPMetricExporter(endpoint=endpoint, headers=headers)
+        metric_reader = PeriodicExportingMetricReader(exporter)
+
+        metrics_api.set_meter_provider(
+            metrics_sdk.MeterProvider(metric_readers=[metric_reader], resource=resource)
+        )
+
+        # 3. init meter
+        self.meter: Meter = metrics.get_meter(name=name)
+
+        # create meter attributes
+        self.llm_invoke_counter = self.meter.create_counter(
+            name="gen_ai.chat.count",
+            description="Number of LLM invocations",
+            unit="count",
+        )
+        self.token_usage = self.meter.create_histogram(
+            name="gen_ai.client.token.usage",
+            description="Token consumption of LLM invocations",
+            unit="count",
+        )
+
+    def record(self, llm_request: LlmRequest, llm_response: LlmResponse) -> None:
+        attributes = {
+            "gen_ai_system": "volcengine",
+            "gen_ai_response_model": llm_request.model,
+            "gen_ai_operation_name": "chat_completions",
+            "stream": "false",
+            "server_address": "api.volcengine.com",
+        }  # required by Volcengine APMPlus
+
+        if llm_response.usage_metadata:
+            # llm invocation number += 1
+            self.llm_invoke_counter.add(1, attributes)
+
+            # upload token usage
+            input_token = llm_response.usage_metadata.prompt_token_count
+            output_token = llm_response.usage_metadata.candidates_token_count
+
+            if input_token:
+                token_attributes = {**attributes, "gen_ai_token_type": "input"}
+                self.token_usage.record(input_token, attributes=token_attributes)
+            if output_token:
+                token_attributes = {**attributes, "gen_ai_token_type": "output"}
+                self.token_usage.record(output_token, attributes=token_attributes)
+
+
 class APMPlusExporterConfig(BaseModel):
     endpoint: str = Field(
         default_factory=lambda: getenv(
@@ -49,50 +119,39 @@ class APMPlusExporterConfig(BaseModel):
     )
 
 
-class APMPlusExporter(BaseModel, BaseExporter):
+class APMPlusExporter(BaseExporter):
     config: APMPlusExporterConfig = Field(default_factory=APMPlusExporterConfig)
 
-    @override
-    def get_processor(self):
-        resource_attributes = {
-            "service.name": self.config.service_name,
-        }
-
+    def model_post_init(self, context: Any) -> None:
         headers = {
             "x-byteapm-appkey": self.config.app_key,
         }
-        exporter = OTLPSpanExporter(
-            endpoint=self.config.endpoint, insecure=True, headers=headers
-        )
-        self._real_exporter = exporter
-        processor = BatchSpanProcessor(exporter)
-        return processor, resource_attributes
-
-    def export(self):
-        self._real_exporter.force_flush()
-        logger.info(
-            f"APMPlusExporter exports data to {self.config.endpoint}, service name: {self.config.service_name}"
-        )
+        self.headers |= headers
 
-    @override
-    def get_meter_context(self) -> MeterContext:
         resource_attributes = {
             "service.name": self.config.service_name,
         }
-        endpoint = self.config.endpoint
-        headers = {
-            "x-byteapm-appkey": self.config.app_key,
-        }
+        self.resource_attributes |= resource_attributes
 
-        resource = Resource.create(resource_attributes)
-        exporter = OTLPMetricExporter(endpoint=endpoint, headers=headers)
-        metric_reader = PeriodicExportingMetricReader(exporter)
-        provider = MeterProvider(metric_readers=[metric_reader], resource=resource)
-        metrics.set_meter_provider(provider)
-        meter = metrics.get_meter("my.meter.name")
-        meter_context = MeterContext(
-            meter=meter,
-            provider=provider,
-            reader=metric_reader,
+        self._exporter = OTLPSpanExporter(
+            endpoint=self.config.endpoint, insecure=True, headers=self.headers
         )
-        return meter_context
+        self.processor = BatchSpanProcessor(self._exporter)
+
+        self.meter_uploader = MeterUploader(
+            name="apmplus_meter",
+            endpoint=self.config.endpoint,
+            headers=self.headers,
+            resource_attributes=self.resource_attributes,
+        )
+
+    @override
+    def export(self) -> None:
+        if self._exporter:
+            self._exporter.force_flush()
+
+            logger.info(
+                f"APMPlusExporter exports data to {self.config.endpoint}, service name: {self.config.service_name}"
+            )
+        else:
+            logger.warning("APMPlusExporter internal exporter is not initialized.")

veadk/tracing/telemetry/exporters/base_exporter.py

@@ -12,20 +12,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from abc import ABC, abstractmethod
-from typing import Any
+from opentelemetry.sdk.trace import SpanProcessor
+from opentelemetry.sdk.trace.export import SpanExporter
+from pydantic import BaseModel, ConfigDict, Field
 
 
-class BaseExporter(ABC):
-    def __init__(self) -> None:
-        pass
+class BaseExporter(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True, extra="allow")
 
-    @abstractmethod
-    def get_processor(self) -> Any:
-        pass
+    resource_attributes: dict = Field(default_factory=dict)
+    headers: dict = Field(default_factory=dict)
 
-    def get_meter_context(self) -> Any:
-        pass
+    _exporter: SpanExporter | None = None
+    processor: SpanProcessor | None = None
 
     def export(self) -> None:
+        """Force export of telemetry data."""
         pass

veadk/tracing/telemetry/exporters/cozeloop_exporter.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,26 +43,31 @@ class CozeloopExporterConfig(BaseModel):
     )
 
 
-class CozeloopExporter(BaseModel, BaseExporter):
+class CozeloopExporter(BaseExporter):
     config: CozeloopExporterConfig = Field(default_factory=CozeloopExporterConfig)
 
-    @override
-    def get_processor(self):
+    def model_post_init(self, context: Any) -> None:
         headers = {
             "cozeloop-workspace-id": self.config.space_id,
             "authorization": f"Bearer {self.config.token}",
         }
-        exporter = OTLPSpanExporter(
+        self.headers |= headers
+
+        self._exporter = OTLPSpanExporter(
             endpoint=self.config.endpoint,
-            headers=headers,
+            headers=self.headers,
             timeout=10,
         )
-        self._real_exporter = exporter
-        processor = BatchSpanProcessor(exporter)
-        return processor, None
 
-    def export(self):
-        self._real_exporter.force_flush()
-        logger.info(
-            f"CozeloopExporter exports data to {self.config.endpoint}, space id: {self.config.space_id}"
-        )
+        self.processor = BatchSpanProcessor(self._exporter)
+
+    @override
+    def export(self) -> None:
+        """Force export of telemetry data."""
+        if self._exporter:
+            self._exporter.force_flush()
+            logger.info(
+                f"CozeloopExporter exports data to {self.config.endpoint}, space id: {self.config.space_id}"
+            )
+        else:
+            logger.warning("CozeloopExporter internal exporter is not initialized.")

veadk/tracing/telemetry/exporters/inmemory_exporter.py

@@ -12,11 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import typing
-from typing import Any
+from typing import Sequence
 
+from opentelemetry.context import (
+    _SUPPRESS_INSTRUMENTATION_KEY,
+    attach,
+    detach,
+    set_value,
+)
 from opentelemetry.sdk.trace import ReadableSpan, export
-from pydantic import BaseModel
 from typing_extensions import override
 
 from veadk.tracing.telemetry.exporters.base_exporter import BaseExporter
@@ -27,20 +31,17 @@ logger = get_logger(__name__)
 
 # ======== Adapted from Google ADK ========
 class _InMemoryExporter(export.SpanExporter):
-    def __init__(self, session_trace_dict):
+    def __init__(self) -> None:
         super().__init__()
         self._spans = []
-        self.session_trace_dict = session_trace_dict
         self.trace_id = ""
-        self.prompt_tokens = []
-        self.completion_tokens = []
+        self.session_trace_dict = {}
 
     @override
-    def export(self, spans: typing.Sequence[ReadableSpan]) -> export.SpanExportResult:
+    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
             else:
                 logger.warning(
                     f"Span context is missing, failed to get `trace_id`. span: {span}"
@@ -48,22 +49,12 @@ class _InMemoryExporter(export.SpanExporter):
 
             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)
+                session_id = attributes.get("gen_ai.session.id", None)
                 if session_id:
                     if session_id not in self.session_trace_dict:
-                        self.session_trace_dict[session_id] = [trace_id]
+                        self.session_trace_dict[session_id] = [self.trace_id]
                     else:
-                        self.session_trace_dict[session_id] += [trace_id]
+                        self.session_trace_dict[session_id] += [self.trace_id]
         self._spans.extend(spans)
         return export.SpanExportResult.SUCCESS
 
@@ -81,14 +72,37 @@ class _InMemoryExporter(export.SpanExporter):
         self._spans.clear()
 
 
-class InMemoryExporter(BaseModel, BaseExporter):
-    name: str = "inmemory_exporter"
+class _InMemorySpanProcessor(export.SimpleSpanProcessor):
+    def __init__(self, exporter: _InMemoryExporter) -> None:
+        super().__init__(exporter)
+        self.spans = []
 
-    def model_post_init(self, __context: Any) -> None:
-        self._exporter = _InMemoryExporter({})
+    def on_start(self, span, parent_context) -> None:
+        if span.context:
+            self.spans.append(span)
 
-    @override
-    def get_processor(self):
-        self._real_exporter = self._exporter
-        processor = export.SimpleSpanProcessor(self._exporter)
-        return processor, None
+    def on_end(self, span: ReadableSpan) -> None:
+        if span.context:
+            if not span.context.trace_flags.sampled:
+                return
+            token = attach(set_value(_SUPPRESS_INSTRUMENTATION_KEY, True))
+            try:
+                self.span_exporter.export((span,))
+            # pylint: disable=broad-exception-caught
+            except Exception:
+                logger.exception("Exception while exporting Span.")
+            detach(token)
+            if span in self.spans:
+                self.spans.remove(span)
+
+
+class InMemoryExporter(BaseExporter):
+    """InMemory Exporter mainly for store spans in memory for debugging / observability purposes."""
+
+    def __init__(self, name: str = "inmemory_exporter") -> None:
+        super().__init__()
+
+        self.name = name
+
+        self._exporter = _InMemoryExporter()
+        self.processor = _InMemorySpanProcessor(self._exporter)

veadk/tracing/telemetry/exporters/tls_exporter.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
@@ -44,28 +46,32 @@ class TLSExporterConfig(BaseModel):
     secret_key: str = Field(default_factory=lambda: getenv("VOLCENGINE_SECRET_KEY"))
 
 
-class TLSExporter(BaseModel, BaseExporter):
+class TLSExporter(BaseExporter):
     config: TLSExporterConfig = Field(default_factory=TLSExporterConfig)
 
-    @override
-    def get_processor(self):
+    def model_post_init(self, context: Any) -> None:
         headers = {
             "x-tls-otel-tracetopic": self.config.topic_id,
             "x-tls-otel-ak": self.config.access_key,
             "x-tls-otel-sk": self.config.secret_key,
             "x-tls-otel-region": self.config.region,
         }
-        exporter = OTLPSpanExporter(
+        self.headers |= headers
+
+        self._exporter = OTLPSpanExporter(
             endpoint=self.config.endpoint,
             headers=headers,
             timeout=10,
         )
-        self._real_exporter = exporter
-        processor = BatchSpanProcessor(exporter)
-        return processor, None
 
-    def export(self):
-        self._real_exporter.force_flush()
-        logger.info(
-            f"TLSExporter exports data to {self.config.endpoint}, topic id: {self.config.topic_id}"
-        )
+        self.processor = BatchSpanProcessor(self._exporter)
+
+    @override
+    def export(self) -> None:
+        if self._exporter:
+            self._exporter.force_flush()
+            logger.info(
+                f"TLSExporter exports data to {self.config.endpoint}, topic id: {self.config.topic_id}"
+            )
+        else:
+            logger.warning("TLSExporter internal exporter is not initialized.")