lmnr 0.6.21__py3-none-any.whl → 0.7.1__py3-none-any.whl

lmnr/opentelemetry_lib/opentelemetry/instrumentation/openai/shared/completion_wrappers.py

@@ -27,6 +27,10 @@ from ..utils import (
     should_emit_events,
     should_send_prompts,
 )
+from lmnr.opentelemetry_lib.tracing.context import (
+    get_current_context,
+    get_event_attributes_from_context,
+)
 from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
 from opentelemetry.semconv.attributes.error_attributes import ERROR_TYPE
 from opentelemetry.semconv_ai import (
@@ -55,6 +59,7 @@ def completion_wrapper(tracer, wrapped, instance, args, kwargs):
         SPAN_NAME,
         kind=SpanKind.CLIENT,
         attributes={SpanAttributes.LLM_REQUEST_TYPE: LLM_REQUEST_TYPE.value},
+        context=get_current_context(),
     )
 
     _handle_request(span, kwargs, instance)
@@ -63,7 +68,8 @@ def completion_wrapper(tracer, wrapped, instance, args, kwargs):
         response = wrapped(*args, **kwargs)
     except Exception as e:
         span.set_attribute(ERROR_TYPE, e.__class__.__name__)
-        span.record_exception(e)
+        attributes = get_event_attributes_from_context()
+        span.record_exception(e, attributes=attributes)
         span.set_status(Status(StatusCode.ERROR, str(e)))
         span.end()
         raise
@@ -89,6 +95,7 @@ async def acompletion_wrapper(tracer, wrapped, instance, args, kwargs):
         name=SPAN_NAME,
         kind=SpanKind.CLIENT,
         attributes={SpanAttributes.LLM_REQUEST_TYPE: LLM_REQUEST_TYPE.value},
+        context=get_current_context(),
     )
 
     _handle_request(span, kwargs, instance)
@@ -97,7 +104,8 @@ async def acompletion_wrapper(tracer, wrapped, instance, args, kwargs):
         response = await wrapped(*args, **kwargs)
     except Exception as e:
         span.set_attribute(ERROR_TYPE, e.__class__.__name__)
-        span.record_exception(e)
+        attributes = get_event_attributes_from_context()
+        span.record_exception(e, attributes=attributes)
         span.set_status(Status(StatusCode.ERROR, str(e)))
         span.end()
         raise

lmnr/opentelemetry_lib/opentelemetry/instrumentation/openai/shared/embeddings_wrappers.py

@@ -3,6 +3,8 @@ import time
 from collections.abc import Iterable
 
 from opentelemetry import context as context_api
