lmnr 0.4.17b0__py2.py3-none-any.whl

lmnr/sdk/laminar.py

@@ -0,0 +1,403 @@
+from lmnr.traceloop_sdk.instruments import Instruments
+from opentelemetry import context
+from opentelemetry.trace import (
+    INVALID_SPAN,
+    get_current_span,
+)
+from opentelemetry.util.types import AttributeValue
+from opentelemetry.context import set_value, attach, detach
+from lmnr.traceloop_sdk import Traceloop
+from lmnr.traceloop_sdk.tracing import get_tracer
+from contextlib import contextmanager
+from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
+
+from pydantic.alias_generators import to_snake
+from typing import Any, Optional, Set, Union
+
+import copy
+import datetime
+import dotenv
+import json
+import logging
+import os
+import requests
+import uuid
+
+from lmnr.traceloop_sdk.tracing.attributes import (
+    SESSION_ID,
+    SPAN_INPUT,
+    SPAN_OUTPUT,
+    SPAN_PATH,
+    TRACE_TYPE,
+    USER_ID,
+)
+from lmnr.traceloop_sdk.tracing.tracing import (
+    get_span_path,
+    set_association_properties,
+    update_association_properties,
+)
+
+from .log import VerboseColorfulFormatter
+
+from .types import (
+    CreateEvaluationResponse,
+    EvaluationResultDatapoint,
+    PipelineRunError,
+    PipelineRunResponse,
+    NodeInput,
+    PipelineRunRequest,
+    TraceType,
+)
+
+
+class Laminar:
+    __base_http_url: str
+    __base_grpc_url: str
+    __project_api_key: Optional[str] = None
+    __env: dict[str, str] = {}
+    __initialized: bool = False
+
+    @classmethod
+    def initialize(
+        cls,
+        project_api_key: Optional[str] = None,
+        env: dict[str, str] = {},
+        base_url: Optional[str] = None,
+        http_port: Optional[int] = None,
+        grpc_port: Optional[int] = None,
+        instruments: Optional[Set[Instruments]] = None,
+    ):
+        """Initialize Laminar context across the application.
+        This method must be called before using any other Laminar methods or
+        decorators.
+
+        Args:
+            project_api_key (Optional[str], optional): Laminar project api key.
+                            You can generate one by going to the projects
+                            settings page on the Laminar dashboard.
+                            If not specified, it will try to read from the
+                            LMNR_PROJECT_API_KEY environment variable
+                            in os.environ or in .env file.
+                            Defaults to None.
+            env (dict[str, str], optional): Default environment passed to
+                            `run` and `evaluate_event` requests, unless
+                            overriden at request time. Usually, model
+                            provider keys are stored here.
+                            Defaults to {}.
+            base_url (Optional[str], optional): Laminar API url.
+                            If not specified, defaults to https://api.lmnr.ai.
+            http_port (Optional[int], optional): Laminar API http port.
+                            If not specified, defaults to 443.
+            grpc_port (Optional[int], optional): Laminar API grpc port.
+                            If not specified, defaults to 8443.
+
+        Raises:
+            ValueError: If project API key is not set
+        """
+        cls.__project_api_key = project_api_key or os.environ.get(
+            "LMNR_PROJECT_API_KEY"
+        )
+        if not cls.__project_api_key:
+            dotenv_path = dotenv.find_dotenv(usecwd=True)
+            cls.__project_api_key = dotenv.get_key(
+                dotenv_path=dotenv_path, key_to_get="LMNR_PROJECT_API_KEY"
+            )
+        if not cls.__project_api_key:
+            raise ValueError(
+                "Please initialize the Laminar object with"
+                " your project API key or set the LMNR_PROJECT_API_KEY"
+                " environment variable in your environment or .env file"
+            )
+
+        cls.__base_http_url = f"{base_url or 'https://api.lmnr.ai'}:{http_port or 443}"
+        cls.__base_grpc_url = f"{base_url or 'https://api.lmnr.ai'}:{grpc_port or 8443}"
+
+        cls.__env = env
+        cls.__initialized = True
+        cls._initialize_logger()
+        Traceloop.init(
+            exporter=OTLPSpanExporter(
+                endpoint=cls.__base_grpc_url,
+                headers={"authorization": f"Bearer {cls.__project_api_key}"},
+            ),
+            instruments=instruments,
+        )
+
+    @classmethod
+    def is_initialized(cls):
+        """Check if Laminar is initialized. A utility to make sure other
+        methods are called after initialization.
+
+        Returns:
+            bool: True if Laminar is initialized, False otherwise
+        """
+        return cls.__initialized
+
+    @classmethod
+    def _initialize_logger(cls):
+        cls.__logger = logging.getLogger(__name__)
+        console_log_handler = logging.StreamHandler()
+        console_log_handler.setFormatter(VerboseColorfulFormatter())
+        cls.__logger.addHandler(console_log_handler)
+
+    @classmethod
+    def run(
+        cls,
+        pipeline: str,
+        inputs: dict[str, NodeInput],
+        env: dict[str, str] = {},
+        metadata: dict[str, str] = {},
+        parent_span_id: Optional[uuid.UUID] = None,
+        trace_id: Optional[uuid.UUID] = None,
+    ) -> PipelineRunResponse:
+        """Runs the pipeline with the given inputs
+
+        Args:
+            pipeline (str): name of the Laminar pipeline.
+                The pipeline must have a target version set.
+            inputs (dict[str, NodeInput]):
+                inputs to the endpoint's target pipeline.
+                Keys in the dictionary must match input node names
+            env (dict[str, str], optional):
+                Environment variables for the pipeline execution.
+                Defaults to {}.
+            metadata (dict[str, str], optional):
+                any custom metadata to be stored
+                with execution trace. Defaults to {}.
+            parent_span_id (Optional[uuid.UUID], optional):
+                parent span id for the resulting span.
+                Defaults to None.
+            trace_id (Optional[uuid.UUID], optional):
+                trace id for the resulting trace.
+                Defaults to None.
+
+        Returns:
+            PipelineRunResponse: response object containing the outputs
+
+        Raises:
+            ValueError: if project API key is not set
+            PipelineRunError: if the endpoint run fails
+        """
+        if cls.__project_api_key is None:
+            raise ValueError(
+                "Please initialize the Laminar object with your project "
+                "API key or set the LMNR_PROJECT_API_KEY environment variable"
+            )
+        try:
+            current_span = get_current_span()
+            if current_span != INVALID_SPAN:
+                parent_span_id = parent_span_id or uuid.UUID(
+                    int=current_span.get_span_context().span_id
+                )
+                trace_id = trace_id or uuid.UUID(
+                    int=current_span.get_span_context().trace_id
+                )
+            request = PipelineRunRequest(
+                inputs=inputs,
+                pipeline=pipeline,
+                env=env,
+                metadata=metadata,
+                parent_span_id=parent_span_id,
+                trace_id=trace_id,
+            )
+        except Exception as e:
+            raise ValueError(f"Invalid request: {e}")
+
+        response = requests.post(
+            cls.__base_http_url + "/v1/pipeline/run",
+            data=json.dumps(request.to_dict()),
+            headers=cls._headers(),
+        )
+        if response.status_code != 200:
+            raise PipelineRunError(response)
+        try:
+            resp_json = response.json()
+            keys = list(resp_json.keys())
+            for key in keys:
+                value = resp_json[key]
+                del resp_json[key]
+                resp_json[to_snake(key)] = value
+            return PipelineRunResponse(**resp_json)
+        except Exception:
+            raise PipelineRunError(response)
+
+    @classmethod
+    def event(
+        cls,
+        name: str,
+        value: Optional[AttributeValue] = None,
+        timestamp: Optional[Union[datetime.datetime, int]] = None,
+    ):
+        """Associate an event with the current span. If event with such name never
+        existed, Laminar will create a new event and infer its type from the value.
+        If the event already exists, Laminar will append the value to the event
+        if and only if the value is of a matching type. Otherwise, the event won't
+        be recorded Supported types are string, numeric, and boolean. If the value
+        is `None`, event is considered a boolean tag with the value of `True`.
+
+        Args:
+            name (str): event name
+            value (Optional[AttributeValue]): event value. Must be a primitive type.
+                            Boolean true is assumed in the backend if value is None.
+                            Defaults to None.
+            timestamp (Optional[Union[datetime.datetime, int]], optional):
+                            If int, must be epoch nanoseconds. If not
+                            specified, relies on the underlying OpenTelemetry
+                            implementation. Defaults to None.
+        """
+        if timestamp and isinstance(timestamp, datetime.datetime):
+            timestamp = int(timestamp.timestamp() * 1e9)
+
+        event = {
+            "lmnr.event.type": "default",
+        }
+        if value is not None:
+            event["lmnr.event.value"] = value
+
+        current_span = get_current_span()
+        if current_span == INVALID_SPAN:
+            cls.__logger.warning(
+                "`Laminar().event()` called outside of span context. "
+                f"Event '{name}' will not be recorded in the trace. "
+                "Make sure to annotate the function with a decorator"
+            )
+            return
+
+        current_span.add_event(name, event, timestamp)
+
+    @classmethod
+    @contextmanager
+    def start_as_current_span(
+        cls,
+        name: str,
+        input: Any = None,
+    ):
+        """Start a new span as the current span. Useful for manual instrumentation.
+        This is the preferred and more stable way to use manual instrumentation.
+
+        Usage example:
+        ```python
+        with Laminar.start_as_current_span("my_span", input="my_input"):
+            await my_async_function()
+        ```
+
+        Args:
+            name (str): name of the span
+            input (Any, optional): input to the span. Will be sent as an
+                attribute, so must be json serializable. Defaults to None.
+        """
+        with get_tracer() as tracer:
+            span_path = get_span_path(name)
+            ctx = set_value("span_path", span_path)
+            ctx_token = attach(set_value("span_path", span_path))
+            with tracer.start_as_current_span(
+                name,
+                context=ctx,
+                attributes={SPAN_PATH: span_path},
+            ) as span:
+                if input is not None:
+                    span.set_attribute(
+                        SPAN_INPUT,
+                        json.dumps(input),
+                    )
+                yield span
+
+            # TODO: Figure out if this is necessary
+            try:
+                detach(ctx_token)
+            except Exception:
+                pass
+
+    @classmethod
+    def set_span_output(cls, output: Any = None):
+        """Set the output of the current span. Useful for manual instrumentation.
+
+        Args:
+            output (Any, optional): output of the span. Will be sent as an
+                attribute, so must be json serializable. Defaults to None.
+        """
+        span = get_current_span()
+        if output is not None and span != INVALID_SPAN:
+            span.set_attribute(SPAN_OUTPUT, json.dumps(output))
+
+    @classmethod
+    def set_session(
+        cls,
+        session_id: Optional[str] = None,
+        user_id: Optional[str] = None,
+    ):
+        """Set the session and user id for the current span and the context
+        (i.e. any children spans created from the current span in the current
+        thread).
+
+        Args:
+            session_id (Optional[str], optional): Custom session id.
+                            Useful to debug and group long-running
+                            sessions/conversations.
+                            Defaults to None.
+            user_id (Optional[str], optional): Custom user id.
+                            Useful for grouping spans or traces by user.
+                            Defaults to None.
+        """
+        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
+        update_association_properties(association_properties)
+
+    @classmethod
+    def _set_trace_type(
+        cls,
+        trace_type: TraceType,
+    ):
+        """Set the trace_type for the current span and the context
+        Args:
+            trace_type (TraceType): Type of the trace
+        """
+        association_properties = {
+            TRACE_TYPE: trace_type.value,
+        }
+        update_association_properties(association_properties)
+
+    @classmethod
+    def clear_session(cls):
+        """Clear the session and user id from  the context"""
+        props: dict = copy.copy(context.get_value("association_properties"))
+        props.pop("session_id", None)
+        props.pop("user_id", None)
+        set_association_properties(props)
+
+    @classmethod
+    def create_evaluation(
+        cls,
+        data: list[EvaluationResultDatapoint],
+        group_id: Optional[str] = None,
+        name: Optional[str] = None,
+    ) -> CreateEvaluationResponse:
+        response = requests.post(
+            cls.__base_http_url + "/v1/evaluations",
+            data=json.dumps(
+                {
+                    "groupId": group_id,
+                    "name": name,
+                    "points": [datapoint.to_dict() for datapoint in data],
+                }
+            ),
+            headers=cls._headers(),
+        )
+        if response.status_code != 200:
+            try:
+                resp_json = response.json()
+                raise ValueError(f"Error creating evaluation {json.dumps(resp_json)}")
+            except Exception:
+                raise ValueError(f"Error creating evaluation {response.text}")
+        return CreateEvaluationResponse.model_validate(response.json())
+
+    @classmethod
+    def _headers(cls):
+        assert cls.__project_api_key is not None, "Project API key is not set"
+        return {
+            "Authorization": "Bearer " + cls.__project_api_key,
+            "Content-Type": "application/json",
+        }

