lmnr 0.4.53.dev0__py3-none-any.whl → 0.7.26__py3-none-any.whl

lmnr/sdk/log.py

@@ -62,11 +62,16 @@ class VerboseFormatter(CustomFormatter):
         return formatter.format(record)
 
 
-def get_default_logger(name: str, level: int = logging.INFO, propagate: bool = False):
+def get_default_logger(
+    name: str, level: int = logging.INFO, propagate: bool = False, verbose: bool = True
+) -> logging.Logger:
     logger = logging.getLogger(name)
     logger.setLevel(level)
     console_log_handler = logging.StreamHandler()
-    console_log_handler.setFormatter(VerboseColorfulFormatter())
+    if verbose:
+        console_log_handler.setFormatter(VerboseColorfulFormatter())
+    else:
+        console_log_handler.setFormatter(ColorfulFormatter())
     logger.addHandler(console_log_handler)
     logger.propagate = propagate
     return logger

lmnr/sdk/types.py

@@ -1,131 +1,56 @@
-import aiohttp
+from __future__ import annotations  # For "Self" | str | ... type hint
+
+import json
+import logging
 import datetime
-from enum import Enum
-import pydantic
-from typing import Any, Awaitable, Callable, Optional, Union
+from pydantic import BaseModel, Field
 import uuid
 
-from .utils import serialize
-
-
-class ChatMessage(pydantic.BaseModel):
-    role: str
-    content: str
+from enum import Enum
+from opentelemetry.trace import SpanContext, TraceFlags
+from typing import Any, Awaitable, Callable, Optional
+from typing_extensions import TypedDict  # compatibility with python < 3.12
 
+from .utils import serialize, json_dumps
 
-class ConditionedValue(pydantic.BaseModel):
-    condition: str
-    value: "NodeInput"
+DEFAULT_DATAPOINT_MAX_DATA_LENGTH = 16_000_000  # 16MB
 
 
-Numeric = Union[int, float]
+Numeric = 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 SemanticSearchRequest(pydantic.BaseModel):
-    query: str
-    dataset_id: uuid.UUID
-    limit: Optional[int] = pydantic.Field(default=None)
-    threshold: Optional[float] = pydantic.Field(default=None, ge=0.0, le=1.0)
-
-    def to_dict(self):
-        res = {
-            "query": self.query,
-            "datasetId": str(self.dataset_id),
-        }
-        if self.limit is not None:
-            res["limit"] = self.limit
-        if self.threshold is not None:
-            res["threshold"] = self.threshold
-        return res
-
-
-class SemanticSearchResult(pydantic.BaseModel):
-    dataset_id: uuid.UUID
-    score: float
-    data: dict[str, Any]
-    content: str
-
-
-class SemanticSearchResponse(pydantic.BaseModel):
-    results: list[SemanticSearchResult]
-
-
-class PipelineRunError(Exception):
-    error_code: str
-    error_message: str
-
-    def __init__(self, response: aiohttp.ClientResponse):
-        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 = Any  # non-null, must be JSON-serializable
-EvaluationDatapointTarget = Optional[Any]  # must be JSON-serializable
-EvaluationDatapointMetadata = Optional[Any]  # must be JSON-serializable
+EvaluationDatapointTarget = Any | None  # must be JSON-serializable
+EvaluationDatapointMetadata = Any | None  # must be JSON-serializable
 
 
 # EvaluationDatapoint is a single data point in the evaluation
-class Datapoint(pydantic.BaseModel):
+class Datapoint(BaseModel):
     # input to the executor function.
     data: EvaluationDatapointData
     # input to the evaluator function (alongside the executor output).
-    target: EvaluationDatapointTarget = pydantic.Field(default=None)
-    metadata: EvaluationDatapointMetadata = pydantic.Field(default=None)
+    target: EvaluationDatapointTarget = Field(default_factory=dict)
+    metadata: EvaluationDatapointMetadata = Field(default_factory=dict)
+    id: uuid.UUID | None = Field(default=None)
+    created_at: datetime.datetime | None = Field(default=None, alias="createdAt")
+
+
+class Dataset(BaseModel):
+    id: uuid.UUID = Field()
+    name: str = Field()
+    created_at: datetime.datetime = Field(alias="createdAt")
+
+
+class PushDatapointsResponse(BaseModel):
+    dataset_id: uuid.UUID = Field(alias="datasetId")
 
 
 ExecutorFunctionReturnType = Any
