lmnr 0.3.7__py3-none-any.whl → 0.4.1__py3-none-any.whl

lmnr/__init__.py

@@ -1,7 +1,4 @@
-from .sdk.client import Laminar
-from .sdk.decorators import observe, lmnr_context, wrap_llm_call
-from .sdk.interface import trace, TraceContext, SpanContext, initialize
-from .sdk.tracing_types import EvaluateEvent
+from .sdk.evaluations import Evaluation
+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,284 +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
-        )
-
-    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
-
-        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, evaluator: str, data: dict):
-        """Evaluate an event with the given name by evaluator based on the given data.
-        Evaluator is the Laminar pipeline name.
-        Data is passed as an input to the the evaluator pipeline, so you need to specify which data you want to evaluate. The prompt
-        of the evaluator will be templated with the keys of the data dictionary.
-
-        Usually, you would want to pass the output of LLM generation, users' messages, and some other surrounding data to 'data'.
-
-        Args:
-            name (str): Name of the event.
-            evaluator (str): Name of the evaluator pipeline.
-            data (dict): Data to be used when evaluating the event.
-        """
-        laminar = LaminarSingleton().get()
-        laminar.evaluate_event(name, evaluator=evaluator, data=data)
+from .laminar import Laminar as L
+from .utils import is_async
 
-    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.
+P = ParamSpec("P")
+R = TypeVar("R")
 
-        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
+    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
-        )
-
-    @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
-            )
-            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)

lmnr/sdk/evaluations.py

@@ -0,0 +1,163 @@
+from typing import Union
+
+from .utils import is_async
+from .types import EvaluatorFunction, ExecutorFunction, EvaluationDatapoint, Numeric
+from .laminar import Laminar as L
+import asyncio
+
+from abc import ABC, abstractmethod
+
+DEFAULT_BATCH_SIZE = 5
+
+
+class EvaluationDataset(ABC):
+    @abstractmethod
+    def __init__(self, *args, **kwargs):
+        pass
+
+    @abstractmethod
+    def __len__(self) -> int:
+        pass
+
+    @abstractmethod
+    def __getitem__(self, idx) -> EvaluationDatapoint:
+        pass
+
+    def slice(self, start: int, end: int):
+        return [self[i] for i in range(max(start, 0), min(end, len(self)))]
+
+
+class Evaluation:
+    def __init__(
+        self,
+        name,
+        data: Union[EvaluationDataset, list[Union[EvaluationDatapoint, dict]]],
+        executor: ExecutorFunction,
+        evaluators: list[EvaluatorFunction],
+        batch_size: int = DEFAULT_BATCH_SIZE,
+        project_api_key: str = "",
+        base_url: str = "https://api.lmnr.ai",
+    ):
+        """
+        Initializes an instance of the Evaluations class.
+        Parameters:
+            name (str): The name of the evaluation.
+            data (Union[List[Union[EvaluationDatapoint, dict]], EvaluationDataset]): List of data points to evaluate or an evaluation dataset.
+                            `data` is the input to the executor function,
+                            `target` is the input to the evaluator function.
+            executor (Callable[..., Any]): The executor function.
+                            Takes the data point + any additional arguments
+                            and returns the output to evaluate.
+            evaluators (List[Callable[..., Any]]): List of evaluator functions.
+                Each evaluator function takes the output of the executor _and_
+                the target data, and returns a score. The score can be a
+                single number or a record of string keys and number values.
+                If the score is a single number, it will be named after the
+                evaluator function. If the function is anonymous, it will be
+                named `evaluator_${index}`, where index is the index of the
+                evaluator function in the list starting from 1.
+            batch_size (int, optional): The batch size for evaluation.
+                            Defaults to DEFAULT_BATCH_SIZE.
+            project_api_key (str, optional): The project API key.
+                            Defaults to an empty string.
+            base_url (str, optional): The base URL for the LMNR API.
+                            Useful if self-hosted elsewhere.
+                            Defaults to "https://api.lmnr.ai".
+        """
+
+        self.name = name
+        self.executor = executor
+        self.evaluators = dict(
+            zip(
+                [
+                    (
+                        e.__name__
+                        if e.__name__ and e.__name__ != "<lambda>"
+                        else f"evaluator_{i+1}"
+                    )
+                    for i, e in enumerate(evaluators)
+                ],
+                evaluators,
+            )
+        )
+        self.evaluator_names = list(self.evaluators.keys())
+        if isinstance(data, list):
+            self.data = [
+                (
+                    EvaluationDatapoint.model_validate(point)
+                    if isinstance(point, dict)
+                    else point
+                )
+                for point in data
+            ]
+        else:
+            self.data = data
+        self.batch_size = batch_size
+        L.initialize(project_api_key=project_api_key, base_url=base_url)
+
+    async def run(self):
+        """Runs the evaluation.
+
+        Creates a new evaluation if no evaluation with such name exists, or
+        adds data to an existing one otherwise. Evaluates data points in
+        batches of `self.batch_size`. The executor
+        function is called on each data point to get the output,
+        and then evaluate it by each evaluator function.
+        """
+        response = L.create_evaluation(self.name)
+        batch_promises = []
+
+        for i in range(0, len(self.data), self.batch_size):
+            batch = (
+                self.data[i : i + self.batch_size]
+                if isinstance(self.data, list)
+                else self.data.slice(i, i + self.batch_size)
+            )
+            batch_promises.append(self._evaluate_batch(batch))
+
+        try:
+            await asyncio.gather(*batch_promises)
+            L.update_evaluation_status(response.name, "Finished")
+            print(f"Evaluation {response.id} complete")
+        except Exception as e:
+            print(f"Error evaluating batch: {e}")
+
+    async def _evaluate_batch(self, batch: list[EvaluationDatapoint]):
+        results = []
+        for datapoint in batch:
+            output = (
+                await self.executor(datapoint.data)
+                if is_async(self.executor)
+                else self.executor(datapoint.data)
+            )
+            target = datapoint.target
+
+            # iterate in order of evaluators
+            scores = {}
+            for evaluator_name in self.evaluator_names:
+                evaluator = self.evaluators[evaluator_name]
+                value = (
+                    await evaluator(output, target)
+                    if is_async(evaluator)
+                    else evaluator(output, target)
+                )
+
+                # if the evaluator returns a single number,
+                # use the evaluator name as the key
+                if isinstance(value, Numeric):
+                    scores[evaluator_name] = value
+                else:
+                    # if the evaluator returns an object,
+                    # use the object keys as the keys
+                    scores.update(value)
+
+            results.append(
+                {
+                    "executorOutput": output,
+                    "data": datapoint.data,
+                    "target": target,
+                    "scores": scores,
+                }
+            )
+
+        return L.post_evaluation_results(self.name, results)