lmnr/sdk/log.py

@@ -0,0 +1,39 @@
+import logging
+
+
+class CustomFormatter(logging.Formatter):
+    grey = "\x1b[38;20m"
+    green = "\x1b[32;20m"
+    yellow = "\x1b[33;20m"
+    red = "\x1b[31;20m"
+    bold_red = "\x1b[31;1m"
+    reset = "\x1b[0m"
+    fmt = "%(asctime)s::%(name)s::%(levelname)s: %(message)s (%(filename)s:%(lineno)d)"
+
+    FORMATS = {
+        logging.DEBUG: grey + fmt + reset,
+        logging.INFO: green + fmt + reset,
+        logging.WARNING: yellow + fmt + reset,
+        logging.ERROR: red + fmt + reset,
+        logging.CRITICAL: bold_red + fmt + reset,
+    }
+
+    def format(self, record: logging.LogRecord):
+        log_fmt = self.FORMATS.get(record.levelno)
+        formatter = logging.Formatter(log_fmt)
+        return formatter.format(record)
+
+
+# For StreamHandlers / console
+class VerboseColorfulFormatter(CustomFormatter):
+    def format(self, record):
+        return super().format(record)
+
+
+# For Verbose FileHandlers / files
+class VerboseFormatter(CustomFormatter):
+    fmt = "%(asctime)s::%(name)s::%(levelname)s| %(message)s (%(filename)s:%(lineno)d)"
+
+    def format(self, record):
+        formatter = logging.Formatter(self.fmt)
+        return formatter.format(record)

