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

lmnr/sdk/datasets/__init__.py

@@ -0,0 +1,94 @@
+from abc import ABC, abstractmethod
+from pathlib import Path
+
+import uuid
+
+from lmnr.sdk.client.synchronous.sync_client import LaminarClient
+from lmnr.sdk.datasets.file_utils import load_from_paths
+from lmnr.sdk.log import get_default_logger
+from lmnr.sdk.types import Datapoint
+
+DEFAULT_FETCH_SIZE = 25
+LOG = get_default_logger(__name__, verbose=False)
+
+
+class EvaluationDataset(ABC):
+    @abstractmethod
+    def __init__(self, *args, **kwargs):
+        pass
+
+    @abstractmethod
+    def __len__(self) -> int:
+        pass
+
+    @abstractmethod
+    def __getitem__(self, idx) -> Datapoint:
+        pass
+
+    def slice(self, start: int, end: int):
+        return [self[i] for i in range(max(start, 0), min(end, len(self)))]
+
+
+class LaminarDataset(EvaluationDataset):
+    client: LaminarClient
+    id: uuid.UUID | None = None
+
+    def __init__(
+        self,
+        name: str | None = None,
+        id: uuid.UUID | None = None,
+        fetch_size: int = DEFAULT_FETCH_SIZE,
+    ):
+        self.name = name
+        self.id = id
+        if name is None and id is None:
+            raise ValueError("Either name or id must be provided")
+        if name is not None and id is not None:
+            raise ValueError("Only one of name or id must be provided")
+        self._len = None
+        self._fetched_items = []
+        self._offset = 0
+        self._fetch_size = fetch_size
+        self._logger = get_default_logger(self.__class__.__name__)
+
+    def _fetch_batch(self):
+        self._logger.debug(
+            f"dataset name: {self.name}, id: {self.id}. Fetching batch from {self._offset} to "
+            + f"{self._offset + self._fetch_size}"
+        )
+        identifier = {"id": self.id} if self.id is not None else {"name": self.name}
+        resp = self.client.datasets.pull(
+            **identifier,
+            offset=self._offset,
+            limit=self._fetch_size,
+        )
+        self._fetched_items += resp.items
+        self._offset = len(self._fetched_items)
+        if self._len is None:
+            self._len = resp.total_count
+
+    def __len__(self) -> int:
+        if self._len is None:
+            self._fetch_batch()
+        return self._len
+
+    def __getitem__(self, idx) -> Datapoint:
+        if idx >= len(self._fetched_items):
+            self._fetch_batch()
+        return self._fetched_items[idx]
+
+    def set_client(self, client: LaminarClient):
+        self.client = client
+
+    def push(self, paths: str | list[str], recursive: bool = False):
+        paths = [paths] if isinstance(paths, str) else paths
+        paths = [Path(path) for path in paths]
+        data = load_from_paths(paths, recursive)
+        if len(data) == 0:
+            LOG.warning("No data to push. Skipping")
+            return
+        identifier = {"id": self.id} if self.id is not None else {"name": self.name}
+        self.client.datasets.push(data, **identifier)
+        LOG.info(
+            f"Successfully pushed {len(data)} datapoints to dataset [{identifier}]"
+        )

lmnr/sdk/datasets/file_utils.py

@@ -0,0 +1,91 @@
+from pathlib import Path
+from typing import Any
+import csv
+import orjson
+
+from lmnr.sdk.log import get_default_logger
+
+LOG = get_default_logger(__name__, verbose=False)
+
+
+def _is_supported_file(file: Path) -> bool:
+    """Check if a file is supported."""
+    return file.suffix in [".json", ".csv", ".jsonl"]
+
+
+def _collect_files(paths: list[Path], recursive: bool = False) -> list[Path]:
+    """
+    Collect all supported files from the given paths.
+
+    Handles both files and directories. If a path is a directory,
+    collects all supported files within it (recursively if specified).
+    """
+    collected_files = []
+
+    for path in paths:
+        if path.is_file():
+            if _is_supported_file(path):
+                collected_files.append(path)
+            else:
+                LOG.warning(f"Skipping unsupported file type: {path}")
+        elif path.is_dir():
+            for item in path.iterdir():
+                if item.is_file() and _is_supported_file(item):
+                    collected_files.append(item)
+                elif recursive and item.is_dir():
+                    # Recursively collect files from subdirectories
+                    collected_files.extend(_collect_files([item], recursive=True))
+        else:
+            LOG.warning(f"Path does not exist or is not accessible: {path}")
+
+    return collected_files
+
+
+def _read_file(file: Path) -> list[dict[str, Any]]:
+    """Read data from a single file and return as a list of dictionaries."""
+    if file.suffix == ".json":
+        result = orjson.loads(file.read_bytes())
+        if isinstance(result, list):
+            return result
+        else:
+            return [result]
+    elif file.suffix == ".csv":
+        return [dict(row) for row in csv.DictReader(file.read_text().splitlines())]
+    elif file.suffix == ".jsonl":
+        return [
+            orjson.loads(line) for line in file.read_text().splitlines() if line.strip()
+        ]
+    else:
+        raise ValueError(f"Unsupported file type: {file.suffix}")
+
+
+def load_from_paths(paths: list[Path], recursive: bool = False) -> list[dict[str, Any]]:
+    """
+    Load data from all files in the specified paths.
+
+    First collects all file paths, then reads each file's data.
+    """
+    files = _collect_files(paths, recursive)
+
+    if not files:
+        LOG.warning("No supported files found in the specified paths")
+        return []
+
+    LOG.info(f"Found {len(files)} file(s) to read")
+
+    result = []
+    for file in files:
+        try:
+            data = _read_file(file)
+            result.extend(data)
+            LOG.info(f"Read {len(data)} record(s) from {file}")
+        except Exception as e:
+            LOG.error(f"Error reading file {file}: {e}")
+            raise
+
+    return result
+
+
+def parse_paths(paths: list[str]) -> list[Path]:
+    """Parse paths."""
+    return [Path(path) for path in paths]

