veadk-python 0.1.0__py3-none-any.whl

veadk/tracing/telemetry/exporters/base_exporter.py

@@ -0,0 +1,28 @@
+# 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 abc import ABC, abstractmethod
+from typing import Any
+
+
+class BaseExporter(ABC):
+    def __init__(self) -> None:
+        pass
+
+    @abstractmethod
+    def get_processor(self) -> Any:
+        pass
+
+    def get_meter_context(self) -> Any:
+        pass

veadk/tracing/telemetry/exporters/cozeloop_exporter.py

@@ -0,0 +1,69 @@
+# 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 opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.sdk.trace.export import BatchSpanProcessor
+from pydantic import BaseModel, Field
+from typing_extensions import override
+
+from veadk.config import getenv
+from veadk.tracing.telemetry.exporters.base_exporter import BaseExporter
+from veadk.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+class CozeloopExporterConfig(BaseModel):
+    endpoint: str = Field(
+        ...,
+        default_factory=lambda: getenv(
+            "OBSERVABILITY_OPENTELEMETRY_COZELOOP_ENDPOINT",
+            "https://api.coze.cn/v1/loop/opentelemetry/v1/traces",
+        ),
+    )
+    space_id: str = Field(
+        ...,
+        default_factory=lambda: getenv(
+            "OBSERVABILITY_OPENTELEMETRY_COZELOOP_SERVICE_NAME"
+        ),
+    )
+    token: str = Field(
+        ...,
+        default_factory=lambda: getenv("OBSERVABILITY_OPENTELEMETRY_COZELOOP_API_KEY"),
+    )
+
+
+class CozeloopExporter(BaseModel, BaseExporter):
+    config: CozeloopExporterConfig = Field(default_factory=CozeloopExporterConfig)
+
+    @override
+    def get_processor(self):
+        headers = {
+            "cozeloop-workspace-id": self.config.space_id,
+            "authorization": f"Bearer {self.config.token}",
+        }
+        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"CozeloopExporter exports data to {self.config.endpoint}, space id: {self.config.space_id}"
+        )

veadk/tracing/telemetry/exporters/inmemory_exporter.py

@@ -0,0 +1,88 @@
+# 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.
+
+import typing
+from typing import Any
+
+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
+
+
+# ======== Adapted from Google ADK ========
+class _InMemoryExporter(export.SpanExporter):
+    def __init__(self, session_trace_dict):
+        super().__init__()
+        self._spans = []
+        self.session_trace_dict = session_trace_dict
+        self.trace_id = ""
+        self.prompt_tokens = []
+        self.completion_tokens = []
+
+    @override
+    def export(self, spans: typing.Sequence[ReadableSpan]) -> export.SpanExportResult:
+        for span in spans:
+            trace_id = span.context.trace_id
+            self.trace_id = trace_id
+            if span.name == "call_llm":
+                attributes = dict(span.attributes)
+                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)
+                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
+
+    @override
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        return True
+
+    def get_finished_spans(self, session_id: str):
+        trace_ids = self.session_trace_dict.get(session_id, None)
+        if trace_ids is None or not trace_ids:
+            return []
+        return [x for x in self._spans if x.context.trace_id in trace_ids]
+
+    def clear(self):
+        self._spans.clear()
+
+
+class InMemoryExporter(BaseModel, BaseExporter):
+    name: str = "inmemory_exporter"
+
+    def model_post_init(self, __context: Any) -> None:
+        self._exporter = _InMemoryExporter({})
+
+    @override
+    def get_processor(self):
+        self._real_exporter = self._exporter
+        processor = export.SimpleSpanProcessor(self._exporter)
+        return processor, None
+
+    def get_meter_context(self) -> Any:
+        return None

veadk/tracing/telemetry/exporters/tls_exporter.py