lmnr/sdk/types.py

@@ -0,0 +1,155 @@
+import datetime
+from enum import Enum
+import pydantic
+import requests
+from typing import Any, Awaitable, Callable, Optional, Union
+import uuid
+
+from .utils import serialize
+
+
+class ChatMessage(pydantic.BaseModel):
+    role: str
+    content: str
+
+
+class ConditionedValue(pydantic.BaseModel):
+    condition: str
+    value: "NodeInput"
+
+
+Numeric = Union[int, float]
+NumericTypes = (int, float)  # for use with isinstance
+
+NodeInput = Union[str, list[ChatMessage], ConditionedValue, Numeric, bool]
+PipelineOutput = Union[NodeInput]
+
+
+class PipelineRunRequest(pydantic.BaseModel):
+    inputs: dict[str, NodeInput]
+    pipeline: str
+    env: dict[str, str] = pydantic.Field(default_factory=dict)
+    metadata: dict[str, str] = pydantic.Field(default_factory=dict)
+    stream: bool = pydantic.Field(default=False)
+    parent_span_id: Optional[uuid.UUID] = pydantic.Field(default=None)
+    trace_id: Optional[uuid.UUID] = pydantic.Field(default=None)
+
+    # uuid is not serializable by default, so we need to convert it to a string
+    def to_dict(self):
+        return {
+            "inputs": {
+                k: v.model_dump() if isinstance(v, pydantic.BaseModel) else serialize(v)
+                for k, v in self.inputs.items()
+            },
+            "pipeline": self.pipeline,
+            "env": self.env,
+            "metadata": self.metadata,
+            "stream": self.stream,
+            "parentSpanId": str(self.parent_span_id) if self.parent_span_id else None,
+            "traceId": str(self.trace_id) if self.trace_id else None,
+        }
+
+
+class PipelineRunResponse(pydantic.BaseModel):
+    outputs: dict[str, dict[str, PipelineOutput]]
+    run_id: str
+
+
+class PipelineRunError(Exception):
+    error_code: str
+    error_message: str
+
+    def __init__(self, response: requests.Response):
+        try:
+            resp_json = response.json()
+            self.error_code = resp_json["error_code"]
+            self.error_message = resp_json["error_message"]
+            super().__init__(self.error_message)
+        except Exception:
+            super().__init__(response.text)
+
+    def __str__(self) -> str:
+        try:
+            return str(
+                {"error_code": self.error_code, "error_message": self.error_message}
+            )
+        except Exception:
+            return super().__str__()
+
+
+EvaluationDatapointData = dict[str, Any]
+EvaluationDatapointTarget = dict[str, Any]
+
+
+# EvaluationDatapoint is a single data point in the evaluation
+class Datapoint(pydantic.BaseModel):
+    # input to the executor function. Must be a dict with string keys
+    data: EvaluationDatapointData
+    # input to the evaluator function (alongside the executor output).
+    # Must be a dict with string keys
+    target: EvaluationDatapointTarget
+
+
+ExecutorFunctionReturnType = Any
+EvaluatorFunctionReturnType = Union[Numeric, dict[str, Numeric]]
+
+ExecutorFunction = Callable[
+    [EvaluationDatapointData, Any, dict[str, Any]],
+    Union[ExecutorFunctionReturnType, Awaitable[ExecutorFunctionReturnType]],
+]
+
+# EvaluatorFunction is a function that 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. The latter is useful for evaluating
+# multiple criteria in one go instead of running multiple evaluators.
+EvaluatorFunction = Callable[
+    [ExecutorFunctionReturnType, Any, dict[str, Any]],
+    Union[EvaluatorFunctionReturnType, Awaitable[EvaluatorFunctionReturnType]],
+]
+
+
+class CreateEvaluationResponse(pydantic.BaseModel):
+    id: uuid.UUID
+    createdAt: datetime.datetime
+    groupId: str
+    name: str
+    projectId: uuid.UUID
+
+
+class EvaluationResultDatapoint(pydantic.BaseModel):
+    data: EvaluationDatapointData
+    target: EvaluationDatapointTarget
+    executor_output: ExecutorFunctionReturnType
+    scores: dict[str, Numeric]
+    trace_id: uuid.UUID
+
+    # uuid is not serializable by default, so we need to convert it to a string
+    def to_dict(self):
+        return {
+            "data": {
+                k: v.model_dump() if isinstance(v, pydantic.BaseModel) else serialize(v)
+                for k, v in self.data.items()
+            },
+            "target": {
+                k: v.model_dump() if isinstance(v, pydantic.BaseModel) else serialize(v)
+                for k, v in self.target.items()
+            },
+            "executorOutput": serialize(self.executor_output),
+            "scores": self.scores,
+            "traceId": str(self.trace_id),
+        }
+
+
+class SpanType(Enum):
+    DEFAULT = "DEFAULT"
+    LLM = "LLM"
+    PIPELINE = "PIPELINE"  # must not be set manually
+    EXECUTOR = "EXECUTOR"
+    EVALUATOR = "EVALUATOR"
+    EVALUATION = "EVALUATION"
+
+
+class TraceType(Enum):
+    DEFAULT = "DEFAULT"
+    EVENT = "EVENT"  # must not be set manually
+    EVALUATION = "EVALUATION"

