lmnr 0.4.33__tar.gz → 0.4.35__tar.gz

{lmnr-0.4.33 → lmnr-0.4.35}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lmnr
-Version: 0.4.33
+Version: 0.4.35
 Summary: Python SDK for Laminar AI
 License: Apache-2.0
 Author: lmnr.ai

{lmnr-0.4.33 → lmnr-0.4.35}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "lmnr"
-version = "0.4.33"
+version = "0.4.35"
 description = "Python SDK for Laminar AI"
 authors = [
   { name = "lmnr.ai", email = "founders@lmnr.ai" }
@@ -11,7 +11,7 @@ license = "Apache-2.0"
 
 [tool.poetry]
 name = "lmnr"
-version = "0.4.33"
+version = "0.4.35"
 description = "Python SDK for Laminar AI"
 authors = ["lmnr.ai"]
 readme = "README.md"

{lmnr-0.4.33 → lmnr-0.4.35}/src/lmnr/sdk/laminar.py

@@ -1,18 +1,19 @@
 from contextlib import contextmanager
 from contextvars import Context
-from opentelemetry import context, trace
-from opentelemetry.util.types import AttributeValue
-from opentelemetry.context import set_value, attach, detach
 from lmnr.traceloop_sdk import Traceloop
 from lmnr.traceloop_sdk.instruments import Instruments
 from lmnr.traceloop_sdk.tracing import get_tracer
 from lmnr.traceloop_sdk.tracing.attributes import (
+    ASSOCIATION_PROPERTIES,
     Attributes,
     SPAN_TYPE,
     OVERRIDE_PARENT_SPAN,
 )
 from lmnr.traceloop_sdk.decorators.base import json_dumps
+from opentelemetry import context, trace
+from opentelemetry.context import attach, detach, set_value
 from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
+from opentelemetry.util.types import AttributeValue
 
 from pydantic.alias_generators import to_snake
 from typing import Any, Literal, Optional, Set, Union
@@ -39,6 +40,7 @@ from lmnr.traceloop_sdk.tracing.attributes import (
 )
 from lmnr.traceloop_sdk.tracing.tracing import (
     get_span_path,
+    remove_association_properties,
     set_association_properties,
     update_association_properties,
 )
@@ -298,6 +300,7 @@ class Laminar:
         span_type: Union[Literal["DEFAULT"], Literal["LLM"]] = "DEFAULT",
         context: Optional[Context] = None,
         trace_id: Optional[uuid.UUID] = None,
+        labels: Optional[dict[str, str]] = None,
     ):
         """Start a new span as the current span. Useful for manual
         instrumentation. If `span_type` is set to `"LLM"`, you should report
@@ -323,6 +326,8 @@ class Laminar:
             trace_id (Optional[uuid.UUID], optional): [EXPERIMENTAL] override\
                 the trace id for the span. If not provided, use the current\
                 trace id. Defaults to None.
+            labels (Optional[dict[str, str]], optional): labels to set for the\
+                span. Defaults to None.
         """
 
         with get_tracer() as tracer:
@@ -345,10 +350,26 @@ class Laminar:
                         " is not a valid UUID"
                     )
             ctx_token = attach(ctx)
+            label_props = {}
+            try:
+                if labels:
+                    label_props = dict(
+                        (f"{ASSOCIATION_PROPERTIES}.label.{k}", json_dumps(v))
+                        for k, v in labels.items()  # noqa: F821
+                    )
+            except Exception:
+                cls.__logger.warning(
+                    f"`start_as_current_span` Could not set labels: {labels}. "
+                    "They will be propagated to the next span."
+                )
             with tracer.start_as_current_span(
                 name,
                 context=ctx,
-                attributes={SPAN_PATH: span_path, SPAN_TYPE: span_type},
+                attributes={
+                    SPAN_PATH: span_path,
+                    SPAN_TYPE: span_type,
+                    **(label_props),
+                },
             ) as span:
                 if trace_id is not None and isinstance(trace_id, uuid.UUID):
                     span.set_attribute(OVERRIDE_PARENT_SPAN, True)
@@ -365,6 +386,41 @@ class Laminar:
             except Exception:
                 pass
 
+    @classmethod
+    @contextmanager
+    def with_labels(cls, labels: dict[str, str], context: Optional[Context] = None):
+        """Set labels for spans within this `with` context. This is useful for
+        adding labels to the spans created in the auto-instrumentations.
+
+        Requirements:
+        - Labels must be created in your project in advance.
+        - Keys must be strings from your label names.
+        - Values must be strings matching the label's allowed values.
+
+        Usage example:
+        ```python
+        with Laminar.with_labels({"sentiment": "positive"}):
+            openai_client.chat.completions.create()
+        ```
+        """
+        with get_tracer():
+            label_props = labels.copy()
+            label_props = dict(
+                (f"label.{k}", json_dumps(v)) for k, v in label_props.items()
+            )
+            update_association_properties(
+                label_props, set_on_current_span=False, context=context
+            )
+            yield
+            try:
+                remove_association_properties(label_props)
+            except Exception:
+                cls.__logger.warning(
+                    f"`with_labels` Could not remove labels: {labels}. They will be "
+                    "propagated to the next span."
+                )
+                pass
+
     @classmethod
     def start_span(
         cls,
@@ -373,6 +429,7 @@ class Laminar:
         span_type: Union[Literal["DEFAULT"], Literal["LLM"]] = "DEFAULT",
         context: Optional[Context] = None,
         trace_id: Optional[uuid.UUID] = None,
+        labels: Optional[dict[str, str]] = None,
     ):
         """Start a new span. Useful for manual instrumentation.
         If `span_type` is set to `"LLM"`, you should report usage and response
@@ -417,6 +474,8 @@ class Laminar:
             trace_id (Optional[uuid.UUID], optional): [EXPERIMENTAL] override\
                 the trace id for the span. If not provided, use the current\
                 trace id. Defaults to None.
+            labels (Optional[dict[str, str]], optional): labels to set for the\
+                span. Defaults to None.
         """
         with get_tracer() as tracer:
             span_path = get_span_path(name)