@@ -0,0 +1,78 @@
+# 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 opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.sdk.trace.export import BatchSpanProcessor
+from pydantic import BaseModel, Field
+from typing_extensions import override
+
+from veadk.config import getenv
+from veadk.tracing.telemetry.exporters.base_exporter import BaseExporter
+from veadk.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+class TLSExporterConfig(BaseModel):
+    endpoint: str = Field(
+        ...,
+        default_factory=lambda: getenv(
+            "OBSERVABILITY_OPENTELEMETRY_TLS_ENDPOINT",
+            "https://tls-cn-beijing.volces.com:4318/v1/traces",
+        ),
+    )
+    region: str = Field(
+        ...,
+        default_factory=lambda: getenv(
+            "OBSERVABILITY_OPENTELEMETRY_TLS_REGION",
+            "cn-beijing",
+        ),
+    )
+    topic_id: str = Field(
+        ...,
+        default_factory=lambda: getenv("OBSERVABILITY_OPENTELEMETRY_TLS_SERVICE_NAME"),
+    )
+    access_key: str = Field(
+        ..., default_factory=lambda: getenv("VOLCENGINE_ACCESS_KEY")
+    )
+    secret_key: str = Field(
+        ..., default_factory=lambda: getenv("VOLCENGINE_SECRET_KEY")
+    )
+
+
+class TLSExporter(BaseModel, BaseExporter):
+    config: TLSExporterConfig = Field(default_factory=TLSExporterConfig)
+
+    @override
+    def get_processor(self):
+        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(
+            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}"
+        )

veadk/tracing/telemetry/metrics/__init__.py

@@ -0,0 +1,13 @@
+# 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.

veadk/tracing/telemetry/metrics/opentelemetry_metrics.py

@@ -0,0 +1,73 @@
+# 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.
+
+import time
+
+from opentelemetry.metrics._internal import Meter
+from opentelemetry.sdk.metrics import MeterProvider
+from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader
+
+from veadk.config import getenv
+
+
+class MeterContext:
+    def __init__(
+        self,
+        meter: Meter,
+        provider: MeterProvider,
+        reader: PeriodicExportingMetricReader,
+    ):
+        self.meter = meter
+        self.provider = provider
+        self.reader = reader
+
+
+class MeterUploader:
+    def __init__(self, meter_context: MeterContext):
+        self.meter = meter_context.meter
+        self.provider = meter_context.provider
+        self.reader = meter_context.reader
+
+        self.base_attributes = {
+            "gen_ai_system": "volcengine",
+            "server_address": "api.volcengine.com",
+            "gen_ai_response_model": getenv("MODEL_AGENT_NAME", "unknown"),
+            "stream": "false",
+            "gen_ai_operation_name": "chat_completions",
+        }
+        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, prompt_tokens: list[int], completion_tokens: list[int]):
+        self.llm_invoke_counter.add(len(completion_tokens), self.base_attributes)
+
+        for prompt_token in prompt_tokens:
+            token_attributes = {**self.base_attributes, "gen_ai_token_type": "input"}
+            self.token_usage.record(prompt_token, attributes=token_attributes)
+        for completion_token in completion_tokens:
+            token_attributes = {**self.base_attributes, "gen_ai_token_type": "output"}
+            self.token_usage.record(completion_token, attributes=token_attributes)
+
+    def close(self):
+        time.sleep(0.05)
+        self.reader.force_flush()
+        self.provider.shutdown()

veadk/tracing/telemetry/opentelemetry_tracer.py