+
+from lmnr.opentelemetry_lib.tracing.context import get_event_attributes_from_context
 from ..shared import (
     OPENAI_LLM_USAGE_TOKEN_TYPES,
     _get_openai_base_url,
@@ -91,7 +93,8 @@ def embeddings_wrapper(
                 exception_counter.add(1, attributes=attributes)
 
             span.set_attribute(ERROR_TYPE, e.__class__.__name__)
-            span.record_exception(e)
+            attributes = get_event_attributes_from_context()
+            span.record_exception(e, attributes=attributes)
             span.set_status(Status(StatusCode.ERROR, str(e)))
             span.end()
 
@@ -156,7 +159,8 @@ async def aembeddings_wrapper(
                 exception_counter.add(1, attributes=attributes)
 
             span.set_attribute(ERROR_TYPE, e.__class__.__name__)
-            span.record_exception(e)
+            attributes = get_event_attributes_from_context()
+            span.record_exception(e, attributes=attributes)
             span.set_status(Status(StatusCode.ERROR, str(e)))
             span.end()
 

lmnr/opentelemetry_lib/opentelemetry/instrumentation/openai/v1/assistant_wrappers.py

@@ -17,6 +17,10 @@ from ..utils import (
     dont_throw,
     should_emit_events,
 )
+from lmnr.opentelemetry_lib.tracing.context import (
+    get_current_context,
+    get_event_attributes_from_context,
+)
 from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
 from opentelemetry.semconv.attributes.error_attributes import ERROR_TYPE
 from opentelemetry.semconv_ai import LLMRequestTypeValues, SpanAttributes
@@ -126,11 +130,12 @@ def messages_list_wrapper(tracer, wrapped, instance, args, kwargs):
         kind=SpanKind.CLIENT,
         attributes={SpanAttributes.LLM_REQUEST_TYPE: LLMRequestTypeValues.CHAT.value},
         start_time=run.get("start_time"),
+        context=get_current_context(),
     )
 
     if exception := run.get("exception"):
         span.set_attribute(ERROR_TYPE, exception.__class__.__name__)
-        span.record_exception(exception)
+        span.record_exception(exception, attributes=get_event_attributes_from_context())
         span.set_status(Status(StatusCode.ERROR, str(exception)))
         span.end(run.get("end_time"))
 
@@ -250,6 +255,7 @@ def runs_create_and_stream_wrapper(tracer, wrapped, instance, args, kwargs):
         "openai.assistant.run_stream",
         kind=SpanKind.CLIENT,
         attributes={SpanAttributes.LLM_REQUEST_TYPE: LLMRequestTypeValues.CHAT.value},
+        context=get_current_context(),
     )
 
     i = 0
@@ -313,7 +319,7 @@ def runs_create_and_stream_wrapper(tracer, wrapped, instance, args, kwargs):
         return response
     except Exception as e:
         span.set_attribute(ERROR_TYPE, e.__class__.__name__)
-        span.record_exception(e)
+        span.record_exception(e, attributes=get_event_attributes_from_context())
         span.set_status(Status(StatusCode.ERROR, str(e)))
         span.end()
         raise

lmnr/opentelemetry_lib/opentelemetry/instrumentation/openai/v1/event_handler_wrapper.py

@@ -1,3 +1,4 @@
+from lmnr.opentelemetry_lib.tracing.context import get_event_attributes_from_context
 from ..shared import _set_span_attribute
 from ..shared.event_emitter import emit_event
 from ..shared.event_models import ChoiceEvent
@@ -69,7 +70,9 @@ class EventHandlerWrapper(AssistantEventHandler):
     @override
     def on_exception(self, exception: Exception):
         self._span.set_attribute(ERROR_TYPE, exception.__class__.__name__)
-        self._span.record_exception(exception)
+        self._span.record_exception(
+            exception, attributes=get_event_attributes_from_context()
+        )
         self._span.set_status(Status(StatusCode.ERROR, str(exception)))
         self._original_handler.on_exception(exception)
 

lmnr/opentelemetry_lib/opentelemetry/instrumentation/openai/v1/responses_wrappers.py

@@ -36,6 +36,10 @@ except ImportError:
     ResponseOutputMessageParam = Dict[str, Any]
     RESPONSES_AVAILABLE = False
 
+from lmnr.opentelemetry_lib.tracing.context import (
+    get_current_context,
+    get_event_attributes_from_context,
+)
 from openai._legacy_response import LegacyAPIResponse
 from opentelemetry import context as context_api
 from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
@@ -429,9 +433,10 @@ def responses_get_or_create_wrapper(tracer: Tracer, wrapped, instance, args, kwa
             start_time=(
                 start_time if traced_data is None else int(traced_data.start_time)
             ),
+            context=get_current_context(),
         )
         span.set_attribute(ERROR_TYPE, e.__class__.__name__)
-        span.record_exception(e)
+        span.record_exception(e, attributes=get_event_attributes_from_context())
         span.set_status(StatusCode.ERROR, str(e))
         if traced_data:
             set_data_attributes(traced_data, span)
@@ -472,6 +477,7 @@ def responses_get_or_create_wrapper(tracer: Tracer, wrapped, instance, args, kwa
             SPAN_NAME,
             kind=SpanKind.CLIENT,
             start_time=int(traced_data.start_time),
+            context=get_current_context(),
         )
         set_data_attributes(traced_data, span)
         span.end()
@@ -523,9 +529,10 @@ async def async_responses_get_or_create_wrapper(
             start_time=(
                 start_time if traced_data is None else int(traced_data.start_time)
             ),
+            context=get_current_context(),
         )
         span.set_attribute(ERROR_TYPE, e.__class__.__name__)
-        span.record_exception(e)
+        span.record_exception(e, attributes=get_event_attributes_from_context())
         span.set_status(StatusCode.ERROR, str(e))
         if traced_data:
             set_data_attributes(traced_data, span)
@@ -566,6 +573,7 @@ async def async_responses_get_or_create_wrapper(
             SPAN_NAME,
             kind=SpanKind.CLIENT,
             start_time=int(traced_data.start_time),
+            context=get_current_context(),
         )
         set_data_attributes(traced_data, span)
         span.end()
@@ -590,8 +598,12 @@ def responses_cancel_wrapper(tracer: Tracer, wrapped, instance, args, kwargs):
             kind=SpanKind.CLIENT,
             start_time=existing_data.start_time,
             record_exception=True,
+            context=get_current_context(),
+        )
+        span.record_exception(
+            Exception("Response cancelled"),
+            attributes=get_event_attributes_from_context(),
         )
