lmnr 0.6.7__py3-none-any.whl → 0.6.9__py3-none-any.whl

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

@@ -0,0 +1,119 @@
+"""OpenTelemetry Langgraph instrumentation"""
+
+import json
+import logging
+from typing import Collection
+
+from .utils import (
+    with_tracer_wrapper,
+)
+
+from langchain_core.runnables.graph import Graph
+from opentelemetry.trace import Tracer
+from wrapt import wrap_function_wrapper
+from opentelemetry.trace import get_tracer, get_current_span
+
+from lmnr.opentelemetry_lib.tracing.context_properties import (
+    update_association_properties,
+)
+
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.instrumentation.utils import unwrap
+
+
+logger = logging.getLogger(__name__)
+
+_instruments = ("langgraph >= 0.1.0",)
+
+
+@with_tracer_wrapper
+def wrap_pregel_stream(tracer: Tracer, to_wrap, wrapped, instance, args, kwargs):
+    graph: Graph = instance.get_graph()
+    nodes = [
+        {
+            "id": node.id,
+            "name": node.name,
+            "metadata": node.metadata,
+        }
+        for node in graph.nodes.values()
+    ]
+    edges = [
+        {
+            "source": edge.source,
+            "target": edge.target,
+            "conditional": edge.conditional,
+        }
+        for edge in graph.edges
+    ]
+    update_association_properties(
+        {
+            "langgraph.edges": json.dumps(edges),
+            "langgraph.nodes": json.dumps(nodes),
+        },
+    )
+    return wrapped(*args, **kwargs)
+
+
+@with_tracer_wrapper
+async def async_wrap_pregel_stream(
+    tracer: Tracer, to_wrap, wrapped, instance, args, kwargs
+):
+    graph: Graph = await instance.aget_graph()
+    nodes = [
+        {
+            "id": node.id,
+            "name": node.name,
+            "metadata": node.metadata,
+        }
+        for node in graph.nodes.values()
+    ]
+    edges = [
+        {
+            "source": edge.source,
+            "target": edge.target,
+            "conditional": edge.conditional,
+        }
+        for edge in graph.edges
+    ]
+    update_association_properties(
+        {
+            "langgraph.edges": json.dumps(edges),
+            "langgraph.nodes": json.dumps(nodes),
+        },
+    )
+    return await wrapped(*args, **kwargs)
+
+
+class LanggraphInstrumentor(BaseInstrumentor):
+    """An instrumentor for Langgraph."""
+
+    def __init__(self):
+        super().__init__()
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return _instruments
+
+    def _instrument(self, **kwargs):
+        tracer_provider = kwargs.get("tracer_provider")
+        tracer = get_tracer(__name__, "0.0.1a0", tracer_provider)
+
+        wrap_function_wrapper(
+            module="langgraph.pregel",
+            name="Pregel.stream",
+            wrapper=wrap_pregel_stream(tracer, "Pregel.stream"),
+        )
+        wrap_function_wrapper(
+            module="langgraph.pregel",
+            name="Pregel.astream",
+            wrapper=async_wrap_pregel_stream(tracer, "Pregel.astream"),
+        )
+
+    def _uninstrument(self, **kwargs):
+        unwrap(
+            module="langgraph.pregel",
+            name="Pregel.stream",
+        )
+        unwrap(
+            module="langgraph.pregel",
+            name="Pregel.astream",
+        )

lmnr/opentelemetry_lib/opentelemetry/instrumentation/langgraph/utils.py