-EvaluatorFunctionReturnType = Union[Numeric, dict[str, Numeric]]
+EvaluatorFunctionReturnType = Numeric | dict[str, Numeric]
 
 ExecutorFunction = Callable[
     [EvaluationDatapointData, Any],
-    Union[ExecutorFunctionReturnType, Awaitable[ExecutorFunctionReturnType]],
+    ExecutorFunctionReturnType | Awaitable[ExecutorFunctionReturnType],
 ]
 
 # EvaluatorFunction is a function that takes the output of the executor and the
@@ -134,18 +59,20 @@ ExecutorFunction = Callable[
 # multiple criteria in one go instead of running multiple evaluators.
 EvaluatorFunction = Callable[
     [ExecutorFunctionReturnType, Any],
-    Union[EvaluatorFunctionReturnType, Awaitable[EvaluatorFunctionReturnType]],
+    EvaluatorFunctionReturnType | Awaitable[EvaluatorFunctionReturnType],
 ]
 
 
-class HumanEvaluator(pydantic.BaseModel):
-    queueName: str
+class HumanEvaluatorOptionsEntry(TypedDict):
+    label: str
+    value: float
 
-    def __init__(self, queue_name: str):
-        super().__init__(queueName=queue_name)
 
+class HumanEvaluator(BaseModel):
+    options: list[HumanEvaluatorOptionsEntry] = Field(default_factory=list)
 
-class CreateEvaluationResponse(pydantic.BaseModel):
+
+class InitEvaluationResponse(BaseModel):
     id: uuid.UUID
     createdAt: datetime.datetime
     groupId: str
@@ -153,33 +80,116 @@ class CreateEvaluationResponse(pydantic.BaseModel):
     projectId: uuid.UUID
 
 
-class EvaluationResultDatapoint(pydantic.BaseModel):
+class EvaluationDatapointDatasetLink(BaseModel):
+    dataset_id: uuid.UUID
+    datapoint_id: uuid.UUID
+    created_at: datetime.datetime
+
+    def to_dict(self):
+        return {
+            "datasetId": str(self.dataset_id),
+            "datapointId": str(self.datapoint_id),
+            "createdAt": self.created_at.isoformat(),
+        }
+
+
+class PartialEvaluationDatapoint(BaseModel):
+    id: uuid.UUID
+    data: EvaluationDatapointData
+    target: EvaluationDatapointTarget
+    index: int
+    trace_id: uuid.UUID
+    executor_span_id: uuid.UUID
+    metadata: EvaluationDatapointMetadata = Field(default=None)
+    dataset_link: EvaluationDatapointDatasetLink | None = Field(default=None)
+
+    # uuid is not serializable by default, so we need to convert it to a string
+    def to_dict(self, max_data_length: int = DEFAULT_DATAPOINT_MAX_DATA_LENGTH):
+        serialized_data = serialize(self.data)
+        serialized_target = serialize(self.target)
+        str_data = json_dumps(serialized_data)
+        str_target = json_dumps(serialized_target)
+        try:
+            return {
+                "id": str(self.id),
+                "data": (
+                    str_data[:max_data_length]
+                    if len(str_data) > max_data_length
+                    else serialized_data
+                ),
+                "target": (
+                    str_target[:max_data_length]
+                    if len(str_target) > max_data_length
+                    else serialized_target
+                ),
+                "index": self.index,
+                "traceId": str(self.trace_id),
+                "executorSpanId": str(self.executor_span_id),
+                "metadata": (
+                    serialize(self.metadata) if self.metadata is not None else {}
+                ),
+                "datasetLink": (
+                    self.dataset_link.to_dict()
+                    if self.dataset_link is not None
+                    else None
+                ),
+            }
+        except Exception as e:
+            raise ValueError(f"Error serializing PartialEvaluationDatapoint: {e}")
+
+
+class EvaluationResultDatapoint(BaseModel):
+    id: uuid.UUID
+    index: int
     data: EvaluationDatapointData
     target: EvaluationDatapointTarget
     executor_output: ExecutorFunctionReturnType
-    scores: dict[str, Numeric]
-    human_evaluators: list[HumanEvaluator] = pydantic.Field(default_factory=list)
+    scores: dict[str, Optional[Numeric]]
     trace_id: uuid.UUID
     executor_span_id: uuid.UUID
+    metadata: EvaluationDatapointMetadata = Field(default=None)
+    dataset_link: EvaluationDatapointDatasetLink | None = Field(default=None)
 
     # uuid is not serializable by default, so we need to convert it to a string
-    def to_dict(self):
+    def to_dict(self, max_data_length: int = DEFAULT_DATAPOINT_MAX_DATA_LENGTH):
         try:
+            serialized_data = serialize(self.data)
+            serialized_target = serialize(self.target)
+            serialized_executor_output = serialize(self.executor_output)
+            str_data = json.dumps(serialized_data)
+            str_target = json.dumps(serialized_target)
+            str_executor_output = json.dumps(serialized_executor_output)
             return {
-                "data": serialize(self.data),
-                "target": serialize(self.target),
-                "executorOutput": serialize(self.executor_output),
+                # preserve only preview of the data, target and executor output
+                # (full data is in trace)
+                "id": str(self.id),
+                "data": (
+                    str_data[:max_data_length]
+                    if len(str_data) > max_data_length
+                    else serialized_data
+                ),
+                "target": (
+                    str_target[:max_data_length]
+                    if len(str_target) > max_data_length
+                    else serialized_target
+                ),
+                "executorOutput": (
+                    str_executor_output[:max_data_length]
+                    if len(str_executor_output) > max_data_length
+                    else serialized_executor_output
+                ),
                 "scores": self.scores,
                 "traceId": str(self.trace_id),
-                "humanEvaluators": [
-                    (
-                        v.model_dump()
-                        if isinstance(v, pydantic.BaseModel)
-                        else serialize(v)
-                    )
-                    for v in self.human_evaluators
-                ],
                 "executorSpanId": str(self.executor_span_id),
+                "index": self.index,
+                "metadata": (
+                    serialize(self.metadata) if self.metadata is not None else {}
+                ),
+                "datasetLink": (
+                    self.dataset_link.to_dict()
+                    if self.dataset_link is not None
+                    else None
+                ),
             }
         except Exception as e:
             raise ValueError(f"Error serializing EvaluationResultDatapoint: {e}")
@@ -191,21 +201,123 @@ class SpanType(Enum):
     PIPELINE = "PIPELINE"  # must not be set manually
     EXECUTOR = "EXECUTOR"
     EVALUATOR = "EVALUATOR"
+    HUMAN_EVALUATOR = "HUMAN_EVALUATOR"
     EVALUATION = "EVALUATION"
 
 
 class TraceType(Enum):
     DEFAULT = "DEFAULT"
-    EVENT = "EVENT"  # deprecated
     EVALUATION = "EVALUATION"
 
 
-class GetDatapointsResponse(pydantic.BaseModel):
+class GetDatapointsResponse(BaseModel):
     items: list[Datapoint]
-    totalCount: int
+    total_count: int = Field(alias="totalCount")
+
+
+class LaminarSpanContext(BaseModel):
+    """
+    A span context that can be used to continue a trace across services. This
+    is a slightly modified version of the OpenTelemetry span context. For
+    usage examples, see `Laminar.serialize_span_context`,
+    `Laminar.get_span_context`, and `Laminar.deserialize_laminar_span_context`.
+
+    The difference between this and the OpenTelemetry span context is that
+    the `trace_id` and `span_id` are stored as UUIDs instead of integers for
+    easier debugging, and the separate trace flags are not currently stored.
+    """
+
+    trace_id: uuid.UUID
+    span_id: uuid.UUID
+    is_remote: bool = Field(default=False)
+    span_path: list[str] = Field(default=[])
+    span_ids_path: list[str] = Field(default=[])  # stringified UUIDs
+    user_id: str | None = Field(default=None)
+    session_id: str | None = Field(default=None)
+    trace_type: TraceType | None = Field(default=None)
+    metadata: dict[str, Any] | None = Field(default=None)
+
+    def __str__(self) -> str:
+        return self.model_dump_json()
+
+    @classmethod
+    def try_to_otel_span_context(
+        cls,
+        span_context: "LaminarSpanContext" | dict[str, Any] | str | SpanContext,
+        logger: logging.Logger | None = None,
+    ) -> SpanContext:
+        if logger is None:
+            logger = logging.getLogger(__name__)
+
+        if isinstance(span_context, LaminarSpanContext):
+            return SpanContext(
+                trace_id=span_context.trace_id.int,
+                span_id=span_context.span_id.int,
+                is_remote=span_context.is_remote,
+                trace_flags=TraceFlags(TraceFlags.SAMPLED),
+            )
+        elif isinstance(span_context, SpanContext) or (
+            isinstance(getattr(span_context, "trace_id", None), int)
+            and isinstance(getattr(span_context, "span_id", None), int)
+        ):
+            logger.warning(
+                "span_context provided"
+                " is likely a raw OpenTelemetry span context. Will try to use it. "
+                "Please use `LaminarSpanContext` instead."
+            )
+            return span_context
+        elif isinstance(span_context, (dict, str)):
+            try:
+                laminar_span_context = cls.deserialize(span_context)
+                return SpanContext(
+                    trace_id=laminar_span_context.trace_id.int,
+                    span_id=laminar_span_context.span_id.int,
+                    is_remote=laminar_span_context.is_remote,
+                    trace_flags=TraceFlags(TraceFlags.SAMPLED),
+                )
+            except Exception:
+                raise ValueError("Invalid span_context provided")
+        else:
+            raise ValueError("Invalid span_context provided")
+
+    @classmethod
+    def deserialize(cls, data: dict[str, Any] | str) -> "LaminarSpanContext":
+        if isinstance(data, dict):
+            # Convert camelCase to snake_case for known fields
+            converted_data = {
+                "trace_id": data.get("trace_id") or data.get("traceId"),
+                "span_id": data.get("span_id") or data.get("spanId"),
+                "is_remote": data.get("is_remote") or data.get("isRemote", False),
+                "span_path": data.get("span_path") or data.get("spanPath", []),
+                "span_ids_path": data.get("span_ids_path")
+                or data.get("spanIdsPath", []),
+                "user_id": data.get("user_id") or data.get("userId"),
+                "session_id": data.get("session_id") or data.get("sessionId"),
+                "trace_type": data.get("trace_type") or data.get("traceType"),
+                "metadata": data.get("metadata") or data.get("metadata", {}),
+            }
+            return cls.model_validate(converted_data)
+        elif isinstance(data, str):
+            return cls.deserialize(json.loads(data))
+        else:
+            raise ValueError("Invalid span_context provided")
+
+
+class ModelProvider(str, Enum):
+    ANTHROPIC = "anthropic"
+    BEDROCK = "bedrock"
+    OPENAI = "openai"
+    GEMINI = "gemini"
+
+
+class MaskInputOptions(TypedDict):
+    textarea: bool | None
+    text: bool | None
+    number: bool | None
+    select: bool | None
+    email: bool | None
+    tel: bool | None
 
 
-class TracingLevel(Enum):
-    OFF = 0
-    META_ONLY = 1
-    ALL = 2
+class SessionRecordingOptions(TypedDict):
+    mask_input_options: MaskInputOptions | None

lmnr/sdk/utils.py

@@ -1,13 +1,19 @@
-import asyncio
 import datetime
 import dataclasses
+import dotenv
 import enum
 import inspect
+import os
+import orjson
 import pydantic
 import queue
 import typing
 import uuid
 
+from lmnr.sdk.log import get_default_logger
+
+logger = get_default_logger(__name__)
+
 
 def is_method(func: typing.Callable) -> bool:
     # inspect.ismethod is True for bound methods only, but in the decorator,
@@ -31,14 +37,13 @@ def is_async(func: typing.Callable) -> bool:
         return False
 
     # Check if the function is asynchronous
-    if asyncio.iscoroutinefunction(func):
+    if inspect.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
+    # This is for cases when a decorator (not ours) did not properly use
     # `functools.wraps` or `functools.update_wrapper`
-    CO_COROUTINE = inspect.CO_COROUTINE
-    return (func.__code__.co_flags & CO_COROUTINE) != 0
+    return (func.__code__.co_flags & inspect.CO_COROUTINE) != 0
 
 
 def is_async_iterator(o: typing.Any) -> bool:
@@ -49,7 +54,7 @@ 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 serialize(obj: typing.Any) -> str | dict[str, typing.Any]:
     def serialize_inner(o: typing.Any):
         if isinstance(o, (datetime.datetime, datetime.date)):
             return o.strftime("%Y-%m-%dT%H:%M:%S.%f%z")
@@ -86,14 +91,153 @@ def get_input_from_func_args(
     is_method: bool = False,
     func_args: list[typing.Any] = [],
     func_kwargs: dict[str, typing.Any] = {},
+    ignore_inputs: list[str] | None = None,
 ) -> dict[str, typing.Any]:
     # Remove implicitly passed "self" or "cls" argument for
     # instance or class methods
-    res = func_kwargs.copy()
+    res = {
+        k: v
+        for k, v in func_kwargs.items()
+        if not (ignore_inputs and k in ignore_inputs)
+    }
     for i, k in enumerate(inspect.signature(func).parameters.keys()):
         if is_method and k in ["self", "cls"]:
             continue
+        if ignore_inputs and k in ignore_inputs:
+            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
+
+
+def from_env(key: str) -> str | None:
+    if val := os.getenv(key):
+        return val
+    dotenv_path = dotenv.find_dotenv(usecwd=True)
+    # use DotEnv directly so we can set verbose to False
+    return dotenv.main.DotEnv(dotenv_path, verbose=False, encoding="utf-8").get(key)
+
+
+def is_otel_attribute_value_type(value: typing.Any) -> bool:
+    def is_primitive_type(value: typing.Any) -> bool:
+        return isinstance(value, (int, float, str, bool))
+
+    if is_primitive_type(value):
+        return True
+    elif isinstance(value, typing.Sequence):
+        if len(value) > 0:
+            return is_primitive_type(value[0]) and all(
+                isinstance(v, type(value[0])) for v in value
+            )
+        return True
+    return False
+
+
+def get_otel_env_var(var_name: str) -> str | None:
+    """Get OTEL environment variable with priority order.
+
+    Checks in order:
+    1. OTEL_EXPORTER_OTLP_TRACES_{var_name}
+    2. OTEL_EXPORTER_OTLP_{var_name}
+    3. OTEL_{var_name}
+
+    Args:
+        var_name: The variable name (e.g., 'ENDPOINT', 'HEADERS', 'TIMEOUT')
+
+    Returns:
+        str | None: The environment variable value or None if not found
+    """
+    candidates = [
+        f"OTEL_EXPORTER_OTLP_TRACES_{var_name}",
+        f"OTEL_EXPORTER_OTLP_{var_name}",
+        f"OTEL_{var_name}",
+    ]
+
+    for candidate in candidates:
+        if value := from_env(candidate):
+            return value
+    return None
+
+
+def parse_otel_headers(headers_str: str | None) -> dict[str, str]:
+    """Parse OTEL headers string into dictionary.
+
+    Format: key1=value1,key2=value2
+    Values are URL-decoded.
+
+    Args:
+        headers_str: Headers string in OTEL format
+
+    Returns:
+        dict[str, str]: Parsed headers dictionary
+    """
+    if not headers_str:
+        return {}
+
+    headers = {}
+    for pair in headers_str.split(","):
+        if "=" in pair:
+            key, value = pair.split("=", 1)
+            import urllib.parse
+
+            headers[key.strip()] = urllib.parse.unquote(value.strip())
+    return headers
+
+
+def format_id(id_value: str | int | uuid.UUID) -> str:
+    """Format trace/span/evaluation ID to a UUID string, or return valid UUID strings as-is.
+
+    Args:
+        id_value: The ID in various formats (UUID, int, or valid UUID string)
+
+    Returns:
+        str: UUID string representation
+
+    Raises:
+        ValueError: If id_value cannot be converted to a valid UUID
+    """
+    if isinstance(id_value, uuid.UUID):
+        return str(id_value)
+    elif isinstance(id_value, int):
+        return str(uuid.UUID(int=id_value))
+    elif isinstance(id_value, str):
+        uuid.UUID(id_value)
+        return id_value
+    else:
+        raise ValueError(f"Invalid ID type: {type(id_value)}")
+
+
+DEFAULT_PLACEHOLDER = {}
+
+
+def default_json(o):
+    if isinstance(o, pydantic.BaseModel):
+        return o.model_dump()
+
+    # Handle various sequence types, but not strings or bytes
+    if isinstance(o, (list, tuple, set, frozenset)):
+        return list(o)
+
+    try:
+        return str(o)
+    except Exception:
+        logger.debug("Failed to serialize data to JSON, inner type: %s", type(o))
+        pass
+    return DEFAULT_PLACEHOLDER
+
+
+def json_dumps(data: dict) -> str:
+    try:
+        return orjson.dumps(
+            data,
+            default=default_json,
+            option=orjson.OPT_SERIALIZE_DATACLASS
+            | orjson.OPT_SERIALIZE_UUID
+            | orjson.OPT_UTC_Z
+            | orjson.OPT_NON_STR_KEYS,
+        ).decode("utf-8")
+    except Exception:
+        # Log the exception and return a placeholder if serialization completely fails
+        logger.info("Failed to serialize data to JSON, type: %s", type(data))
+        return "{}"  # Return an empty JSON object as a fallback