Move span methods to LaminarSpan, propagate association properties down to the context

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",

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",

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

src/lmnr/opentelemetry_lib/decorators/__init__.py

@@ -1,75 +1,40 @@
 from functools import wraps
-import pydantic
-import orjson
 import types
 from typing import Any, AsyncGenerator, Callable, Generator, Literal, TypeVar
 
 from opentelemetry import context as context_api
 from opentelemetry.trace import Span, Status, StatusCode
 
 from lmnr.opentelemetry_lib.tracing.context import (
-    CONTEXT_SESSION_ID_KEY,
-    CONTEXT_USER_ID_KEY,
+    CONTEXT_METADATA_KEY,
     attach_context,
     detach_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 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 +45,17 @@ 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 +79,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 +115,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 +145,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,17 +161,20 @@ 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,
             )
+
+            # 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
@@ -255,6 +227,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,17 +243,20 @@ 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,
             )
+
+            # 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

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__)
 

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

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
 

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

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,

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

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
 
 

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

src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/openhands_ai/__init__.py

@@ -3,15 +3,14 @@
 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,
     USER_ID,
 )
 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__