-        span.record_exception(Exception("Response cancelled"))
         set_data_attributes(existing_data, span)
         span.end()
     return response
@@ -616,8 +628,12 @@ async def async_responses_cancel_wrapper(
             kind=SpanKind.CLIENT,
             start_time=existing_data.start_time,
             record_exception=True,
+            context=get_current_context(),
+        )
+        span.record_exception(
+            Exception("Response cancelled"),
+            attributes=get_event_attributes_from_context(),
         )
-        span.record_exception(Exception("Response cancelled"))
         set_data_attributes(existing_data, span)
         span.end()
     return response

lmnr/opentelemetry_lib/opentelemetry/instrumentation/threading/__init__.py

@@ -0,0 +1,190 @@
+# Copyright The OpenTelemetry Authors
+#
+# 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.
+"""
+Instrument threading to propagate OpenTelemetry context.
+
+Copied from opentelemetry-instrumentation-threading at commit:
+ad2fe813abb2ab0b6e25bedeebef5041ca3189f7
+https://github.com/open-telemetry/opentelemetry-python-contrib/blob/ad2fe813abb2ab0b6e25bedeebef5041ca3189f7/instrumentation/opentelemetry-instrumentation-threading/src/opentelemetry/instrumentation/threading/__init__.py
+
+Modified to use the Laminar isolated context.
+
+Usage
+-----
+
+.. code-block:: python
+
+    from opentelemetry.instrumentation.threading import ThreadingInstrumentor
+
+    ThreadingInstrumentor().instrument()
+
+This library provides instrumentation for the `threading` module to ensure that
+the OpenTelemetry context is propagated across threads. It is important to note
+that this instrumentation does not produce any telemetry data on its own. It
+merely ensures that the context is correctly propagated when threads are used.
+
+
+When instrumented, new threads created using threading.Thread, threading.Timer,
+or within futures.ThreadPoolExecutor will have the current OpenTelemetry
+context attached, and this context will be re-activated in the thread's
+run method or the executor's worker thread."
+"""
+
+from __future__ import annotations
+
+import threading
+from concurrent import futures
+from typing import TYPE_CHECKING, Any, Callable, Collection
+
+from wrapt import (
+    wrap_function_wrapper,  # type: ignore[reportUnknownVariableType]
+)
+
+from lmnr.opentelemetry_lib.tracing.context import (
+    get_current_context,
+    attach_context,
+    detach_context,
+)
+from opentelemetry import context
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.instrumentation.utils import unwrap
+
+_instruments = ()
+
+if TYPE_CHECKING:
+    from typing import Protocol, TypeVar
+
+    R = TypeVar("R")
+
+    class HasOtelContext(Protocol):
+        _otel_context: context.Context
+
+
+class ThreadingInstrumentor(BaseInstrumentor):
+    __WRAPPER_START_METHOD = "start"
+    __WRAPPER_RUN_METHOD = "run"
+    __WRAPPER_SUBMIT_METHOD = "submit"
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return _instruments
+
+    def _instrument(self, **kwargs: Any):
+        self._instrument_thread()
+        self._instrument_timer()
+        self._instrument_thread_pool()
+
+    def _uninstrument(self, **kwargs: Any):
+        self._uninstrument_thread()
+        self._uninstrument_timer()
+        self._uninstrument_thread_pool()
+
+    @staticmethod
+    def _instrument_thread():
+        wrap_function_wrapper(
+            threading.Thread,
+            ThreadingInstrumentor.__WRAPPER_START_METHOD,
+            ThreadingInstrumentor.__wrap_threading_start,
+        )
+        wrap_function_wrapper(
+            threading.Thread,
+            ThreadingInstrumentor.__WRAPPER_RUN_METHOD,
+            ThreadingInstrumentor.__wrap_threading_run,
+        )
+
+    @staticmethod
+    def _instrument_timer():
+        wrap_function_wrapper(
+            threading.Timer,
+            ThreadingInstrumentor.__WRAPPER_START_METHOD,
+            ThreadingInstrumentor.__wrap_threading_start,
+        )
+        wrap_function_wrapper(
+            threading.Timer,
+            ThreadingInstrumentor.__WRAPPER_RUN_METHOD,
+            ThreadingInstrumentor.__wrap_threading_run,
+        )
+
+    @staticmethod
+    def _instrument_thread_pool():
+        wrap_function_wrapper(
+            futures.ThreadPoolExecutor,
+            ThreadingInstrumentor.__WRAPPER_SUBMIT_METHOD,
+            ThreadingInstrumentor.__wrap_thread_pool_submit,
+        )
+
+    @staticmethod
+    def _uninstrument_thread():
+        unwrap(threading.Thread, ThreadingInstrumentor.__WRAPPER_START_METHOD)
+        unwrap(threading.Thread, ThreadingInstrumentor.__WRAPPER_RUN_METHOD)
+
+    @staticmethod
+    def _uninstrument_timer():
+        unwrap(threading.Timer, ThreadingInstrumentor.__WRAPPER_START_METHOD)
+        unwrap(threading.Timer, ThreadingInstrumentor.__WRAPPER_RUN_METHOD)
+
+    @staticmethod
+    def _uninstrument_thread_pool():
+        unwrap(
+            futures.ThreadPoolExecutor,
+            ThreadingInstrumentor.__WRAPPER_SUBMIT_METHOD,
+        )
+
+    @staticmethod
+    def __wrap_threading_start(
+        call_wrapped: Callable[[], None],
+        instance: HasOtelContext,
+        args: tuple[()],
+        kwargs: dict[str, Any],
+    ) -> None:
+        instance._otel_context = get_current_context()
+        return call_wrapped(*args, **kwargs)
+
+    @staticmethod
+    def __wrap_threading_run(
+        call_wrapped: Callable[..., R],
+        instance: HasOtelContext,
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any],
+    ) -> R:
+        token = None
+        try:
+            token = attach_context(instance._otel_context)
+            return call_wrapped(*args, **kwargs)
+        finally:
+            if token is not None:
+                detach_context(token)
+
+    @staticmethod
+    def __wrap_thread_pool_submit(
+        call_wrapped: Callable[..., R],
+        instance: futures.ThreadPoolExecutor,
+        args: tuple[Callable[..., Any], ...],
+        kwargs: dict[str, Any],
+    ) -> R:
+        # obtain the original function and wrapped kwargs
+        original_func = args[0]
+        otel_context = get_current_context()
+
+        def wrapped_func(*func_args: Any, **func_kwargs: Any) -> R:
+            token = None
+            try:
+                token = attach_context(otel_context)
+                return original_func(*func_args, **func_kwargs)
+            finally:
+                if token is not None:
+                    detach_context(token)
+
+        # replace the original function with the wrapped function
+        new_args: tuple[Callable[..., Any], ...] = (wrapped_func,) + args[1:]
+        return call_wrapped(*new_args, **kwargs)