lmnr/sdk/utils.py

@@ -0,0 +1,99 @@
+import asyncio
+import datetime
+import dataclasses
+import enum
+import inspect
+import pydantic
+import queue
+import typing
+import uuid
+
+
+def is_method(func: typing.Callable) -> bool:
+    # inspect.ismethod is True for bound methods only, but in the decorator,
+    # the method is not bound yet, so we need to check if the first parameter
+    # is either 'self' or 'cls'. This only relies on naming conventions
+
+    # `signature._parameters` is an OrderedDict,
+    # so the order of insertion is preserved
+    params = list(inspect.signature(func).parameters)
+    return len(params) > 0 and params[0] in ["self", "cls"]
+
+
+def is_async(func: typing.Callable) -> bool:
+    # `__wrapped__` is set automatically by `functools.wraps` and
+    # `functools.update_wrapper`
+    # so we can use it to get the original function
+    while hasattr(func, "__wrapped__"):
+        func = func.__wrapped__
+
+    if not inspect.isfunction(func):
+        return False
+
+    # Check if the function is asynchronous
+    if asyncio.iscoroutinefunction(func):
+        return True
+
+    # Fallback: check if the function's code object contains 'async'.
+    # This is for cases when a decorator did not properly use
+    # `functools.wraps` or `functools.update_wrapper`
+    CO_COROUTINE = inspect.CO_COROUTINE
+    return (func.__code__.co_flags & CO_COROUTINE) != 0
+
+
+def is_async_iterator(o: typing.Any) -> bool:
+    return hasattr(o, "__aiter__") and hasattr(o, "__anext__")
+
+
+def is_iterator(o: typing.Any) -> bool:
+    return hasattr(o, "__iter__") and hasattr(o, "__next__")
+
+
+def serialize(obj: typing.Any) -> dict[str, typing.Any]:
+    def to_dict_inner(o: typing.Any):
+        if isinstance(o, (datetime.datetime, datetime.date)):
+            return o.strftime("%Y-%m-%dT%H:%M:%S.%f%z")
+        elif o is None:
+            return None
+        elif isinstance(o, (int, float, str, bool)):
+            return o
+        elif isinstance(o, uuid.UUID):
+            return str(o)  # same as in final return, but explicit
+        elif isinstance(o, enum.Enum):
+            return o.value
+        elif dataclasses.is_dataclass(o):
+            return dataclasses.asdict(o)
+        elif isinstance(o, bytes):
+            return o.decode("utf-8")
+        elif isinstance(o, pydantic.BaseModel):
+            return o.model_dump()
+        elif isinstance(o, (tuple, set, frozenset)):
+            return [to_dict_inner(item) for item in o]
+        elif isinstance(o, list):
+            return [to_dict_inner(item) for item in o]
+        elif isinstance(o, dict):
+            return {to_dict_inner(k): to_dict_inner(v) for k, v in o.items()}
+        elif isinstance(o, queue.Queue):
+            return type(o).__name__
+
+        return str(o)
+
+    return to_dict_inner(obj)
+
+
+def get_input_from_func_args(
+    func: typing.Callable,
+    is_method: bool = False,
+    func_args: list[typing.Any] = [],
+    func_kwargs: dict[str, typing.Any] = {},
+) -> dict[str, typing.Any]:
+    # Remove implicitly passed "self" or "cls" argument for
+    # instance or class methods
+    res = func_kwargs.copy()
+    for i, k in enumerate(inspect.signature(func).parameters.keys()):
+        if is_method and k in ["self", "cls"]:
+            continue
+        # If param has default value, then it's not present in func args
+        if i < len(func_args):
+            res[k] = func_args[i]
+    return res