From 24dee361e1de3746958d7d318fb55209045afbf0 Mon Sep 17 00:00:00 2001 From: Vyacheslav Morov Date: Sun, 31 Aug 2025 16:32:47 +0200 Subject: [PATCH 1/7] Initial version of interceptors. --- examples/interceptors_examples.py | 22 +++++++ examples/openai_manual_tracing.py | 10 +-- tracely/src/tracely/__init__.py | 8 ++- tracely/src/tracely/_context.py | 83 +++++++++++++++++++++++++ tracely/src/tracely/_tracer_provider.py | 67 ++++---------------- tracely/src/tracely/context.py | 13 ++-- tracely/src/tracely/decorators.py | 40 +++++++++--- tracely/src/tracely/interceptors.py | 21 +++++++ tracely/src/tracely/proxy.py | 4 +- tracely/tests/test_interceptors.py | 52 ++++++++++++++++ 10 files changed, 244 insertions(+), 76 deletions(-) create mode 100644 examples/interceptors_examples.py create mode 100644 tracely/src/tracely/_context.py create mode 100644 tracely/src/tracely/interceptors.py create mode 100644 tracely/tests/test_interceptors.py diff --git a/examples/interceptors_examples.py b/examples/interceptors_examples.py new file mode 100644 index 0000000..cee1102 --- /dev/null +++ b/examples/interceptors_examples.py @@ -0,0 +1,22 @@ +from tracely import Interceptor +from tracely import init_tracing +from tracely import trace_event + + +class ExampleInterceptor(Interceptor): + def before_call(self, span, context, *args, **kwargs): + span.set_attribute("custom_attribute", 1) + + def after_call(self, span, context, *args, **kwargs): + pass + + def on_exception(self, span, context, exception) -> bool: + pass + + +init_tracing(exporter_type="console", interceptors=[ExampleInterceptor()]) + + +@trace_event() +def func(data: str) -> str: + return data diff --git a/examples/openai_manual_tracing.py b/examples/openai_manual_tracing.py index 597a7b6..341faaf 100644 --- a/examples/openai_manual_tracing.py +++ b/examples/openai_manual_tracing.py @@ -77,9 +77,9 @@ def call_with_user_id_explicit(input: str): ) print(call_openai("What is LLM?")) - print(call_openai_with_helper("What is LLM?")) - print(call_openai_with_context("What is LLM?")) - print(multiple_calls_openai("What is LLM?")) - print(call_with_user_id_param("What is LLM?", "user_id")) - print(call_with_user_id_explicit("What is LLM?")) + # print(call_openai_with_helper("What is LLM?")) + # print(call_openai_with_context("What is LLM?")) + # print(multiple_calls_openai("What is LLM?")) + # print(call_with_user_id_param("What is LLM?", "user_id")) + # print(call_with_user_id_explicit("What is LLM?")) sleep(1) diff --git a/tracely/src/tracely/__init__.py b/tracely/src/tracely/__init__.py index f34b4e3..e2d45a4 100644 --- a/tracely/src/tracely/__init__.py +++ b/tracely/src/tracely/__init__.py @@ -1,10 +1,13 @@ from ._tracer_provider import UsageDetails from ._tracer_provider import init_tracing -from ._tracer_provider import get_info +from ._context import get_info +from ._context import get_interceptors +from ._context import get_tracer from .decorators import trace_event from .context import create_trace_event from .context import bind_to_trace from .proxy import get_current_span +from .proxy import SpanObject from ._version import __version__ @@ -13,8 +16,11 @@ "create_trace_event", "get_current_span", "get_info", + "get_tracer", + "get_interceptors", "init_tracing", "bind_to_trace", "trace_event", + "SpanObject", "__version__", ] diff --git a/tracely/src/tracely/_context.py b/tracely/src/tracely/_context.py new file mode 100644 index 0000000..f316078 --- /dev/null +++ b/tracely/src/tracely/_context.py @@ -0,0 +1,83 @@ +import dataclasses +import typing +import uuid +from typing import Dict +from typing import List +from typing import Optional +from typing import Union + +import opentelemetry +from opentelemetry import trace +from opentelemetry.context import Context +from opentelemetry.trace import NonRecordingSpan +from opentelemetry.trace import SpanContext +from opentelemetry.trace import TraceFlags + +if typing.TYPE_CHECKING: + from .interceptors import Interceptor + + +@dataclasses.dataclass +class UsageDetails: + cost_per_token: Dict[str, float] + + +class DataContext: + export_id: Union[str, uuid.UUID] + project_id: Union[str, uuid.UUID] + default_usage_details: Optional[UsageDetails] + usage_details_by_model_id: Optional[Dict[str, UsageDetails]] + interceptors: List["Interceptor"] + + def __init__( + self, + export_id: str, + project_id: str, + default_usage_details: Optional[UsageDetails] = None, + usage_details_by_model_id: Optional[Dict[str, UsageDetails]] = None, + interceptors: Optional[List["Interceptor"]] = None, + ): + self.export_id = export_id + self.project_id = project_id + self.default_usage_details = default_usage_details + self.usage_details_by_model_id = usage_details_by_model_id + self.interceptors = interceptors or [] + + def get_model_usage_details(self, model_id: str) -> Optional[UsageDetails]: + if self.usage_details_by_model_id is None: + return self.default_usage_details + return self.usage_details_by_model_id.get(model_id, self.default_usage_details) + + +_tracer: Optional[trace.Tracer] = None +_context: Optional[Context] = None +_data_context: DataContext = DataContext("", "") + + +def set_tracer(new_tracer: trace.Tracer) -> None: + global _tracer + _tracer = new_tracer + + +def get_tracer() -> Optional[trace.Tracer]: + return _tracer + + +def get_info(): + return { + "export_id": _data_context.export_id, + "project_id": _data_context.project_id, + } + + +def get_interceptors() -> List["Interceptor"]: + return _data_context.interceptors + + +def create_context(trace_id: int, parent_span_id: Optional[int]): + if parent_span_id is None: + generator = opentelemetry.sdk.trace.RandomIdGenerator() + parent_span_id = generator.generate_span_id() + span_context = SpanContext(trace_id=trace_id, span_id=parent_span_id, is_remote=True, trace_flags=TraceFlags(0x01)) + context = opentelemetry.trace.set_span_in_context(NonRecordingSpan(span_context)) + return context diff --git a/tracely/src/tracely/_tracer_provider.py b/tracely/src/tracely/_tracer_provider.py index d5fd180..90b1bfd 100644 --- a/tracely/src/tracely/_tracer_provider.py +++ b/tracely/src/tracely/_tracer_provider.py @@ -1,14 +1,13 @@ -import dataclasses import urllib.parse import uuid from typing import Dict +from typing import List from typing import Optional from typing import Union import opentelemetry.trace import requests from opentelemetry import trace -from opentelemetry.context import Context from opentelemetry.sdk.resources import Resource from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import BatchSpanProcessor, SpanExporter, SimpleSpanProcessor @@ -16,6 +15,9 @@ from opentelemetry.trace import SpanContext from opentelemetry.trace import TraceFlags +from ._context import UsageDetails +from ._context import _data_context +from ._context import set_tracer from ._env import ( _EVIDENTLY_API_KEY, _TRACE_COLLECTOR_ADDRESS, @@ -25,40 +27,7 @@ _TRACE_COLLECTOR_PROJECT_ID, ) from .evidently_cloud_client import EvidentlyCloudClient - - -@dataclasses.dataclass -class UsageDetails: - cost_per_token: Dict[str, float] - - -class DataContext: - export_id: Union[str, uuid.UUID] - project_id: Union[str, uuid.UUID] - default_usage_details: Optional[UsageDetails] - usage_details_by_model_id: Optional[Dict[str, UsageDetails]] - - def __init__( - self, - export_id: str, - project_id: str, - default_usage_details: Optional[UsageDetails] = None, - usage_details_by_model_id: Optional[Dict[str, UsageDetails]] = None, - ): - self.export_id = export_id - self.project_id = project_id - self.default_usage_details = default_usage_details - self.usage_details_by_model_id = usage_details_by_model_id - - def get_model_usage_details(self, model_id: str) -> Optional[UsageDetails]: - if self.usage_details_by_model_id is None: - return self.default_usage_details - return self.usage_details_by_model_id.get(model_id, self.default_usage_details) - - -_tracer: Optional[trace.Tracer] = None -_context: Optional[Context] = None -_data_context: DataContext = DataContext("", "") +from .interceptors import Interceptor def _create_tracer_provider( @@ -70,6 +39,7 @@ def _create_tracer_provider( export_name: Optional[str] = None, default_usage_details: Optional[UsageDetails] = None, usage_details_by_model_id: Optional[Dict[str, UsageDetails]] = None, + interceptors: Optional[List[Interceptor]] = None, ) -> trace.TracerProvider: """ Creates Evidently telemetry tracer provider which would be used for sending traces. @@ -80,7 +50,6 @@ def _create_tracer_provider( project_id: id of project in Evidently Cloud export_name: string name of exported data, all data with same id would be grouped into single dataset """ - global _tracer # noqa: PLW0603 _address = address or _TRACE_COLLECTOR_ADDRESS if len(_address) == 0: @@ -140,6 +109,7 @@ def _create_tracer_provider( _data_context.project_id = uuid.UUID(_project_id) _data_context.default_usage_details = default_usage_details _data_context.usage_details_by_model_id = usage_details_by_model_id + _data_context.interceptors = interceptors or [] tracer_provider = TracerProvider( resource=Resource.create( @@ -181,7 +151,7 @@ def _create_tracer_provider( tracer_provider.add_span_processor(SimpleSpanProcessor(exporter)) else: raise ValueError(f"Unexpected processor type: {processor_type}. Expected values: batch or simple") - _tracer = tracer_provider.get_tracer("evidently") + set_tracer(tracer_provider.get_tracer("evidently")) return tracer_provider @@ -196,6 +166,7 @@ def init_tracing( processor_type: str = "batch", default_usage_details: Optional[UsageDetails] = None, usage_details_by_model_id: Optional[Dict[str, UsageDetails]] = None, + interceptors: Optional[List[Interceptor]] = None, ) -> trace.TracerProvider: """ Initialize Evidently tracing @@ -213,9 +184,9 @@ def init_tracing( 'simple' - upload traces synchronously as it is reported, can cause performance issues. default_usage_details: usage data for tokens usage_details_by_model_id: usage data for tokens by model id (if provided) + interceptors: list of interceptors to use """ - global _tracer # noqa: PLW0603 provider = _create_tracer_provider( address, exporter_type, @@ -225,22 +196,17 @@ def init_tracing( export_name, default_usage_details, usage_details_by_model_id, + interceptors, ) if as_global: trace.set_tracer_provider(provider) - _tracer = trace.get_tracer("evidently") + set_tracer(trace.get_tracer("evidently")) else: - _tracer = provider.get_tracer("evidently") + set_tracer(provider.get_tracer("evidently")) return provider -def get_tracer() -> trace.Tracer: - if _tracer is None: - raise ValueError("TracerProvider not initialized, use init_tracer()") - return _tracer - - def create_context(trace_id: int, parent_span_id: Optional[int]): if parent_span_id is None: generator = opentelemetry.sdk.trace.RandomIdGenerator() @@ -248,10 +214,3 @@ def create_context(trace_id: int, parent_span_id: Optional[int]): span_context = SpanContext(trace_id=trace_id, span_id=parent_span_id, is_remote=True, trace_flags=TraceFlags(0x01)) context = opentelemetry.trace.set_span_in_context(NonRecordingSpan(span_context)) return context - - -def get_info(): - return { - "export_id": _data_context.export_id, - "project_id": _data_context.project_id, - } diff --git a/tracely/src/tracely/context.py b/tracely/src/tracely/context.py index e4dd6e8..46b885e 100644 --- a/tracely/src/tracely/context.py +++ b/tracely/src/tracely/context.py @@ -4,12 +4,13 @@ import opentelemetry.sdk.trace -from . import _tracer_provider -from .proxy import _ProxySpanObject +from ._context import get_tracer +from ._context import create_context +from .proxy import SpanObject @contextmanager -def create_trace_event(name: str, **params) -> Generator[_ProxySpanObject, None, None]: +def create_trace_event(name: str, **params) -> Generator[SpanObject, None, None]: """ Create a span with given name. @@ -20,9 +21,9 @@ def create_trace_event(name: str, **params) -> Generator[_ProxySpanObject, None, Returns: span object to work with """ - _tracer = _tracer_provider.get_tracer() + _tracer = get_tracer() with _tracer.start_as_current_span(f"{name}") as span: - obj = _ProxySpanObject(span) + obj = SpanObject(span) try: yield obj finally: @@ -32,7 +33,7 @@ def create_trace_event(name: str, **params) -> Generator[_ProxySpanObject, None, @contextmanager def bind_to_trace(trace_id: int, parent_span_id: Optional[int] = None): - context = _tracer_provider.create_context(trace_id, parent_span_id) + context = create_context(trace_id, parent_span_id) token = opentelemetry.sdk.trace.context_api.attach(context) try: yield diff --git a/tracely/src/tracely/decorators.py b/tracely/src/tracely/decorators.py index 110da41..3066842 100644 --- a/tracely/src/tracely/decorators.py +++ b/tracely/src/tracely/decorators.py @@ -5,9 +5,11 @@ from opentelemetry.trace import StatusCode import tracely -from . import _tracer_provider +from ._context import get_interceptors +from ._context import get_tracer +from .interceptors import InterceptorContext from .proxy import set_result -from .proxy import _ProxySpanObject +from .proxy import SpanObject def _fill_span_from_signature( @@ -15,7 +17,7 @@ def _fill_span_from_signature( ignore_args: Optional[List[str]], sign: Signature, bind: BoundArguments, - span: _ProxySpanObject, + span: SpanObject, ): final_args = track_args if final_args is None: @@ -60,16 +62,27 @@ async def func(*args, **kwargs): sign = inspect.signature(f) bind = sign.bind(*args, **kwargs) + interceptor_context = InterceptorContext() with tracely.create_trace_event(f"{span_name or f.__name__}", parse_output) as span: _fill_span_from_signature(track_args, ignore_args, bind.signature, bind, span) + for interceptor in get_interceptors(): + interceptor.before_call(span, interceptor_context, *args, **kwargs) try: result = await f(*args, **kwargs) if result is not None and track_output: span.set_result(result) + for interceptor in get_interceptors(): + interceptor.after_call(span, interceptor_context, result) span.set_status(StatusCode.OK) except Exception as e: - span.set_attribute("exception", str(e)) - span.set_status(StatusCode.ERROR) + processed = False + for interceptor in get_interceptors(): + processed = processed or interceptor.on_exception( + span, interceptor_context, *args, **kwargs + ) + if not processed: + span.set_attribute("exception", str(e)) + span.set_status(StatusCode.ERROR) raise return result @@ -80,19 +93,30 @@ async def func(*args, **kwargs): def func(*args, **kwargs): import inspect - _tracer = _tracer_provider.get_tracer() + _tracer = get_tracer() sign = inspect.signature(f) bind = sign.bind(*args, **kwargs) + interceptor_context = InterceptorContext() with _tracer.start_as_current_span(f"{span_name or f.__name__}") as span: _fill_span_from_signature(track_args, ignore_args, bind.signature, bind, span) + for interceptor in get_interceptors(): + interceptor.before_call(span, interceptor_context, *args, **kwargs) try: result = f(*args, **kwargs) if result is not None and track_output: set_result(span, result, parse_output) + for interceptor in get_interceptors(): + interceptor.after_call(span, interceptor_context, result) span.set_status(StatusCode.OK) except Exception as e: - span.set_attribute("exception", str(e)) - span.set_status(StatusCode.ERROR) + processed = False + for interceptor in get_interceptors(): + processed = processed or interceptor.on_exception( + span, interceptor_context, *args, **kwargs + ) + if not processed: + span.set_attribute("exception", str(e)) + span.set_status(StatusCode.ERROR) raise return result diff --git a/tracely/src/tracely/interceptors.py b/tracely/src/tracely/interceptors.py new file mode 100644 index 0000000..d294e93 --- /dev/null +++ b/tracely/src/tracely/interceptors.py @@ -0,0 +1,21 @@ +import abc + +from tracely.proxy import SpanObject + + +class InterceptorContext: + pass + + +class Interceptor: + @abc.abstractmethod + def before_call(self, span: SpanObject, context: InterceptorContext, *args, **kwargs): + pass + + @abc.abstractmethod + def after_call(self, span: SpanObject, context: InterceptorContext, return_value): + pass + + @abc.abstractmethod + def on_exception(self, span: SpanObject, context: InterceptorContext, ex: Exception) -> bool: + pass diff --git a/tracely/src/tracely/proxy.py b/tracely/src/tracely/proxy.py index 4660731..41ec616 100644 --- a/tracely/src/tracely/proxy.py +++ b/tracely/src/tracely/proxy.py @@ -10,7 +10,7 @@ from openai.types.responses import ResponseUsage -class _ProxySpanObject: +class SpanObject: def __init__(self, span: Optional[opentelemetry.trace.Span] = None): if span is None: self.span = opentelemetry.trace.get_current_span() @@ -75,7 +75,7 @@ def _update_usage_openai(self, usage: "ResponseUsage"): def get_current_span(): - return _ProxySpanObject() + return SpanObject() def set_result(span, result, parse_output: bool): diff --git a/tracely/tests/test_interceptors.py b/tracely/tests/test_interceptors.py new file mode 100644 index 0000000..21cebc2 --- /dev/null +++ b/tracely/tests/test_interceptors.py @@ -0,0 +1,52 @@ +from uuid import UUID + +import pytest +import opentelemetry +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from tracely import init_tracing +from tracely import trace_event +from tracely.interceptors import Interceptor + + +class AddAttributeBeforeCallInterceptor(Interceptor): + def before_call(self, span, context, *args, **kwargs): + span.set_attribute("custom_attribute", 1) + + def after_call(self, span, context, *args, **kwargs): + pass + + def on_exception(self, span, context, exception) -> bool: + pass + + +@pytest.fixture +def exporter(): + provider = init_tracing( + exporter_type="console", + project_id=UUID(int=0), + export_name="test", + as_global=False, + interceptors=[AddAttributeBeforeCallInterceptor()], + ) + exporter = InMemorySpanExporter() + if isinstance(provider, opentelemetry.sdk.trace.TracerProvider): + provider.add_span_processor(SimpleSpanProcessor(exporter)) + return exporter + + +@trace_event(track_output=True) +def trace_func_with_output(): + return 100 + + +def test_trace_func_with_output(exporter): + result = trace_func_with_output() + + spans = exporter.get_finished_spans() + assert len(spans) == 1 + span = spans[0] + assert span.name == "trace_func_with_output" + assert span.attributes["custom_attribute"] == 1 + assert span.attributes["result"] == result From 86cc658311d89a996ddeae2bca573029ca496baa Mon Sep 17 00:00:00 2001 From: Vyacheslav Morov Date: Mon, 1 Sep 2025 10:32:39 +0100 Subject: [PATCH 2/7] Improve interceptors. --- tracely/src/tracely/__init__.py | 5 +- tracely/src/tracely/_runtime_context.py | 39 +++++++++++++++ tracely/src/tracely/decorators.py | 21 +++++---- tracely/src/tracely/interceptors.py | 13 ++++- tracely/src/tracely/proxy.py | 15 ++++-- tracely/tests/test_interceptors.py | 63 +++++++++++++++++++++++-- 6 files changed, 138 insertions(+), 18 deletions(-) create mode 100644 tracely/src/tracely/_runtime_context.py diff --git a/tracely/src/tracely/__init__.py b/tracely/src/tracely/__init__.py index e2d45a4..503614d 100644 --- a/tracely/src/tracely/__init__.py +++ b/tracely/src/tracely/__init__.py @@ -6,8 +6,10 @@ from .decorators import trace_event from .context import create_trace_event from .context import bind_to_trace -from .proxy import get_current_span +from .interceptors import Interceptor from .proxy import SpanObject +from ._runtime_context import get_current_span +from ._runtime_context import RuntimeContext from ._version import __version__ @@ -20,6 +22,7 @@ "get_interceptors", "init_tracing", "bind_to_trace", + "Interceptor", "trace_event", "SpanObject", "__version__", diff --git a/tracely/src/tracely/_runtime_context.py b/tracely/src/tracely/_runtime_context.py new file mode 100644 index 0000000..c774766 --- /dev/null +++ b/tracely/src/tracely/_runtime_context.py @@ -0,0 +1,39 @@ +from typing import Optional + +from .proxy import SpanObject + + +class RuntimeContext: + def __init__(self): + self.span = None + + def set_current_span(self, span: SpanObject): + self.span = span + + def get_current_span(self) -> Optional[SpanObject]: + return self.span + + def reset_span(self): + self.span = None + + +_DEFAULT_CONTEXT = RuntimeContext() + +def get_current_span(context: Optional[RuntimeContext] = None) -> Optional[SpanObject]: + if context is None: + return _DEFAULT_CONTEXT.get_current_span() + return context.get_current_span() + + +def set_current_span(span: SpanObject, context: Optional[RuntimeContext] = None): + if context is None: + _DEFAULT_CONTEXT.set_current_span(span) + else: + context.set_current_span(span) + + +def reset_span(context: Optional[RuntimeContext] = None): + if context is None: + _DEFAULT_CONTEXT.reset_span() + else: + context.reset_span() diff --git a/tracely/src/tracely/decorators.py b/tracely/src/tracely/decorators.py index 3066842..d0fa42a 100644 --- a/tracely/src/tracely/decorators.py +++ b/tracely/src/tracely/decorators.py @@ -7,9 +7,11 @@ import tracely from ._context import get_interceptors from ._context import get_tracer -from .interceptors import InterceptorContext -from .proxy import set_result from .proxy import SpanObject +from .proxy import set_result +from ._runtime_context import get_current_span +from ._runtime_context import set_current_span +from .interceptors import InterceptorContext def _fill_span_from_signature( @@ -77,9 +79,7 @@ async def func(*args, **kwargs): except Exception as e: processed = False for interceptor in get_interceptors(): - processed = processed or interceptor.on_exception( - span, interceptor_context, *args, **kwargs - ) + processed = processed or interceptor.on_exception(span, interceptor_context, e) if not processed: span.set_attribute("exception", str(e)) span.set_status(StatusCode.ERROR) @@ -97,7 +97,10 @@ def func(*args, **kwargs): sign = inspect.signature(f) bind = sign.bind(*args, **kwargs) interceptor_context = InterceptorContext() - with _tracer.start_as_current_span(f"{span_name or f.__name__}") as span: + with _tracer.start_as_current_span(f"{span_name or f.__name__}") as otel_span: + prev_span = get_current_span() + span = SpanObject(otel_span) + set_current_span(span) _fill_span_from_signature(track_args, ignore_args, bind.signature, bind, span) for interceptor in get_interceptors(): interceptor.before_call(span, interceptor_context, *args, **kwargs) @@ -111,13 +114,13 @@ def func(*args, **kwargs): except Exception as e: processed = False for interceptor in get_interceptors(): - processed = processed or interceptor.on_exception( - span, interceptor_context, *args, **kwargs - ) + processed = processed or interceptor.on_exception(span, interceptor_context, e) if not processed: span.set_attribute("exception", str(e)) span.set_status(StatusCode.ERROR) raise + finally: + set_current_span(prev_span) return result return func diff --git a/tracely/src/tracely/interceptors.py b/tracely/src/tracely/interceptors.py index d294e93..8f44079 100644 --- a/tracely/src/tracely/interceptors.py +++ b/tracely/src/tracely/interceptors.py @@ -1,10 +1,21 @@ import abc +from typing import Any +from typing import Dict from tracely.proxy import SpanObject class InterceptorContext: - pass + data: Dict[str, Any] + + def __init__(self): + self.data = {} + + def set(self, key: str, value): + self.data[key] = value + + def get(self, key: str): + return self.data.get(key) class Interceptor: diff --git a/tracely/src/tracely/proxy.py b/tracely/src/tracely/proxy.py index 41ec616..261119e 100644 --- a/tracely/src/tracely/proxy.py +++ b/tracely/src/tracely/proxy.py @@ -12,6 +12,7 @@ class SpanObject: def __init__(self, span: Optional[opentelemetry.trace.Span] = None): + self.context = {} if span is None: self.span = opentelemetry.trace.get_current_span() else: @@ -45,6 +46,15 @@ def update_usage( raise ValueError("Must specify either tokens or usage") self._update_usage(tokens=tokens, costs=costs) + def set_context_value(self, key, value): + self.context[key] = value + + def get_context_value(self, key): + return self.context.get(key) + + def get_context(self): + return self.context + def _update_usage( self, *, @@ -73,9 +83,8 @@ def _update_usage_openai(self, usage: "ResponseUsage"): }, ) - -def get_current_span(): - return SpanObject() + def set_status(self, status): + self.span.set_status(status) def set_result(span, result, parse_output: bool): diff --git a/tracely/tests/test_interceptors.py b/tracely/tests/test_interceptors.py index 21cebc2..c15cdde 100644 --- a/tracely/tests/test_interceptors.py +++ b/tracely/tests/test_interceptors.py @@ -1,3 +1,4 @@ +from functools import wraps from uuid import UUID import pytest @@ -7,24 +8,43 @@ from tracely import init_tracing from tracely import trace_event +from tracely import get_current_span from tracely.interceptors import Interceptor +class MyException(Exception): + message = "Exception message" + + def __init__(self, value): + self.message = value + + class AddAttributeBeforeCallInterceptor(Interceptor): def before_call(self, span, context, *args, **kwargs): - span.set_attribute("custom_attribute", 1) + pass def after_call(self, span, context, *args, **kwargs): - pass + value = span.get_context_value("deco") + if value is not None: + span.set_attribute("deco", value) + span.set_attribute("status", "passed") def on_exception(self, span, context, exception) -> bool: - pass + value = span.get_context_value("deco") + if value is not None: + if isinstance(exception, MyException): + span.set_attribute("deco", value) + span.set_attribute("status", "failed") + span.set_attribute("error", exception.message) + return True + return False @pytest.fixture def exporter(): provider = init_tracing( exporter_type="console", + processor_type="simple", project_id=UUID(int=0), export_name="test", as_global=False, @@ -36,11 +56,32 @@ def exporter(): return exporter +def stub_deco(fail, message): + def decorator(func): + @wraps(func) + def wrapper(*args, **kwargs): + span = get_current_span() + if span: + span.set_context_value("deco", "test_deco") + if fail: + raise MyException(message) + return func(*args, **kwargs) + return wrapper + return decorator + + @trace_event(track_output=True) +@stub_deco(False, "message") def trace_func_with_output(): return 100 +@trace_event(track_output=True) +@stub_deco(True, "message") +def trace_func_with_output_failed_deco(): + return 100 + + def test_trace_func_with_output(exporter): result = trace_func_with_output() @@ -48,5 +89,19 @@ def test_trace_func_with_output(exporter): assert len(spans) == 1 span = spans[0] assert span.name == "trace_func_with_output" - assert span.attributes["custom_attribute"] == 1 + assert span.attributes["deco"] == "test_deco" + assert span.attributes["status"] == "passed" assert span.attributes["result"] == result + + +def test_trace_func_with_output_failed_deco(exporter): + with pytest.raises(MyException): + result = trace_func_with_output_failed_deco() + + spans = exporter.get_finished_spans() + assert len(spans) == 1 + span = spans[0] + assert span.name == "trace_func_with_output_failed_deco" + assert span.attributes["deco"] == "test_deco" + assert span.attributes["status"] == "failed" + assert span.attributes["error"] == "message" From e23a35245a2cb432c751c321bd68af30e0f05589 Mon Sep 17 00:00:00 2001 From: Vyacheslav Morov Date: Mon, 1 Sep 2025 12:02:41 +0100 Subject: [PATCH 3/7] Fix linter. --- tracely/src/tracely/_runtime_context.py | 1 + tracely/tests/test_interceptors.py | 2 ++ 2 files changed, 3 insertions(+) diff --git a/tracely/src/tracely/_runtime_context.py b/tracely/src/tracely/_runtime_context.py index c774766..585d49a 100644 --- a/tracely/src/tracely/_runtime_context.py +++ b/tracely/src/tracely/_runtime_context.py @@ -19,6 +19,7 @@ def reset_span(self): _DEFAULT_CONTEXT = RuntimeContext() + def get_current_span(context: Optional[RuntimeContext] = None) -> Optional[SpanObject]: if context is None: return _DEFAULT_CONTEXT.get_current_span() diff --git a/tracely/tests/test_interceptors.py b/tracely/tests/test_interceptors.py index c15cdde..edf8fcb 100644 --- a/tracely/tests/test_interceptors.py +++ b/tracely/tests/test_interceptors.py @@ -66,7 +66,9 @@ def wrapper(*args, **kwargs): if fail: raise MyException(message) return func(*args, **kwargs) + return wrapper + return decorator From 2de1e9f964246285869cbeed4cb1d94692b53fee Mon Sep 17 00:00:00 2001 From: Vyacheslav Morov Date: Mon, 1 Sep 2025 12:05:04 +0100 Subject: [PATCH 4/7] Fix linter. --- tracely/src/tracely/__init__.py | 1 + tracely/tests/test_interceptors.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/tracely/src/tracely/__init__.py b/tracely/src/tracely/__init__.py index 503614d..e663b02 100644 --- a/tracely/src/tracely/__init__.py +++ b/tracely/src/tracely/__init__.py @@ -25,5 +25,6 @@ "Interceptor", "trace_event", "SpanObject", + "RuntimeContext", "__version__", ] diff --git a/tracely/tests/test_interceptors.py b/tracely/tests/test_interceptors.py index edf8fcb..e0d9e15 100644 --- a/tracely/tests/test_interceptors.py +++ b/tracely/tests/test_interceptors.py @@ -98,7 +98,7 @@ def test_trace_func_with_output(exporter): def test_trace_func_with_output_failed_deco(exporter): with pytest.raises(MyException): - result = trace_func_with_output_failed_deco() + trace_func_with_output_failed_deco() spans = exporter.get_finished_spans() assert len(spans) == 1 From 46a90ac677ef66b59162a714f81c49fc4b2b60cb Mon Sep 17 00:00:00 2001 From: Vyacheslav Morov Date: Tue, 2 Sep 2025 11:33:35 +0100 Subject: [PATCH 5/7] Add tests. --- tracely/tests/test_spans.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/tracely/tests/test_spans.py b/tracely/tests/test_spans.py index d21e294..1fc7e34 100644 --- a/tracely/tests/test_spans.py +++ b/tracely/tests/test_spans.py @@ -23,6 +23,10 @@ def trace_func_with_output_struct(): "f2": 100, } +@trace_event() +def trace_func_with_inner_trace(): + return trace_func_with_output() + @trace_event(track_output=True, parse_output=True) def trace_func_with_tokens(): @@ -155,3 +159,12 @@ def test_trace_func_with_tokens(exporter): assert span.attributes["tokens.output"] == 200 assert span.attributes["cost.input"] == 0.1 assert span.attributes["cost.output"] == 1.0 + + +def test_trace_func_with_inner_trace(exporter): + trace_func_with_inner_trace() + + spans = exporter.get_finished_spans() + assert len(spans) == 2 + assert spans[0].name == "trace_func_with_output" + assert spans[1].name == "trace_func_with_inner_trace" From c6122ed56010bf4b5b114c3ef4323d39dadb1163 Mon Sep 17 00:00:00 2001 From: Vyacheslav Morov Date: Mon, 8 Sep 2025 11:47:00 +0100 Subject: [PATCH 6/7] Add new spans into create_trace_event. --- tracely/src/tracely/context.py | 6 ++++++ tracely/tests/test_spans.py | 1 + 2 files changed, 7 insertions(+) diff --git a/tracely/src/tracely/context.py b/tracely/src/tracely/context.py index 46b885e..7c16fd2 100644 --- a/tracely/src/tracely/context.py +++ b/tracely/src/tracely/context.py @@ -6,6 +6,8 @@ from ._context import get_tracer from ._context import create_context +from ._runtime_context import get_current_span +from ._runtime_context import set_current_span from .proxy import SpanObject @@ -24,11 +26,15 @@ def create_trace_event(name: str, **params) -> Generator[SpanObject, None, None] _tracer = get_tracer() with _tracer.start_as_current_span(f"{name}") as span: obj = SpanObject(span) + prev_span = get_current_span() + set_current_span(obj) + try: yield obj finally: for attr, value in params.items(): span.set_attribute(attr, value) + set_current_span(prev_span) @contextmanager diff --git a/tracely/tests/test_spans.py b/tracely/tests/test_spans.py index 1fc7e34..30b8f56 100644 --- a/tracely/tests/test_spans.py +++ b/tracely/tests/test_spans.py @@ -23,6 +23,7 @@ def trace_func_with_output_struct(): "f2": 100, } + @trace_event() def trace_func_with_inner_trace(): return trace_func_with_output() From 75e79b81c6bb87ae76bd55c1a9ce4bbedbd0c7c2 Mon Sep 17 00:00:00 2001 From: Vyacheslav Morov Date: Mon, 8 Sep 2025 12:00:10 +0100 Subject: [PATCH 7/7] Fix typing errors. --- tracely/src/tracely/_runtime_context.py | 4 ++-- tracely/src/tracely/context.py | 2 ++ tracely/src/tracely/proxy.py | 2 ++ 3 files changed, 6 insertions(+), 2 deletions(-) diff --git a/tracely/src/tracely/_runtime_context.py b/tracely/src/tracely/_runtime_context.py index 585d49a..cdfa364 100644 --- a/tracely/src/tracely/_runtime_context.py +++ b/tracely/src/tracely/_runtime_context.py @@ -7,7 +7,7 @@ class RuntimeContext: def __init__(self): self.span = None - def set_current_span(self, span: SpanObject): + def set_current_span(self, span: Optional[SpanObject]): self.span = span def get_current_span(self) -> Optional[SpanObject]: @@ -26,7 +26,7 @@ def get_current_span(context: Optional[RuntimeContext] = None) -> Optional[SpanO return context.get_current_span() -def set_current_span(span: SpanObject, context: Optional[RuntimeContext] = None): +def set_current_span(span: Optional[SpanObject], context: Optional[RuntimeContext] = None): if context is None: _DEFAULT_CONTEXT.set_current_span(span) else: diff --git a/tracely/src/tracely/context.py b/tracely/src/tracely/context.py index 7c16fd2..4e1f8dd 100644 --- a/tracely/src/tracely/context.py +++ b/tracely/src/tracely/context.py @@ -24,6 +24,8 @@ def create_trace_event(name: str, **params) -> Generator[SpanObject, None, None] span object to work with """ _tracer = get_tracer() + if _tracer is None: + raise ValueError("tracer not initialized") with _tracer.start_as_current_span(f"{name}") as span: obj = SpanObject(span) prev_span = get_current_span() diff --git a/tracely/src/tracely/proxy.py b/tracely/src/tracely/proxy.py index 261119e..061d00f 100644 --- a/tracely/src/tracely/proxy.py +++ b/tracely/src/tracely/proxy.py @@ -11,6 +11,8 @@ class SpanObject: + context: Dict[str, typing.Any] + def __init__(self, span: Optional[opentelemetry.trace.Span] = None): self.context = {} if span is None: