add span path to LaminarSpanContext (#171)

* add span path to LaminarSpanContext

* improve tests: checks parent id

* remove debug print statements

* remove breaking openai import

* remove path context

* bump version to 0.7.4

Author: dinmukhamedm

pyproject.toml

@@ -6,7 +6,7 @@
 
 [project]
 name = "lmnr"
-version = "0.7.3"
+version = "0.7.4"
 description = "Python SDK for Laminar"
 authors = [
   { name = "lmnr.ai", email = "[email protected]" }
@@ -124,7 +124,7 @@ dev = [
   "pytest-asyncio>=0.26.0",
   "playwright>=1.52.0",
   "vcrpy>=7.0.0",
-  "openai>=1.99.3",
+  "openai>=1.99.7",
   "pytest-recording>=0.13.4",
   "patchright>=1.52.3",
   "google-genai>=1.19.0",

src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/openai/shared/chat_wrappers.py

@@ -56,8 +56,8 @@
 from opentelemetry.trace.status import Status, StatusCode
 from wrapt import ObjectProxy
 
-from openai.types.chat import ChatCompletionMessageToolCall
 from openai.types.chat.chat_completion_message import FunctionCall
+import pydantic
 
 SPAN_NAME = "openai.chat"
 PROMPT_FILTER_KEY = "prompt_filter_results"
@@ -995,7 +995,7 @@ async def _abuild_from_streaming_response(
 
 
 def _parse_tool_calls(
-    tool_calls: Optional[List[Union[dict, ChatCompletionMessageToolCall]]],
+    tool_calls: Optional[List[Union[dict, pydantic.BaseModel]]],
 ) -> Union[List[ToolCall], None]:
     """
     Util to correctly parse the tool calls data from the OpenAI API to this module's

src/lmnr/opentelemetry_lib/tracing/__init__.py

@@ -74,7 +74,7 @@ def __new__(
             if not hasattr(cls, "instance"):
                 cls._initialize_logger(cls)
                 obj = super(TracerWrapper, cls).__new__(cls)
-                
+
                 # Store session recording options
                 cls.session_recording_options = session_recording_options or {}
 
@@ -249,13 +249,13 @@ def flush(self):
             self._logger.warning("TracerWrapper not fully initialized, cannot flush")
             return False
         return self._span_processor.force_flush()
-    
+
     @classmethod
     def get_session_recording_options(cls) -> SessionRecordingOptions:
         """Get the session recording options set during initialization."""
         return cls.session_recording_options
 
-    def get_tracer(self):
+    def get_tracer(self) -> trace.Tracer:
         if self._tracer_provider is None:
             return trace.get_tracer_provider().get_tracer(TRACER_NAME)
         return self._tracer_provider.get_tracer(TRACER_NAME)

src/lmnr/opentelemetry_lib/tracing/attributes.py

@@ -14,6 +14,8 @@
 SPAN_TYPE = "lmnr.span.type"
 SPAN_PATH = "lmnr.span.path"
 SPAN_IDS_PATH = "lmnr.span.ids_path"
+PARENT_SPAN_PATH = "lmnr.span.parent_path"
+PARENT_SPAN_IDS_PATH = "lmnr.span.parent_ids_path"
 SPAN_INSTRUMENTATION_SOURCE = "lmnr.span.instrumentation_source"
 SPAN_SDK_VERSION = "lmnr.span.sdk_version"
 SPAN_LANGUAGE_VERSION = "lmnr.span.language_version"

src/lmnr/opentelemetry_lib/tracing/processor.py

@@ -7,9 +7,11 @@
     SimpleSpanProcessor,
 )
 from opentelemetry.sdk.trace import Span
-from opentelemetry.context import Context, get_value, get_current, set_value
+from opentelemetry.context import Context, get_value
 
 from lmnr.opentelemetry_lib.tracing.attributes import (
+    PARENT_SPAN_IDS_PATH,
+    PARENT_SPAN_PATH,
     SPAN_IDS_PATH,
     SPAN_INSTRUMENTATION_SOURCE,
     SPAN_LANGUAGE_VERSION,
@@ -52,20 +54,18 @@ def __init__(
         )
 
     def on_start(self, span: Span, parent_context: Context | None = None):
-        span_path_in_context = get_value("span_path", parent_context or get_current())
-        parent_span_path = span_path_in_context or (
+        parent_span_path = list(span.attributes.get(PARENT_SPAN_PATH, tuple())) or (
             self.__span_id_to_path.get(span.parent.span_id) if span.parent else None
         )
-        parent_span_ids_path = (
-            self.__span_id_lists.get(span.parent.span_id, []) if span.parent else []
-        )
+        parent_span_ids_path = list(
+            span.attributes.get(PARENT_SPAN_IDS_PATH, tuple())
+        ) or (self.__span_id_lists.get(span.parent.span_id, []) if span.parent else [])
         span_path = parent_span_path + [span.name] if parent_span_path else [span.name]
         span_ids_path = parent_span_ids_path + [
             str(uuid.UUID(int=span.get_span_context().span_id))
         ]
         span.set_attribute(SPAN_PATH, span_path)
         span.set_attribute(SPAN_IDS_PATH, span_ids_path)
-        set_value("span_path", span_path, get_current())
         self.__span_id_to_path[span.get_span_context().span_id] = span_path
         self.__span_id_lists[span.get_span_context().span_id] = span_ids_path
 

src/lmnr/sdk/laminar.py

@@ -13,6 +13,10 @@
 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,
@@ -316,9 +320,29 @@ def start_as_current_span(
 
         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(
-                    parent_span_context, cls.__logger
+                    laminar_span_context, cls.__logger
                 )
                 ctx = trace.set_span_in_context(
                     trace.NonRecordingSpan(span_context), ctx
@@ -352,6 +376,8 @@ def start_as_current_span(
                 context=ctx,
                 attributes={
                     SPAN_TYPE: span_type,
+                    PARENT_SPAN_PATH: path,
+                    PARENT_SPAN_IDS_PATH: span_ids_path,
                     **(label_props),
                     **(tag_props),
                 },
@@ -454,9 +480,29 @@ 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(
-                    parent_span_context, cls.__logger
+                    laminar_span_context, cls.__logger
                 )
                 ctx = trace.set_span_in_context(
                     trace.NonRecordingSpan(span_context), ctx
@@ -491,6 +537,8 @@ def bar():
                 context=ctx,
                 attributes={
                     SPAN_TYPE: span_type,
+                    PARENT_SPAN_PATH: path,
+                    PARENT_SPAN_IDS_PATH: span_ids_path,
                     **(label_props),
                     **(tag_props),
                 },
@@ -662,6 +710,8 @@ def get_laminar_span_context(
             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, []),
         )
 
     @classmethod

src/lmnr/sdk/types.py

@@ -169,6 +169,8 @@ class LaminarSpanContext(pydantic.BaseModel):
     trace_id: uuid.UUID
     span_id: uuid.UUID
     is_remote: bool = pydantic.Field(default=False)
+    span_path: list[str] = pydantic.Field(default=[])
+    span_ids_path: list[str] = pydantic.Field(default=[])  # stringified UUIDs
 
     def __str__(self) -> str:
         return self.model_dump_json()
@@ -199,7 +201,7 @@ def try_to_otel_span_context(
                 "Please use `LaminarSpanContext` instead."
             )
             return span_context
-        elif isinstance(span_context, dict) or isinstance(span_context, str):
+        elif isinstance(span_context, (dict, str)):
             try:
                 laminar_span_context = cls.deserialize(span_context)
                 return SpanContext(
@@ -221,6 +223,9 @@ def deserialize(cls, data: dict[str, Any] | str) -> "LaminarSpanContext":
                 "trace_id": data.get("trace_id") or data.get("traceId"),
                 "span_id": data.get("span_id") or data.get("spanId"),
                 "is_remote": data.get("is_remote") or data.get("isRemote", False),
+                "span_path": data.get("span_path") or data.get("spanPath", []),
+                "span_ids_path": data.get("span_ids_path")
+                or data.get("spanIdsPath", []),
             }
             return cls.model_validate(converted_data)
         elif isinstance(data, str):
@@ -356,5 +361,6 @@ class MaskInputOptions(TypedDict):
     email: bool | None
     tel: bool | None
 
+
 class SessionRecordingOptions(TypedDict):
-    mask_input_options: MaskInputOptions | None
+    mask_input_options: MaskInputOptions | None

src/lmnr/version.py

@@ -3,7 +3,7 @@
 from packaging import version
 
 
-__version__ = "0.7.3"
+__version__ = "0.7.4"
 PYTHON_VERSION = f"{sys.version_info.major}.{sys.version_info.minor}"