lmnr/sdk/decorators.py

@@ -1,46 +1,132 @@
-from lmnr.openllmetry_sdk.decorators.base import (
-    entity_method,
-    aentity_method,
+from lmnr.opentelemetry_lib.decorators import (
+    observe_base,
+    async_observe_base,
 )
 from opentelemetry.trace import INVALID_SPAN, get_current_span
 
-from typing import Callable, Optional, TypeVar, cast
+from typing import Any, Callable, Coroutine, Literal, TypeVar, overload
 from typing_extensions import ParamSpec
 
-from lmnr.openllmetry_sdk.tracing.attributes import SESSION_ID
-from lmnr.openllmetry_sdk.tracing.tracing import update_association_properties
+from lmnr.opentelemetry_lib.tracing.attributes import SESSION_ID
+from lmnr.sdk.log import get_default_logger
+from lmnr.sdk.types import TraceType
 
 from .utils import is_async
 
+logger = get_default_logger(__name__)
 
 P = ParamSpec("P")
 R = TypeVar("R")
 
 
+# Overload for synchronous functions
+@overload
 def observe(
     *,
-    name: Optional[str] = None,
-    session_id: Optional[str] = None,
-) -> Callable[[Callable[P, R]], Callable[P, R]]:
+    name: str | None = None,
+    session_id: str | None = None,
+    user_id: str | None = None,
+    ignore_input: bool = False,
+    ignore_output: bool = False,
+    span_type: Literal["DEFAULT", "LLM", "TOOL"] = "DEFAULT",
+    ignore_inputs: list[str] | None = None,
+    input_formatter: Callable[..., str] | None = None,
+    output_formatter: Callable[..., str] | None = None,
+    metadata: dict[str, Any] | None = None,
+    tags: list[str] | None = None,
+    preserve_global_context: bool = False,
+) -> Callable[[Callable[P, R]], Callable[P, R]]: ...
+
+
+# Overload for asynchronous functions
+@overload
+def observe(
+    *,
+    name: str | None = None,
+    session_id: str | None = None,
+    user_id: str | None = None,
+    ignore_input: bool = False,
+    ignore_output: bool = False,
+    span_type: Literal["DEFAULT", "LLM", "TOOL"] = "DEFAULT",
+    ignore_inputs: list[str] | None = None,
+    input_formatter: Callable[..., str] | None = None,
+    output_formatter: Callable[..., str] | None = None,
+    metadata: dict[str, Any] | None = None,
+    tags: list[str] | None = None,
+    preserve_global_context: bool = False,
+) -> Callable[
+    [Callable[P, Coroutine[Any, Any, R]]], Callable[P, Coroutine[Any, Any, R]]
+]: ...
+
+
+# Implementation
+def observe(
+    *,
+    name: str | None = None,
+    session_id: str | None = None,
+    user_id: str | None = None,
+    ignore_input: bool = False,
+    ignore_output: bool = False,
+    span_type: Literal["DEFAULT", "LLM", "TOOL"] = "DEFAULT",
+    ignore_inputs: list[str] | None = None,
+    input_formatter: Callable[..., str] | None = None,
+    output_formatter: Callable[..., str] | None = None,
+    metadata: dict[str, Any] | None = None,
+    tags: list[str] | None = None,
+    preserve_global_context: bool = False,
+):
+    # Return type is determined by overloads above
     """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.
-        session_id (Optional[str], optional): Session ID to associate with the
-                        span and the following context. Defaults to None.
-
+        name (str | None, optional): Name of the span. Function name is used if\
+            not specified. Defaults to None.
+        session_id (str | None, optional): Session ID to associate with the\
+            span and the following context. Defaults to None.
+        user_id (str | None, optional): User ID to associate with the span and\
+            the following context. This is different from ID of a Laminar user.
+            Defaults to None.
+        ignore_input (bool, optional): Whether to ignore ALL input of the\
+            wrapped function. Defaults to False.
+        ignore_output (bool, optional): Whether to ignore ALL output of the\
+            wrapped function. Defaults to False.
+        span_type (Literal["DEFAULT", "LLM", "TOOL"], optional): Type of the span.
+            Defaults to "DEFAULT".
+        ignore_inputs (list[str] | None, optional): List of input keys to\
+            ignore. For example, if the wrapped function takes three arguments\
+            def foo(a, b, `sensitive_data`), and you want to ignore the\
+            `sensitive_data` argument, you can pass ["sensitive_data"] to\
+            this argument. Defaults to None.
+        input_formatter (Callable[P, str] | None, optional): A custom function\
+            to format the input of the wrapped function. This function should\
+            accept the same parameters as the wrapped function and return a string.\
+            All function arguments are passed to this function. Ignored if\
+            `ignore_input` is True. Does not respect `ignore_inputs` argument.
+            Defaults to None.
+        output_formatter (Callable[[R], str] | None, optional): A custom function\
+            to format the output of the wrapped function. This function should\
+            accept a single parameter (the return value of the wrapped function)\
+            and return a string. Ignored if `ignore_output` is True.\
+            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.
+        preserve_global_context (bool, optional): Whether to preserve the global\
+            OpenTelemetry context. If set to True, Laminar spans will continue\
+            traces started in the global context. Defaults to False.
     Raises:
-        Exception: re-raises the exception if the wrapped function raises
-                   an exception
+        Exception: re-raises the exception if the wrapped function raises an\
+            exception
 
     Returns:
         R: Returns the result of the wrapped function
     """
 