@@ -0,0 +1,167 @@
+# 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 __future__ import annotations
+
+import json
+import time
+from typing import Any
+
+from openinference.instrumentation.google_adk import GoogleADKInstrumentor
+from opentelemetry.sdk import trace as trace_sdk
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace import TracerProvider
+from pydantic import BaseModel, ConfigDict, Field
+from typing_extensions import override
+
+from veadk.tracing.base_tracer import BaseTracer
+from veadk.tracing.telemetry.exporters.apiserver_exporter import ApiServerExporter
+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.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+DEFAULT_VEADK_TRACER_NAME = "veadk_global_tracer"
+
+
+def update_resource_attributions(provider: TracerProvider, resource_attributes: dict):
+    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."
+    )
+
+    def model_post_init(self, context: Any, /) -> None:
+        self._processors = []
+        self._inmemory_exporter: InMemoryExporter = None
+
+        # Inmemory & APIServer are the default exporters
+        have_inmemory_exporter = False
+        have_apiserver_exporter = False
+        for exporter in self.exporters:
+            if isinstance(exporter, InMemoryExporter):
+                have_inmemory_exporter = True
+            elif isinstance(exporter, ApiServerExporter):
+                have_apiserver_exporter = True
+
+        if not have_inmemory_exporter:
+            inmemory_exporter = InMemoryExporter()
+            self.exporters.append(inmemory_exporter)
+            self._inmemory_exporter = inmemory_exporter
+        if not have_apiserver_exporter:
+            apiserver_exporter = ApiServerExporter()
+            self.exporters.append(apiserver_exporter)
+            self._apiserver_exporter = apiserver_exporter
+
+        tracer_provider = trace_sdk.TracerProvider()
+        for exporter in self.exporters:
+            processor, resource_attributes = exporter.get_processor()
+            if resource_attributes is not None:
+                update_resource_attributions(tracer_provider, resource_attributes)
+            tracer_provider.add_span_processor(processor)
+
+            logger.debug(f"Add exporter `{exporter.__class__.__name__}` to tracing.")
+            self._processors.append(processor)
+        logger.debug(f"Init OpentelemetryTracer with {len(self.exporters)} exporters.")
+
+        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)
+
+        for meter_context in self._meter_contexts:
+            meter_uploader = MeterUploader(meter_context)
+            self._meter_uploaders.append(meter_uploader)
+
+        # just for debug
+        self._trace_file_path = ""
+
+        # patch this before starting instrumentation
+        # enable_veadk_tracing(self.dump)
+
+        GoogleADKInstrumentor().instrument(tracer_provider=tracer_provider)
+
+    @override
+    def dump(
+        self,
+        user_id: str,
+        session_id: str,
+        path: str = "/tmp",
+    ) -> str:
+        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(
+            session_id=session_id
+        )
+        if not spans:
+            data = []
+        else:
+            data = [
+                {
+                    "name": s.name,
+                    "span_id": s.context.span_id,
+                    "trace_id": s.context.trace_id,
+                    "start_time": s.start_time,
+                    "end_time": s.end_time,
+                    "attributes": dict(s.attributes),
+                    "parent_span_id": s.parent.span_id if s.parent else None,
+                }
+                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"
+        with open(file_path, "w") as f:
+            json.dump(data, f, indent=4)
+
+        self._trace_file_path = file_path
+
+        for exporter in self.exporters:
+            if not isinstance(exporter, InMemoryExporter) and not isinstance(
+                exporter, ApiServerExporter
+            ):
+                exporter.export()
+        logger.info(
+            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

veadk/types.py

@@ -0,0 +1,23 @@
+# 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 pydantic import BaseModel
+
+
+class MediaMessage(BaseModel):
+    text: str
+    """Text-based prompt"""
+
+    media: str
+    """Media file (e.g., `.pdf`, `.docx`, `.png`, `.jpg`, `.jpeg`, `.mp4`, `.mp3`, `.wav`, `.txt`) path"""

veadk/utils/__init__.py

@@ -0,0 +1,13 @@
+# 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.

veadk/utils/logger.py

@@ -0,0 +1,59 @@
+# 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.
+
+import sys
+
+from loguru import logger
+
+from veadk.config import getenv
+
+
+def filter_log():
+    import logging
+    import warnings
+
+    from urllib3.exceptions import InsecureRequestWarning
+
+    # ignore all warnings
+    warnings.filterwarnings("ignore")
+
+    # ignore UserWarning
+    warnings.filterwarnings(
+        "ignore", category=UserWarning, module="opensearchpy.connection.http_urllib3"
+    )
+
+    # ignore InsecureRequestWarning
+    warnings.filterwarnings("ignore", category=InsecureRequestWarning)
+
+    # disable logs
+    logging.basicConfig(level=logging.ERROR)
+
+
+def setup_logger():
+    logger.remove()
+    logger.add(
+        sys.stdout,
+        format="<green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{level}</level> | <cyan>{file}:{line}</cyan> - {message}",
+        colorize=True,
+        level=getenv("LOGGING_LEVEL", "DEBUG"),
+    )
+    return logger
+
+
+filter_log()
+setup_logger()
+
+
+def get_logger(name: str):
+    return logger.bind(name=name)

veadk/utils/misc.py

@@ -0,0 +1,33 @@
+# 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.
+
+import time
+
+
+def read_file(file_path):
+    with open(file_path, "r", encoding="utf-8") as f:
+        data = f.readlines()
+    data = [x.strip() for x in data]
+    return data
+
+
+def formatted_timestamp():
+    # YYYYMMDDHHMMSS
+    return time.strftime("%Y%m%d%H%M%S", time.localtime())
+
+
+def read_png_to_bytes(png_path: str) -> bytes:
+    with open(png_path, "rb") as f:
+        data = f.read()
+    return data