From bf22d8aed7ee9f763f54daa9e60f961d92073d25 Mon Sep 17 00:00:00 2001 From: Din Date: Tue, 2 Dec 2025 18:05:03 +0000 Subject: [PATCH 01/10] start span params, global metadata, span object, add_tags - Add user id and session id as params to start span for consistency with JS - Global metadata set for all traces at Laminar.initialize - Return LaminarSpan object from start trace - Expose most control methods on LaminarSpan and use them in Laminar - Add an ability to add tags --- pyproject.toml | 3 +- src/lmnr/__init__.py | 2 + .../opentelemetry_lib/decorators/__init__.py | 100 +++---- .../opentelemetry_lib/litellm/__init__.py | 2 +- .../instrumentation/claude_agent/__init__.py | 3 +- .../instrumentation/cua_agent/__init__.py | 2 +- .../instrumentation/cua_computer/__init__.py | 3 +- .../instrumentation/google_genai/__init__.py | 2 +- .../instrumentation/kernel/__init__.py | 3 +- .../instrumentation/kernel/utils.py | 2 +- .../openai/v1/responses_wrappers.py | 2 +- .../instrumentation/openhands_ai/__init__.py | 3 +- .../instrumentation/skyvern/__init__.py | 120 ++++---- .../opentelemetry_lib/tracing/attributes.py | 1 + src/lmnr/opentelemetry_lib/tracing/context.py | 3 +- src/lmnr/opentelemetry_lib/tracing/span.py | 258 +++++++++++++---- src/lmnr/sdk/browser/browser_use_otel.py | 3 +- src/lmnr/sdk/decorators.py | 12 +- src/lmnr/sdk/evaluations.py | 3 +- src/lmnr/sdk/laminar.py | 270 ++++++++++++------ src/lmnr/sdk/types.py | 8 +- src/lmnr/sdk/utils.py | 40 +++ tests/test_global_metadata.py | 149 ++++++++++ tests/test_observe.py | 35 +++ tests/test_tracing.py | 70 +++++ tests/test_utils.py | 3 +- 26 files changed, 796 insertions(+), 306 deletions(-) create mode 100644 tests/test_global_metadata.py diff --git a/pyproject.toml b/pyproject.toml index 3140307f..5ccfdb3b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -67,7 +67,7 @@ lmnr = "lmnr.cli:cli" alephalpha=["opentelemetry-instrumentation-alephalpha>=0.47.1"] bedrock=["opentelemetry-instrumentation-bedrock>=0.47.1"] chromadb=["opentelemetry-instrumentation-chromadb>=0.47.1"] -claude-agent-sdk=["lmnr-claude-code-proxy>=0.1.0a5"] +claude-agent-sdk=["lmnr-claude-code-proxy>=0.1.3"] cohere=["opentelemetry-instrumentation-cohere>=0.47.1"] crewai=["opentelemetry-instrumentation-crewai>=0.47.1"] haystack=["opentelemetry-instrumentation-haystack>=0.47.1"] @@ -93,6 +93,7 @@ weaviate=["opentelemetry-instrumentation-weaviate>=0.47.1"] # we suggest using package-manager-specific commands instead, # like `uv add lmnr --all-extras` all = [ + "lmnr-claude-code-proxy>=0.1.3", "opentelemetry-instrumentation-alephalpha>=0.47.1", "opentelemetry-instrumentation-bedrock>=0.47.1", "opentelemetry-instrumentation-chromadb>=0.47.1", diff --git a/src/lmnr/__init__.py b/src/lmnr/__init__.py index 23fccd7a..f95f82ce 100644 --- a/src/lmnr/__init__.py +++ b/src/lmnr/__init__.py @@ -11,6 +11,7 @@ from .opentelemetry_lib.tracing.attributes import Attributes from .opentelemetry_lib.tracing.instruments import Instruments from .opentelemetry_lib.tracing.processor import LaminarSpanProcessor +from .opentelemetry_lib.tracing.span import LaminarSpan from .opentelemetry_lib.tracing.tracer import get_laminar_tracer_provider, get_tracer __all__ = [ @@ -25,6 +26,7 @@ "LaminarLiteLLMCallback", "LaminarSpanContext", "LaminarSpanProcessor", + "LaminarSpan", "get_laminar_tracer_provider", "get_tracer", "evaluate", diff --git a/src/lmnr/opentelemetry_lib/decorators/__init__.py b/src/lmnr/opentelemetry_lib/decorators/__init__.py index fb7da05e..ba1cb0fc 100644 --- a/src/lmnr/opentelemetry_lib/decorators/__init__.py +++ b/src/lmnr/opentelemetry_lib/decorators/__init__.py @@ -1,6 +1,4 @@ from functools import wraps -import pydantic -import orjson import types from typing import Any, AsyncGenerator, Callable, Generator, Literal, TypeVar @@ -8,68 +6,36 @@ from opentelemetry.trace import Span, Status, StatusCode from lmnr.opentelemetry_lib.tracing.context import ( + CONTEXT_METADATA_KEY, CONTEXT_SESSION_ID_KEY, CONTEXT_USER_ID_KEY, attach_context, detach_context, get_event_attributes_from_context, ) +from lmnr.opentelemetry_lib.tracing.span import LaminarSpan from lmnr.sdk.utils import get_input_from_func_args, is_method -from lmnr.opentelemetry_lib import MAX_MANUAL_SPAN_PAYLOAD_SIZE from lmnr.opentelemetry_lib.tracing.tracer import get_tracer_with_context from lmnr.opentelemetry_lib.tracing.attributes import ( ASSOCIATION_PROPERTIES, - SPAN_INPUT, - SPAN_OUTPUT, + METADATA, SPAN_TYPE, ) from lmnr.opentelemetry_lib.tracing import TracerWrapper from lmnr.sdk.log import get_default_logger +from lmnr.sdk.utils import is_otel_attribute_value_type, json_dumps logger = get_default_logger(__name__) F = TypeVar("F", bound=Callable[..., Any]) -DEFAULT_PLACEHOLDER = {} - - -def default_json(o): - if isinstance(o, pydantic.BaseModel): - return o.model_dump() - - # Handle various sequence types, but not strings or bytes - if isinstance(o, (list, tuple, set, frozenset)): - return list(o) - - try: - return str(o) - except Exception: - logger.debug("Failed to serialize data to JSON, inner type: %s", type(o)) - pass - return DEFAULT_PLACEHOLDER - - -def json_dumps(data: dict) -> str: - try: - return orjson.dumps( - data, - default=default_json, - option=orjson.OPT_SERIALIZE_DATACLASS - | orjson.OPT_SERIALIZE_UUID - | orjson.OPT_UTC_Z - | orjson.OPT_NON_STR_KEYS, - ).decode("utf-8") - except Exception: - # Log the exception and return a placeholder if serialization completely fails - logger.info("Failed to serialize data to JSON, type: %s", type(data)) - return "{}" # Return an empty JSON object as a fallback - def _setup_span( span_name: str, span_type: str, association_properties: dict[str, Any] | None, preserve_global_context: bool = False, + metadata: dict[str, Any] | None = None, ): """Set up a span with the given name, type, and association properties.""" with get_tracer_with_context() as (tracer, isolated_context): @@ -80,6 +46,14 @@ def _setup_span( attributes={SPAN_TYPE: span_type}, ) + ctx_metadata = context_api.get_value(CONTEXT_METADATA_KEY, isolated_context) + merged_metadata = {**(ctx_metadata or {}), **(metadata or {})} + for key, value in merged_metadata.items(): + span.set_attribute( + f"{ASSOCIATION_PROPERTIES}.{METADATA}.{key}", + (value if is_otel_attribute_value_type(value) else json_dumps(value)), + ) + if association_properties is not None: for key, value in association_properties.items(): span.set_attribute(f"{ASSOCIATION_PROPERTIES}.{key}", value) @@ -103,23 +77,18 @@ def _process_input( try: if input_formatter is not None: inp = input_formatter(*args, **kwargs) - if not isinstance(inp, str): - inp = json_dumps(inp) else: - inp = json_dumps( - get_input_from_func_args( - fn, - is_method=is_method(fn), - func_args=args, - func_kwargs=kwargs, - ignore_inputs=ignore_inputs, - ) + inp = get_input_from_func_args( + fn, + is_method=is_method(fn), + func_args=args, + func_kwargs=kwargs, + ignore_inputs=ignore_inputs, ) - if len(inp) > MAX_MANUAL_SPAN_PAYLOAD_SIZE: - span.set_attribute(SPAN_INPUT, "Laminar: input too large to record") - else: - span.set_attribute(SPAN_INPUT, inp) + if not isinstance(span, LaminarSpan): + span = LaminarSpan(span) + span.set_input(inp) except Exception: msg = "Failed to process input, ignoring" if input_formatter is not None: @@ -144,15 +113,12 @@ def _process_output( try: if output_formatter is not None: output = output_formatter(result) - if not isinstance(output, str): - output = json_dumps(output) else: - output = json_dumps(result) + output = result - if len(output) > MAX_MANUAL_SPAN_PAYLOAD_SIZE: - span.set_attribute(SPAN_OUTPUT, "Laminar: output too large to record") - else: - span.set_attribute(SPAN_OUTPUT, output) + if not isinstance(span, LaminarSpan): + span = LaminarSpan(span) + span.set_output(output) except Exception: msg = "Failed to process output, ignoring" if output_formatter is not None: @@ -177,6 +143,7 @@ def observe_base( ignore_inputs: list[str] | None = None, ignore_output: bool = False, span_type: Literal["DEFAULT", "LLM", "TOOL"] = "DEFAULT", + metadata: dict[str, Any] | None = None, association_properties: dict[str, Any] | None = None, input_formatter: Callable[..., str] | None = None, output_formatter: Callable[..., str] | None = None, @@ -192,7 +159,11 @@ def wrap(*args, **kwargs): wrapper = TracerWrapper() span = _setup_span( - span_name, span_type, association_properties, preserve_global_context + span_name, + span_type, + association_properties, + preserve_global_context, + metadata, ) new_context = wrapper.push_span_context(span) if session_id := association_properties.get("session_id"): @@ -255,6 +226,7 @@ def async_observe_base( ignore_inputs: list[str] | None = None, ignore_output: bool = False, span_type: Literal["DEFAULT", "LLM", "TOOL"] = "DEFAULT", + metadata: dict[str, Any] | None = None, association_properties: dict[str, Any] | None = None, input_formatter: Callable[..., str] | None = None, output_formatter: Callable[..., str] | None = None, @@ -270,7 +242,11 @@ async def wrap(*args, **kwargs): wrapper = TracerWrapper() span = _setup_span( - span_name, span_type, association_properties, preserve_global_context + span_name, + span_type, + association_properties, + preserve_global_context, + metadata, ) new_context = wrapper.push_span_context(span) if session_id := association_properties.get("session_id"): diff --git a/src/lmnr/opentelemetry_lib/litellm/__init__.py b/src/lmnr/opentelemetry_lib/litellm/__init__.py index 9973bea4..4ff3b731 100644 --- a/src/lmnr/opentelemetry_lib/litellm/__init__.py +++ b/src/lmnr/opentelemetry_lib/litellm/__init__.py @@ -5,7 +5,6 @@ from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import GEN_AI_PROMPT from opentelemetry.trace import SpanKind, Status, StatusCode, Tracer -from lmnr.opentelemetry_lib.decorators import json_dumps from lmnr.opentelemetry_lib.litellm.utils import ( get_tool_definition, is_validator_iterator, @@ -21,6 +20,7 @@ from lmnr.opentelemetry_lib.tracing.attributes import ASSOCIATION_PROPERTIES from lmnr.opentelemetry_lib.utils.package_check import is_package_installed from lmnr.sdk.log import get_default_logger +from lmnr.sdk.utils import json_dumps logger = get_default_logger(__name__) diff --git a/src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/claude_agent/__init__.py b/src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/claude_agent/__init__.py index 0942a1b8..4084d413 100644 --- a/src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/claude_agent/__init__.py +++ b/src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/claude_agent/__init__.py @@ -3,11 +3,10 @@ import sys from lmnr import Laminar -from lmnr.opentelemetry_lib.decorators import json_dumps from lmnr.opentelemetry_lib.tracing import get_current_context from lmnr.opentelemetry_lib.tracing.attributes import SPAN_IDS_PATH, SPAN_PATH from lmnr.sdk.log import get_default_logger -from lmnr.sdk.utils import get_input_from_func_args, is_method +from lmnr.sdk.utils import get_input_from_func_args, is_method, json_dumps from opentelemetry import trace from opentelemetry.sdk.trace import ReadableSpan diff --git a/src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/cua_agent/__init__.py b/src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/cua_agent/__init__.py index 0a6c7925..56c6ea6e 100644 --- a/src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/cua_agent/__init__.py +++ b/src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/cua_agent/__init__.py @@ -3,8 +3,8 @@ import logging from typing import Any, AsyncGenerator, Collection -from lmnr.opentelemetry_lib.decorators import json_dumps from lmnr import Laminar +from lmnr.sdk.utils import json_dumps from opentelemetry.instrumentation.instrumentor import BaseInstrumentor from opentelemetry.instrumentation.utils import unwrap diff --git a/src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/cua_computer/__init__.py b/src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/cua_computer/__init__.py index 21279bb2..4d62e49e 100644 --- a/src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/cua_computer/__init__.py +++ b/src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/cua_computer/__init__.py @@ -3,8 +3,7 @@ import logging from typing import Collection -from lmnr.opentelemetry_lib.decorators import json_dumps -from lmnr.sdk.utils import get_input_from_func_args +from lmnr.sdk.utils import get_input_from_func_args, json_dumps from lmnr import Laminar from lmnr.opentelemetry_lib.tracing.context import get_current_context from opentelemetry.instrumentation.instrumentor import BaseInstrumentor diff --git a/src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/google_genai/__init__.py b/src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/google_genai/__init__.py index 2731b304..3f4ade5e 100644 --- a/src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/google_genai/__init__.py +++ b/src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/google_genai/__init__.py @@ -8,11 +8,11 @@ from google.genai import types -from lmnr.opentelemetry_lib.decorators import json_dumps from lmnr.opentelemetry_lib.tracing.context import ( get_current_context, get_event_attributes_from_context, ) +from lmnr.sdk.utils import json_dumps from .config import ( Config, diff --git a/src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/kernel/__init__.py b/src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/kernel/__init__.py index e5693b4c..cb065cfa 100644 --- a/src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/kernel/__init__.py +++ b/src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/kernel/__init__.py @@ -3,13 +3,12 @@ import functools from typing import Collection -from lmnr.opentelemetry_lib.decorators import json_dumps from lmnr.opentelemetry_lib.opentelemetry.instrumentation.kernel.utils import ( process_tool_output_formatter, screenshot_tool_output_formatter, ) from lmnr.sdk.decorators import observe -from lmnr.sdk.utils import get_input_from_func_args, is_async +from lmnr.sdk.utils import get_input_from_func_args, is_async, json_dumps from lmnr import Laminar from opentelemetry.instrumentation.instrumentor import BaseInstrumentor from opentelemetry.instrumentation.utils import unwrap diff --git a/src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/kernel/utils.py b/src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/kernel/utils.py index 69c888a1..25d6ea12 100644 --- a/src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/kernel/utils.py +++ b/src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/kernel/utils.py @@ -3,7 +3,7 @@ from copy import deepcopy from typing import Any -from lmnr.opentelemetry_lib.decorators import json_dumps +from lmnr.sdk.utils import json_dumps from pydantic import BaseModel diff --git a/src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/openai/v1/responses_wrappers.py b/src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/openai/v1/responses_wrappers.py index 5b1ff2fb..d0dc11dc 100644 --- a/src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/openai/v1/responses_wrappers.py +++ b/src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/openai/v1/responses_wrappers.py @@ -36,11 +36,11 @@ ResponseOutputMessageParam = Dict[str, Any] RESPONSES_AVAILABLE = False -from lmnr.opentelemetry_lib.decorators import json_dumps from lmnr.opentelemetry_lib.tracing.context import ( get_current_context, get_event_attributes_from_context, ) +from lmnr.sdk.utils import json_dumps from openai._legacy_response import LegacyAPIResponse from opentelemetry import context as context_api from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY diff --git a/src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/openhands_ai/__init__.py b/src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/openhands_ai/__init__.py index e6d482d8..bc0532a8 100644 --- a/src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/openhands_ai/__init__.py +++ b/src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/openhands_ai/__init__.py @@ -3,7 +3,6 @@ import sys from typing import Collection -from lmnr.opentelemetry_lib.decorators import json_dumps from lmnr.opentelemetry_lib.tracing.attributes import ( ASSOCIATION_PROPERTIES, SESSION_ID, @@ -11,7 +10,7 @@ ) from lmnr.opentelemetry_lib.utils.wrappers import _with_tracer_wrapper from lmnr.sdk.log import get_default_logger -from lmnr.sdk.utils import get_input_from_func_args +from lmnr.sdk.utils import get_input_from_func_args, json_dumps from lmnr import Laminar from lmnr.version import __version__ diff --git a/src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/skyvern/__init__.py b/src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/skyvern/__init__.py index 7d8e17fd..b45fa167 100644 --- a/src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/skyvern/__init__.py +++ b/src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/skyvern/__init__.py @@ -1,6 +1,5 @@ -from lmnr.opentelemetry_lib.decorators import json_dumps from lmnr.sdk.browser.utils import with_tracer_wrapper -from lmnr.sdk.utils import get_input_from_func_args +from lmnr.sdk.utils import get_input_from_func_args, json_dumps from lmnr.version import __version__ from opentelemetry.instrumentation.instrumentor import BaseInstrumentor @@ -25,14 +24,14 @@ { "package": "skyvern.library.skyvern", "object": "Skyvern", # Class name - "method": "run_task", # Method name + "method": "run_task", # Method name "span_name": "Skyvern.run_task", "span_type": "DEFAULT", }, { "package": "skyvern.webeye.scraper.scraper", # No "object" field for module-level functions - "method": "get_interactable_element_tree", # Function name + "method": "get_interactable_element_tree", # Function name "span_name": "get_interactable_element_tree", "span_type": "DEFAULT", }, @@ -43,31 +42,31 @@ "span_name": "ForgeAgent.execute_step", "span_type": "DEFAULT", }, - { - "package": "skyvern.services.task_v2_service", - "method": "initialize_task_v2", - "span_name": "initialize_task_v2", - "span_type": "DEFAULT", + { + "package": "skyvern.services.task_v2_service", + "method": "initialize_task_v2", + "span_name": "initialize_task_v2", + "span_type": "DEFAULT", }, - { - "package": "skyvern.services.task_v2_service", - "method": "run_task_v2_helper", - "span_name": "run_task_v2_helper", - "span_type": "DEFAULT", + { + "package": "skyvern.services.task_v2_service", + "method": "run_task_v2_helper", + "span_name": "run_task_v2_helper", + "span_type": "DEFAULT", }, - { + { "package": "skyvern.forge.sdk.workflow.models.block", "object": "Block", - "method": "_generate_workflow_run_block_description", - "span_name": "Block._generate_workflow_run_block_description", - "span_type": "DEFAULT", - }, - { - "package": "skyvern.webeye.actions.handler", - "method": "extract_information_for_navigation_goal", - "span_name": "extract_information_for_navigation_goal", - "span_type": "DEFAULT", - }, + "method": "_generate_workflow_run_block_description", + "span_name": "Block._generate_workflow_run_block_description", + "span_type": "DEFAULT", + }, + { + "package": "skyvern.webeye.actions.handler", + "method": "extract_information_for_navigation_goal", + "span_name": "extract_information_for_navigation_goal", + "span_type": "DEFAULT", + }, ] @@ -77,36 +76,36 @@ async def _wrap(tracer: Tracer, to_wrap, wrapped, instance, args, kwargs): attributes = { "lmnr.span.type": to_wrap.get("span_type"), } - + attributes["lmnr.span.input"] = json_dumps( get_input_from_func_args(wrapped, True, args, kwargs) ) - + with tracer.start_as_current_span(span_name, attributes=attributes) as span: - try: - result = await wrapped(*args, **kwargs) - - to_serialize = result - serialized = ( - to_serialize.model_dump_json() - if isinstance(to_serialize, pydantic.BaseModel) - else json_dumps(to_serialize) - ) - span.set_attribute("lmnr.span.output", serialized) - return result - - except Exception as e: - span.record_exception(e) + try: + result = await wrapped(*args, **kwargs) + + to_serialize = result + serialized = ( + to_serialize.model_dump_json() + if isinstance(to_serialize, pydantic.BaseModel) + else json_dumps(to_serialize) + ) + span.set_attribute("lmnr.span.output", serialized) + return result + + except Exception as e: + span.record_exception(e) raise -def instrument_llm_handler(tracer: Tracer): - from skyvern.forge import app - - # Store the original handler - original_handler = app.LLM_API_HANDLER - - async def wrapped_llm_handler(*args, **kwargs): +def instrument_llm_handler(tracer: Tracer): + from skyvern.forge import app + + # Store the original handler + original_handler = app.LLM_API_HANDLER + + async def wrapped_llm_handler(*args, **kwargs): prompt_name = kwargs.get("prompt_name", "") @@ -115,13 +114,13 @@ async def wrapped_llm_handler(*args, **kwargs): else: span_name = "app.LLM_API_HANDLER" - attributes = { - "lmnr.span.type": "DEFAULT", - } + attributes = { + "lmnr.span.type": "DEFAULT", + } - with tracer.start_as_current_span(span_name, attributes=attributes) as span: - try: - result = await original_handler(*args, **kwargs) + with tracer.start_as_current_span(span_name, attributes=attributes) as span: + try: + result = await original_handler(*args, **kwargs) to_serialize = result serialized = ( @@ -129,13 +128,13 @@ async def wrapped_llm_handler(*args, **kwargs): if isinstance(to_serialize, pydantic.BaseModel) else json_dumps(to_serialize) ) - span.set_attribute("lmnr.span.output", serialized) - return result - except Exception as e: - span.record_exception(e) + span.set_attribute("lmnr.span.output", serialized) + return result + except Exception as e: + span.record_exception(e) raise - - # Replace the global handler + + # Replace the global handler app.LLM_API_HANDLER = wrapped_llm_handler @@ -190,4 +189,3 @@ def _uninstrument(self, **kwargs): module_path = wrap_package unwrap(module_path, wrap_method) - diff --git a/src/lmnr/opentelemetry_lib/tracing/attributes.py b/src/lmnr/opentelemetry_lib/tracing/attributes.py index 7214ef12..579c0516 100644 --- a/src/lmnr/opentelemetry_lib/tracing/attributes.py +++ b/src/lmnr/opentelemetry_lib/tracing/attributes.py @@ -24,6 +24,7 @@ ASSOCIATION_PROPERTIES = "lmnr.association.properties" SESSION_ID = "session_id" USER_ID = "user_id" +METADATA = "metadata" TRACE_TYPE = "trace_type" TRACING_LEVEL = "tracing_level" diff --git a/src/lmnr/opentelemetry_lib/tracing/context.py b/src/lmnr/opentelemetry_lib/tracing/context.py index 85b2668d..365d3f16 100644 --- a/src/lmnr/opentelemetry_lib/tracing/context.py +++ b/src/lmnr/opentelemetry_lib/tracing/context.py @@ -4,7 +4,7 @@ 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 +from lmnr.opentelemetry_lib.tracing.attributes import METADATA, SESSION_ID, USER_ID class _IsolatedRuntimeContext(ABC): @@ -113,6 +113,7 @@ def detach_context(token: Token[Context]) -> None: CONTEXT_USER_ID_KEY = create_key(f"lmnr.{USER_ID}") CONTEXT_SESSION_ID_KEY = create_key(f"lmnr.{SESSION_ID}") +CONTEXT_METADATA_KEY = create_key(f"lmnr.{METADATA}") def get_event_attributes_from_context(context: Context | None = None) -> dict[str, str]: diff --git a/src/lmnr/opentelemetry_lib/tracing/span.py b/src/lmnr/opentelemetry_lib/tracing/span.py index 99ff95d8..49b337e9 100644 --- a/src/lmnr/opentelemetry_lib/tracing/span.py +++ b/src/lmnr/opentelemetry_lib/tracing/span.py @@ -1,4 +1,8 @@ +from logging import Logger from inspect import Traceback +from typing import Any, Literal +import uuid + from opentelemetry.sdk.resources import Resource from opentelemetry.sdk.trace import Event, ReadableSpan, Span as SDKSpan from opentelemetry.sdk.util.instrumentation import ( @@ -9,42 +13,175 @@ from opentelemetry.util.types import AttributeValue from opentelemetry.context import detach +from lmnr.opentelemetry_lib import MAX_MANUAL_SPAN_PAYLOAD_SIZE +from lmnr.opentelemetry_lib.tracing.attributes import ( + ASSOCIATION_PROPERTIES, + SPAN_IDS_PATH, + SPAN_INPUT, + SPAN_OUTPUT, + SPAN_PATH, +) from lmnr.opentelemetry_lib.tracing.context import ( pop_span_context, ) from lmnr.sdk.log import get_default_logger +from lmnr.sdk.types import LaminarSpanContext +from lmnr.sdk.utils import is_otel_attribute_value_type, json_dumps -class LaminarSpan(Span, ReadableSpan): - # We wrap the SDK span in a LaminarSpan instead of inheriting from it, - # because OpenTelemetry discourages direct initialization of SdkSpan objects. - # Instead, we rely on the tracer to create the span for us, and then we - # wrap it in a LaminarSpan. - span: SDKSpan - # Whether the span has been popped from the context stack to prevent - # double popping if span.end() is called multiple times. - _popped: bool = False +class LaminarSpanInterfaceMixin: + """Mixin providing Laminar-specific span methods and properties.""" - def __init__(self, span: SDKSpan): - self.span = span - self.logger = get_default_logger(__name__) + span: SDKSpan + logger: Logger + + def set_trace_session_id(self, session_id: str | None = None) -> None: + """Set the session id for the current trace. Must be called at most once per trace. + + Args: + session_id (str | None): Session id to set for the span. + """ + if session_id is not None: + self.set_attribute(f"{ASSOCIATION_PROPERTIES}.session_id", session_id) + + def set_trace_user_id(self, user_id: str | None = None) -> None: + """Set the user id for the current trace. Must be called at most once per trace. + + Args: + user_id (str | None): User id to set for the span. + """ + if user_id is not None: + self.span.set_attribute(f"{ASSOCIATION_PROPERTIES}.user_id", user_id) + + def set_trace_metadata(self, metadata: dict[str, AttributeValue]) -> None: + """Set the metadata for the current trace, merging with any global metadata. + Must be called at most once per trace. + + Args: + metadata (dict[str, AttributeValue]): Metadata to set for the trace. + """ + formatted_metadata = {} + for key, value in metadata.items(): + if is_otel_attribute_value_type(value): + formatted_metadata[f"{ASSOCIATION_PROPERTIES}.metadata.{key}"] = value + else: + formatted_metadata[f"{ASSOCIATION_PROPERTIES}.metadata.{key}"] = ( + json_dumps(value) + ) + self.span.set_attributes(formatted_metadata) + + def get_laminar_span_context(self) -> LaminarSpanContext: + span_path = [] + span_ids_path = [] + if hasattr(self.span, "attributes"): + span_path = list(self.span.attributes.get(SPAN_PATH, tuple())) + span_ids_path = list(self.span.attributes.get(SPAN_IDS_PATH, tuple())) + else: + self.logger.warning( + "Attributes object is not available. Most likely the span is not a LaminarSpan ", + "and not an OpenTelemetry default SDK span. Span path and ids path will be empty.", + ) + return LaminarSpanContext( + trace_id=uuid.UUID(int=self.span.get_span_context().trace_id), + span_id=uuid.UUID(int=self.span.get_span_context().span_id), + is_remote=self.span.get_span_context().is_remote, + span_path=span_path, + span_ids_path=span_ids_path, + ) + + def span_id(self, format: Literal["int", "uuid"] = "int") -> int | uuid.UUID: + if format == "int": + return self.span.get_span_context().span_id + elif format == "uuid": + return uuid.UUID(int=self.span.get_span_context().span_id) + self.logger.warning(f"Invalid format: {format}. Returning int.") + return self.span.get_span_context().span_id + + def trace_id(self, format: Literal["int", "uuid"] = "int") -> int | uuid.UUID: + if format == "int": + return self.span.get_span_context().trace_id + elif format == "uuid": + return uuid.UUID(int=self.span.get_span_context().trace_id) + self.logger.warning(f"Invalid format: {format}. Returning int.") + return self.span.get_span_context().trace_id + + def parent_span_id( + self, format: Literal["int", "uuid"] = "int" + ) -> int | uuid.UUID | None: + parent_span_id = self.span.parent.span_id if self.span.parent else None + if parent_span_id is None: + return None + if format == "int": + return parent_span_id + elif format == "uuid": + return uuid.UUID(int=parent_span_id) + self.logger.warning(f"Invalid format: {format}. Returning int.") + return parent_span_id + + def set_output(self, output: Any = None) -> None: + if output is not None: + serialized_output = json_dumps(output) + if len(serialized_output) > MAX_MANUAL_SPAN_PAYLOAD_SIZE: + self.span.set_attribute( + SPAN_OUTPUT, + "Laminar: output too large to record", + ) + else: + self.span.set_attribute(SPAN_OUTPUT, serialized_output) + + def set_input(self, input: Any = None) -> None: + if input is not None: + serialized_input = json_dumps(input) + if len(serialized_input) > MAX_MANUAL_SPAN_PAYLOAD_SIZE: + self.span.set_attribute( + SPAN_INPUT, + "Laminar: input too large to record", + ) + else: + self.span.set_attribute(SPAN_INPUT, serialized_input) + + def add_tags(self, tags: list[str]) -> None: + if not isinstance(tags, list) or not all(isinstance(tag, str) for tag in tags): + self.logger.warning("Tags must be a list of strings. Tags will be ignored.") + return + current_tags = self.tags + if current_tags is None: + current_tags = [] + current_tags.extend(tags) + self.span.set_attribute( + f"{ASSOCIATION_PROPERTIES}.tags", list(set(current_tags)) + ) + + def set_tags(self, tags: list[str]) -> None: + """Set the tags for the current span. + + Args: + tags (list[str]): Tags to set for the span. + """ + if not isinstance(tags, list) or not all(isinstance(tag, str) for tag in tags): + self.logger.warning("Tags must be a list of strings. Tags will be ignored.") + return + self.span.set_attribute(f"{ASSOCIATION_PROPERTIES}.tags", list(set(tags))) - def end(self, end_time: int | None = None) -> None: - self.span.end(end_time) - if hasattr(self, "_lmnr_ctx_token") and not self._popped: - try: - pop_span_context() - # Internally handles and logs the error - detach(self._lmnr_ctx_token) - self._popped = True - except Exception: - pass + @property + def tags(self) -> list[str]: + self.logger.debug( + "[LaminarSpan.tags] WARNING. Reading current span attributes to get tags. " + "This works in default OTel SDK, but is not guaranteed by the " + "OpenTelemetry API specification. If you are using a custom SDK, " + "you may need to implement your own way to get the tags. In such case, " + "this function returns an empty list." + ) + try: + return list(self.span.attributes.get(f"{ASSOCIATION_PROPERTIES}.tags", [])) + except Exception: + return [] + + +class SpanDelegationMixin: + """Mixin providing delegation to the wrapped SDK span for standard OpenTelemetry methods.""" - ### ======================================================================== - # The below methods are just passthrough of abstract Span methods - # to the SDK span. If you need to override them, or add additional logic, - # move them above this section. - # ========================================================================== + span: SDKSpan def get_span_context(self) -> SpanContext: return self.span.get_span_context() @@ -86,17 +223,6 @@ def record_exception( ) -> None: self.span.record_exception(exception, attributes, timestamp, escaped) - def __enter__(self) -> "LaminarSpan": - return self - - def __exit__( - self, - exc_type: type[BaseException] | None, - exc_value: BaseException | None, - exc_tb: Traceback | None, - ) -> None: - self.end() - def _readable_span(self) -> ReadableSpan: return self.span._readable_span() @@ -109,11 +235,11 @@ def context(self) -> SpanContext: return self.span.context @property - def start_time(self) -> int: + def start_time(self) -> int | None: return self.span.start_time @property - def end_time(self) -> int: + def end_time(self) -> int | None: return self.span.end_time @property @@ -144,18 +270,6 @@ def links(self) -> list[Link]: def status(self) -> Status: return self.span.status - @property - def parent_span_context(self) -> SpanContext: - return self.span.parent_span_context - - @property - def span_context(self) -> SpanContext: - return self.span.span_context - - @property - def span_id(self) -> str: - return self.span.span_id - @property def kind(self) -> SpanKind: return self.span.kind @@ -174,3 +288,43 @@ def instrumentation_info(self) -> InstrumentationInfo: def to_json(self) -> str: return self.span.to_json() + + +class LaminarSpan(LaminarSpanInterfaceMixin, SpanDelegationMixin, Span, ReadableSpan): + """ + Laminar's span wrapper that complies with OpenTelemetry's Span and ReadableSpan interfaces. + + We wrap the SDK span instead of inheriting from it, because OpenTelemetry discourages + direct initialization of SdkSpan objects. Instead, we rely on the tracer to create + the span for us, and then we wrap it in a LaminarSpan. + """ + + span: SDKSpan + _popped: bool = False + + def __init__(self, span: SDKSpan): + if isinstance(span, LaminarSpan): + span = span.span + self.logger = get_default_logger(__name__) + self.span = span + + def end(self, end_time: int | None = None) -> None: + self.span.end(end_time) + if hasattr(self, "_lmnr_ctx_token") and not self._popped: + try: + pop_span_context() + detach(self._lmnr_ctx_token) + self._popped = True + except Exception: + pass + + def __enter__(self) -> "LaminarSpan": + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + exc_tb: Traceback | None, + ) -> None: + self.end() diff --git a/src/lmnr/sdk/browser/browser_use_otel.py b/src/lmnr/sdk/browser/browser_use_otel.py index d6506898..4ee3013c 100644 --- a/src/lmnr/sdk/browser/browser_use_otel.py +++ b/src/lmnr/sdk/browser/browser_use_otel.py @@ -1,7 +1,6 @@ -from lmnr.opentelemetry_lib.decorators import json_dumps from lmnr import Laminar from lmnr.sdk.browser.utils import with_tracer_wrapper -from lmnr.sdk.utils import get_input_from_func_args +from lmnr.sdk.utils import get_input_from_func_args, json_dumps from lmnr.version import __version__ from opentelemetry.instrumentation.instrumentor import BaseInstrumentor diff --git a/src/lmnr/sdk/decorators.py b/src/lmnr/sdk/decorators.py index ae700a0e..ee3aa5c8 100644 --- a/src/lmnr/sdk/decorators.py +++ b/src/lmnr/sdk/decorators.py @@ -1,7 +1,6 @@ from lmnr.opentelemetry_lib.decorators import ( observe_base, async_observe_base, - json_dumps, ) from opentelemetry.trace import INVALID_SPAN, get_current_span @@ -136,15 +135,6 @@ def decorator( association_properties["session_id"] = session_id if user_id is not None: association_properties["user_id"] = user_id - if metadata is not None: - association_properties.update( - { - f"metadata.{k}": ( - v if isinstance(v, (str, int, float, bool)) else json_dumps(v) - ) - for k, v in metadata.items() - } - ) if tags is not None: if not isinstance(tags, list) or not all( isinstance(tag, str) for tag in tags @@ -178,6 +168,7 @@ def decorator( ignore_input=ignore_input, ignore_output=ignore_output, span_type=span_type, + metadata=metadata, ignore_inputs=ignore_inputs, input_formatter=input_formatter, output_formatter=output_formatter, @@ -190,6 +181,7 @@ def decorator( ignore_input=ignore_input, ignore_output=ignore_output, span_type=span_type, + metadata=metadata, ignore_inputs=ignore_inputs, input_formatter=input_formatter, output_formatter=output_formatter, diff --git a/src/lmnr/sdk/evaluations.py b/src/lmnr/sdk/evaluations.py index 7b570065..ea395c35 100644 --- a/src/lmnr/sdk/evaluations.py +++ b/src/lmnr/sdk/evaluations.py @@ -7,7 +7,6 @@ from tqdm import tqdm -from lmnr.opentelemetry_lib.decorators import json_dumps from lmnr.opentelemetry_lib.tracing.instruments import Instruments from lmnr.opentelemetry_lib.tracing.attributes import HUMAN_EVALUATOR_OPTIONS, SPAN_TYPE @@ -30,7 +29,7 @@ SpanType, TraceType, ) -from lmnr.sdk.utils import from_env, is_async +from lmnr.sdk.utils import from_env, is_async, json_dumps DEFAULT_BATCH_SIZE = 5 MAX_EXPORT_BATCH_SIZE = 64 diff --git a/src/lmnr/sdk/laminar.py b/src/lmnr/sdk/laminar.py index 624d441e..a7b3eb6c 100644 --- a/src/lmnr/sdk/laminar.py +++ b/src/lmnr/sdk/laminar.py @@ -4,6 +4,7 @@ from lmnr.opentelemetry_lib import TracerManager from lmnr.opentelemetry_lib.tracing import TracerWrapper, get_current_context from lmnr.opentelemetry_lib.tracing.context import ( + CONTEXT_METADATA_KEY, CONTEXT_SESSION_ID_KEY, CONTEXT_USER_ID_KEY, attach_context, @@ -12,19 +13,16 @@ ) from lmnr.opentelemetry_lib.tracing.instruments import Instruments from lmnr.opentelemetry_lib.tracing.processor import LaminarSpanProcessor +from lmnr.opentelemetry_lib.tracing.span import LaminarSpan from lmnr.opentelemetry_lib.tracing.tracer import get_tracer_with_context from lmnr.opentelemetry_lib.tracing.attributes import ( ASSOCIATION_PROPERTIES, PARENT_SPAN_IDS_PATH, PARENT_SPAN_PATH, - SPAN_IDS_PATH, - SPAN_PATH, USER_ID, Attributes, SPAN_TYPE, ) -from lmnr.opentelemetry_lib import MAX_MANUAL_SPAN_PAYLOAD_SIZE -from lmnr.opentelemetry_lib.decorators import json_dumps from lmnr.sdk.utils import get_otel_env_var from opentelemetry import trace @@ -41,13 +39,9 @@ import re import uuid -from lmnr.opentelemetry_lib.tracing.attributes import ( - SESSION_ID, - SPAN_INPUT, - SPAN_OUTPUT, - TRACE_TYPE, -) -from lmnr.sdk.utils import from_env, is_otel_attribute_value_type +from lmnr.opentelemetry_lib.tracing.attributes import SESSION_ID, TRACE_TYPE + +from lmnr.sdk.utils import from_env, is_otel_attribute_value_type, json_dumps from .log import VerboseColorfulFormatter @@ -62,6 +56,7 @@ class Laminar: __project_api_key: str | None = None __initialized: bool = False __base_http_url: str | None = None + __global_metadata: dict[str, AttributeValue] = {} @classmethod def initialize( @@ -84,6 +79,7 @@ def initialize( otel_logger_level: int = logging.ERROR, session_recording_options: SessionRecordingOptions | None = None, force_http: bool = False, + metadata: dict[str, AttributeValue] | None = None, ): """Initialize Laminar context across the application. This method must be called before using any other Laminar methods or @@ -181,6 +177,7 @@ def initialize( cls.__initialized = True cls.__base_http_url = f"{http_url}:{http_port or 443}" + cls.__global_metadata = metadata or {} if not os.getenv("OTEL_ATTRIBUTE_COUNT_LIMIT"): # each message is at least 2 attributes: role and content, @@ -204,6 +201,11 @@ def initialize( session_recording_options=session_recording_options, force_http=force_http, ) + with get_tracer_with_context() as (tracer, isolated_context): + new_ctx = context_api.set_value( + CONTEXT_METADATA_KEY, cls.__global_metadata, isolated_context + ) + attach_context(new_ctx) cls._initialize_context_from_env() @@ -321,7 +323,10 @@ def start_as_current_span( labels: list[str] | None = None, parent_span_context: LaminarSpanContext | None = None, tags: list[str] | None = None, - ) -> Iterator[Span]: + user_id: str | None = None, + session_id: str | None = None, + metadata: dict[str, AttributeValue] | None = None, + ) -> Iterator[LaminarSpan]: """Start a new span as the current span. Useful for manual instrumentation. If `span_type` is set to `"LLM"`, you should report usage and response attributes manually. See `Laminar.set_span_attributes` @@ -357,6 +362,12 @@ def start_as_current_span( instead. Labels to set for the span. Defaults to None. tags (list[str] | None, optional): tags to set for the span. Defaults to None. + user_id (str | None, optional): user id to set for the trace. + Defaults to None. + session_id (str | None, optional): session id to set for the trace. + Defaults to None. + metadata (dict[str, AttributeValue] | None, optional): metadata to\ + set for the trace. Defaults to None. """ if not cls.is_initialized(): @@ -400,6 +411,14 @@ def start_as_current_span( ctx = trace.set_span_in_context( trace.NonRecordingSpan(span_context), ctx ) + if user_id is not None: + ctx = cls._set_association_prop_context( + user_id=user_id, context=ctx, attach=False + ) + if session_id is not None: + ctx = cls._set_association_prop_context( + session_id=session_id, context=ctx, attach=False + ) ctx_token = context_api.attach(ctx) label_props = {} try: @@ -424,6 +443,12 @@ def start_as_current_span( "Tags will be ignored." ) + association_props = cls._get_association_prop_attributes( + user_id=user_id, + session_id=session_id, + metadata=metadata, + ) + with tracer.start_as_current_span( name, context=ctx, @@ -433,21 +458,13 @@ def start_as_current_span( PARENT_SPAN_IDS_PATH: span_ids_path, **(label_props), **(tag_props), + **(association_props), }, ) as span: + if not isinstance(span, LaminarSpan): + span = LaminarSpan(span) + span.set_input(input) wrapper.push_span_context(span) - if input is not None: - serialized_input = json_dumps(input) - if len(serialized_input) > MAX_MANUAL_SPAN_PAYLOAD_SIZE: - span.set_attribute( - SPAN_INPUT, - "Laminar: input too large to record", - ) - else: - span.set_attribute( - SPAN_INPUT, - serialized_input, - ) yield span wrapper.pop_span_context() @@ -466,7 +483,10 @@ def start_span( parent_span_context: LaminarSpanContext | None = None, labels: dict[str, str] | None = None, tags: list[str] | None = None, - ): + user_id: str | None = None, + session_id: str | None = None, + metadata: dict[str, AttributeValue] | None = None, + ) -> LaminarSpan | Span: """Start a new span. Useful for manual instrumentation. If `span_type` is set to `"LLM"`, you should report usage and response attributes manually. See `Laminar.set_span_attributes` for more @@ -529,6 +549,12 @@ def bar(): Defaults to None. labels (dict[str, str] | None, optional): [DEPRECATED] Use tags\ instead. Labels to set for the span. Defaults to None. + user_id (str | None, optional): user id to set for the trace. + Defaults to None. + session_id (str | None, optional): session id to set for the trace. + Defaults to None. + metadata (dict[str, AttributeValue] | None, optional): metadata to\ + set for the trace. Defaults to None. """ if not cls.is_initialized(): return trace.NonRecordingSpan( @@ -593,6 +619,12 @@ def bar(): + "Tags will be ignored." ) + association_props = cls._get_association_prop_attributes( + user_id=user_id, + session_id=session_id, + metadata=metadata, + ) + span = tracer.start_span( name, context=ctx, @@ -602,21 +634,13 @@ def bar(): PARENT_SPAN_IDS_PATH: span_ids_path, **(label_props), **(tag_props), + **(association_props), }, ) - if input is not None: - serialized_input = json_dumps(input) - if len(serialized_input) > MAX_MANUAL_SPAN_PAYLOAD_SIZE: - span.set_attribute( - SPAN_INPUT, - "Laminar: input too large to record", - ) - else: - span.set_attribute( - SPAN_INPUT, - serialized_input, - ) + if not isinstance(span, LaminarSpan): + span = LaminarSpan(span) + span.set_input(input) return span @classmethod @@ -627,7 +651,7 @@ def use_span( end_on_exit: bool = False, record_exception: bool = True, set_status_on_exception: bool = True, - ) -> Iterator[Span]: + ) -> Iterator[LaminarSpan | Span]: """Use a span as the current span. Useful for manual instrumentation. Fully copies the implementation of `use_span` from opentelemetry.trace @@ -662,7 +686,10 @@ def use_span( # span and trace_id. context_token = context_api.attach(context) try: - yield span + if isinstance(span, LaminarSpan): + yield span + else: + yield LaminarSpan(span) finally: context_api.detach(context_token) wrapper.pop_span_context() @@ -706,7 +733,10 @@ def start_active_span( context: Context | None = None, parent_span_context: LaminarSpanContext | None = None, tags: list[str] | None = None, - ) -> Span: + user_id: str | None = None, + session_id: str | None = None, + metadata: dict[str, AttributeValue] | None = None, + ) -> LaminarSpan | Span: """Start a span and mark it as active within the current context. All spans started after this one will be children of this span. Useful for manual instrumentation. Must be ended manually. @@ -736,17 +766,17 @@ def start_active_span( def foo(): with Laminar.start_as_current_span("foo_inner"): some_function() - + @observe() def bar(): openai_client.chat.completions.create() - + span = Laminar.start_active_span("outer") foo() bar() # IMPORTANT: End the span manually span.end() - + # Results in: # | outer # | | foo @@ -776,6 +806,12 @@ def bar(): Defaults to None. tags (list[str] | None, optional): tags to set for the span. Defaults to None. + user_id (str | None, optional): user id to set for the trace. + Defaults to None. + session_id (str | None, optional): session id to set for the trace. + Defaults to None. + metadata (dict[str, AttributeValue] | None, optional): metadata to\ + set for the trace. Defaults to None. """ span = cls.start_span( name=name, @@ -784,6 +820,9 @@ def bar(): context=context, parent_span_context=parent_span_context, tags=tags, + user_id=user_id, + session_id=session_id, + metadata=metadata, ) if not cls.is_initialized(): return span @@ -791,7 +830,10 @@ def bar(): context = wrapper.push_span_context(span) context_token = context_api.attach(context) span._lmnr_ctx_token = context_token - return span + if isinstance(span, LaminarSpan): + return span + else: + return LaminarSpan(span) @classmethod def set_span_output(cls, output: Any = None): @@ -802,16 +844,10 @@ def set_span_output(cls, output: Any = None): output (Any, optional): output of the span. Will be sent as an\ attribute, so must be json serializable. Defaults to None. """ - span = trace.get_current_span(context=get_current_context()) - if output is not None and span != trace.INVALID_SPAN: - serialized_output = json_dumps(output) - if len(serialized_output) > MAX_MANUAL_SPAN_PAYLOAD_SIZE: - span.set_attribute( - SPAN_OUTPUT, - "Laminar: output too large to record", - ) - else: - span.set_attribute(SPAN_OUTPUT, serialized_output) + span = cls.current_span() + if span is None: + return + span.set_output(output) @classmethod def set_span_attributes( @@ -840,7 +876,7 @@ def set_span_attributes( Args: attributes (dict[Attributes | str, Any]): attributes to set for the span """ - span = trace.get_current_span(context=get_current_context()) + span = cls.current_span() if span == trace.INVALID_SPAN: return @@ -862,16 +898,12 @@ def get_laminar_span_context( if not cls.is_initialized(): return None - span = span or trace.get_current_span(context=get_current_context()) + span = span or cls.current_span() if span == trace.INVALID_SPAN: return None - return LaminarSpanContext( - trace_id=uuid.UUID(int=span.get_span_context().trace_id), - span_id=uuid.UUID(int=span.get_span_context().span_id), - is_remote=span.get_span_context().is_remote, - span_path=span.attributes.get(SPAN_PATH, []), - span_ids_path=span.attributes.get(SPAN_IDS_PATH, []), - ) + if not isinstance(span, LaminarSpan): + span = LaminarSpan(span) + return span.get_laminar_span_context() @classmethod def get_laminar_span_context_dict( @@ -920,6 +952,17 @@ def call_service_b(request, headers): def deserialize_span_context(cls, span_context: dict | str) -> LaminarSpanContext: return LaminarSpanContext.deserialize(span_context) + @classmethod + def current_span(cls, context: Context | None = None) -> LaminarSpan | None: + context = context or get_current_context() + span = trace.get_current_span(context=context) + if span == trace.INVALID_SPAN: + return None + if isinstance(span, LaminarSpan): + return span + else: + return LaminarSpan(span) + @classmethod def flush(cls) -> bool: """Flush the internal tracer. @@ -964,17 +1007,18 @@ def set_span_tags(cls, tags: list[str]): if not cls.is_initialized(): return - span = trace.get_current_span(context=get_current_context()) - if span == trace.INVALID_SPAN: - cls.__logger.warning("No active span to set tags on") + span = cls.current_span() + if span is None: return - if not isinstance(tags, list) or not all(isinstance(tag, str) for tag in tags): - cls.__logger.warning( - "Tags must be a list of strings. Tags will be ignored." - ) + span.set_tags(tags) + + @classmethod + def add_span_tags(cls, tags: list[str]): + """Add tags to the current span.""" + span = cls.current_span() + if span is None: return - # list(set(tags)) to deduplicate tags - span.set_attribute(f"{ASSOCIATION_PROPERTIES}.tags", list(set(tags))) + span.add_tags(tags) @classmethod def set_trace_session_id(cls, session_id: str | None = None): @@ -987,16 +1031,13 @@ def set_trace_session_id(cls, session_id: str | None = None): if not cls.is_initialized(): return - context = get_current_context() - context = context_api.set_value(CONTEXT_SESSION_ID_KEY, session_id, context) - attach_context(context) + context = cls._set_association_prop_context(session_id=session_id, attach=True) - span = trace.get_current_span(context=context) - if span == trace.INVALID_SPAN: + span = cls.current_span(context=context) + if span is None: cls.__logger.warning("No active span to set session id on") return - if session_id is not None: - span.set_attribute(f"{ASSOCIATION_PROPERTIES}.{SESSION_ID}", session_id) + span.set_trace_session_id(session_id) @classmethod def set_trace_user_id(cls, user_id: str | None = None): @@ -1009,16 +1050,13 @@ def set_trace_user_id(cls, user_id: str | None = None): if not cls.is_initialized(): return - context = get_current_context() - context = context_api.set_value(CONTEXT_USER_ID_KEY, user_id, context) - attach_context(context) + context = cls._set_association_prop_context(user_id=user_id, attach=True) - span = trace.get_current_span(context=context) - if span == trace.INVALID_SPAN: + span = cls.current_span(context=context) + if span is None: cls.__logger.warning("No active span to set user id on") return - if user_id is not None: - span.set_attribute(f"{ASSOCIATION_PROPERTIES}.{USER_ID}", user_id) + span.set_trace_user_id(user_id) @classmethod def set_trace_metadata(cls, metadata: dict[str, AttributeValue]): @@ -1030,17 +1068,13 @@ def set_trace_metadata(cls, metadata: dict[str, AttributeValue]): if not cls.is_initialized(): return - span = trace.get_current_span(context=get_current_context()) - if span == trace.INVALID_SPAN: + merged_metadata = {**cls.__global_metadata, **(metadata or {})} + + span = cls.current_span() + if span is None: cls.__logger.warning("No active span to set metadata on") return - for key, value in metadata.items(): - if is_otel_attribute_value_type(value): - span.set_attribute(f"{ASSOCIATION_PROPERTIES}.metadata.{key}", value) - else: - span.set_attribute( - f"{ASSOCIATION_PROPERTIES}.metadata.{key}", json_dumps(value) - ) + span.set_trace_metadata(merged_metadata) @classmethod def get_base_http_url(cls): @@ -1093,3 +1127,49 @@ def _set_trace_type( cls.__logger.warning("No active span to set trace type on") return span.set_attribute(f"{ASSOCIATION_PROPERTIES}.{TRACE_TYPE}", trace_type.value) + + @classmethod + def _get_association_prop_attributes( + cls, + user_id: str | None = None, + session_id: str | None = None, + metadata: dict[str, AttributeValue] | None = None, + ) -> dict[str, AttributeValue]: + association_properties = {} + if user_id is not None: + association_properties[f"{ASSOCIATION_PROPERTIES}.{USER_ID}"] = user_id + if session_id is not None: + association_properties[f"{ASSOCIATION_PROPERTIES}.{SESSION_ID}"] = ( + session_id + ) + merged_metadata = {**cls.__global_metadata, **(metadata or {})} + association_properties.update( + { + f"{ASSOCIATION_PROPERTIES}.metadata.{k}": ( + v if is_otel_attribute_value_type(v) else json_dumps(v) + ) + for k, v in merged_metadata.items() + } + ) + return association_properties + + @classmethod + def _set_association_prop_context( + cls, + user_id: str | None = None, + session_id: str | None = None, + context: Context | None = None, + attach: bool = True, + ) -> Context: + context = context or get_current_context() + if user_id is not None: + context = context_api.set_value(CONTEXT_USER_ID_KEY, user_id, context) + if session_id is not None: + context = context_api.set_value(CONTEXT_SESSION_ID_KEY, session_id, context) + if attach: + attach_context(context) + return context + + @property + def global_metadata(cls) -> dict[str, AttributeValue]: + return cls.__global_metadata diff --git a/src/lmnr/sdk/types.py b/src/lmnr/sdk/types.py index 43b67e9f..12a120ac 100644 --- a/src/lmnr/sdk/types.py +++ b/src/lmnr/sdk/types.py @@ -11,7 +11,7 @@ from typing import Any, Awaitable, Callable, Optional from typing_extensions import TypedDict # compatibility with python < 3.12 -from .utils import serialize +from .utils import serialize, json_dumps DEFAULT_DATAPOINT_MAX_DATA_LENGTH = 16_000_000 # 16MB @@ -107,10 +107,8 @@ class PartialEvaluationDatapoint(BaseModel): def to_dict(self, max_data_length: int = DEFAULT_DATAPOINT_MAX_DATA_LENGTH): serialized_data = serialize(self.data) serialized_target = serialize(self.target) - # TODO: use json_dumps instead of json.dumps once we - # move it to utils so we can avoid circular imports - str_data = json.dumps(serialized_data) - str_target = json.dumps(serialized_target) + str_data = json_dumps(serialized_data) + str_target = json_dumps(serialized_target) try: return { "id": str(self.id), diff --git a/src/lmnr/sdk/utils.py b/src/lmnr/sdk/utils.py index 14b572e2..248250b3 100644 --- a/src/lmnr/sdk/utils.py +++ b/src/lmnr/sdk/utils.py @@ -4,11 +4,16 @@ import enum import inspect import os +import orjson import pydantic import queue import typing import uuid +from lmnr.sdk.log import get_default_logger + +logger = get_default_logger(__name__) + def is_method(func: typing.Callable) -> bool: # inspect.ismethod is True for bound methods only, but in the decorator, @@ -201,3 +206,38 @@ def format_id(id_value: str | int | uuid.UUID) -> str: return id_value else: raise ValueError(f"Invalid ID type: {type(id_value)}") + + +DEFAULT_PLACEHOLDER = {} + + +def default_json(o): + if isinstance(o, pydantic.BaseModel): + return o.model_dump() + + # Handle various sequence types, but not strings or bytes + if isinstance(o, (list, tuple, set, frozenset)): + return list(o) + + try: + return str(o) + except Exception: + logger.debug("Failed to serialize data to JSON, inner type: %s", type(o)) + pass + return DEFAULT_PLACEHOLDER + + +def json_dumps(data: dict) -> str: + try: + return orjson.dumps( + data, + default=default_json, + option=orjson.OPT_SERIALIZE_DATACLASS + | orjson.OPT_SERIALIZE_UUID + | orjson.OPT_UTC_Z + | orjson.OPT_NON_STR_KEYS, + ).decode("utf-8") + except Exception: + # Log the exception and return a placeholder if serialization completely fails + logger.info("Failed to serialize data to JSON, type: %s", type(data)) + return "{}" # Return an empty JSON object as a fallback diff --git a/tests/test_global_metadata.py b/tests/test_global_metadata.py new file mode 100644 index 00000000..ffd8f79c --- /dev/null +++ b/tests/test_global_metadata.py @@ -0,0 +1,149 @@ +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +import pytest +from lmnr.sdk.laminar import Laminar +from lmnr.sdk.decorators import observe + + +@pytest.fixture(autouse=True) +def setup_and_teardown(): + """Reset Laminar state before each test.""" + # Save the current state + original_initialized = Laminar._Laminar__initialized + original_base_http_url = Laminar._Laminar__base_http_url + original_project_api_key = Laminar._Laminar__project_api_key + + # Reset the initialized state for the test + Laminar._Laminar__initialized = False + Laminar._Laminar__base_http_url = None + Laminar._Laminar__project_api_key = None + + yield + + # Restore the original state after test + Laminar._Laminar__initialized = original_initialized + Laminar._Laminar__base_http_url = original_base_http_url + Laminar._Laminar__project_api_key = original_project_api_key + + +def test_global_metadata_no_trace_metadata(span_exporter: InMemorySpanExporter): + Laminar.initialize(project_api_key="test_key", metadata={"foo": "bar"}) + span = Laminar.start_span("test") + span.end() + + spans = span_exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].attributes["lmnr.association.properties.metadata.foo"] == "bar" + assert spans[0].name == "test" + + +def test_global_metadata_no_trace_metadata_start_span_merge( + span_exporter: InMemorySpanExporter, +): + Laminar.initialize( + project_api_key="test_key", metadata={"foo": "bar", "replace": "me"} + ) + span = Laminar.start_span("test", metadata={"baz": "qux", "replace": "new"}) + span.end() + + spans = span_exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].attributes["lmnr.association.properties.metadata.foo"] == "bar" + assert spans[0].attributes["lmnr.association.properties.metadata.baz"] == "qux" + assert spans[0].attributes["lmnr.association.properties.metadata.replace"] == "new" + assert spans[0].name == "test" + + +def test_global_metadata_no_trace_metadata_start_as_current_span_merge( + span_exporter: InMemorySpanExporter, +): + Laminar.initialize( + project_api_key="test_key", metadata={"foo": "bar", "replace": "me"} + ) + with Laminar.start_as_current_span( + "test", metadata={"baz": "qux", "replace": "new"} + ): + pass + + spans = span_exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].attributes["lmnr.association.properties.metadata.foo"] == "bar" + assert spans[0].attributes["lmnr.association.properties.metadata.baz"] == "qux" + assert spans[0].attributes["lmnr.association.properties.metadata.replace"] == "new" + assert spans[0].name == "test" + + +def test_global_metadata_no_trace_metadata_start_active_span_merge( + span_exporter: InMemorySpanExporter, +): + Laminar.initialize( + project_api_key="test_key", metadata={"foo": "bar", "replace": "me"} + ) + span = Laminar.start_active_span("test", metadata={"baz": "qux", "replace": "new"}) + span.end() + + spans = span_exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].name == "test" + assert spans[0].attributes["lmnr.association.properties.metadata.foo"] == "bar" + assert spans[0].attributes["lmnr.association.properties.metadata.baz"] == "qux" + assert spans[0].attributes["lmnr.association.properties.metadata.replace"] == "new" + + +def test_global_metadata_no_trace_metadata_observe(span_exporter: InMemorySpanExporter): + Laminar.initialize(project_api_key="test_key", metadata={"foo": "bar"}) + + @observe() + def test(): + return "test" + + test() + + spans = span_exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].attributes["lmnr.association.properties.metadata.foo"] == "bar" + assert spans[0].name == "test" + + +@pytest.mark.asyncio +async def test_global_metadata_no_trace_metadata_observe_async( + span_exporter: InMemorySpanExporter, +): + Laminar.initialize(project_api_key="test_key", metadata={"foo": "bar"}) + + @observe() + async def test(): + return "test" + + await test() + + spans = span_exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].attributes["lmnr.association.properties.metadata.foo"] == "bar" + assert spans[0].name == "test" + + +def test_global_metadata_no_trace_metadata_two_traces( + span_exporter: InMemorySpanExporter, +): + Laminar.initialize(project_api_key="test_key", metadata={"foo": "bar"}) + + actual_span = Laminar.start_span("test", metadata={"baz": "qux"}) + actual_span.end() + + actual_span2 = Laminar.start_span("test2") + actual_span2.end() + + spans = span_exporter.get_finished_spans() + assert len(spans) == 2 + span = [s for s in spans if s.name == "test"][0] + span2 = [s for s in spans if s.name == "test2"][0] + + assert span.attributes["lmnr.association.properties.metadata.foo"] == "bar" + assert span.attributes["lmnr.association.properties.metadata.baz"] == "qux" + + assert span2.attributes["lmnr.association.properties.metadata.foo"] == "bar" + assert span2.attributes.get("lmnr.association.properties.metadata.baz") is None + + assert span.parent is None or span.parent.span_id == 0 + assert span2.parent is None or span2.parent.span_id == 0 + assert span.get_span_context().trace_id != span2.get_span_context().trace_id diff --git a/tests/test_observe.py b/tests/test_observe.py index b32ef746..2778da60 100644 --- a/tests/test_observe.py +++ b/tests/test_observe.py @@ -1609,3 +1609,38 @@ def test(): os.environ["LMNR_SPAN_CONTEXT"] = old_val else: os.environ.pop("LMNR_SPAN_CONTEXT", None) + + +def test_add_span_tags(span_exporter: InMemorySpanExporter): + @observe(tags=["foo"]) + def test(): + Laminar.add_span_tags(["bar", "baz", "foo"]) + Laminar.add_span_tags(["qux", "bar"]) + + test() + + spans = span_exporter.get_finished_spans() + assert len(spans) == 1 + assert sorted(spans[0].attributes["lmnr.association.properties.tags"]) == [ + "bar", + "baz", + "foo", + "qux", + ] + + +def test_set_span_tags_add_span_tags(span_exporter: InMemorySpanExporter): + @observe(tags=["foo"]) + def test(): + Laminar.set_span_tags(["bar", "baz"]) + Laminar.add_span_tags(["qux", "bar"]) + + test() + + spans = span_exporter.get_finished_spans() + assert len(spans) == 1 + assert sorted(spans[0].attributes["lmnr.association.properties.tags"]) == [ + "bar", + "baz", + "qux", + ] diff --git a/tests/test_tracing.py b/tests/test_tracing.py index c6a805d6..ee0eb489 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -289,6 +289,47 @@ def test_session_id(span_exporter: InMemorySpanExporter): assert spans[0].attributes["lmnr.span.path"] == ("test",) +def test_session_id_start_span(span_exporter: InMemorySpanExporter): + span = Laminar.start_span( + "test", session_id="123", user_id="123", metadata={"foo": "bar"} + ) + span.end() + + spans = span_exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].name == "test" + assert spans[0].attributes["lmnr.association.properties.session_id"] == "123" + assert spans[0].attributes["lmnr.association.properties.metadata.foo"] == "bar" + assert spans[0].attributes["lmnr.association.properties.user_id"] == "123" + + +def test_session_id_start_as_current_span(span_exporter: InMemorySpanExporter): + with Laminar.start_as_current_span( + "test", session_id="123", user_id="123", metadata={"foo": "bar"} + ): + pass + + spans = span_exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].attributes["lmnr.association.properties.session_id"] == "123" + assert spans[0].name == "test" + assert spans[0].attributes["lmnr.association.properties.metadata.foo"] == "bar" + assert spans[0].attributes["lmnr.association.properties.user_id"] == "123" + + +def test_session_id_start_active_span(span_exporter: InMemorySpanExporter): + span = Laminar.start_active_span( + "test", session_id="123", user_id="123", metadata={"foo": "bar"} + ) + span.end() + + spans = span_exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].attributes["lmnr.association.properties.session_id"] == "123" + assert spans[0].name == "test" + assert spans[0].attributes["lmnr.association.properties.metadata.foo"] == "bar" + + def test_session_id_doesnt_leak(span_exporter: InMemorySpanExporter): with Laminar.start_as_current_span("no_session"): pass @@ -974,3 +1015,32 @@ async def nested_level1(): assert level0.attributes["lmnr.span.path"] == ("level0",) assert level1.attributes["lmnr.span.path"] == ("level0", "level1") assert level2.attributes["lmnr.span.path"] == ("level0", "level1", "level2") + + +def test_add_span_tags(span_exporter: InMemorySpanExporter): + with Laminar.start_as_current_span("test"): + Laminar.add_span_tags(["bar", "baz", "foo"]) + Laminar.add_span_tags(["qux", "bar"]) + + spans = span_exporter.get_finished_spans() + assert len(spans) == 1 + assert sorted(spans[0].attributes["lmnr.association.properties.tags"]) == [ + "bar", + "baz", + "foo", + "qux", + ] + + +def test_set_span_tags_add_span_tags(span_exporter: InMemorySpanExporter): + with Laminar.start_as_current_span("test"): + Laminar.set_span_tags(["bar", "baz"]) + Laminar.add_span_tags(["qux", "bar"]) + + spans = span_exporter.get_finished_spans() + assert len(spans) == 1 + assert sorted(spans[0].attributes["lmnr.association.properties.tags"]) == [ + "bar", + "baz", + "qux", + ] diff --git a/tests/test_utils.py b/tests/test_utils.py index 7ae799fc..5a24be9a 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -6,8 +6,7 @@ from typing import Any, Dict, List from pydantic import BaseModel -from lmnr.opentelemetry_lib.decorators import json_dumps -from lmnr.sdk.utils import is_otel_attribute_value_type, format_id +from lmnr.sdk.utils import is_otel_attribute_value_type, format_id, json_dumps class SimplePydanticModel(BaseModel): From 959a29c84e1e2c5111374d7f333b887b9f347f12 Mon Sep 17 00:00:00 2001 From: Din Date: Sun, 7 Dec 2025 18:04:57 +0000 Subject: [PATCH 02/10] properly propagate association properties --- src/lmnr/opentelemetry_lib/__init__.py | 2 - .../opentelemetry_lib/decorators/__init__.py | 86 +++++-- .../opentelemetry_lib/tracing/__init__.py | 2 - src/lmnr/opentelemetry_lib/tracing/context.py | 36 ++- .../opentelemetry_lib/tracing/processor.py | 41 ++++ src/lmnr/opentelemetry_lib/tracing/span.py | 53 ++++- src/lmnr/opentelemetry_lib/tracing/tracer.py | 7 +- src/lmnr/sdk/decorators.py | 3 + src/lmnr/sdk/laminar.py | 186 +++++++++++---- tests/test_context_propagation.py | 213 ++++++++++++++++++ 10 files changed, 552 insertions(+), 77 deletions(-) create mode 100644 tests/test_context_propagation.py diff --git a/src/lmnr/opentelemetry_lib/__init__.py b/src/lmnr/opentelemetry_lib/__init__.py index 52455449..602b3df2 100644 --- a/src/lmnr/opentelemetry_lib/__init__.py +++ b/src/lmnr/opentelemetry_lib/__init__.py @@ -8,8 +8,6 @@ from lmnr.opentelemetry_lib.tracing import TracerWrapper from lmnr.sdk.types import SessionRecordingOptions -MAX_MANUAL_SPAN_PAYLOAD_SIZE = 1024 * 1024 * 10 # 10MB - class TracerManager: __tracer_wrapper: TracerWrapper diff --git a/src/lmnr/opentelemetry_lib/decorators/__init__.py b/src/lmnr/opentelemetry_lib/decorators/__init__.py index ba1cb0fc..2f5bf854 100644 --- a/src/lmnr/opentelemetry_lib/decorators/__init__.py +++ b/src/lmnr/opentelemetry_lib/decorators/__init__.py @@ -8,9 +8,11 @@ from lmnr.opentelemetry_lib.tracing.context import ( CONTEXT_METADATA_KEY, CONTEXT_SESSION_ID_KEY, + CONTEXT_TRACE_TYPE_KEY, CONTEXT_USER_ID_KEY, attach_context, detach_context, + get_current_context, get_event_attributes_from_context, ) from lmnr.opentelemetry_lib.tracing.span import LaminarSpan @@ -20,10 +22,14 @@ ASSOCIATION_PROPERTIES, METADATA, SPAN_TYPE, + USER_ID, + SESSION_ID, + TRACE_TYPE, ) from lmnr.opentelemetry_lib.tracing import TracerWrapper from lmnr.sdk.log import get_default_logger from lmnr.sdk.utils import is_otel_attribute_value_type, json_dumps +from opentelemetry.context import set_value logger = get_default_logger(__name__) @@ -47,7 +53,10 @@ def _setup_span( ) ctx_metadata = context_api.get_value(CONTEXT_METADATA_KEY, isolated_context) - merged_metadata = {**(ctx_metadata or {}), **(metadata or {})} + merged_metadata = { + **(ctx_metadata or {}), + **(metadata or {}), + } for key, value in merged_metadata.items(): span.set_attribute( f"{ASSOCIATION_PROPERTIES}.{METADATA}.{key}", @@ -136,6 +145,51 @@ def _cleanup_span(span: Span, wrapper: TracerWrapper): wrapper.pop_span_context() +def _set_association_props_in_context(span: Span): + """Set association properties from span in context before push_span_context. + + Returns the token that needs to be detached when the span ends. + """ + if not isinstance(span, LaminarSpan): + return None + + props = span.laminar_association_properties + user_id_key = f"{ASSOCIATION_PROPERTIES}.{USER_ID}" + session_id_key = f"{ASSOCIATION_PROPERTIES}.{SESSION_ID}" + trace_type_key = f"{ASSOCIATION_PROPERTIES}.{TRACE_TYPE}" + + # Extract values from props + extracted_user_id = props.get(user_id_key) + extracted_session_id = props.get(session_id_key) + extracted_trace_type = props.get(trace_type_key) + + # Extract metadata from props (keys without ASSOCIATION_PROPERTIES prefix) + metadata_dict = {} + for key, value in props.items(): + if not key.startswith(f"{ASSOCIATION_PROPERTIES}."): + metadata_dict[key] = value + + # Set context with association props + current_ctx = get_current_context() + ctx_with_props = current_ctx + if extracted_user_id: + ctx_with_props = set_value( + CONTEXT_USER_ID_KEY, extracted_user_id, ctx_with_props + ) + if extracted_session_id: + ctx_with_props = set_value( + CONTEXT_SESSION_ID_KEY, extracted_session_id, ctx_with_props + ) + if extracted_trace_type: + ctx_with_props = set_value( + CONTEXT_TRACE_TYPE_KEY, extracted_trace_type, ctx_with_props + ) + if metadata_dict: + ctx_with_props = set_value(CONTEXT_METADATA_KEY, metadata_dict, ctx_with_props) + + return attach_context(ctx_with_props) + + def observe_base( *, name: str | None = None, @@ -165,15 +219,14 @@ def wrap(*args, **kwargs): preserve_global_context, metadata, ) + + # Set association props in context before push_span_context + # so child spans inherit them + assoc_props_token = _set_association_props_in_context(span) + if assoc_props_token and isinstance(span, LaminarSpan): + span._lmnr_assoc_props_token = assoc_props_token + new_context = wrapper.push_span_context(span) - if session_id := association_properties.get("session_id"): - new_context = context_api.set_value( - CONTEXT_SESSION_ID_KEY, session_id, new_context - ) - if user_id := association_properties.get("user_id"): - new_context = context_api.set_value( - CONTEXT_USER_ID_KEY, user_id, new_context - ) # Some auto-instrumentations are not under our control, so they # don't have access to our isolated context. We attach the context # to the OTEL global context, so that spans know their parent @@ -248,15 +301,14 @@ async def wrap(*args, **kwargs): preserve_global_context, metadata, ) + + # Set association props in context before push_span_context + # so child spans inherit them + assoc_props_token = _set_association_props_in_context(span) + if assoc_props_token and isinstance(span, LaminarSpan): + span._lmnr_assoc_props_token = assoc_props_token + new_context = wrapper.push_span_context(span) - if session_id := association_properties.get("session_id"): - new_context = context_api.set_value( - CONTEXT_SESSION_ID_KEY, session_id, new_context - ) - if user_id := association_properties.get("user_id"): - new_context = context_api.set_value( - CONTEXT_USER_ID_KEY, user_id, new_context - ) # Some auto-instrumentations are not under our control, so they # don't have access to our isolated context. We attach the context # to the OTEL global context, so that spans know their parent diff --git a/src/lmnr/opentelemetry_lib/tracing/__init__.py b/src/lmnr/opentelemetry_lib/tracing/__init__.py index e53d8259..ad5aeeb8 100644 --- a/src/lmnr/opentelemetry_lib/tracing/__init__.py +++ b/src/lmnr/opentelemetry_lib/tracing/__init__.py @@ -188,9 +188,7 @@ 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) - # Store the token for later detachment - tokens are much lighter than contexts ctx_push_span_context(new_context) - return new_context def pop_span_context(self) -> None: diff --git a/src/lmnr/opentelemetry_lib/tracing/context.py b/src/lmnr/opentelemetry_lib/tracing/context.py index 365d3f16..93f97afd 100644 --- a/src/lmnr/opentelemetry_lib/tracing/context.py +++ b/src/lmnr/opentelemetry_lib/tracing/context.py @@ -2,9 +2,16 @@ 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 METADATA, SESSION_ID, USER_ID +from typing import Any +from opentelemetry.context import Context, Token, create_key, get_value, set_value + +from lmnr.opentelemetry_lib.tracing.attributes import ( + METADATA, + SESSION_ID, + TRACE_TYPE, + USER_ID, +) +from lmnr.sdk.types import TraceType class _IsolatedRuntimeContext(ABC): @@ -114,6 +121,7 @@ def detach_context(token: Token[Context]) -> None: CONTEXT_USER_ID_KEY = create_key(f"lmnr.{USER_ID}") CONTEXT_SESSION_ID_KEY = create_key(f"lmnr.{SESSION_ID}") CONTEXT_METADATA_KEY = create_key(f"lmnr.{METADATA}") +CONTEXT_TRACE_TYPE_KEY = create_key(f"lmnr.{TRACE_TYPE}") def get_event_attributes_from_context(context: Context | None = None) -> dict[str, str]: @@ -127,6 +135,28 @@ def get_event_attributes_from_context(context: Context | None = None) -> dict[st return attributes +def set_association_prop_context( + user_id: str | None = None, + session_id: str | None = None, + trace_type: TraceType | None = None, + context: Context | None = None, + metadata: dict[str, Any] | None = None, + attach: bool = True, +) -> Context: + context = context or get_current_context() + if user_id is not None: + context = set_value(CONTEXT_USER_ID_KEY, user_id, context) + if session_id is not None: + context = set_value(CONTEXT_SESSION_ID_KEY, session_id, context) + if trace_type is not None: + context = set_value(CONTEXT_TRACE_TYPE_KEY, trace_type.value, context) + if metadata is not None: + context = set_value(CONTEXT_METADATA_KEY, metadata, context) + if attach: + attach_context(context) + return context + + def pop_span_context() -> None: """Pop the current span context from the stack.""" current_stack = get_token_stack().copy() diff --git a/src/lmnr/opentelemetry_lib/tracing/processor.py b/src/lmnr/opentelemetry_lib/tracing/processor.py index 7d0826ca..bf7a300a 100644 --- a/src/lmnr/opentelemetry_lib/tracing/processor.py +++ b/src/lmnr/opentelemetry_lib/tracing/processor.py @@ -12,16 +12,27 @@ from opentelemetry.context import Context, get_value from lmnr.opentelemetry_lib.tracing.attributes import ( + ASSOCIATION_PROPERTIES, PARENT_SPAN_IDS_PATH, PARENT_SPAN_PATH, + SESSION_ID, SPAN_IDS_PATH, SPAN_INSTRUMENTATION_SOURCE, SPAN_LANGUAGE_VERSION, SPAN_PATH, SPAN_SDK_VERSION, + TRACE_TYPE, + USER_ID, +) +from lmnr.opentelemetry_lib.tracing.context import ( + CONTEXT_METADATA_KEY, + CONTEXT_SESSION_ID_KEY, + CONTEXT_TRACE_TYPE_KEY, + CONTEXT_USER_ID_KEY, ) from lmnr.opentelemetry_lib.tracing.exporter import LaminarSpanExporter from lmnr.sdk.log import get_default_logger +from lmnr.sdk.utils import is_otel_attribute_value_type, json_dumps from lmnr.version import PYTHON_VERSION, __version__ @@ -89,6 +100,36 @@ def on_start(self, span: Span, parent_context: Context | None = None): span.set_attribute(SPAN_SDK_VERSION, __version__) span.set_attribute(SPAN_LANGUAGE_VERSION, f"python@{PYTHON_VERSION}") + if parent_context: + trace_type = get_value(CONTEXT_TRACE_TYPE_KEY, parent_context) + if trace_type: + span.set_attribute(f"{ASSOCIATION_PROPERTIES}.{TRACE_TYPE}", trace_type) + user_id = get_value(CONTEXT_USER_ID_KEY, parent_context) + if user_id: + span.set_attribute(f"{ASSOCIATION_PROPERTIES}.{USER_ID}", user_id) + session_id = get_value(CONTEXT_SESSION_ID_KEY, parent_context) + if session_id: + span.set_attribute(f"{ASSOCIATION_PROPERTIES}.{SESSION_ID}", session_id) + ctx_metadata = get_value(CONTEXT_METADATA_KEY, parent_context) + if ctx_metadata and isinstance(ctx_metadata, dict): + span_metadata = {} + if hasattr(span, "attributes") and hasattr(span.attributes, "items"): + for key, value in span.attributes.items(): + if key.startswith(f"{ASSOCIATION_PROPERTIES}.metadata."): + span_metadata[ + key.replace(f"{ASSOCIATION_PROPERTIES}.metadata.", "") + ] = value + + for key, value in {**ctx_metadata, **span_metadata}.items(): + span.set_attribute( + f"{ASSOCIATION_PROPERTIES}.metadata.{key}", + ( + value + if is_otel_attribute_value_type(value) + else json_dumps(value) + ), + ) + if span.name == "LangGraph.workflow": graph_context = get_value("lmnr.langgraph.graph") or {} for key, value in graph_context.items(): diff --git a/src/lmnr/opentelemetry_lib/tracing/span.py b/src/lmnr/opentelemetry_lib/tracing/span.py index 49b337e9..800726ac 100644 --- a/src/lmnr/opentelemetry_lib/tracing/span.py +++ b/src/lmnr/opentelemetry_lib/tracing/span.py @@ -1,6 +1,7 @@ from logging import Logger from inspect import Traceback from typing import Any, Literal +import orjson import uuid from opentelemetry.sdk.resources import Resource @@ -13,7 +14,6 @@ from opentelemetry.util.types import AttributeValue from opentelemetry.context import detach -from lmnr.opentelemetry_lib import MAX_MANUAL_SPAN_PAYLOAD_SIZE from lmnr.opentelemetry_lib.tracing.attributes import ( ASSOCIATION_PROPERTIES, SPAN_IDS_PATH, @@ -22,12 +22,15 @@ SPAN_PATH, ) from lmnr.opentelemetry_lib.tracing.context import ( + detach_context, pop_span_context, ) from lmnr.sdk.log import get_default_logger from lmnr.sdk.types import LaminarSpanContext from lmnr.sdk.utils import is_otel_attribute_value_type, json_dumps +MAX_MANUAL_SPAN_PAYLOAD_SIZE = 1024 * 1024 * 10 # 10MB + class LaminarSpanInterfaceMixin: """Mixin providing Laminar-specific span methods and properties.""" @@ -165,18 +168,48 @@ def set_tags(self, tags: list[str]) -> None: @property def tags(self) -> list[str]: - self.logger.debug( - "[LaminarSpan.tags] WARNING. Reading current span attributes to get tags. " - "This works in default OTel SDK, but is not guaranteed by the " - "OpenTelemetry API specification. If you are using a custom SDK, " - "you may need to implement your own way to get the tags. In such case, " - "this function returns an empty list." - ) + if not hasattr(self.span, "attributes"): + self.logger.debug( + "[LaminarSpan.tags] WARNING. Current span does not have attributes object. " + "Perhaps, the span was created with a custom OTel SDK. Returning an empty list." + "Help: OpenTelemetry API does not guarantee reading attributes from a span, but OTel SDK ", + "allows it by default. Laminar SDK allows to read attributes too.", + ) + return [] try: return list(self.span.attributes.get(f"{ASSOCIATION_PROPERTIES}.tags", [])) except Exception: return [] + @property + def laminar_association_properties(self) -> dict[str, Any]: + if not hasattr(self.span, "attributes"): + self.logger.debug( + "[LaminarSpan.laminar_association_properties] WARNING. Current span " + "does not have attributes object. Perhaps, the span was created with a " + "custom OTel SDK. Returning an empty dictionary." + "Help: OpenTelemetry API does not guarantee reading attributes from a span, but OTel SDK " + "allows it by default. Laminar SDK allows to read attributes too.", + ) + return {} + try: + values = {} + for key, value in self.span.attributes.items(): + if key.startswith(f"{ASSOCIATION_PROPERTIES}."): + if key.startswith(f"{ASSOCIATION_PROPERTIES}.metadata."): + meta_key = key.replace( + f"{ASSOCIATION_PROPERTIES}.metadata.", "" + ) + try: + values[meta_key] = orjson.loads(value) + except Exception: + values[meta_key] = value + else: + values[key] = value + return values + except Exception: + return {} + class SpanDelegationMixin: """Mixin providing delegation to the wrapped SDK span for standard OpenTelemetry methods.""" @@ -317,6 +350,10 @@ def end(self, end_time: int | None = None) -> None: self._popped = True except Exception: pass + if hasattr(self, "_lmnr_isolated_ctx_token"): + detach_context(self._lmnr_isolated_ctx_token) + if hasattr(self, "_lmnr_assoc_props_token") and self._lmnr_assoc_props_token: + detach_context(self._lmnr_assoc_props_token) def __enter__(self) -> "LaminarSpan": return self diff --git a/src/lmnr/opentelemetry_lib/tracing/tracer.py b/src/lmnr/opentelemetry_lib/tracing/tracer.py index 433c726c..e25bb8fc 100644 --- a/src/lmnr/opentelemetry_lib/tracing/tracer.py +++ b/src/lmnr/opentelemetry_lib/tracing/tracer.py @@ -48,5 +48,10 @@ def start_span(self, *args, **kwargs) -> trace.Span: @contextmanager def start_as_current_span(self, *args, **kwargs) -> Iterator[trace.Span]: + wrapper = TracerWrapper() with self._instance.start_as_current_span(*args, **kwargs) as span: - yield LaminarSpan(span) + wrapper.push_span_context(span) + try: + yield LaminarSpan(span) + finally: + wrapper.pop_span_context() diff --git a/src/lmnr/sdk/decorators.py b/src/lmnr/sdk/decorators.py index ee3aa5c8..a4036eee 100644 --- a/src/lmnr/sdk/decorators.py +++ b/src/lmnr/sdk/decorators.py @@ -9,6 +9,7 @@ from lmnr.opentelemetry_lib.tracing.attributes import SESSION_ID from lmnr.sdk.log import get_default_logger +from lmnr.sdk.types import TraceType from .utils import is_async @@ -135,6 +136,8 @@ def decorator( association_properties["session_id"] = session_id if user_id is not None: association_properties["user_id"] = user_id + if span_type in ["EVALUATION", "EXECUTOR", "EVALUATOR"]: + association_properties["trace_type"] = TraceType.EVALUATION.value if tags is not None: if not isinstance(tags, list) or not all( isinstance(tag, str) for tag in tags diff --git a/src/lmnr/sdk/laminar.py b/src/lmnr/sdk/laminar.py index a7b3eb6c..7a31cf45 100644 --- a/src/lmnr/sdk/laminar.py +++ b/src/lmnr/sdk/laminar.py @@ -1,16 +1,20 @@ from contextlib import contextmanager -from contextvars import Context +from contextvars import Context, Token import warnings from lmnr.opentelemetry_lib import TracerManager from lmnr.opentelemetry_lib.tracing import TracerWrapper, get_current_context from lmnr.opentelemetry_lib.tracing.context import ( CONTEXT_METADATA_KEY, CONTEXT_SESSION_ID_KEY, + CONTEXT_TRACE_TYPE_KEY, CONTEXT_USER_ID_KEY, attach_context, + detach_context, get_event_attributes_from_context, push_span_context, + set_association_prop_context, ) +from opentelemetry.context import get_value, set_value from lmnr.opentelemetry_lib.tracing.instruments import Instruments from lmnr.opentelemetry_lib.tracing.processor import LaminarSpanProcessor from lmnr.opentelemetry_lib.tracing.span import LaminarSpan @@ -52,6 +56,54 @@ ) +def _set_span_association_props_in_context(span: Span) -> Token[Context] | None: + """Set association properties from span in context. + + Extracts association properties from span attributes and sets them in the + isolated context so child spans can inherit them. + + Returns the token that needs to be stored on the span for cleanup. + """ + if not isinstance(span, LaminarSpan): + return None + + props = span.laminar_association_properties + user_id_key = f"{ASSOCIATION_PROPERTIES}.{USER_ID}" + session_id_key = f"{ASSOCIATION_PROPERTIES}.{SESSION_ID}" + trace_type_key = f"{ASSOCIATION_PROPERTIES}.{TRACE_TYPE}" + + # Extract values from props + extracted_user_id = props.get(user_id_key) + extracted_session_id = props.get(session_id_key) + extracted_trace_type = props.get(trace_type_key) + + # Extract metadata from props (keys without ASSOCIATION_PROPERTIES prefix) + metadata_dict = {} + for key, value in props.items(): + if not key.startswith(f"{ASSOCIATION_PROPERTIES}."): + metadata_dict[key] = value + + # Set context with association props + current_ctx = get_current_context() + ctx_with_props = current_ctx + if extracted_user_id: + ctx_with_props = set_value( + CONTEXT_USER_ID_KEY, extracted_user_id, ctx_with_props + ) + if extracted_session_id: + ctx_with_props = set_value( + CONTEXT_SESSION_ID_KEY, extracted_session_id, ctx_with_props + ) + if extracted_trace_type: + ctx_with_props = set_value( + CONTEXT_TRACE_TYPE_KEY, extracted_trace_type, ctx_with_props + ) + if metadata_dict: + ctx_with_props = set_value(CONTEXT_METADATA_KEY, metadata_dict, ctx_with_props) + + return attach_context(ctx_with_props) + + class Laminar: __project_api_key: str | None = None __initialized: bool = False @@ -302,10 +354,9 @@ def event( current_span = trace.get_current_span(context=get_current_context()) if current_span == trace.INVALID_SPAN: - with cls.start_as_current_span(name) as span: - span.add_event( - name, {**(attributes or {}), **extra_attributes}, timestamp - ) + span = cls.start_span(name) + span.add_event(name, {**(attributes or {}), **extra_attributes}, timestamp) + span.end() return current_span.add_event( @@ -411,15 +462,32 @@ def start_as_current_span( ctx = trace.set_span_in_context( trace.NonRecordingSpan(span_context), ctx ) - if user_id is not None: - ctx = cls._set_association_prop_context( - user_id=user_id, context=ctx, attach=False - ) - if session_id is not None: - ctx = cls._set_association_prop_context( - session_id=session_id, context=ctx, attach=False - ) + trace_type = None + if span_type in ["EVALUATION", "EXECUTOR", "EVALUATOR"]: + trace_type = TraceType.EVALUATION + + # Merge metadata: context (inherited) + global + explicit (explicit wins) + # Get metadata from context if it exists + ctx_metadata = get_value(CONTEXT_METADATA_KEY, ctx) or {} + # Merge: context metadata + global metadata + explicit metadata + # Later keys override earlier ones + merged_metadata = { + **ctx_metadata, + **cls.__global_metadata, + **(metadata or {}), + } + + ctx = set_association_prop_context( + trace_type=trace_type, + user_id=user_id, + session_id=session_id, + metadata=merged_metadata if merged_metadata else None, + context=ctx, + # we need a token separately, so we manually attach the context + attach=False, + ) ctx_token = context_api.attach(ctx) + isolated_context_token = attach_context(ctx) label_props = {} try: if labels: @@ -443,12 +511,6 @@ def start_as_current_span( "Tags will be ignored." ) - association_props = cls._get_association_prop_attributes( - user_id=user_id, - session_id=session_id, - metadata=metadata, - ) - with tracer.start_as_current_span( name, context=ctx, @@ -458,17 +520,17 @@ def start_as_current_span( PARENT_SPAN_IDS_PATH: span_ids_path, **(label_props), **(tag_props), - **(association_props), + # Association properties are attached to context above + # and the relevant attributes are populated in the processor }, ) as span: if not isinstance(span, LaminarSpan): span = LaminarSpan(span) span.set_input(input) - wrapper.push_span_context(span) yield span - wrapper.pop_span_context() try: + detach_context(isolated_context_token) context_api.detach(ctx_token) except Exception: pass @@ -594,6 +656,18 @@ def bar(): ctx = trace.set_span_in_context( trace.NonRecordingSpan(span_context), ctx ) + + # Get association props from context (fallback values) + ctx_user_id = ( + get_value(CONTEXT_USER_ID_KEY, ctx) if user_id is None else None + ) + ctx_session_id = ( + get_value(CONTEXT_SESSION_ID_KEY, ctx) if session_id is None else None + ) + ctx_metadata = ( + get_value(CONTEXT_METADATA_KEY, ctx) if metadata is None else None + ) + label_props = {} try: if labels: @@ -619,10 +693,25 @@ def bar(): + "Tags will be ignored." ) + trace_type = None + if span_type in ["EVALUATION", "EXECUTOR", "EVALUATOR"]: + trace_type = TraceType.EVALUATION + # Get trace_type from context if not set explicitly + ctx_trace_type = ( + get_value(CONTEXT_TRACE_TYPE_KEY, ctx) if trace_type is None else None + ) + if ctx_trace_type: + try: + trace_type = TraceType(ctx_trace_type) + except (ValueError, TypeError): + pass + + # Build association_props using explicit params or context fallbacks association_props = cls._get_association_prop_attributes( - user_id=user_id, - session_id=session_id, - metadata=metadata, + user_id=user_id or ctx_user_id, + session_id=session_id or ctx_session_id, + metadata=metadata or ctx_metadata, + trace_type=trace_type, ) span = tracer.start_span( @@ -679,11 +768,18 @@ def use_span( wrapper = TracerWrapper() try: + # Set association props in context before push_span_context + # so child spans inherit them + assoc_props_token = _set_span_association_props_in_context(span) + if assoc_props_token and isinstance(span, LaminarSpan): + span._lmnr_assoc_props_token = assoc_props_token + context = wrapper.push_span_context(span) # Some auto-instrumentations are not under our control, so they # don't have access to our isolated context. We attach the context # to the OTEL global context, so that spans know their parent # span and trace_id. + isolated_context_token = attach_context(context) context_token = context_api.attach(context) try: if isinstance(span, LaminarSpan): @@ -692,6 +788,7 @@ def use_span( yield LaminarSpan(span) finally: context_api.detach(context_token) + detach_context(isolated_context_token) wrapper.pop_span_context() # Record only exceptions that inherit Exception class but not BaseException, because @@ -827,9 +924,18 @@ def bar(): if not cls.is_initialized(): return span wrapper = TracerWrapper() + + # Set association props in context before push_span_context + # so child spans inherit them + assoc_props_token = _set_span_association_props_in_context(span) + if assoc_props_token and isinstance(span, LaminarSpan): + span._lmnr_assoc_props_token = assoc_props_token + context = wrapper.push_span_context(span) context_token = context_api.attach(context) + isolated_context_token = attach_context(context) span._lmnr_ctx_token = context_token + span._lmnr_isolated_ctx_token = isolated_context_token if isinstance(span, LaminarSpan): return span else: @@ -1031,7 +1137,7 @@ def set_trace_session_id(cls, session_id: str | None = None): if not cls.is_initialized(): return - context = cls._set_association_prop_context(session_id=session_id, attach=True) + context = set_association_prop_context(session_id=session_id, attach=True) span = cls.current_span(context=context) if span is None: @@ -1050,7 +1156,7 @@ def set_trace_user_id(cls, user_id: str | None = None): if not cls.is_initialized(): return - context = cls._set_association_prop_context(user_id=user_id, attach=True) + context = set_association_prop_context(user_id=user_id, attach=True) span = cls.current_span(context=context) if span is None: @@ -1133,6 +1239,7 @@ def _get_association_prop_attributes( cls, user_id: str | None = None, session_id: str | None = None, + trace_type: TraceType | None = None, metadata: dict[str, AttributeValue] | None = None, ) -> dict[str, AttributeValue]: association_properties = {} @@ -1142,6 +1249,14 @@ def _get_association_prop_attributes( association_properties[f"{ASSOCIATION_PROPERTIES}.{SESSION_ID}"] = ( session_id ) + if trace_type is not None: + trace_type_val = ( + trace_type.value if isinstance(trace_type, TraceType) else trace_type + ) + association_properties[f"{ASSOCIATION_PROPERTIES}.{TRACE_TYPE}"] = ( + trace_type_val + ) + merged_metadata = {**cls.__global_metadata, **(metadata or {})} association_properties.update( { @@ -1153,23 +1268,6 @@ def _get_association_prop_attributes( ) return association_properties - @classmethod - def _set_association_prop_context( - cls, - user_id: str | None = None, - session_id: str | None = None, - context: Context | None = None, - attach: bool = True, - ) -> Context: - context = context or get_current_context() - if user_id is not None: - context = context_api.set_value(CONTEXT_USER_ID_KEY, user_id, context) - if session_id is not None: - context = context_api.set_value(CONTEXT_SESSION_ID_KEY, session_id, context) - if attach: - attach_context(context) - return context - @property def global_metadata(cls) -> dict[str, AttributeValue]: return cls.__global_metadata diff --git a/tests/test_context_propagation.py b/tests/test_context_propagation.py new file mode 100644 index 00000000..229a20a1 --- /dev/null +++ b/tests/test_context_propagation.py @@ -0,0 +1,213 @@ +from opentelemetry.sdk.trace import ReadableSpan +from lmnr import Laminar, observe +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + + +def _assert_assoiation_properties( + span: ReadableSpan, + user_id: str, + session_id: str, + metadata: dict[str, str], + trace_type: str, +): + assert span.attributes["lmnr.association.properties.user_id"] == user_id + assert span.attributes["lmnr.association.properties.session_id"] == session_id + assert span.attributes["lmnr.association.properties.trace_type"] == trace_type + for key, value in metadata.items(): + assert span.attributes[f"lmnr.association.properties.metadata.{key}"] == value + + +def test_ctx_prop_parent_sc_child_s(span_exporter: InMemorySpanExporter): + with Laminar.start_as_current_span( + "parent", + user_id="user_id", + session_id="session_id", + span_type="EVALUATION", + metadata={"foo": "bar"}, + ): + inner_span = Laminar.start_span("child") + inner_span.end() + + spans = span_exporter.get_finished_spans() + assert len(spans) == 2 + parent_span = [s for s in spans if s.name == "parent"][0] + child_span = [s for s in spans if s.name == "child"][0] + _assert_assoiation_properties( + parent_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" + ) + _assert_assoiation_properties( + child_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" + ) + + +def test_ctx_prop_parent_sc_child_sc(span_exporter: InMemorySpanExporter): + with Laminar.start_as_current_span( + "parent", + user_id="user_id", + session_id="session_id", + span_type="EVALUATION", + metadata={"foo": "bar"}, + ): + with Laminar.start_as_current_span("child"): + pass + + spans = span_exporter.get_finished_spans() + assert len(spans) == 2 + parent_span = [s for s in spans if s.name == "parent"][0] + child_span = [s for s in spans if s.name == "child"][0] + _assert_assoiation_properties( + parent_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" + ) + _assert_assoiation_properties( + child_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" + ) + + +def test_ctx_prop_parent_sc_child_obs(span_exporter: InMemorySpanExporter): + @observe() + def child(): + pass + + with Laminar.start_as_current_span( + "parent", + user_id="user_id", + session_id="session_id", + span_type="EVALUATION", + metadata={"foo": "bar"}, + ): + child() + + spans = span_exporter.get_finished_spans() + assert len(spans) == 2 + parent_span = [s for s in spans if s.name == "parent"][0] + child_span = [s for s in spans if s.name == "child"][0] + _assert_assoiation_properties( + parent_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" + ) + _assert_assoiation_properties( + child_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" + ) + + +def test_ctx_prop_parent_sa_child_s(span_exporter: InMemorySpanExporter): + span = Laminar.start_active_span( + "parent", + user_id="user_id", + session_id="session_id", + span_type="EVALUATION", + metadata={"foo": "bar"}, + ) + inner_span = Laminar.start_span("child") + inner_span.end() + span.end() + + spans = span_exporter.get_finished_spans() + assert len(spans) == 2 + parent_span = [s for s in spans if s.name == "parent"][0] + child_span = [s for s in spans if s.name == "child"][0] + _assert_assoiation_properties( + parent_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" + ) + _assert_assoiation_properties( + child_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" + ) + + +def test_ctx_prop_parent_sa_child_sc(span_exporter: InMemorySpanExporter): + span = Laminar.start_active_span( + "parent", + user_id="user_id", + session_id="session_id", + span_type="EVALUATION", + metadata={"foo": "bar"}, + ) + inner_span = Laminar.start_span("child") + inner_span.end() + span.end() + + spans = span_exporter.get_finished_spans() + assert len(spans) == 2 + parent_span = [s for s in spans if s.name == "parent"][0] + child_span = [s for s in spans if s.name == "child"][0] + _assert_assoiation_properties( + parent_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" + ) + _assert_assoiation_properties( + child_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" + ) + + +def test_ctx_prop_parent_sa_child_obs(span_exporter: InMemorySpanExporter): + @observe() + def child(): + pass + + span = Laminar.start_active_span( + "parent", + user_id="user_id", + session_id="session_id", + span_type="EVALUATION", + metadata={"foo": "bar"}, + ) + child() + span.end() + + spans = span_exporter.get_finished_spans() + assert len(spans) == 2 + parent_span = [s for s in spans if s.name == "parent"][0] + child_span = [s for s in spans if s.name == "child"][0] + _assert_assoiation_properties( + parent_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" + ) + _assert_assoiation_properties( + child_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" + ) + + +def test_ctx_prop_parent_obs_child_s(span_exporter: InMemorySpanExporter): + @observe( + user_id="user_id", + session_id="session_id", + span_type="EVALUATION", + metadata={"foo": "bar"}, + ) + def parent(): + inner_span = Laminar.start_span("child") + inner_span.end() + + parent() + + spans = span_exporter.get_finished_spans() + assert len(spans) == 2 + parent_span = [s for s in spans if s.name == "parent"][0] + child_span = [s for s in spans if s.name == "child"][0] + _assert_assoiation_properties( + parent_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" + ) + _assert_assoiation_properties( + child_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" + ) + + +def test_ctx_prop_parent_use_span(span_exporter: InMemorySpanExporter): + span = Laminar.start_span( + "parent", + user_id="user_id", + session_id="session_id", + span_type="EVALUATION", + metadata={"foo": "bar"}, + ) + with Laminar.use_span(span, end_on_exit=True): + inner_span = Laminar.start_span("child") + inner_span.end() + + spans = span_exporter.get_finished_spans() + assert len(spans) == 2 + parent_span = [s for s in spans if s.name == "parent"][0] + child_span = [s for s in spans if s.name == "child"][0] + _assert_assoiation_properties( + parent_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" + ) + _assert_assoiation_properties( + child_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" + ) From 8b738bfbb191f46d4118f506150f864103e9ac4e Mon Sep 17 00:00:00 2001 From: Din Date: Mon, 8 Dec 2025 15:12:29 +0000 Subject: [PATCH 03/10] pass other association props in LaminarSpanContext --- src/lmnr/opentelemetry_lib/tracing/span.py | 31 ++- src/lmnr/sdk/laminar.py | 301 +++++++++++++++------ src/lmnr/sdk/types.py | 13 +- tests/test_context_propagation.py | 76 ++++-- 4 files changed, 311 insertions(+), 110 deletions(-) diff --git a/src/lmnr/opentelemetry_lib/tracing/span.py b/src/lmnr/opentelemetry_lib/tracing/span.py index 800726ac..b3a9b88f 100644 --- a/src/lmnr/opentelemetry_lib/tracing/span.py +++ b/src/lmnr/opentelemetry_lib/tracing/span.py @@ -16,10 +16,14 @@ from lmnr.opentelemetry_lib.tracing.attributes import ( ASSOCIATION_PROPERTIES, + METADATA, + SESSION_ID, SPAN_IDS_PATH, SPAN_INPUT, SPAN_OUTPUT, SPAN_PATH, + TRACE_TYPE, + USER_ID, ) from lmnr.opentelemetry_lib.tracing.context import ( detach_context, @@ -79,9 +83,28 @@ def get_laminar_span_context(self) -> LaminarSpanContext: if hasattr(self.span, "attributes"): span_path = list(self.span.attributes.get(SPAN_PATH, tuple())) span_ids_path = list(self.span.attributes.get(SPAN_IDS_PATH, tuple())) + user_id = self.span.attributes.get( + f"{ASSOCIATION_PROPERTIES}.{USER_ID}", None + ) + session_id = self.span.attributes.get( + f"{ASSOCIATION_PROPERTIES}.{SESSION_ID}", None + ) + trace_type = self.span.attributes.get( + f"{ASSOCIATION_PROPERTIES}.{TRACE_TYPE}", None + ) + metadata = { + k.replace(f"{ASSOCIATION_PROPERTIES}.{METADATA}.", ""): v + for k, v in self.span.attributes.items() + if k.startswith(f"{ASSOCIATION_PROPERTIES}.{METADATA}.") + } + for k, v in metadata.items(): + try: + metadata[k] = orjson.loads(v) + except Exception: + metadata[k] = v else: self.logger.warning( - "Attributes object is not available. Most likely the span is not a LaminarSpan ", + "Attributes object is not available. Most likely the span is not a LaminarSpan " "and not an OpenTelemetry default SDK span. Span path and ids path will be empty.", ) return LaminarSpanContext( @@ -90,6 +113,10 @@ def get_laminar_span_context(self) -> LaminarSpanContext: is_remote=self.span.get_span_context().is_remote, span_path=span_path, span_ids_path=span_ids_path, + user_id=user_id, + session_id=session_id, + trace_type=trace_type, + metadata=metadata, ) def span_id(self, format: Literal["int", "uuid"] = "int") -> int | uuid.UUID: @@ -172,7 +199,7 @@ def tags(self) -> list[str]: self.logger.debug( "[LaminarSpan.tags] WARNING. Current span does not have attributes object. " "Perhaps, the span was created with a custom OTel SDK. Returning an empty list." - "Help: OpenTelemetry API does not guarantee reading attributes from a span, but OTel SDK ", + "Help: OpenTelemetry API does not guarantee reading attributes from a span, but OTel SDK " "allows it by default. Laminar SDK allows to read attributes too.", ) return [] diff --git a/src/lmnr/sdk/laminar.py b/src/lmnr/sdk/laminar.py index 7a31cf45..c0832083 100644 --- a/src/lmnr/sdk/laminar.py +++ b/src/lmnr/sdk/laminar.py @@ -36,6 +36,7 @@ from opentelemetry.util.types import AttributeValue from typing import Any, Iterator, Literal +from typing_extensions import TypedDict import datetime import logging @@ -56,6 +57,18 @@ ) +class ParsedParentSpanContext(TypedDict): + """Parsed information from a parent span context.""" + + otel_span_context: trace.SpanContext | None + path: list[str] + span_ids_path: list[str] + user_id: str | None + session_id: str | None + trace_type: TraceType | None + metadata: dict[str, Any] | None + + def _set_span_association_props_in_context(span: Span) -> Token[Context] | None: """Set association properties from span in context. @@ -104,6 +117,97 @@ def _set_span_association_props_in_context(span: Span) -> Token[Context] | None: return attach_context(ctx_with_props) +def _parse_parent_span_context( + parent_span_context: LaminarSpanContext | dict | str | None, + logger: logging.Logger, +) -> ParsedParentSpanContext: + """Parse parent_span_context and extract all relevant information. + + Args: + parent_span_context: Parent span context to parse + logger: Logger for warnings + + Returns: + ParsedParentSpanContext with otel_span_context, path, span_ids_path, + user_id, session_id, trace_type, and metadata + """ + if parent_span_context is None: + return ParsedParentSpanContext( + otel_span_context=None, + path=[], + span_ids_path=[], + user_id=None, + session_id=None, + trace_type=None, + metadata=None, + ) + + path = [] + span_ids_path = [] + user_id = None + session_id = None + trace_type = None + metadata = None + laminar_span_context = None + + # Try to deserialize if dict or str + if isinstance(parent_span_context, (dict, str)): + try: + laminar_span_context = LaminarSpanContext.deserialize(parent_span_context) + except Exception: + logger.warning( + f"Could not deserialize parent_span_context: {parent_span_context}. " + "Will use it as is." + ) + laminar_span_context = parent_span_context + else: + laminar_span_context = parent_span_context + + # Extract path and association props from LaminarSpanContext + if isinstance(laminar_span_context, LaminarSpanContext): + path = laminar_span_context.span_path + span_ids_path = laminar_span_context.span_ids_path + user_id = laminar_span_context.user_id + session_id = laminar_span_context.session_id + if laminar_span_context.trace_type is not None: + try: + trace_type = ( + TraceType(laminar_span_context.trace_type) + if isinstance(laminar_span_context.trace_type, str) + else laminar_span_context.trace_type + ) + except (ValueError, TypeError): + pass + metadata = laminar_span_context.metadata + + # Convert to OTEL span context + try: + otel_span_context = LaminarSpanContext.try_to_otel_span_context( + laminar_span_context, logger + ) + except ValueError as exc: + logger.warning(f"Invalid span context provided: {exc}") + return ParsedParentSpanContext( + otel_span_context=None, + path=path, + span_ids_path=span_ids_path, + user_id=user_id, + session_id=session_id, + trace_type=trace_type, + metadata=metadata, + ) + + return ParsedParentSpanContext( + otel_span_context=otel_span_context, + path=path, + span_ids_path=span_ids_path, + user_id=user_id, + session_id=session_id, + trace_type=trace_type, + metadata=metadata, + ) + + class Laminar: __project_api_key: str | None = None __initialized: bool = False @@ -406,8 +510,7 @@ def start_as_current_span( obtained from `Laminar.get_laminar_span_context_dict()` or\ `Laminar.get_laminar_span_context_str()` respectively, it will be\ converted to a `LaminarSpanContext` if possible. See also\ - `Laminar.get_span_context`, `Laminar.get_span_context_dict` and\ - `Laminar.get_span_context_str` for more information. + `Laminar.serialize_span_context` for more information. Defaults to None. labels (list[str] | None, optional): [DEPRECATED] Use tags\ instead. Labels to set for the span. Defaults to None. @@ -431,56 +534,62 @@ def start_as_current_span( ) return - wrapper = TracerWrapper() - with get_tracer_with_context() as (tracer, isolated_context): ctx = context or isolated_context - path = [] - span_ids_path = [] - if parent_span_context is not None: - if isinstance(parent_span_context, (dict, str)): - try: - laminar_span_context = LaminarSpanContext.deserialize( - parent_span_context - ) - path = laminar_span_context.span_path - span_ids_path = laminar_span_context.span_ids_path - except Exception: - cls.__logger.warning( - f"`start_as_current_span` Could not deserialize parent_span_context: {parent_span_context}. " - "Will use it as is." - ) - laminar_span_context = parent_span_context - else: - laminar_span_context = parent_span_context - if isinstance(laminar_span_context, LaminarSpanContext): - path = laminar_span_context.span_path - span_ids_path = laminar_span_context.span_ids_path - span_context = LaminarSpanContext.try_to_otel_span_context( - laminar_span_context, cls.__logger - ) + + # Parse parent_span_context and extract all info + parsed = _parse_parent_span_context(parent_span_context, cls.__logger) + + # Set parent span in context if present + if parsed["otel_span_context"] is not None: ctx = trace.set_span_in_context( - trace.NonRecordingSpan(span_context), ctx + trace.NonRecordingSpan(parsed["otel_span_context"]), ctx ) + + # Determine trace_type with proper priority trace_type = None if span_type in ["EVALUATION", "EXECUTOR", "EVALUATOR"]: trace_type = TraceType.EVALUATION + elif parsed["trace_type"] is not None: + trace_type = parsed["trace_type"] - # Merge metadata: context (inherited) + global + explicit (explicit wins) + # Merge metadata: context (inherited) + global + parent + explicit (explicit wins) # Get metadata from context if it exists ctx_metadata = get_value(CONTEXT_METADATA_KEY, ctx) or {} - # Merge: context metadata + global metadata + explicit metadata - # Later keys override earlier ones + # Merge with priority: global < context < parent < explicit merged_metadata = { - **ctx_metadata, - **cls.__global_metadata, + **(cls.__global_metadata or {}), + **(ctx_metadata or {}), + **(parsed["metadata"] or {}), **(metadata or {}), } + # Get association props from context (fallback values) + ctx_user_id = get_value(CONTEXT_USER_ID_KEY, ctx) + ctx_session_id = get_value(CONTEXT_SESSION_ID_KEY, ctx) + + # Merge user_id and session_id with priority: context < parent < explicit + final_user_id = ( + user_id + if user_id is not None + else ( + parsed["user_id"] if parsed["user_id"] is not None else ctx_user_id + ) + ) + final_session_id = ( + session_id + if session_id is not None + else ( + parsed["session_id"] + if parsed["session_id"] is not None + else ctx_session_id + ) + ) + ctx = set_association_prop_context( trace_type=trace_type, - user_id=user_id, - session_id=session_id, + user_id=final_user_id, + session_id=final_session_id, metadata=merged_metadata if merged_metadata else None, context=ctx, # we need a token separately, so we manually attach the context @@ -516,8 +625,8 @@ def start_as_current_span( context=ctx, attributes={ SPAN_TYPE: span_type, - PARENT_SPAN_PATH: path, - PARENT_SPAN_IDS_PATH: span_ids_path, + PARENT_SPAN_PATH: parsed["path"], + PARENT_SPAN_IDS_PATH: parsed["span_ids_path"], **(label_props), **(tag_props), # Association properties are attached to context above @@ -629,44 +738,20 @@ def bar(): with get_tracer_with_context() as (tracer, isolated_context): ctx = context or isolated_context - path = [] - span_ids_path = [] - if parent_span_context is not None: - if isinstance(parent_span_context, (dict, str)): - try: - laminar_span_context = LaminarSpanContext.deserialize( - parent_span_context - ) - path = laminar_span_context.span_path - span_ids_path = laminar_span_context.span_ids_path - except Exception: - cls.__logger.warning( - f"`start_span` Could not deserialize parent_span_context: {parent_span_context}. " - "Will use it as is." - ) - laminar_span_context = parent_span_context - else: - laminar_span_context = parent_span_context - if isinstance(laminar_span_context, LaminarSpanContext): - path = laminar_span_context.span_path - span_ids_path = laminar_span_context.span_ids_path - span_context = LaminarSpanContext.try_to_otel_span_context( - laminar_span_context, cls.__logger - ) + + # Parse parent_span_context and extract all info + parsed = _parse_parent_span_context(parent_span_context, cls.__logger) + + # Set parent span in context if present + if parsed["otel_span_context"] is not None: ctx = trace.set_span_in_context( - trace.NonRecordingSpan(span_context), ctx + trace.NonRecordingSpan(parsed["otel_span_context"]), ctx ) # Get association props from context (fallback values) - ctx_user_id = ( - get_value(CONTEXT_USER_ID_KEY, ctx) if user_id is None else None - ) - ctx_session_id = ( - get_value(CONTEXT_SESSION_ID_KEY, ctx) if session_id is None else None - ) - ctx_metadata = ( - get_value(CONTEXT_METADATA_KEY, ctx) if metadata is None else None - ) + ctx_user_id = get_value(CONTEXT_USER_ID_KEY, ctx) + ctx_session_id = get_value(CONTEXT_SESSION_ID_KEY, ctx) + ctx_metadata = get_value(CONTEXT_METADATA_KEY, ctx) label_props = {} try: @@ -693,24 +778,52 @@ def bar(): + "Tags will be ignored." ) + # Determine trace_type with proper priority: explicit > parent > context trace_type = None if span_type in ["EVALUATION", "EXECUTOR", "EVALUATOR"]: trace_type = TraceType.EVALUATION - # Get trace_type from context if not set explicitly - ctx_trace_type = ( - get_value(CONTEXT_TRACE_TYPE_KEY, ctx) if trace_type is None else None + elif parsed["trace_type"] is not None: + trace_type = parsed["trace_type"] + else: + # Get trace_type from context if not set explicitly or from parent + ctx_trace_type = get_value(CONTEXT_TRACE_TYPE_KEY, ctx) + if ctx_trace_type: + try: + trace_type = TraceType(ctx_trace_type) + except (ValueError, TypeError): + pass + + # Merge with priority: global < context < parent < explicit + merged_metadata = { + **(cls.__global_metadata or {}), + **(ctx_metadata or {}), + **(parsed["metadata"] or {}), + **(metadata or {}), + } + + # Merge user_id and session_id with priority: context < parent < explicit + final_user_id = ( + user_id + if user_id is not None + else ( + parsed["user_id"] if parsed["user_id"] is not None else ctx_user_id + ) + ) + final_session_id = ( + session_id + if session_id is not None + else ( + parsed["session_id"] + if parsed["session_id"] is not None + else ctx_session_id + ) ) - if ctx_trace_type: - try: - trace_type = TraceType(ctx_trace_type) - except (ValueError, TypeError): - pass - # Build association_props using explicit params or context fallbacks + # Build association_props using merged values association_props = cls._get_association_prop_attributes( - user_id=user_id or ctx_user_id, - session_id=session_id or ctx_session_id, - metadata=metadata or ctx_metadata, + user_id=final_user_id, + session_id=final_session_id, + metadata=merged_metadata if merged_metadata else None, trace_type=trace_type, ) @@ -719,8 +832,8 @@ def bar(): context=ctx, attributes={ SPAN_TYPE: span_type, - PARENT_SPAN_PATH: path, - PARENT_SPAN_IDS_PATH: span_ids_path, + PARENT_SPAN_PATH: parsed["path"], + PARENT_SPAN_IDS_PATH: parsed["span_ids_path"], **(label_props), **(tag_props), **(association_props), @@ -983,7 +1096,7 @@ def set_span_attributes( attributes (dict[Attributes | str, Any]): attributes to set for the span """ span = cls.current_span() - if span == trace.INVALID_SPAN: + if span == trace.INVALID_SPAN or span is None: return for key, value in attributes.items(): @@ -1005,7 +1118,7 @@ def get_laminar_span_context( return None span = span or cls.current_span() - if span == trace.INVALID_SPAN: + if span == trace.INVALID_SPAN or span is None: return None if not isinstance(span, LaminarSpan): span = LaminarSpan(span) @@ -1060,6 +1173,18 @@ def deserialize_span_context(cls, span_context: dict | str) -> LaminarSpanContex @classmethod def current_span(cls, context: Context | None = None) -> LaminarSpan | None: + """Get the current active span. If a context is provided, the span will + be retrieved from that context. + + Args: + context (Context | None, optional): The context to get the span\ + from. If not provided, the current context will be used. + Defaults to None. + + Returns: + LaminarSpan | None: The current active span, or None if there is no\ + active span. + """ context = context or get_current_context() span = trace.get_current_span(context=context) if span == trace.INVALID_SPAN: diff --git a/src/lmnr/sdk/types.py b/src/lmnr/sdk/types.py index 12a120ac..65b3dfff 100644 --- a/src/lmnr/sdk/types.py +++ b/src/lmnr/sdk/types.py @@ -219,9 +219,8 @@ class LaminarSpanContext(BaseModel): """ A span context that can be used to continue a trace across services. This is a slightly modified version of the OpenTelemetry span context. For - usage examples, see `Laminar.get_laminar_span_context_dict`, - `Laminar.get_laminar_span_context_str`, `Laminar.get_span_context`, and - `Laminar.deserialize_laminar_span_context`. + usage examples, see `Laminar.serialize_span_context`, + `Laminar.get_span_context`, and `Laminar.deserialize_laminar_span_context`. The difference between this and the OpenTelemetry span context is that the `trace_id` and `span_id` are stored as UUIDs instead of integers for @@ -233,6 +232,10 @@ class LaminarSpanContext(BaseModel): is_remote: bool = Field(default=False) span_path: list[str] = Field(default=[]) span_ids_path: list[str] = Field(default=[]) # stringified UUIDs + user_id: str | None = Field(default=None) + session_id: str | None = Field(default=None) + trace_type: TraceType | None = Field(default=None) + metadata: dict[str, Any] | None = Field(default=None) def __str__(self) -> str: return self.model_dump_json() @@ -288,6 +291,10 @@ def deserialize(cls, data: dict[str, Any] | str) -> "LaminarSpanContext": "span_path": data.get("span_path") or data.get("spanPath", []), "span_ids_path": data.get("span_ids_path") or data.get("spanIdsPath", []), + "user_id": data.get("user_id") or data.get("userId"), + "session_id": data.get("session_id") or data.get("sessionId"), + "trace_type": data.get("trace_type") or data.get("traceType"), + "metadata": data.get("metadata") or data.get("metadata", {}), } return cls.model_validate(converted_data) elif isinstance(data, str): diff --git a/tests/test_context_propagation.py b/tests/test_context_propagation.py index 229a20a1..faf167dc 100644 --- a/tests/test_context_propagation.py +++ b/tests/test_context_propagation.py @@ -3,7 +3,7 @@ from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter -def _assert_assoiation_properties( +def _assert_association_properties( span: ReadableSpan, user_id: str, session_id: str, @@ -32,10 +32,10 @@ def test_ctx_prop_parent_sc_child_s(span_exporter: InMemorySpanExporter): assert len(spans) == 2 parent_span = [s for s in spans if s.name == "parent"][0] child_span = [s for s in spans if s.name == "child"][0] - _assert_assoiation_properties( + _assert_association_properties( parent_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" ) - _assert_assoiation_properties( + _assert_association_properties( child_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" ) @@ -55,10 +55,10 @@ def test_ctx_prop_parent_sc_child_sc(span_exporter: InMemorySpanExporter): assert len(spans) == 2 parent_span = [s for s in spans if s.name == "parent"][0] child_span = [s for s in spans if s.name == "child"][0] - _assert_assoiation_properties( + _assert_association_properties( parent_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" ) - _assert_assoiation_properties( + _assert_association_properties( child_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" ) @@ -81,10 +81,10 @@ def child(): assert len(spans) == 2 parent_span = [s for s in spans if s.name == "parent"][0] child_span = [s for s in spans if s.name == "child"][0] - _assert_assoiation_properties( + _assert_association_properties( parent_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" ) - _assert_assoiation_properties( + _assert_association_properties( child_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" ) @@ -105,10 +105,10 @@ def test_ctx_prop_parent_sa_child_s(span_exporter: InMemorySpanExporter): assert len(spans) == 2 parent_span = [s for s in spans if s.name == "parent"][0] child_span = [s for s in spans if s.name == "child"][0] - _assert_assoiation_properties( + _assert_association_properties( parent_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" ) - _assert_assoiation_properties( + _assert_association_properties( child_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" ) @@ -129,10 +129,10 @@ def test_ctx_prop_parent_sa_child_sc(span_exporter: InMemorySpanExporter): assert len(spans) == 2 parent_span = [s for s in spans if s.name == "parent"][0] child_span = [s for s in spans if s.name == "child"][0] - _assert_assoiation_properties( + _assert_association_properties( parent_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" ) - _assert_assoiation_properties( + _assert_association_properties( child_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" ) @@ -156,10 +156,10 @@ def child(): assert len(spans) == 2 parent_span = [s for s in spans if s.name == "parent"][0] child_span = [s for s in spans if s.name == "child"][0] - _assert_assoiation_properties( + _assert_association_properties( parent_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" ) - _assert_assoiation_properties( + _assert_association_properties( child_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" ) @@ -181,10 +181,10 @@ def parent(): assert len(spans) == 2 parent_span = [s for s in spans if s.name == "parent"][0] child_span = [s for s in spans if s.name == "child"][0] - _assert_assoiation_properties( + _assert_association_properties( parent_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" ) - _assert_assoiation_properties( + _assert_association_properties( child_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" ) @@ -205,9 +205,51 @@ def test_ctx_prop_parent_use_span(span_exporter: InMemorySpanExporter): assert len(spans) == 2 parent_span = [s for s in spans if s.name == "parent"][0] child_span = [s for s in spans if s.name == "child"][0] - _assert_assoiation_properties( + _assert_association_properties( parent_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" ) - _assert_assoiation_properties( + _assert_association_properties( + child_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" + ) + + +def test_ctx_prop_laminar_span_context(span_exporter: InMemorySpanExporter): + span = Laminar.start_span( + "parent", + user_id="user_id", + session_id="session_id", + span_type="EVALUATION", + metadata={"foo": "bar"}, + ) + span_context = Laminar.get_laminar_span_context(span) + passed_span_context = Laminar.serialize_span_context(span) + span2 = Laminar.start_span( + "child", + parent_span_context=Laminar.deserialize_span_context(passed_span_context), + ) + span2.end() + span.end() + + assert span_context is not None + assert span_context.user_id == "user_id" + assert span_context.session_id == "session_id" + assert span_context.trace_type.value == "EVALUATION" + assert span_context.metadata == {"foo": "bar"} + + spans = span_exporter.get_finished_spans() + assert len(spans) == 2 + parent_span = [s for s in spans if s.name == "parent"][0] + child_span = [s for s in spans if s.name == "child"][0] + + assert ( + parent_span.get_span_context().trace_id + == child_span.get_span_context().trace_id + ) + assert child_span.parent.span_id == parent_span.get_span_context().span_id + + _assert_association_properties( + parent_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" + ) + _assert_association_properties( child_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" ) From 08c348a9a8e0a8c87483de42ecbbd32c520e1db1 Mon Sep 17 00:00:00 2001 From: Din Date: Mon, 8 Dec 2025 16:27:02 +0000 Subject: [PATCH 04/10] fix laminar span context when attributes not available --- src/lmnr/opentelemetry_lib/tracing/span.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/src/lmnr/opentelemetry_lib/tracing/span.py b/src/lmnr/opentelemetry_lib/tracing/span.py index b3a9b88f..324e68d7 100644 --- a/src/lmnr/opentelemetry_lib/tracing/span.py +++ b/src/lmnr/opentelemetry_lib/tracing/span.py @@ -107,6 +107,13 @@ def get_laminar_span_context(self) -> LaminarSpanContext: "Attributes object is not available. Most likely the span is not a LaminarSpan " "and not an OpenTelemetry default SDK span. Span path and ids path will be empty.", ) + return LaminarSpanContext( + trace_id=uuid.UUID(int=self.span.get_span_context().trace_id), + span_id=uuid.UUID(int=self.span.get_span_context().span_id), + is_remote=self.span.get_span_context().is_remote, + span_path=span_path, + span_ids_path=span_ids_path, + ) return LaminarSpanContext( trace_id=uuid.UUID(int=self.span.get_span_context().trace_id), span_id=uuid.UUID(int=self.span.get_span_context().span_id), From c957918daa2b0f965bab51537a2d6d57ecd91a11 Mon Sep 17 00:00:00 2001 From: Din Date: Mon, 8 Dec 2025 16:27:33 +0000 Subject: [PATCH 05/10] tiny refactor --- src/lmnr/opentelemetry_lib/tracing/span.py | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/src/lmnr/opentelemetry_lib/tracing/span.py b/src/lmnr/opentelemetry_lib/tracing/span.py index 324e68d7..32355b58 100644 --- a/src/lmnr/opentelemetry_lib/tracing/span.py +++ b/src/lmnr/opentelemetry_lib/tracing/span.py @@ -80,6 +80,10 @@ def set_trace_metadata(self, metadata: dict[str, AttributeValue]) -> None: def get_laminar_span_context(self) -> LaminarSpanContext: span_path = [] span_ids_path = [] + user_id = None + session_id = None + trace_type = None + metadata = {} if hasattr(self.span, "attributes"): span_path = list(self.span.attributes.get(SPAN_PATH, tuple())) span_ids_path = list(self.span.attributes.get(SPAN_IDS_PATH, tuple())) @@ -107,13 +111,6 @@ def get_laminar_span_context(self) -> LaminarSpanContext: "Attributes object is not available. Most likely the span is not a LaminarSpan " "and not an OpenTelemetry default SDK span. Span path and ids path will be empty.", ) - return LaminarSpanContext( - trace_id=uuid.UUID(int=self.span.get_span_context().trace_id), - span_id=uuid.UUID(int=self.span.get_span_context().span_id), - is_remote=self.span.get_span_context().is_remote, - span_path=span_path, - span_ids_path=span_ids_path, - ) return LaminarSpanContext( trace_id=uuid.UUID(int=self.span.get_span_context().trace_id), span_id=uuid.UUID(int=self.span.get_span_context().span_id), From 3dca2c419bd0feec195c126d7207422112ce4f8d Mon Sep 17 00:00:00 2001 From: Din Date: Tue, 9 Dec 2025 11:04:33 +0000 Subject: [PATCH 06/10] remove global metadata accessor --- src/lmnr/sdk/laminar.py | 21 ++++++++------------- 1 file changed, 8 insertions(+), 13 deletions(-) diff --git a/src/lmnr/sdk/laminar.py b/src/lmnr/sdk/laminar.py index c0832083..6ecde5be 100644 --- a/src/lmnr/sdk/laminar.py +++ b/src/lmnr/sdk/laminar.py @@ -894,15 +894,14 @@ def use_span( # span and trace_id. isolated_context_token = attach_context(context) context_token = context_api.attach(context) - try: - if isinstance(span, LaminarSpan): - yield span - else: - yield LaminarSpan(span) - finally: - context_api.detach(context_token) - detach_context(isolated_context_token) - wrapper.pop_span_context() + if isinstance(span, LaminarSpan): + yield span + else: + yield LaminarSpan(span) + + context_api.detach(context_token) + detach_context(isolated_context_token) + wrapper.pop_span_context() # Record only exceptions that inherit Exception class but not BaseException, because # classes that directly inherit BaseException are not technically errors, e.g. GeneratorExit. @@ -1392,7 +1391,3 @@ def _get_association_prop_attributes( } ) return association_properties - - @property - def global_metadata(cls) -> dict[str, AttributeValue]: - return cls.__global_metadata From ed85a5d561bb38b6386eb1c8f3adfdcc8d144d43 Mon Sep 17 00:00:00 2001 From: Din Date: Tue, 9 Dec 2025 18:11:45 +0000 Subject: [PATCH 07/10] refactor one function, rename get_current_span, more span nesting in tests --- .../opentelemetry_lib/decorators/__init__.py | 58 +------ src/lmnr/opentelemetry_lib/tracing/utils.py | 62 +++++++ src/lmnr/sdk/laminar.py | 83 ++-------- tests/test_context_propagation.py | 155 +++++++++++++----- 4 files changed, 199 insertions(+), 159 deletions(-) create mode 100644 src/lmnr/opentelemetry_lib/tracing/utils.py diff --git a/src/lmnr/opentelemetry_lib/decorators/__init__.py b/src/lmnr/opentelemetry_lib/decorators/__init__.py index 2f5bf854..ca6a2047 100644 --- a/src/lmnr/opentelemetry_lib/decorators/__init__.py +++ b/src/lmnr/opentelemetry_lib/decorators/__init__.py @@ -7,29 +7,22 @@ from lmnr.opentelemetry_lib.tracing.context import ( CONTEXT_METADATA_KEY, - CONTEXT_SESSION_ID_KEY, - CONTEXT_TRACE_TYPE_KEY, - CONTEXT_USER_ID_KEY, attach_context, detach_context, - get_current_context, get_event_attributes_from_context, ) from lmnr.opentelemetry_lib.tracing.span import LaminarSpan +from lmnr.opentelemetry_lib.tracing.utils import set_association_props_in_context from lmnr.sdk.utils import get_input_from_func_args, is_method from lmnr.opentelemetry_lib.tracing.tracer import get_tracer_with_context from lmnr.opentelemetry_lib.tracing.attributes import ( ASSOCIATION_PROPERTIES, METADATA, SPAN_TYPE, - USER_ID, - SESSION_ID, - TRACE_TYPE, ) from lmnr.opentelemetry_lib.tracing import TracerWrapper from lmnr.sdk.log import get_default_logger from lmnr.sdk.utils import is_otel_attribute_value_type, json_dumps -from opentelemetry.context import set_value logger = get_default_logger(__name__) @@ -145,51 +138,6 @@ def _cleanup_span(span: Span, wrapper: TracerWrapper): wrapper.pop_span_context() -def _set_association_props_in_context(span: Span): - """Set association properties from span in context before push_span_context. - - Returns the token that needs to be detached when the span ends. - """ - if not isinstance(span, LaminarSpan): - return None - - props = span.laminar_association_properties - user_id_key = f"{ASSOCIATION_PROPERTIES}.{USER_ID}" - session_id_key = f"{ASSOCIATION_PROPERTIES}.{SESSION_ID}" - trace_type_key = f"{ASSOCIATION_PROPERTIES}.{TRACE_TYPE}" - - # Extract values from props - extracted_user_id = props.get(user_id_key) - extracted_session_id = props.get(session_id_key) - extracted_trace_type = props.get(trace_type_key) - - # Extract metadata from props (keys without ASSOCIATION_PROPERTIES prefix) - metadata_dict = {} - for key, value in props.items(): - if not key.startswith(f"{ASSOCIATION_PROPERTIES}."): - metadata_dict[key] = value - - # Set context with association props - current_ctx = get_current_context() - ctx_with_props = current_ctx - if extracted_user_id: - ctx_with_props = set_value( - CONTEXT_USER_ID_KEY, extracted_user_id, ctx_with_props - ) - if extracted_session_id: - ctx_with_props = set_value( - CONTEXT_SESSION_ID_KEY, extracted_session_id, ctx_with_props - ) - if extracted_trace_type: - ctx_with_props = set_value( - CONTEXT_TRACE_TYPE_KEY, extracted_trace_type, ctx_with_props - ) - if metadata_dict: - ctx_with_props = set_value(CONTEXT_METADATA_KEY, metadata_dict, ctx_with_props) - - return attach_context(ctx_with_props) - - def observe_base( *, name: str | None = None, @@ -222,7 +170,7 @@ def wrap(*args, **kwargs): # Set association props in context before push_span_context # so child spans inherit them - assoc_props_token = _set_association_props_in_context(span) + assoc_props_token = set_association_props_in_context(span) if assoc_props_token and isinstance(span, LaminarSpan): span._lmnr_assoc_props_token = assoc_props_token @@ -304,7 +252,7 @@ async def wrap(*args, **kwargs): # Set association props in context before push_span_context # so child spans inherit them - assoc_props_token = _set_association_props_in_context(span) + assoc_props_token = set_association_props_in_context(span) if assoc_props_token and isinstance(span, LaminarSpan): span._lmnr_assoc_props_token = assoc_props_token diff --git a/src/lmnr/opentelemetry_lib/tracing/utils.py b/src/lmnr/opentelemetry_lib/tracing/utils.py new file mode 100644 index 00000000..f8060381 --- /dev/null +++ b/src/lmnr/opentelemetry_lib/tracing/utils.py @@ -0,0 +1,62 @@ +from opentelemetry.trace import Span +from lmnr.opentelemetry_lib.tracing.span import LaminarSpan +from lmnr.opentelemetry_lib.tracing.attributes import ( + ASSOCIATION_PROPERTIES, + USER_ID, + SESSION_ID, + TRACE_TYPE, +) +from lmnr.opentelemetry_lib.tracing.context import ( + get_current_context, + attach_context, + set_value, + CONTEXT_USER_ID_KEY, + CONTEXT_SESSION_ID_KEY, + CONTEXT_TRACE_TYPE_KEY, + CONTEXT_METADATA_KEY, +) + + +def set_association_props_in_context(span: Span): + """Set association properties from span in context before push_span_context. + + Returns the token that needs to be detached when the span ends. + """ + if not isinstance(span, LaminarSpan): + return None + + props = span.laminar_association_properties + user_id_key = f"{ASSOCIATION_PROPERTIES}.{USER_ID}" + session_id_key = f"{ASSOCIATION_PROPERTIES}.{SESSION_ID}" + trace_type_key = f"{ASSOCIATION_PROPERTIES}.{TRACE_TYPE}" + + # Extract values from props + extracted_user_id = props.get(user_id_key) + extracted_session_id = props.get(session_id_key) + extracted_trace_type = props.get(trace_type_key) + + # Extract metadata from props (keys without ASSOCIATION_PROPERTIES prefix) + metadata_dict = {} + for key, value in props.items(): + if not key.startswith(f"{ASSOCIATION_PROPERTIES}."): + metadata_dict[key] = value + + # Set context with association props + current_ctx = get_current_context() + ctx_with_props = current_ctx + if extracted_user_id: + ctx_with_props = set_value( + CONTEXT_USER_ID_KEY, extracted_user_id, ctx_with_props + ) + if extracted_session_id: + ctx_with_props = set_value( + CONTEXT_SESSION_ID_KEY, extracted_session_id, ctx_with_props + ) + if extracted_trace_type: + ctx_with_props = set_value( + CONTEXT_TRACE_TYPE_KEY, extracted_trace_type, ctx_with_props + ) + if metadata_dict: + ctx_with_props = set_value(CONTEXT_METADATA_KEY, metadata_dict, ctx_with_props) + + return attach_context(ctx_with_props) diff --git a/src/lmnr/sdk/laminar.py b/src/lmnr/sdk/laminar.py index 6ecde5be..8dacf6c5 100644 --- a/src/lmnr/sdk/laminar.py +++ b/src/lmnr/sdk/laminar.py @@ -1,5 +1,5 @@ from contextlib import contextmanager -from contextvars import Context, Token +from contextvars import Context import warnings from lmnr.opentelemetry_lib import TracerManager from lmnr.opentelemetry_lib.tracing import TracerWrapper, get_current_context @@ -14,11 +14,7 @@ push_span_context, set_association_prop_context, ) -from opentelemetry.context import get_value, set_value -from lmnr.opentelemetry_lib.tracing.instruments import Instruments -from lmnr.opentelemetry_lib.tracing.processor import LaminarSpanProcessor -from lmnr.opentelemetry_lib.tracing.span import LaminarSpan -from lmnr.opentelemetry_lib.tracing.tracer import get_tracer_with_context +from opentelemetry.context import get_value from lmnr.opentelemetry_lib.tracing.attributes import ( ASSOCIATION_PROPERTIES, PARENT_SPAN_IDS_PATH, @@ -27,6 +23,11 @@ Attributes, SPAN_TYPE, ) +from lmnr.opentelemetry_lib.tracing.instruments import Instruments +from lmnr.opentelemetry_lib.tracing.processor import LaminarSpanProcessor +from lmnr.opentelemetry_lib.tracing.span import LaminarSpan +from lmnr.opentelemetry_lib.tracing.tracer import get_tracer_with_context +from lmnr.opentelemetry_lib.tracing.utils import set_association_props_in_context from lmnr.sdk.utils import get_otel_env_var from opentelemetry import trace @@ -69,54 +70,6 @@ class ParsedParentSpanContext(TypedDict): metadata: dict[str, Any] | None -def _set_span_association_props_in_context(span: Span) -> Token[Context] | None: - """Set association properties from span in context. - - Extracts association properties from span attributes and sets them in the - isolated context so child spans can inherit them. - - Returns the token that needs to be stored on the span for cleanup. - """ - if not isinstance(span, LaminarSpan): - return None - - props = span.laminar_association_properties - user_id_key = f"{ASSOCIATION_PROPERTIES}.{USER_ID}" - session_id_key = f"{ASSOCIATION_PROPERTIES}.{SESSION_ID}" - trace_type_key = f"{ASSOCIATION_PROPERTIES}.{TRACE_TYPE}" - - # Extract values from props - extracted_user_id = props.get(user_id_key) - extracted_session_id = props.get(session_id_key) - extracted_trace_type = props.get(trace_type_key) - - # Extract metadata from props (keys without ASSOCIATION_PROPERTIES prefix) - metadata_dict = {} - for key, value in props.items(): - if not key.startswith(f"{ASSOCIATION_PROPERTIES}."): - metadata_dict[key] = value - - # Set context with association props - current_ctx = get_current_context() - ctx_with_props = current_ctx - if extracted_user_id: - ctx_with_props = set_value( - CONTEXT_USER_ID_KEY, extracted_user_id, ctx_with_props - ) - if extracted_session_id: - ctx_with_props = set_value( - CONTEXT_SESSION_ID_KEY, extracted_session_id, ctx_with_props - ) - if extracted_trace_type: - ctx_with_props = set_value( - CONTEXT_TRACE_TYPE_KEY, extracted_trace_type, ctx_with_props - ) - if metadata_dict: - ctx_with_props = set_value(CONTEXT_METADATA_KEY, metadata_dict, ctx_with_props) - - return attach_context(ctx_with_props) - - def _parse_parent_span_context( parent_span_context: LaminarSpanContext | dict | str | None, logger: logging.Logger, @@ -883,7 +836,7 @@ def use_span( try: # Set association props in context before push_span_context # so child spans inherit them - assoc_props_token = _set_span_association_props_in_context(span) + assoc_props_token = set_association_props_in_context(span) if assoc_props_token and isinstance(span, LaminarSpan): span._lmnr_assoc_props_token = assoc_props_token @@ -1039,7 +992,7 @@ def bar(): # Set association props in context before push_span_context # so child spans inherit them - assoc_props_token = _set_span_association_props_in_context(span) + assoc_props_token = set_association_props_in_context(span) if assoc_props_token and isinstance(span, LaminarSpan): span._lmnr_assoc_props_token = assoc_props_token @@ -1062,7 +1015,7 @@ def set_span_output(cls, output: Any = None): output (Any, optional): output of the span. Will be sent as an\ attribute, so must be json serializable. Defaults to None. """ - span = cls.current_span() + span = cls.get_current_span() if span is None: return span.set_output(output) @@ -1094,7 +1047,7 @@ def set_span_attributes( Args: attributes (dict[Attributes | str, Any]): attributes to set for the span """ - span = cls.current_span() + span = cls.get_current_span() if span == trace.INVALID_SPAN or span is None: return @@ -1116,7 +1069,7 @@ def get_laminar_span_context( if not cls.is_initialized(): return None - span = span or cls.current_span() + span = span or cls.get_current_span() if span == trace.INVALID_SPAN or span is None: return None if not isinstance(span, LaminarSpan): @@ -1171,7 +1124,7 @@ def deserialize_span_context(cls, span_context: dict | str) -> LaminarSpanContex return LaminarSpanContext.deserialize(span_context) @classmethod - def current_span(cls, context: Context | None = None) -> LaminarSpan | None: + def get_current_span(cls, context: Context | None = None) -> LaminarSpan | None: """Get the current active span. If a context is provided, the span will be retrieved from that context. @@ -1237,7 +1190,7 @@ def set_span_tags(cls, tags: list[str]): if not cls.is_initialized(): return - span = cls.current_span() + span = cls.get_current_span() if span is None: return span.set_tags(tags) @@ -1245,7 +1198,7 @@ def set_span_tags(cls, tags: list[str]): @classmethod def add_span_tags(cls, tags: list[str]): """Add tags to the current span.""" - span = cls.current_span() + span = cls.get_current_span() if span is None: return span.add_tags(tags) @@ -1263,7 +1216,7 @@ def set_trace_session_id(cls, session_id: str | None = None): context = set_association_prop_context(session_id=session_id, attach=True) - span = cls.current_span(context=context) + span = cls.get_current_span(context=context) if span is None: cls.__logger.warning("No active span to set session id on") return @@ -1282,7 +1235,7 @@ def set_trace_user_id(cls, user_id: str | None = None): context = set_association_prop_context(user_id=user_id, attach=True) - span = cls.current_span(context=context) + span = cls.get_current_span(context=context) if span is None: cls.__logger.warning("No active span to set user id on") return @@ -1300,7 +1253,7 @@ def set_trace_metadata(cls, metadata: dict[str, AttributeValue]): merged_metadata = {**cls.__global_metadata, **(metadata or {})} - span = cls.current_span() + span = cls.get_current_span() if span is None: cls.__logger.warning("No active span to set metadata on") return diff --git a/tests/test_context_propagation.py b/tests/test_context_propagation.py index faf167dc..7da0949b 100644 --- a/tests/test_context_propagation.py +++ b/tests/test_context_propagation.py @@ -1,4 +1,5 @@ from opentelemetry.sdk.trace import ReadableSpan +from opentelemetry.trace import INVALID_SPAN_ID from lmnr import Laminar, observe from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter @@ -17,50 +18,83 @@ def _assert_association_properties( assert span.attributes[f"lmnr.association.properties.metadata.{key}"] == value +def _assert_same_trace_and_inheritance( + spans: list[ReadableSpan], expected_parent_span_id: str | None = None +): + trace_ids = [span.get_span_context().trace_id for span in spans] + assert len(set(trace_ids)) == 1 + if expected_parent_span_id is not None: + assert spans[0].parent.span_id == expected_parent_span_id + else: + assert ( + spans[0].parent is None + or spans[0].parent.span_id is None + or spans[0].parent.span_id == INVALID_SPAN_ID + ) + assert spans[0].get_span_context().trace_id == trace_ids[0] + + for i, span in enumerate(spans[1:]): + assert span.get_span_context().trace_id == trace_ids[0] + assert span.parent.span_id == spans[i].get_span_context().span_id + + def test_ctx_prop_parent_sc_child_s(span_exporter: InMemorySpanExporter): with Laminar.start_as_current_span( - "parent", + "grandparent", user_id="user_id", session_id="session_id", span_type="EVALUATION", metadata={"foo": "bar"}, ): - inner_span = Laminar.start_span("child") - inner_span.end() + p_span = Laminar.start_span("parent") + with Laminar.use_span(p_span, end_on_exit=True): + c_span = Laminar.start_span("child") + c_span.end() spans = span_exporter.get_finished_spans() - assert len(spans) == 2 + assert len(spans) == 3 + grandparent_span = [s for s in spans if s.name == "grandparent"][0] parent_span = [s for s in spans if s.name == "parent"][0] child_span = [s for s in spans if s.name == "child"][0] + _assert_association_properties( + grandparent_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" + ) _assert_association_properties( parent_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" ) _assert_association_properties( child_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" ) + _assert_same_trace_and_inheritance([grandparent_span, parent_span, child_span]) def test_ctx_prop_parent_sc_child_sc(span_exporter: InMemorySpanExporter): with Laminar.start_as_current_span( - "parent", + "grandparent", user_id="user_id", session_id="session_id", span_type="EVALUATION", metadata={"foo": "bar"}, ): - with Laminar.start_as_current_span("child"): - pass + with Laminar.start_as_current_span("parent"): + with Laminar.start_as_current_span("child"): + pass spans = span_exporter.get_finished_spans() - assert len(spans) == 2 + assert len(spans) == 3 + grandparent_span = [s for s in spans if s.name == "grandparent"][0] parent_span = [s for s in spans if s.name == "parent"][0] child_span = [s for s in spans if s.name == "child"][0] + _assert_association_properties( + grandparent_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" + ) _assert_association_properties( parent_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" ) _assert_association_properties( child_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" ) + _assert_same_trace_and_inheritance([grandparent_span, parent_span, child_span]) def test_ctx_prop_parent_sc_child_obs(span_exporter: InMemorySpanExporter): @@ -68,65 +102,84 @@ def test_ctx_prop_parent_sc_child_obs(span_exporter: InMemorySpanExporter): def child(): pass + @observe() + def parent(): + child() + with Laminar.start_as_current_span( - "parent", + "grandparent", user_id="user_id", session_id="session_id", span_type="EVALUATION", metadata={"foo": "bar"}, ): - child() + parent() spans = span_exporter.get_finished_spans() - assert len(spans) == 2 + assert len(spans) == 3 + grandparent_span = [s for s in spans if s.name == "grandparent"][0] parent_span = [s for s in spans if s.name == "parent"][0] child_span = [s for s in spans if s.name == "child"][0] + _assert_association_properties( + grandparent_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" + ) _assert_association_properties( parent_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" ) _assert_association_properties( child_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" ) + _assert_same_trace_and_inheritance([grandparent_span, parent_span, child_span]) def test_ctx_prop_parent_sa_child_s(span_exporter: InMemorySpanExporter): - span = Laminar.start_active_span( - "parent", + g_span = Laminar.start_active_span( + "grandparent", user_id="user_id", session_id="session_id", span_type="EVALUATION", metadata={"foo": "bar"}, ) - inner_span = Laminar.start_span("child") - inner_span.end() - span.end() + p_span = Laminar.start_span("parent") + with Laminar.use_span(p_span, end_on_exit=True): + c_span = Laminar.start_span("child") + c_span.end() + g_span.end() spans = span_exporter.get_finished_spans() - assert len(spans) == 2 + assert len(spans) == 3 + grandparent_span = [s for s in spans if s.name == "grandparent"][0] parent_span = [s for s in spans if s.name == "parent"][0] child_span = [s for s in spans if s.name == "child"][0] + _assert_association_properties( + grandparent_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" + ) _assert_association_properties( parent_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" ) _assert_association_properties( child_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" ) + _assert_same_trace_and_inheritance([grandparent_span, parent_span, child_span]) def test_ctx_prop_parent_sa_child_sc(span_exporter: InMemorySpanExporter): - span = Laminar.start_active_span( - "parent", + grandparent_span = Laminar.start_active_span( + "grandparent", user_id="user_id", session_id="session_id", span_type="EVALUATION", metadata={"foo": "bar"}, ) - inner_span = Laminar.start_span("child") - inner_span.end() - span.end() + + with Laminar.start_as_current_span("parent"): + with Laminar.start_as_current_span("child"): + pass + grandparent_span.end() spans = span_exporter.get_finished_spans() - assert len(spans) == 2 + assert len(spans) == 3 + grandparent_span = [s for s in spans if s.name == "grandparent"][0] parent_span = [s for s in spans if s.name == "parent"][0] child_span = [s for s in spans if s.name == "child"][0] _assert_association_properties( @@ -135,6 +188,7 @@ def test_ctx_prop_parent_sa_child_sc(span_exporter: InMemorySpanExporter): _assert_association_properties( child_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" ) + _assert_same_trace_and_inheritance([grandparent_span, parent_span, child_span]) def test_ctx_prop_parent_sa_child_obs(span_exporter: InMemorySpanExporter): @@ -142,26 +196,35 @@ def test_ctx_prop_parent_sa_child_obs(span_exporter: InMemorySpanExporter): def child(): pass - span = Laminar.start_active_span( - "parent", + @observe() + def parent(): + child() + + g_span = Laminar.start_active_span( + "grandparent", user_id="user_id", session_id="session_id", span_type="EVALUATION", metadata={"foo": "bar"}, ) - child() - span.end() + parent() + g_span.end() spans = span_exporter.get_finished_spans() - assert len(spans) == 2 + assert len(spans) == 3 + grandparent_span = [s for s in spans if s.name == "grandparent"][0] parent_span = [s for s in spans if s.name == "parent"][0] child_span = [s for s in spans if s.name == "child"][0] + _assert_association_properties( + grandparent_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" + ) _assert_association_properties( parent_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" ) _assert_association_properties( child_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" ) + _assert_same_trace_and_inheritance([grandparent_span, parent_span, child_span]) def test_ctx_prop_parent_obs_child_s(span_exporter: InMemorySpanExporter): @@ -171,46 +234,60 @@ def test_ctx_prop_parent_obs_child_s(span_exporter: InMemorySpanExporter): span_type="EVALUATION", metadata={"foo": "bar"}, ) - def parent(): - inner_span = Laminar.start_span("child") - inner_span.end() + def grandparent(): + p_span = Laminar.start_span("parent") + with Laminar.use_span(p_span, end_on_exit=True): + c_span = Laminar.start_span("child") + c_span.end() - parent() + grandparent() spans = span_exporter.get_finished_spans() - assert len(spans) == 2 + assert len(spans) == 3 + grandparent_span = [s for s in spans if s.name == "grandparent"][0] parent_span = [s for s in spans if s.name == "parent"][0] child_span = [s for s in spans if s.name == "child"][0] + _assert_association_properties( + grandparent_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" + ) _assert_association_properties( parent_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" ) _assert_association_properties( child_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" ) + _assert_same_trace_and_inheritance([grandparent_span, parent_span, child_span]) def test_ctx_prop_parent_use_span(span_exporter: InMemorySpanExporter): - span = Laminar.start_span( - "parent", + g_span = Laminar.start_span( + "grandparent", user_id="user_id", session_id="session_id", span_type="EVALUATION", metadata={"foo": "bar"}, ) - with Laminar.use_span(span, end_on_exit=True): - inner_span = Laminar.start_span("child") - inner_span.end() + with Laminar.use_span(g_span, end_on_exit=True): + p_span = Laminar.start_span("parent") + with Laminar.use_span(p_span, end_on_exit=True): + c_span = Laminar.start_span("child") + c_span.end() spans = span_exporter.get_finished_spans() - assert len(spans) == 2 + assert len(spans) == 3 + grandparent_span = [s for s in spans if s.name == "grandparent"][0] parent_span = [s for s in spans if s.name == "parent"][0] child_span = [s for s in spans if s.name == "child"][0] + _assert_association_properties( + grandparent_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" + ) _assert_association_properties( parent_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" ) _assert_association_properties( child_span, "user_id", "session_id", {"foo": "bar"}, "EVALUATION" ) + _assert_same_trace_and_inheritance([grandparent_span, parent_span, child_span]) def test_ctx_prop_laminar_span_context(span_exporter: InMemorySpanExporter): From ff9b8ed7d4c54d26240287d14719629508f495eb Mon Sep 17 00:00:00 2001 From: Din Date: Wed, 10 Dec 2025 16:45:50 +0000 Subject: [PATCH 08/10] fix content detachment if inner use_span fails --- src/lmnr/sdk/laminar.py | 52 ++++++++++++++++++++--------------------- 1 file changed, 26 insertions(+), 26 deletions(-) diff --git a/src/lmnr/sdk/laminar.py b/src/lmnr/sdk/laminar.py index 8dacf6c5..fd9cfacf 100644 --- a/src/lmnr/sdk/laminar.py +++ b/src/lmnr/sdk/laminar.py @@ -573,29 +573,30 @@ def start_as_current_span( "Tags will be ignored." ) - with tracer.start_as_current_span( - name, - context=ctx, - attributes={ - SPAN_TYPE: span_type, - PARENT_SPAN_PATH: parsed["path"], - PARENT_SPAN_IDS_PATH: parsed["span_ids_path"], - **(label_props), - **(tag_props), - # Association properties are attached to context above - # and the relevant attributes are populated in the processor - }, - ) as span: - if not isinstance(span, LaminarSpan): - span = LaminarSpan(span) - span.set_input(input) - yield span - try: - detach_context(isolated_context_token) - context_api.detach(ctx_token) - except Exception: - pass + with tracer.start_as_current_span( + name, + context=ctx, + attributes={ + SPAN_TYPE: span_type, + PARENT_SPAN_PATH: parsed["path"], + PARENT_SPAN_IDS_PATH: parsed["span_ids_path"], + **(label_props), + **(tag_props), + # Association properties are attached to context above + # and the relevant attributes are populated in the processor + }, + ) as span: + if not isinstance(span, LaminarSpan): + span = LaminarSpan(span) + span.set_input(input) + yield span + finally: + try: + detach_context(isolated_context_token) + context_api.detach(ctx_token) + except Exception: + pass @classmethod def start_span( @@ -852,10 +853,6 @@ def use_span( else: yield LaminarSpan(span) - context_api.detach(context_token) - detach_context(isolated_context_token) - wrapper.pop_span_context() - # Record only exceptions that inherit Exception class but not BaseException, because # classes that directly inherit BaseException are not technically errors, e.g. GeneratorExit. # See https://github.com/open-telemetry/opentelemetry-python/issues/4484 @@ -883,6 +880,9 @@ def use_span( raise finally: + context_api.detach(context_token) + detach_context(isolated_context_token) + wrapper.pop_span_context() if end_on_exit: span.end() From 503ee3916c94054b54be718842916871f2732ef7 Mon Sep 17 00:00:00 2001 From: Din Date: Wed, 10 Dec 2025 19:05:33 +0000 Subject: [PATCH 09/10] fix small bugs --- src/lmnr/opentelemetry_lib/decorators/__init__.py | 4 ++++ src/lmnr/sdk/laminar.py | 12 +++++++----- 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/src/lmnr/opentelemetry_lib/decorators/__init__.py b/src/lmnr/opentelemetry_lib/decorators/__init__.py index ca6a2047..a23172a4 100644 --- a/src/lmnr/opentelemetry_lib/decorators/__init__.py +++ b/src/lmnr/opentelemetry_lib/decorators/__init__.py @@ -79,6 +79,8 @@ def _process_input( try: if input_formatter is not None: inp = input_formatter(*args, **kwargs) + if not isinstance(inp, str): + inp = json_dumps(inp) else: inp = get_input_from_func_args( fn, @@ -115,6 +117,8 @@ def _process_output( try: if output_formatter is not None: output = output_formatter(result) + if not isinstance(output, str): + output = json_dumps(output) else: output = result diff --git a/src/lmnr/sdk/laminar.py b/src/lmnr/sdk/laminar.py index fd9cfacf..98396612 100644 --- a/src/lmnr/sdk/laminar.py +++ b/src/lmnr/sdk/laminar.py @@ -880,11 +880,13 @@ def use_span( raise finally: - context_api.detach(context_token) - detach_context(isolated_context_token) - wrapper.pop_span_context() - if end_on_exit: - span.end() + try: + context_api.detach(context_token) + detach_context(isolated_context_token) + wrapper.pop_span_context() + finally: + if end_on_exit: + span.end() @classmethod def start_active_span( From 4d9a9574f56bbbd6f4e33526e699c2666c8683dc Mon Sep 17 00:00:00 2001 From: Din Date: Thu, 11 Dec 2025 14:24:27 +0000 Subject: [PATCH 10/10] revert json_dumps change in observe (caused double dumps) issues --- src/lmnr/opentelemetry_lib/decorators/__init__.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/src/lmnr/opentelemetry_lib/decorators/__init__.py b/src/lmnr/opentelemetry_lib/decorators/__init__.py index a23172a4..ca6a2047 100644 --- a/src/lmnr/opentelemetry_lib/decorators/__init__.py +++ b/src/lmnr/opentelemetry_lib/decorators/__init__.py @@ -79,8 +79,6 @@ def _process_input( try: if input_formatter is not None: inp = input_formatter(*args, **kwargs) - if not isinstance(inp, str): - inp = json_dumps(inp) else: inp = get_input_from_func_args( fn, @@ -117,8 +115,6 @@ def _process_output( try: if output_formatter is not None: output = output_formatter(result) - if not isinstance(output, str): - output = json_dumps(output) else: output = result