@@ -437,10 +496,26 @@ class Laminar:
                         "trace_id provided to `Laminar.start_span`"
                         " is not a valid UUID"
                     )
+            label_props = {}
+            try:
+                if labels:
+                    label_props = dict(
+                        (f"{ASSOCIATION_PROPERTIES}.label.{k}", json_dumps(v))
+                        for k, v in labels.items()  # noqa: F821
+                    )
+            except Exception:
+                cls.__logger.warning(
+                    f"`start_span` Could not set labels: {labels}. They will be "
+                    "propagated to the next span."
+                )
             span = tracer.start_span(
                 name,
                 context=ctx,
-                attributes={SPAN_PATH: span_path, SPAN_TYPE: span_type},
+                attributes={
+                    SPAN_PATH: span_path,
+                    SPAN_TYPE: span_type,
+                    **(label_props),
+                },
             )
             if trace_id is not None and isinstance(trace_id, uuid.UUID):
                 span.set_attribute(OVERRIDE_PARENT_SPAN, True)

{lmnr-0.4.33 → lmnr-0.4.35}/src/lmnr/traceloop_sdk/tracing/tracing.py

@@ -1,8 +1,18 @@
 import atexit
+import copy
 import logging
-import os
 
 
+from contextvars import Context
+from lmnr.traceloop_sdk.instruments import Instruments
+from lmnr.traceloop_sdk.tracing.attributes import (
+    ASSOCIATION_PROPERTIES,
+    SPAN_INSTRUMENTATION_SOURCE,
+    SPAN_PATH,
+)
+from lmnr.traceloop_sdk.tracing.content_allow_list import ContentAllowList
+from lmnr.traceloop_sdk.utils import is_notebook
+from lmnr.traceloop_sdk.utils.package_check import is_package_installed
 from opentelemetry import trace
 from opentelemetry.exporter.otlp.proto.http.trace_exporter import (
     OTLPSpanExporter as HTTPExporter,
@@ -10,28 +20,19 @@ from opentelemetry.exporter.otlp.proto.http.trace_exporter import (
 from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import (
     OTLPSpanExporter as GRPCExporter,
 )
+from opentelemetry.instrumentation.threading import ThreadingInstrumentor
+from opentelemetry.context import get_value, attach, set_value
+from opentelemetry.propagate import set_global_textmap
+from opentelemetry.propagators.textmap import TextMapPropagator
 from opentelemetry.sdk.resources import Resource
 from opentelemetry.sdk.trace import TracerProvider, SpanProcessor
-from opentelemetry.propagators.textmap import TextMapPropagator
-from opentelemetry.propagate import set_global_textmap
 from opentelemetry.sdk.trace.export import (
     SpanExporter,
     SimpleSpanProcessor,
     BatchSpanProcessor,
 )
 from opentelemetry.trace import get_tracer_provider, ProxyTracerProvider
-from opentelemetry.context import get_value, attach, set_value
-from opentelemetry.instrumentation.threading import ThreadingInstrumentor
 
-from lmnr.traceloop_sdk.instruments import Instruments
-from lmnr.traceloop_sdk.tracing.attributes import (
-    ASSOCIATION_PROPERTIES,
-    SPAN_INSTRUMENTATION_SOURCE,
-    SPAN_PATH,
-)
-from lmnr.traceloop_sdk.tracing.content_allow_list import ContentAllowList
-from lmnr.traceloop_sdk.utils import is_notebook
-from lmnr.traceloop_sdk.utils.package_check import is_package_installed
 from typing import Dict, Optional, Set
 
 
@@ -188,15 +189,27 @@ def set_association_properties(properties: dict) -> None:
     _set_association_properties_attributes(span, properties)
 
 
-def update_association_properties(properties: dict) -> None:
+def update_association_properties(
+    properties: dict,
+    set_on_current_span: bool = True,
+    context: Optional[Context] = None,
+) -> None:
     """Only adds or updates properties that are not already present"""
-    association_properties = get_value("association_properties") or {}
+    association_properties = get_value("association_properties", context) or {}
     association_properties.update(properties)
 
-    attach(set_value("association_properties", association_properties))
+    attach(set_value("association_properties", association_properties, context))
 
-    span = trace.get_current_span()
-    _set_association_properties_attributes(span, properties)
+    if set_on_current_span:
+        span = trace.get_current_span()
+        _set_association_properties_attributes(span, properties)
+
+
+def remove_association_properties(properties: dict) -> None:
+    props: dict = copy.copy(get_value("association_properties") or {})
+    for k in properties.keys():
+        props.pop(k, None)
+    set_association_properties(props)
 
 
 def _set_association_properties_attributes(span, properties: dict) -> None:

lmnr-0.4.35/src/lmnr/traceloop_sdk/utils/package_check.py

@@ -0,0 +1,6 @@
+from importlib.metadata import distributions
+
+installed_packages = {dist.metadata["Name"].lower() for dist in distributions()}
+
+def is_package_installed(package_name: str) -> bool:
+    return package_name.lower() in installed_packages

lmnr-0.4.33/src/lmnr/traceloop_sdk/utils/package_check.py

@@ -1,8 +0,0 @@
-import pkg_resources
-
-installed_packages = {p.key for p in pkg_resources.working_set}
-
-
-def is_package_installed(package_name: str) -> bool:
-    # return importlib.util.find_spec(package_name) is not None
-    return package_name in installed_packages