lmnr/opentelemetry_lib/tracing/__init__.py

@@ -10,9 +10,22 @@ from lmnr.opentelemetry_lib.tracing.instruments import (
     Instruments,
     init_instrumentations,
 )
+from lmnr.opentelemetry_lib.tracing.context import (
+    attach_context,
+    detach_context,
+    get_current_context,
+    get_token_stack,
+    _isolated_token_stack,
+    _isolated_token_stack_storage,
+    set_token_stack,
+)
 
 from opentelemetry import trace
-from opentelemetry.instrumentation.threading import ThreadingInstrumentor
+from opentelemetry.context import Context
+
+# instead of importing from opentelemetry.instrumentation.threading,
+# we import from our modified copy to use Laminar's isolated context.
+from ..opentelemetry.instrumentation.threading import ThreadingInstrumentor
 from opentelemetry.sdk.resources import Resource
 from opentelemetry.sdk.trace import TracerProvider, SpanProcessor
 from opentelemetry.sdk.trace.export import SpanExporter
@@ -32,6 +45,7 @@ class TracerWrapper(object):
     _async_client: AsyncLaminarClient
     _resource: Resource
     _span_processor: SpanProcessor
+    _original_thread_init = None
 
     def __new__(
         cls,
@@ -91,6 +105,9 @@ class TracerWrapper(object):
 
                 obj._tracer_provider.add_span_processor(obj._span_processor)
 
+                # Setup threading context inheritance
+                obj._setup_threading_inheritance()
+
                 # This is not a real instrumentation and does not generate telemetry
                 # data, but it is required to ensure that OpenTelemetry context
                 # propagation is enabled.
@@ -113,6 +130,43 @@ class TracerWrapper(object):
 
             return cls.instance
 
+    def _setup_threading_inheritance(self):
+        """Setup threading inheritance for isolated context."""
+        if TracerWrapper._original_thread_init is None:
+            # Monkey patch Thread.__init__ to capture context inheritance
+            TracerWrapper._original_thread_init = threading.Thread.__init__
+
+            def patched_thread_init(thread_self, *args, **kwargs):
+                # Capture current isolated context and token stack for inheritance
+                current_context = get_current_context()
+                current_token_stack = get_token_stack().copy()
+
+                # Get the original target function
+                original_target = kwargs.get("target")
+                if not original_target and args:
+                    original_target = args[0]
+
+                # Only inherit if we have a target function
+                if original_target:
+                    # Create a wrapper function that sets up context
+                    def thread_wrapper(*target_args, **target_kwargs):
+                        # Set inherited context and token stack in the new thread
+                        attach_context(current_context)
+                        set_token_stack(current_token_stack)
+                        # Run original target
+                        return original_target(*target_args, **target_kwargs)
+
+                    # Replace the target with our wrapper
+                    if "target" in kwargs:
+                        kwargs["target"] = thread_wrapper
+                    elif args:
+                        args = (thread_wrapper,) + args[1:]
+
+                # Call original init
+                TracerWrapper._original_thread_init(thread_self, *args, **kwargs)
+
+            threading.Thread.__init__ = patched_thread_init
+
     def exit_handler(self):
         if isinstance(self._span_processor, LaminarSpanProcessor):
             self._span_processor.clear()
@@ -124,6 +178,31 @@ class TracerWrapper(object):
         console_log_handler.setFormatter(VerboseColorfulFormatter())
         self._logger.addHandler(console_log_handler)
 
+    def get_isolated_context(self) -> Context:
+        """Get the current isolated context."""
+        return get_current_context()
+
+    def push_span_context(self, span: trace.Span) -> Context:
+        """Push a new context with the given span onto the stack."""
+        current_ctx = get_current_context()
+        new_context = trace.set_span_in_context(span, current_ctx)
+        token = attach_context(new_context)
+
+        # Store the token for later detachment - tokens are much lighter than contexts
+        current_stack = get_token_stack().copy()
+        current_stack.append(token)
+        set_token_stack(current_stack)
+
+        return new_context
+
+    def pop_span_context(self) -> None:
+        """Pop the current span context from the stack."""
+        current_stack = get_token_stack().copy()
+        if current_stack:
+            token = current_stack.pop()
+            set_token_stack(current_stack)
+            detach_context(token)
+
     @staticmethod
     def set_static_params(
         resource_attributes: dict,
@@ -144,6 +223,15 @@ class TracerWrapper(object):
         # Any state cleanup. Now used in between tests
         if isinstance(cls.instance._span_processor, LaminarSpanProcessor):
             cls.instance._span_processor.clear()
+        # Clear the isolated context state for clean test state
+        try:
+            _isolated_token_stack.set([])
+        except LookupError:
+            pass
+        if hasattr(_isolated_token_stack_storage, "token_stack"):
+            _isolated_token_stack_storage.token_stack = []
+        # Reset the isolated context to a fresh state
+        attach_context(Context())
 
     def shutdown(self):
         if self._tracer_provider is None:

lmnr/opentelemetry_lib/tracing/context.py

@@ -0,0 +1,126 @@
+import threading
+
+from abc import ABC, abstractmethod
+from contextvars import ContextVar
+from opentelemetry.context import Context, Token, create_key, get_value
+
+from lmnr.opentelemetry_lib.tracing.attributes import SESSION_ID, USER_ID
+
+
+class _IsolatedRuntimeContext(ABC):
+    """The isolated RuntimeContext interface, identical to OpenTelemetry's _RuntimeContext
+    but isolated from the global context.
+    """
+
+    @abstractmethod
+    def attach(self, context: Context) -> Token[Context]:
+        """Sets the current `Context` object. Returns a
+        token that can be used to reset to the previous `Context`.
+
+        Args:
+            context: The Context to set.
+        """
+
+    @abstractmethod
+    def get_current(self) -> Context:
+        """Returns the current `Context` object."""
+
+    @abstractmethod
+    def detach(self, token: Token[Context]) -> None:
+        """Resets Context to a previous value
+
+        Args:
+            token: A reference to a previous Context.
+        """
+
+
+class IsolatedContextVarsRuntimeContext(_IsolatedRuntimeContext):
+    """An isolated implementation of the RuntimeContext interface which wraps ContextVar
+    but uses its own ContextVar instead of the global one.
+    """
+
+    def __init__(self) -> None:
+        self._current_context = ContextVar(
+            "isolated_current_context", default=Context()
+        )
+
+    def attach(self, context: Context) -> Token[Context]:
+        """Sets the current `Context` object. Returns a
+        token that can be used to reset to the previous `Context`.
+
+        Args:
+            context: The Context to set.
+        """
+        return self._current_context.set(context)
+
+    def get_current(self) -> Context:
+        """Returns the current `Context` object."""
+        return self._current_context.get()
+
+    def detach(self, token: Token[Context]) -> None:
+        """Resets Context to a previous value
+
+        Args:
+            token: A reference to a previous Context.
+        """
+        self._current_context.reset(token)
+
+
+# Create the isolated runtime context
+_ISOLATED_RUNTIME_CONTEXT = IsolatedContextVarsRuntimeContext()
+
+# Token stack for push/pop API compatibility - much lighter than copying contexts
+_isolated_token_stack: ContextVar[list[Token[Context]]] = ContextVar(
+    "isolated_token_stack", default=[]
+)
+
+# Thread-local storage for threading support
+_isolated_token_stack_storage = threading.local()
+
+
+def get_token_stack() -> list[Token[Context]]:
+    """Get the token stack, supporting both asyncio and threading."""
+    try:
+        return _isolated_token_stack.get()
+    except LookupError:
+        if not hasattr(_isolated_token_stack_storage, "token_stack"):
+            _isolated_token_stack_storage.token_stack = []
+        return _isolated_token_stack_storage.token_stack
+
+
+def set_token_stack(stack: list[Token[Context]]) -> None:
+    """Set the token stack, supporting both asyncio and threading."""
+    try:
+        _isolated_token_stack.set(stack)
+    except LookupError:
+        _isolated_token_stack_storage.token_stack = stack
+
+
+def get_current_context() -> Context:
+    """Get the current isolated context."""
+    return _ISOLATED_RUNTIME_CONTEXT.get_current()
+
+
+def attach_context(context: Context) -> Token[Context]:
+    """Attach a context to the isolated runtime context."""
+    return _ISOLATED_RUNTIME_CONTEXT.attach(context)
+
+
+def detach_context(token: Token[Context]) -> None:
+    """Detach a context from the isolated runtime context."""
+    _ISOLATED_RUNTIME_CONTEXT.detach(token)
+
+
+CONTEXT_USER_ID_KEY = create_key(f"lmnr.{USER_ID}")
+CONTEXT_SESSION_ID_KEY = create_key(f"lmnr.{SESSION_ID}")
+
+
+def get_event_attributes_from_context(context: Context | None = None) -> dict[str, str]:
+    """Get the event attributes from the context."""
+    context = context or get_current_context()
+    attributes = {}
+    if session_id := get_value(CONTEXT_SESSION_ID_KEY, context):
+        attributes["lmnr.event.session_id"] = session_id
+    if user_id := get_value(CONTEXT_USER_ID_KEY, context):
+        attributes["lmnr.event.user_id"] = user_id
+    return attributes

lmnr/opentelemetry_lib/tracing/processor.py

@@ -17,9 +17,6 @@ from lmnr.opentelemetry_lib.tracing.attributes import (
     SPAN_SDK_VERSION,
 )
 from lmnr.opentelemetry_lib.tracing.exporter import LaminarSpanExporter
-from lmnr.opentelemetry_lib.tracing.context_properties import (
-    _set_association_properties_attributes,
-)
 from lmnr.version import PYTHON_VERSION, __version__
 
 
@@ -76,9 +73,11 @@ class LaminarSpanProcessor(SpanProcessor):
         span.set_attribute(SPAN_SDK_VERSION, __version__)
         span.set_attribute(SPAN_LANGUAGE_VERSION, f"python@{PYTHON_VERSION}")
 
-        association_properties = get_value("association_properties")
-        if association_properties is not None:
-            _set_association_properties_attributes(span, association_properties)
+        if span.name == "LangGraph.workflow":
+            graph_context = get_value("lmnr.langgraph.graph") or {}
+            for key, value in graph_context.items():
+                span.set_attribute(f"lmnr.association.properties.{key}", value)
+
         self.instance.on_start(span, parent_context)
 
     def on_end(self, span: Span):