lmnr 0.3.6__py3-none-any.whl → 0.4.0__py3-none-any.whl

lmnr/__init__.py

@@ -1,7 +1,3 @@
-from .sdk.client import Laminar
-from .sdk.decorators import observe, lmnr_context, wrap_llm_call
-from .sdk.interface import trace, TraceContext, SpanContext
-from .sdk.tracing_types import EvaluateEvent
+from .sdk.laminar import Laminar
 from .sdk.types import ChatMessage, PipelineRunError, PipelineRunResponse, NodeInput
-
-from .semantic_conventions import *
+from .sdk.decorators import observe

lmnr/sdk/decorators.py

@@ -1,281 +1,72 @@
-import datetime
-import functools
-from typing import Any, Callable, Literal, Optional, Union
-
-
-from .context import LaminarSingleton
-from .providers.fallback import FallbackProvider
-from ..semantic_conventions.gen_ai_spans import PROVIDER
-from .types import NodeInput, PipelineRunResponse
-from .utils import (
-    PROVIDER_NAME_TO_OBJECT,
-    get_input_from_func_args,
-    is_async,
-    is_method,
+from traceloop.sdk.decorators.base import (
+    entity_method,
+    aentity_method,
 )
+from opentelemetry.trace import INVALID_SPAN, get_current_span
+from traceloop.sdk import Traceloop
 
+from typing import Callable, Optional, ParamSpec, TypeVar, cast
 
-class LaminarDecorator:
-    def observe(
-        self,
-        *,
-        name: Optional[str] = None,
-        span_type: Optional[Literal["DEFAULT", "LLM"]] = "DEFAULT",
-        capture_input: bool = True,
-        capture_output: bool = True,
-        release: Optional[str] = None,
-        user_id: Optional[str] = None,
-        session_id: Optional[str] = None,
-    ):
-        """The main decorator entrypoint for Laminar. This is used to wrap functions and methods to create spans.
-
-        Args:
-            name (Optional[str], optional): Name of the span. Function name is used if not specified. Defaults to None.
-            span_type (Literal["DEFAULT", "LLM"], optional): Type of this span. Prefer `wrap_llm_call` instead of specifying
-                                                                                 this as "LLM" . Defaults to "DEFAULT".
-            capture_input (bool, optional): Whether to capture input parameters to the function. Defaults to True.
-            capture_output (bool, optional): Whether to capture returned type from the function. Defaults to True.
-            release (Optional[str], optional): Release version of your app. Useful for further grouping and analytics. Defaults to None.
-            user_id (Optional[str], optional): Custom user_id of your user. Useful for grouping and further analytics. Defaults to None.
-            session_id (Optional[str], optional): Custom session_id for your session. Random UUID is generated on Laminar side, if not specified.
-                                                  Defaults to None.
-
-        Raises:
-            Exception: re-raises the exception if the wrapped function raises an exception
-
-        Returns:
-            Any: Returns the result of the wrapped function
-        """
-        context_manager = LaminarSingleton().get()
-
-        def decorator(func: Callable):
-            @functools.wraps(func)
-            def wrapper(*args, **kwargs):
-                span = context_manager.observe_start(
-                    name=name or func.__name__,
-                    span_type=span_type,
-                    input=(
-                        get_input_from_func_args(func, is_method(func), args, kwargs)
-                        if capture_input
-                        else None
-                    ),
-                    user_id=user_id,
-                    session_id=session_id,
-                    release=release,
-                )
-                try:
-                    result = func(*args, **kwargs)
-                except Exception as e:
-                    context_manager.observe_end(result=None, span=span, error=e)
-                    raise e
-                context_manager.observe_end(
-                    result=result if capture_output else None, span=span
-                )
-                return result
-
-            @functools.wraps(func)
-            async def async_wrapper(*args, **kwargs):
-                span = context_manager.observe_start(
-                    name=name or func.__name__,
-                    span_type=span_type,
-                    input=(
-                        get_input_from_func_args(func, is_method(func), args, kwargs)
-                        if capture_input
-                        else None
-                    ),
-                    user_id=user_id,
-                    session_id=session_id,
-                    release=release,
-                )
-                try:
-                    result = await func(*args, **kwargs)
-                except Exception as e:
-                    context_manager.observe_end(result=None, span=span, error=e)
-                    raise e
-                context_manager.observe_end(
-                    result=result if capture_output else None, span=span
-                )
-                return result
-
-            return async_wrapper if is_async(func) else wrapper
-
-        return decorator
-
-    def update_current_span(
-        self,
-        metadata: Optional[dict[str, Any]] = None,
-        attributes: Optional[dict[str, Any]] = None,
-        override: bool = False,
-    ):
-        """Update the current span with any optional metadata.
-
-        Args:
-            metadata (Optional[dict[str, Any]], optional): metadata to the span. Defaults to None.
-            override (bool, optional): Whether to override the existing metadata. If False, metadata is merged with the existing metadata. Defaults to False.
-        """
-        laminar = LaminarSingleton().get()
-        laminar.update_current_span(
-            metadata=metadata, attributes=attributes, override=override
-        )
-
-    def update_current_trace(
-        self,
-        user_id: Optional[str] = None,
-        session_id: Optional[str] = None,
-        release: Optional[str] = None,
-        metadata: Optional[dict[str, Any]] = None,
-    ):
-        """Update the current trace with any optional metadata.
-
-        Args:
-            user_id (Optional[str], optional): Custom user_id of your user. Useful for grouping and further analytics. Defaults to None.
-            session_id (Optional[str], optional): Custom session_id for your session. Random UUID is generated on Laminar side, if not specified.
-                                                  Defaults to None.
-            release (Optional[str], optional): Release version of your app. Useful for further grouping and analytics. Defaults to None.
-            metadata (Optional[dict[str, Any]], optional): metadata to the trace. Defaults to None.
-        """
-        laminar = LaminarSingleton().get()
-        laminar.update_current_trace(
-            user_id=user_id, session_id=session_id, release=release, metadata=metadata
-        )
+from .laminar import Laminar as L
+from .utils import is_async
 
-    def event(
-        self,
-        name: str,
-        value: Optional[Union[str, int, float, bool]] = None,
-        timestamp: Optional[datetime.datetime] = None,
-    ):
-        """Associate an event with the current span
+P = ParamSpec("P")
+R = TypeVar("R")
 
-        Args:
-            name (str): name of the event. Must be predefined in the Laminar events page.
-            value (Optional[Union[str, int, float, bool]], optional): value of the event. Must match range definition in Laminar events page. Defaults to None.
-            timestamp (Optional[datetime.datetime], optional): If you need custom timestamp. If not specified, current time is used. Defaults to None.
-        """
-        laminar = LaminarSingleton().get()
-        laminar.event(name, value=value, timestamp=timestamp)
 
-    def evaluate_event(self, name: str, data: str):
-        """Evaluate an event with the given name and data. The event value will be assessed by the Laminar evaluation engine.
-        Data is passed as an input to the agent, so you need to specify which data you want to evaluate. Most of the times,
-        this is an output of the LLM generation, but sometimes, you may want to evaluate the input or both. In the latter case,
-        concatenate the input and output annotating with natural language.
-
-        Args:
-            name (str): Name of the event. Must be predefined in the Laminar events page.
-            data (str): Data to be evaluated. Typically the output of the LLM generation.
-        """
-        laminar = LaminarSingleton().get()
-        laminar.evaluate_event(name, data)
-
-    def run(
-        self,
-        pipeline: str,
-        inputs: dict[str, NodeInput],
-        env: dict[str, str] = {},
-        metadata: dict[str, str] = {},
-    ) -> PipelineRunResponse:
-        """Run the laminar pipeline with the given inputs. Pipeline must be defined in the Laminar UI and have a target version.
-
-        Args:
-            pipeline (str): pipeline name
-            inputs (dict[str, NodeInput]): Map from input node name to input value
-            env (dict[str, str], optional): Environment variables for the pipeline executions. Typically contains API keys. Defaults to None.
-            metadata (dict[str, str], optional): Any additional data to associate with the resulting span. Defaults to None.
-
-        Returns:
-            PipelineRunResponse: Response from the pipeline execution
-        """
-        laminar = LaminarSingleton().get()
-        return laminar.run_pipeline(pipeline, inputs, env, metadata)
-
-
-def wrap_llm_call(func: Callable, name: str = None, provider: str = None) -> Callable:
-    """Wrap an LLM call with Laminar observability. This is a convenience function that does the same as `@observe()`, plus
-    a few utilities around LLM-specific things, such as counting tokens and recording model params.
-
-    Example usage:
-    ```python
-    wrap_llm_call(client.chat.completions.create)(
-        model="gpt-4o-mini",
-        messages=[
-            {"role": "system", "content": "You are a helpful assistant."},
-            {"role": "user", "content": "Hello"},
-        ],
-        stream=True,
-    )
-    ```
+def observe(
+    *,
+    name: Optional[str] = None,
+    user_id: Optional[str] = None,
+    session_id: Optional[str] = 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.
 
     Args:
-        func (Callable): The function to wrap
-        name (str, optional): Name of the resulting span. Default "{provider name} completion" if not specified. Defaults to None.
-        provider (str, optional): LLM model provider, e.g. openai, anthropic. This is needed to help us correctly parse
-                                  things like token usage. If not specified, we infer it from the name of the package,
-                                  where the function is imported from. Defaults to None.
+        name (Optional[str], optional): Name of the span. Function
+                        name is used if not specified.
+                        Defaults to None.
+        user_id (Optional[str], optional): User ID to associate
+                        with the span and the following context.
+                        Defaults to None.
+        session_id (Optional[str], optional): Session ID to associate with the
+                        span and the following context. Defaults to None.
 
     Raises:
-        Exctption: re-raises the exception if the wrapped function raises an exception
+        Exception: re-raises the exception if the wrapped function raises
+                   an exception
 
     Returns:
-        Callable: the wrapped function
+        R: Returns the result of the wrapped function
     """
-    laminar = LaminarSingleton().get()
-    # Simple heuristic to determine the package from where the LLM call is imported.
-    # This works for major providers, but will likely make no sense for custom providers.
-    provider_name = (
-        provider.lower().strip() if provider else func.__module__.split(".")[0]
-    )
-    provider_module = PROVIDER_NAME_TO_OBJECT.get(provider_name, FallbackProvider())
-    name = name or f"{provider_module.display_name()} completion"
-
-    @functools.wraps(func)
-    def wrapper(*args, **kwargs):
-        inp = kwargs.get("messages")
-        attributes = (
-            provider_module.extract_llm_attributes_from_args(args, kwargs)
-            if provider_module
-            else {}
-        )
-        attributes[PROVIDER] = provider_name
-        span = laminar.observe_start(
-            name=name, span_type="LLM", input=inp, attributes=attributes
-        )
-        try:
-            result = func(*args, **kwargs)
-        except Exception as e:
-            laminar.observe_end(
-                result=None, span=span, error=e, provider_name=provider_name
-            )
-            raise e
-        return laminar.observe_end(
-            result=result, span=span, provider_name=provider_name
-        )
 
-    @functools.wraps(func)
-    async def async_wrapper(*args, **kwargs):
-        inp = kwargs.get("messages")
-        attributes = (
-            provider_module.extract_llm_attributes_from_args(args, kwargs)
-            if provider_module
-            else {}
-        )
-        attributes[PROVIDER] = provider_name
-        span = laminar.observe_start(
-            name=name, span_type="LLM", input=inp, attributes=attributes
-        )
-        try:
-            result = await func(*args, **kwargs)
-        except Exception as e:
-            laminar.observe_end(
-                result=None, span=span, error=e, provider_name=provider_name
+    def decorator(func: Callable[P, R]) -> Callable[P, R]:
+        if not L.is_initialized():
+            raise Exception(
+                "Laminar is not initialized. Please "
+                + "call Laminar.initialize() first."
             )
-            raise e
-        return laminar.observe_end(
-            result=result, span=span, provider_name=provider_name
+        current_span = get_current_span()
+        if current_span != INVALID_SPAN:
+            if session_id is not None:
+                current_span.set_attribute(
+                    "traceloop.association.properties.session_id", session_id
+                )
+            if user_id is not None:
+                current_span.set_attribute(
+                    "traceloop.association.properties.user_id", user_id
+                )
+        association_properties = {}
+        if session_id is not None:
+            association_properties["session_id"] = session_id
+        if user_id is not None:
+            association_properties["user_id"] = user_id
+        Traceloop.set_association_properties(association_properties)
+        return (
+            aentity_method(name=name)(func)
+            if is_async(func)
+            else entity_method(name=name)(func)
         )
 
-    return async_wrapper if is_async(func) else wrapper
-
-
-lmnr_context = LaminarDecorator()
-observe = lmnr_context.observe
+    return cast(Callable[P, R], decorator)