-    def decorator(func: Callable) -> Callable:
+    def decorator(
+        func: Callable[P, R] | Callable[P, Coroutine[Any, Any, R]],
+    ) -> Callable[P, R] | Callable[P, Coroutine[Any, Any, R]]:
         current_span = get_current_span()
         if current_span != INVALID_SPAN:
             if session_id is not None:
@@ -48,11 +134,62 @@ def observe(
         association_properties = {}
         if session_id is not None:
             association_properties["session_id"] = session_id
-        update_association_properties(association_properties)
-        return (
-            aentity_method(name=name)(func)
-            if is_async(func)
-            else entity_method(name=name)(func)
-        )
-
-    return cast(Callable, decorator)
+        if user_id is not None:
+            association_properties["user_id"] = user_id
+        if span_type in ["EVALUATION", "EXECUTOR", "EVALUATOR"]:
+            association_properties["trace_type"] = TraceType.EVALUATION.value
+        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:
+                # list(set(tags)) to deduplicate tags
+                association_properties["tags"] = list(set(tags))
+        if input_formatter is not None and ignore_input:
+            logger.warning(
+                f"observe, function {func.__name__}: Input formatter"
+                " is ignored because `ignore_input` is True. Specify only one of"
+                " `ignore_input` or `input_formatter`."
+            )
+        if input_formatter is not None and ignore_inputs is not None:
+            logger.warning(
+                f"observe, function {func.__name__}: Both input formatter and"
+                " `ignore_inputs` are specified. Input formatter"
+                " will pass all arguments to the formatter regardless of"
+                " `ignore_inputs`."
+            )
+        if output_formatter is not None and ignore_output:
+            logger.warning(
+                f"observe, function {func.__name__}: Output formatter"
+                " is ignored because `ignore_output` is True. Specify only one of"
+                " `ignore_output` or `output_formatter`."
+            )
+        if is_async(func):
+            return async_observe_base(
+                name=name,
+                ignore_input=ignore_input,
+                ignore_output=ignore_output,
+                span_type=span_type,
+                metadata=metadata,
+                ignore_inputs=ignore_inputs,
+                input_formatter=input_formatter,
+                output_formatter=output_formatter,
+                association_properties=association_properties,
+                preserve_global_context=preserve_global_context,
+            )(func)
+        else:
+            return observe_base(
+                name=name,
+                ignore_input=ignore_input,
+                ignore_output=ignore_output,
+                span_type=span_type,
+                metadata=metadata,
+                ignore_inputs=ignore_inputs,
+                input_formatter=input_formatter,
+                output_formatter=output_formatter,
+                association_properties=association_properties,
+                preserve_global_context=preserve_global_context,
+            )(func)
+
+    return decorator

lmnr/sdk/eval_control.py

@@ -1,4 +1,5 @@
 from contextvars import ContextVar
 
-PREPARE_ONLY = ContextVar("__lmnr_prepare_only", default=False)
-EVALUATION_INSTANCE = ContextVar("__lmnr_evaluation_instance", default=None)
+
+PREPARE_ONLY: ContextVar[bool] = ContextVar("__lmnr_prepare_only", default=False)
+EVALUATION_INSTANCES = ContextVar("__lmnr_evaluation_instances")