@@ -0,0 +1,60 @@
+import logging
+import traceback
+
+import pydantic
+from opentelemetry.trace import Span
+from typing import Any
+
+
+def set_span_attribute(span: Span, name: str, value: str):
+    if value is not None:
+        if value != "":
+            span.set_attribute(name, value)
+    return
+
+
+def dont_throw(func):
+    """
+    A decorator that wraps the passed in function and logs exceptions instead of throwing them.
+
+    @param func: The function to wrap
+    @return: The wrapper function
+    """
+    # Obtain a logger specific to the function's module
+    logger = logging.getLogger(func.__module__)
+
+    def wrapper(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except Exception:
+            logger.debug(
+                "Laminar failed to trace in %s, error: %s",
+                func.__name__,
+                traceback.format_exc(),
+            )
+
+    return wrapper
+
+
+def to_dict(obj: pydantic.BaseModel | dict) -> dict[str, Any]:
+    try:
+        if isinstance(obj, pydantic.BaseModel):
+            return obj.model_dump()
+        elif isinstance(obj, dict):
+            return obj
+        else:
+            return dict(obj)
+    except Exception:
+        return dict(obj)
+
+
+def with_tracer_wrapper(func):
+    """Helper for providing tracer for wrapper functions."""
+
+    def _with_tracer(tracer, to_wrap):
+        def wrapper(wrapped, instance, args, kwargs):
+            return func(tracer, to_wrap, wrapped, instance, args, kwargs)
+
+        return wrapper
+
+    return _with_tracer

lmnr/opentelemetry_lib/tracing/_instrument_initializers.py

@@ -171,6 +171,18 @@ class LangchainInstrumentorInitializer(InstrumentorInitializer):
         return LangchainInstrumentor()
 
 
+class LanggraphInstrumentorInitializer(InstrumentorInitializer):
+    def init_instrumentor(self, *args, **kwargs) -> BaseInstrumentor | None:
+        if not is_package_installed("langgraph"):
+            return None
+        if not is_package_installed("langchain-core"):
+            return None
+
+        from ..opentelemetry.instrumentation.langgraph import LanggraphInstrumentor
+
+        return LanggraphInstrumentor()
+
+
 class LlamaIndexInstrumentorInitializer(InstrumentorInitializer):
     def init_instrumentor(self, *args, **kwargs) -> BaseInstrumentor | None:
         if not (

lmnr/opentelemetry_lib/tracing/context_properties.py

@@ -6,10 +6,12 @@ from lmnr.opentelemetry_lib.tracing.attributes import (
 )
 
 from opentelemetry.context import Context, attach, set_value, get_value
-from opentelemetry.trace import Span
+from opentelemetry.sdk.trace import Span
 from opentelemetry import trace
 
 
+# TODO: delete this once deprecated Laminar.with_labels is removed. The logic
+# should be moved into Laminar.set_tracing_level
 def set_association_properties(properties: dict) -> None:
     attach(set_value("association_properties", properties))
 
@@ -17,10 +19,13 @@ def set_association_properties(properties: dict) -> None:
     _set_association_properties_attributes(span, properties)
 
 
+# TODO: delete this once deprecated Laminar.with_labels is removed
 def get_association_properties(context: Context | None = None) -> dict:
     return get_value("association_properties", context) or {}
 
 
+# TODO: delete this once deprecated Laminar.with_labels is removed. The logic
+# should be moved into Laminar.set_tracing_level
 def update_association_properties(
     properties: dict,
     set_on_current_span: bool = True,
@@ -37,6 +42,7 @@ def update_association_properties(
         _set_association_properties_attributes(span, properties)
 
 
+# TODO: this logic should be moved into Laminar.set_tracing_level
 def remove_association_properties(properties: dict) -> None:
     props: dict = copy.copy(get_value("association_properties") or {})
     for k in properties.keys():
@@ -45,8 +51,15 @@ def remove_association_properties(properties: dict) -> None:
 
 
 def _set_association_properties_attributes(span: Span, properties: dict) -> None:
+    if not span.is_recording():
+        return
     for key, value in properties.items():
         if key == TRACING_LEVEL:
             span.set_attribute(f"lmnr.internal.{TRACING_LEVEL}", value)
             continue
+        if (
+            key in ["langgraph.edges", "langgraph.nodes"]
+            and span.name != "LangGraph.workflow"
+        ):
+            continue
         span.set_attribute(f"{ASSOCIATION_PROPERTIES}.{key}", value)

lmnr/opentelemetry_lib/tracing/instruments.py

@@ -26,6 +26,7 @@ class Instruments(Enum):
     HAYSTACK = "haystack"
     LANCEDB = "lancedb"
     LANGCHAIN = "langchain"
+    LANGGRAPH = "langgraph"
     LLAMA_INDEX = "llama_index"
     MARQO = "marqo"
     MCP = "mcp"
@@ -62,6 +63,7 @@ INSTRUMENTATION_INITIALIZERS: dict[
     Instruments.HAYSTACK: initializers.HaystackInstrumentorInitializer(),
     Instruments.LANCEDB: initializers.LanceDBInstrumentorInitializer(),
     Instruments.LANGCHAIN: initializers.LangchainInstrumentorInitializer(),
+    Instruments.LANGGRAPH: initializers.LanggraphInstrumentorInitializer(),
     Instruments.LLAMA_INDEX: initializers.LlamaIndexInstrumentorInitializer(),
     Instruments.MARQO: initializers.MarqoInstrumentorInitializer(),
     Instruments.MCP: initializers.MCPInstrumentorInitializer(),

lmnr/opentelemetry_lib/tracing/processor.py

@@ -6,7 +6,7 @@ from opentelemetry.sdk.trace.export import (
     BatchSpanProcessor,
     SimpleSpanProcessor,
 )
-from opentelemetry.trace import Span
+from opentelemetry.sdk.trace import Span
 from opentelemetry.context import Context, get_value, get_current, set_value
 
 from lmnr.opentelemetry_lib.tracing.attributes import (

lmnr/sdk/decorators.py

@@ -9,9 +9,11 @@ from typing import Any, Callable, Literal, TypeVar, cast
 from typing_extensions import ParamSpec
 
 from lmnr.opentelemetry_lib.tracing.attributes import SESSION_ID
+from lmnr.sdk.log import get_default_logger
 
 from .utils import is_async
 
+logger = get_default_logger(__name__)
 
 P = ParamSpec("P")
 R = TypeVar("R")
@@ -27,6 +29,7 @@ def observe(
     span_type: Literal["DEFAULT", "LLM", "TOOL"] = "DEFAULT",
     ignore_inputs: list[str] | None = None,
     metadata: dict[str, Any] | None = None,
+    tags: list[str] | None = None,
 ) -> Callable[[Callable[P, R]], Callable[P, R]]:
     """The main decorator entrypoint for Laminar. This is used to wrap
     functions and methods to create spans.
@@ -52,6 +55,8 @@ def observe(
             this argument. Defaults to None.
         metadata (dict[str, Any] | None, optional): Metadata to associate with\
             the trace. Must be JSON serializable. Defaults to None.
+        tags (list[str] | None, optional): Tags to associate with the trace.
+            Defaults to None.
     Raises:
         Exception: re-raises the exception if the wrapped function raises an\
             exception
@@ -79,6 +84,13 @@ def observe(
                     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
+            ):
+                logger.warning("Tags must be a list of strings. Tags will be ignored.")
+            else:
+                association_properties["tags"] = tags
         result = (
             aentity_method(
                 name=name,

lmnr/sdk/evaluations.py

@@ -233,10 +233,10 @@ class Evaluation:
             export_timeout_seconds=trace_export_timeout_seconds,
         )
 
-    async def run(self) -> Awaitable[None]:
+    async def run(self) -> Awaitable[dict[str, int | float]]:
         return await self._run()
 
-    async def _run(self) -> None:
+    async def _run(self) -> dict[str, int | float]:
         if isinstance(self.data, LaminarDataset):
             self.data.set_client(
                 LaminarClient(
@@ -261,11 +261,12 @@ class Evaluation:
         except Exception as e:
             self.reporter.stopWithError(e)
             await self._shutdown()
-            return
+            return {}
 
         average_scores = get_average_scores(result_datapoints)
         self.reporter.stop(average_scores, evaluation.projectId, evaluation.id)
         await self._shutdown()
+        return average_scores
 
     async def _shutdown(self):
         # We use flush() instead of shutdown() because multiple evaluations

lmnr/sdk/laminar.py

@@ -1,10 +1,13 @@
 from contextlib import contextmanager
 from contextvars import Context
+import warnings
+from typing_extensions import deprecated
 from lmnr.opentelemetry_lib import TracerManager
 from lmnr.opentelemetry_lib.tracing.instruments import Instruments
 from lmnr.opentelemetry_lib.tracing.tracer import get_tracer
 from lmnr.opentelemetry_lib.tracing.attributes import (
     ASSOCIATION_PROPERTIES,
+    USER_ID,
     Attributes,
     SPAN_TYPE,
 )
@@ -18,7 +21,6 @@ from opentelemetry.util.types import AttributeValue
 
 from typing import Any, Literal
 
-import copy
 import datetime
 import logging
 import os
@@ -37,7 +39,7 @@ from lmnr.opentelemetry_lib.tracing.context_properties import (
     set_association_properties,
     update_association_properties,
 )
-from lmnr.sdk.utils import from_env
+from lmnr.sdk.utils import from_env, is_otel_attribute_value_type
 
 from .log import VerboseColorfulFormatter
 
@@ -224,6 +226,7 @@ class Laminar:
         context: Context | None = None,
         labels: list[str] | None = None,
         parent_span_context: LaminarSpanContext | None = None,
+        tags: list[str] | None = None,
     ):
         """Start a new span as the current span. Useful for manual
         instrumentation. If `span_type` is set to `"LLM"`, you should report
@@ -256,8 +259,10 @@ class Laminar:
                 `Laminar.get_span_context`, `Laminar.get_span_context_dict` and\
                 `Laminar.get_span_context_str` for more information.
                 Defaults to None.
-            labels (list[str] | None, optional): labels to set for the\
-                span. Defaults to None.
+            labels (list[str] | None, optional): [DEPRECATED] Use tags\
+                instead. Labels to set for the span. Defaults to None.
+            tags (list[str] | None, optional): tags to set for the span.
+                Defaults to None.
         """
 
         if not cls.is_initialized():
@@ -283,18 +288,32 @@ class Laminar:
             label_props = {}
             try:
                 if labels:
+                    warnings.warn(
+                        "`Laminar.start_as_current_span` `labels` is deprecated. Use `tags` instead.",
+                        DeprecationWarning,
+                    )
                     label_props = {f"{ASSOCIATION_PROPERTIES}.labels": labels}
             except Exception:
                 cls.__logger.warning(
                     f"`start_as_current_span` Could not set labels: {labels}. "
                     "They will be propagated to the next span."
                 )
+            tag_props = {}
+            if tags:
+                if isinstance(tags, list) and all(isinstance(tag, str) for tag in tags):
+                    tag_props = {f"{ASSOCIATION_PROPERTIES}.tags": tags}
+                else:
+                    cls.__logger.warning(
+                        f"`start_as_current_span` Could not set tags: {tags}. Tags must be a list of strings. "
+                        "Tags will be ignored."
+                    )
             with tracer.start_as_current_span(
                 name,
                 context=ctx,
                 attributes={
                     SPAN_TYPE: span_type,
                     **(label_props),
+                    **(tag_props),
                 },
             ) as span:
                 if input is not None:
@@ -319,6 +338,10 @@ class Laminar:
 
     @classmethod
     @contextmanager
+    @deprecated(
+        "Use `Laminar.set_span_tags` or the `tags` argument of "
+        "`Laminar.start_as_current_span` or `Laminar.start_span` instead"
+    )
     def with_labels(cls, labels: list[str], context: Context | None = None):
         """Set labels for spans within this `with` context. This is useful for
         adding labels to the spans created in the auto-instrumentations.
@@ -334,6 +357,11 @@ class Laminar:
             openai_client.chat.completions.create()
         ```
         """
+        warnings.warn(
+            "`Laminar.with_labels` is deprecated. Use `Laminar.set_span_tags` or the `tags` argument of "
+            "`Laminar.start_as_current_span` or `Laminar.start_span` instead",
+            DeprecationWarning,
+        )
         if not cls.is_initialized():
             yield
             return
@@ -365,6 +393,7 @@ class Laminar:
         context: Context | None = None,
         parent_span_context: LaminarSpanContext | None = None,
         labels: dict[str, str] | None = None,
+        tags: list[str] | None = None,
     ):
         """Start a new span. Useful for manual instrumentation.
         If `span_type` is set to `"LLM"`, you should report usage and response
@@ -416,8 +445,10 @@ class Laminar:
                 `Laminar.get_span_context`, `Laminar.get_span_context_dict` and\
                 `Laminar.get_span_context_str` for more information.
                 Defaults to None.
-            labels (dict[str, str] | None, optional): labels to set for the\
-                span. Defaults to None.
+            tags (list[str] | None, optional): tags to set for the span.
+                Defaults to None.
+            labels (dict[str, str] | None, optional): [DEPRECATED] Use tags\
+                instead. Labels to set for the span. Defaults to None.
         """
         if not cls.is_initialized():
             return trace.NonRecordingSpan(
@@ -440,6 +471,10 @@ class Laminar:
             label_props = {}
             try:
                 if labels:
+                    warnings.warn(
+                        "`Laminar.start_span` `labels` is deprecated. Use `tags` instead.",
+                        DeprecationWarning,
+                    )
                     label_props = {
                         f"{ASSOCIATION_PROPERTIES}.labels": json_dumps(labels)
                     }
@@ -448,12 +483,22 @@ class Laminar:
                     f"`start_span` Could not set labels: {labels}. They will be "
                     "propagated to the next span."
                 )
+            tag_props = {}
+            if tags:
+                if isinstance(tags, list) and all(isinstance(tag, str) for tag in tags):
+                    tag_props = {f"{ASSOCIATION_PROPERTIES}.tags": tags}
+                else:
+                    cls.__logger.warning(
+                        f"`start_span` Could not set tags: {tags}. Tags must be a list of strings. "
+                        + "Tags will be ignored."
+                    )
             span = tracer.start_span(
                 name,
                 context=ctx,
                 attributes={
                     SPAN_TYPE: span_type,
                     **(label_props),
+                    **(tag_props),
                 },
             )
             if input is not None:
@@ -562,7 +607,6 @@ class Laminar:
                 cls.__logger.warning(
                     f"Attribute {key} is not a valid Laminar attribute."
                 )
-                continue
             if not isinstance(value, (str, int, float, bool)):
                 span.set_attribute(key.value, json_dumps(value))
             else:
@@ -649,11 +693,30 @@ class Laminar:
             TracerManager.shutdown()
 
     @classmethod
+    def set_span_tags(cls, tags: list[str]):
+        """Set the tags for the current span.
+
+        Args:
+            tags (list[str]): Tags to set for the span.
+        """
+        span = trace.get_current_span()
+        if span == trace.INVALID_SPAN:
+            cls.__logger.warning("No active span to set tags on")
+            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."
+            )
+            return
+        span.set_attribute(f"{ASSOCIATION_PROPERTIES}.tags", tags)
+
+    @classmethod
+    @deprecated("Use `Laminar.set_trace_session_id` instead")
     def set_session(
         cls,
         session_id: str | None = None,
     ):
-        """Set the session and user id for the current span and the context
+        """Set the session id for the current span and the context
         (i.e. any children spans created from the current span in the current
         thread).
 
@@ -663,12 +726,53 @@ class Laminar:
                             sessions/conversations.
                             Defaults to None.
         """
+        warnings.warn(
+            "`Laminar.set_session` is deprecated. Use `Laminar.set_trace_session_id` instead",
+            DeprecationWarning,
+        )
         association_properties = {}
         if session_id is not None:
             association_properties[SESSION_ID] = session_id
-        update_association_properties(association_properties)
+        # update_association_properties(association_properties)
+        span = trace.get_current_span()
+        if span == trace.INVALID_SPAN:
+            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)
+
+    @classmethod
+    def set_trace_session_id(cls, session_id: str | None = None):
+        """Set the session id for the current trace.
+        Overrides any existing session id.
+
+        Args:
+            session_id (str | None, optional): Custom session id. Defaults to None.
+        """
+        span = trace.get_current_span()
+        if span == trace.INVALID_SPAN:
+            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)
 
     @classmethod
+    def set_trace_user_id(cls, user_id: str | None = None):
+        """Set the user id for the current trace.
+        Overrides any existing user id.
+
+        Args:
+            user_id (str | None, optional): Custom user id. Defaults to None.
+        """
+        span = trace.get_current_span()
+        if span == trace.INVALID_SPAN:
+            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)
+
+    @classmethod
+    @deprecated("Use `Laminar.set_trace_metadata` instead")
     def set_metadata(cls, metadata: dict[str, str]):
         """Set the metadata for the current trace.
 
@@ -676,25 +780,37 @@ class Laminar:
             metadata (dict[str, str]): Metadata to set for the trace. Will be\
                 sent as attributes, so must be json serializable.
         """
+        warnings.warn(
+            "`Laminar.set_metadata` is deprecated. Use `Laminar.set_trace_metadata` instead",
+            DeprecationWarning,
+        )
         props = {f"metadata.{k}": json_dumps(v) for k, v in metadata.items()}
-        update_association_properties(props)
+        # update_association_properties(props)
+        span = trace.get_current_span()
+        if span == trace.INVALID_SPAN:
+            cls.__logger.warning("No active span to set metadata on")
+            return
+        for key, value in props.items():
+            span.set_attribute(key, value)
 
     @classmethod
-    def clear_metadata(cls):
-        """Clear the metadata from the context"""
-        props: dict = copy.copy(context_api.get_value("association_properties"))
-        metadata_keys = [k for k in props.keys() if k.startswith("metadata.")]
-        for k in metadata_keys:
-            props.pop(k)
-        set_association_properties(props)
+    def set_trace_metadata(cls, metadata: dict[str, AttributeValue]):
+        """Set the metadata for the current trace.
 
-    @classmethod
-    def clear_session(cls):
-        """Clear the session and user id from  the context"""
-        props: dict = copy.copy(context_api.get_value("association_properties"))
-        props.pop("session_id", None)
-        props.pop("user_id", None)
-        set_association_properties(props)
+        Args:
+            metadata (dict[str, AttributeValue]): Metadata to set for the trace.
+        """
+        span = trace.get_current_span()
+        if span == trace.INVALID_SPAN:
+            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)
+                )
 
     @classmethod
     def get_base_http_url(cls):
@@ -735,7 +851,8 @@ class Laminar:
         Args:
             trace_type (TraceType): Type of the trace
         """
-        association_properties = {
-            TRACE_TYPE: trace_type.value,
-        }
-        update_association_properties(association_properties)
+        span = trace.get_current_span()
+        if span == trace.INVALID_SPAN:
+            cls.__logger.warning("No active span to set trace type on")
+            return
+        span.set_attribute(f"{ASSOCIATION_PROPERTIES}.{TRACE_TYPE}", trace_type.value)

lmnr/sdk/types.py

@@ -81,7 +81,7 @@ class PartialEvaluationDatapoint(pydantic.BaseModel):
                 "traceId": str(self.trace_id),
                 "executorSpanId": str(self.executor_span_id),
                 "metadata": (
-                    serialize(self.metadata) if self.metadata is not None else None
+                    serialize(self.metadata) if self.metadata is not None else {}
                 ),
             }
         except Exception as e:
@@ -123,7 +123,7 @@ class EvaluationResultDatapoint(pydantic.BaseModel):
                 "executorSpanId": str(self.executor_span_id),
                 "index": self.index,
                 "metadata": (
-                    serialize(self.metadata) if self.metadata is not None else None
+                    serialize(self.metadata) if self.metadata is not None else {}
                 ),
             }
         except Exception as e: