lmnr 0.6.16__py3-none-any.whl → 0.7.26__py3-none-any.whl

lmnr/sdk/client/synchronous/resources/datasets.py

@@ -0,0 +1,131 @@
+"""Datasets resource for interacting with Laminar datasets API."""
+
+import math
+import uuid
+
+from lmnr.sdk.client.synchronous.resources.base import BaseResource
+from lmnr.sdk.log import get_default_logger
+from lmnr.sdk.types import (
+    Datapoint,
+    Dataset,
+    GetDatapointsResponse,
+    PushDatapointsResponse,
+)
+from lmnr.sdk.utils import serialize
+
+logger = get_default_logger(__name__)
+
+DEFAULT_DATASET_PULL_LIMIT = 100
+DEFAULT_DATASET_PUSH_BATCH_SIZE = 100
+
+
+class Datasets(BaseResource):
+    """Resource for interacting with Laminar datasets API."""
+
+    def list_datasets(self) -> list[Dataset]:
+        """List all datasets."""
+        response = self._client.get(
+            f"{self._base_url}/v1/datasets",
+            headers=self._headers(),
+        )
+        if response.status_code != 200:
+            raise ValueError(
+                f"Error listing datasets: [{response.status_code}] {response.text}"
+            )
+        return [Dataset.model_validate(dataset) for dataset in response.json()]
+
+    def get_dataset_by_name(self, name: str) -> list[Dataset]:
+        """Get a dataset by name."""
+        response = self._client.get(
+            f"{self._base_url}/v1/datasets",
+            params={"name": name},
+            headers=self._headers(),
+        )
+        if response.status_code != 200:
+            raise ValueError(
+                f"Error getting dataset: [{response.status_code}] {response.text}"
+            )
+        return [Dataset.model_validate(dataset) for dataset in response.json()]
+
+    def push(
+        self,
+        points: list[Datapoint],
+        name: str | None = None,
+        id: uuid.UUID | None = None,
+        batch_size: int = DEFAULT_DATASET_PUSH_BATCH_SIZE,
+        create_dataset: bool = False,
+    ) -> PushDatapointsResponse | None:
+        """Push data to a dataset."""
+
+        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")
+
+        if create_dataset and name is None:
+            raise ValueError("Name must be provided when creating a new dataset")
+
+        identifier = {"name": name} if name is not None else {"datasetId": id}
+
+        batch_num = 0
+        total_batches = math.ceil(len(points) / batch_size)
+        response = None
+        for i in range(0, len(points), batch_size):
+            batch_num += 1
+            logger.debug(f"Pushing batch {batch_num} of {total_batches}")
+            batch = points[i : i + batch_size]
+            response = self._client.post(
+                f"{self._base_url}/v1/datasets/datapoints",
+                json={
+                    **identifier,
+                    "datapoints": [serialize(point) for point in batch],
+                    "createDataset": create_dataset,
+                },
+                headers=self._headers(),
+            )
+
+            # 201 when creating a new dataset
+            if response.status_code not in [200, 201]:
+                raise ValueError(
+                    f"Error pushing data to dataset: [{response.status_code}] {response.text}"
+                )
+
+            response = PushDatapointsResponse.model_validate(response.json())
+        # Currently, the response only contains the dataset ID,
+        # so it's safe to return the last response only.
+        return response
+
+    def pull(
+        self,
+        name: str | None = None,
+        id: uuid.UUID | None = None,
+        # TODO: move const to one file, import in CLI
+        limit: int = DEFAULT_DATASET_PULL_LIMIT,
+        offset: int = 0,
+    ) -> GetDatapointsResponse:
+        """Pull data from a dataset."""
+
+        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")
+
+        identifier = {"name": name} if name is not None else {"datasetId": id}
+
+        params = {
+            **identifier,
+            "offset": offset,
+            "limit": limit,
+        }
+        response = self._client.get(
+            f"{self._base_url}/v1/datasets/datapoints",
+            params=params,
+            headers=self._headers(),
+        )
+        if response.status_code != 200:
+            raise ValueError(
+                f"Error pulling data from dataset: [{response.status_code}] {response.text}"
+            )
+        return GetDatapointsResponse.model_validate(response.json())

lmnr/sdk/client/synchronous/resources/evals.py

@@ -1,23 +1,32 @@
 """Evals resource for interacting with Laminar evaluations API."""
 
 import uuid
-import urllib.parse
+import warnings
+
 from typing import Any
 
 from lmnr.sdk.client.synchronous.resources.base import BaseResource
+from lmnr.sdk.log import get_default_logger
 from lmnr.sdk.types import (
-    InitEvaluationResponse,
+    GetDatapointsResponse,
     EvaluationResultDatapoint,
+    InitEvaluationResponse,
     PartialEvaluationDatapoint,
-    GetDatapointsResponse,
 )
+from lmnr.sdk.utils import serialize
+
+INITIAL_EVALUATION_DATAPOINT_MAX_DATA_LENGTH = 16_000_000  # 16MB
+logger = get_default_logger(__name__)
 
 
 class Evals(BaseResource):
     """Resource for interacting with Laminar evaluations API."""
 
     def init(
-        self, name: str | None = None, group_name: str | None = None, metadata: dict[str, Any] | None = None
+        self,
+        name: str | None = None,
+        group_name: str | None = None,
+        metadata: dict[str, Any] | None = None,
     ) -> InitEvaluationResponse:
         """Initialize a new evaluation.
 
@@ -53,7 +62,7 @@ class Evals(BaseResource):
     ) -> uuid.UUID:
         """
         Create a new evaluation and return its ID.
-        
+
         Parameters:
             name (str | None, optional): Optional name of the evaluation.
             group_name (str | None, optional): An identifier to group evaluations.
@@ -76,7 +85,7 @@ class Evals(BaseResource):
     ) -> uuid.UUID:
         """
         Create a datapoint for an evaluation.
-        
+
         Parameters:
             eval_id (uuid.UUID): The evaluation ID.
             data: The input data for the executor.
@@ -84,13 +93,13 @@ class Evals(BaseResource):
             metadata (dict[str, Any] | None, optional): Optional metadata.
             index (int | None, optional): Optional index of the datapoint.
             trace_id (uuid.UUID | None, optional): Optional trace ID.
-        
+
         Returns:
             uuid.UUID: The datapoint ID.
         """
-        
+
         datapoint_id = uuid.uuid4()
-        
+
         # Create a minimal datapoint first
         partial_datapoint = PartialEvaluationDatapoint(
             id=datapoint_id,
@@ -101,7 +110,7 @@ class Evals(BaseResource):
             executor_span_id=uuid.uuid4(),  # Will be updated when executor runs
             metadata=metadata,
         )
-        
+
         self.save_datapoints(eval_id, [partial_datapoint])
         return datapoint_id
 
@@ -121,16 +130,24 @@ class Evals(BaseResource):
         Raises:
             ValueError: If there's an error saving the datapoints.
         """
+        length = INITIAL_EVALUATION_DATAPOINT_MAX_DATA_LENGTH
+        points = [datapoint.to_dict(max_data_length=length) for datapoint in datapoints]
         response = self._client.post(
             self._base_url + f"/v1/evals/{eval_id}/datapoints",
             json={
-                "points": [datapoint.to_dict() for datapoint in datapoints],
+                "points": points,
                 "groupName": group_name,
             },
             headers=self._headers(),
         )
+        if response.status_code == 413:
+            self._retry_save_datapoints(eval_id, datapoints, group_name)
+            return
+
         if response.status_code != 200:
-            raise ValueError(f"Error saving evaluation datapoints: {response.text}")
+            raise ValueError(
+                f"Error saving evaluation datapoints: [{response.status_code}] {response.text}"
+            )
 
     def update_datapoint(
         self,
@@ -147,11 +164,17 @@ class Evals(BaseResource):
             executor_output (Any): The executor output.
             scores (dict[str, float | int] | None, optional): The scores. Defaults to None.
         """
-        
+
         response = self._client.post(
             self._base_url + f"/v1/evals/{eval_id}/datapoints/{datapoint_id}",
             json={
-                "executorOutput": executor_output,
+                "executorOutput": (
+                    str(serialize(executor_output))[
+                        :INITIAL_EVALUATION_DATAPOINT_MAX_DATA_LENGTH
+                    ]
+                    if executor_output is not None
+                    else None
+                ),
                 "scores": scores,
             },
             headers=self._headers(),
@@ -179,11 +202,18 @@ class Evals(BaseResource):
         Raises:
             ValueError: If there's an error fetching the datapoints.
         """
+
+        warnings.warn(
+            "Use client.datasets.pull instead",
+            DeprecationWarning,
+        )
+
         params = {"name": dataset_name, "offset": offset, "limit": limit}
-        url = (
-            self._base_url + "/v1/datasets/datapoints?" + urllib.parse.urlencode(params)
+        response = self._client.get(
+            self._base_url + "/v1/datasets/datapoints",
+            params=params,
+            headers=self._headers(),
         )
-        response = self._client.get(url, headers=self._headers())
         if response.status_code != 200:
             try:
                 resp_json = response.json()
@@ -195,3 +225,39 @@ class Evals(BaseResource):
                     f"Error fetching datapoints: [{response.status_code}] {response.text}"
                 )
         return GetDatapointsResponse.model_validate(response.json())
+
+    def _retry_save_datapoints(
+        self,
+        eval_id: uuid.UUID,
+        datapoints: list[EvaluationResultDatapoint | PartialEvaluationDatapoint],
+        group_name: str | None = None,
+        initial_length: int = INITIAL_EVALUATION_DATAPOINT_MAX_DATA_LENGTH,
+        max_retries: int = 20,
+    ):
+        retry = 0
+        length = initial_length
+        while retry < max_retries:
+            retry += 1
+            length = length // 2
+            logger.debug(
+                f"Retrying save datapoints: {retry} of {max_retries}, length: {length}"
+            )
+            if length == 0:
+                raise ValueError("Error saving evaluation datapoints")
+            points = [
+                datapoint.to_dict(max_data_length=length) for datapoint in datapoints
+            ]
+            response = self._client.post(
+                self._base_url + f"/v1/evals/{eval_id}/datapoints",
+                json={
+                    "points": points,
+                    "groupName": group_name,
+                },
+                headers=self._headers(),
+            )
+            if response.status_code != 413:
+                break
+        if response.status_code != 200:
+            raise ValueError(
+                f"Error saving evaluation datapoints: [{response.status_code}] {response.text}"
+            )

lmnr/sdk/client/synchronous/resources/evaluators.py

@@ -0,0 +1,85 @@
+"""Evaluators resource for creating evaluator scores."""
+
+import uuid
+from typing import Any
+
+from lmnr.sdk.client.synchronous.resources.base import BaseResource
+from lmnr.sdk.utils import format_id
+
+
+class Evaluators(BaseResource):
+    """Resource for creating evaluator scores."""
+
+    def score(
+        self,
+        *,
+        name: str,
+        trace_id: str | int | uuid.UUID | None = None,
+        span_id: str | int | uuid.UUID | None = None,
+        metadata: dict[str, Any] | None = None,
+        score: float,
+    ) -> None:
+        """Create a score for a span.
+
+        Args:
+            name (str): Name of the score
+            trace_id (str | int | uuid.UUID | None, optional): The trace ID to score (will be attached to root span)
+            span_id (str | int | uuid.UUID | None, optional): The span ID to score
+            metadata (dict[str, Any] | None, optional): Additional metadata. Defaults to None.
+            score (float): The score value (float)
+
+        Raises:
+            ValueError: If there's an error creating the score.
+
+        Example:
+            Score by trace ID (will attach to root span):
+
+            >>> laminar_client.evaluators.score(
+            ...     name="quality",
+            ...     trace_id="trace-id-here",
+            ...     score=0.95,
+            ...     metadata={"model": "gpt-4"}
+            ... )
+
+            Score by span ID:
+
+            >>> laminar_client.evaluators.score(
+            ...     name="relevance",
+            ...     span_id="span-id-here", 
+            ...     score=0.87
+            ... )
+        """
+        if trace_id is not None and span_id is not None:
+            raise ValueError("Cannot provide both trace_id and span_id. Please provide only one.")
+        if trace_id is None and span_id is None:
+            raise ValueError("Either 'trace_id' or 'span_id' must be provided.")
+
+        if trace_id is not None:
+            formatted_trace_id = format_id(trace_id)
+            payload = {
+                "name": name,
+                "traceId": formatted_trace_id,
+                "metadata": metadata,
+                "score": score,
+                "source": "Code",
+            }
+        else:
+            formatted_span_id = format_id(span_id)
+            payload = {
+                "name": name,
+                "spanId": formatted_span_id,
+                "metadata": metadata,
+                "score": score,
+                "source": "Code",
+            }
+
+        response = self._client.post(
+            self._base_url + "/v1/evaluators/score",
+            json=payload,
+            headers=self._headers(),
+        )
+
+        if response.status_code != 200:
+            if response.status_code == 401:
+                raise ValueError("Unauthorized. Please check your project API key.")
+            raise ValueError(f"Error creating evaluator score: {response.text}")

lmnr/sdk/client/synchronous/resources/tags.py

@@ -5,6 +5,7 @@ import uuid
 
 from lmnr.sdk.client.synchronous.resources.base import BaseResource
 from lmnr.sdk.log import get_default_logger
+from lmnr.sdk.utils import format_id
 
 logger = get_default_logger(__name__)
 
@@ -54,18 +55,11 @@ class Tags(BaseResource):
         ```
         """
         trace_tags = tags if isinstance(tags, list) else [tags]
-        if isinstance(trace_id, uuid.UUID):
-            trace_id = str(trace_id)
-        elif isinstance(trace_id, int):
-            trace_id = str(uuid.UUID(int=trace_id))
-        elif isinstance(trace_id, str):
-            uuid.UUID(trace_id)
-        else:
-            raise ValueError(f"Invalid trace id: {trace_id}")
+        formatted_trace_id = format_id(trace_id)
 
         url = self._base_url + "/v1/tag"
         payload = {
-            "traceId": trace_id,
+            "traceId": formatted_trace_id,
             "names": trace_tags,
         }
         response = self._client.post(
@@ -78,7 +72,7 @@ class Tags(BaseResource):
 
         if response.status_code == 404:
             logger.warning(
-                f"Trace {trace_id} not found. The trace may have not been ended yet."
+                f"Trace {formatted_trace_id} not found. The trace may have not been ended yet."
             )
             return []
 

lmnr/sdk/client/synchronous/sync_client.py

@@ -8,11 +8,12 @@ from typing import TypeVar
 from types import TracebackType
 
 from lmnr.sdk.client.synchronous.resources import (
-    Agent,
     BrowserEvents,
     Evals,
+    Evaluators,
     Tags,
 )
+from lmnr.sdk.client.synchronous.resources.datasets import Datasets
 from lmnr.sdk.utils import from_env
 
 _T = TypeVar("_T", bound="LaminarClient")
@@ -23,10 +24,9 @@ class LaminarClient:
     __project_api_key: str
     __client: httpx.Client = None
 
-    # Resource properties
-    __agent: Agent | None = None
     __evals: Evals | None = None
     __tags: Tags | None = None
+    __evaluators: Evaluators | None = None
 
     def __init__(
         self,
@@ -69,24 +69,40 @@ class LaminarClient:
         self.__client = httpx.Client(
             headers=self._headers(),
             timeout=timeout,
+            # Context: If the server responds with a 413, the connection becomes
+            # poisoned and freezes on subsequent requests, and there is no way
+            # to recover or recycle such connection.
+            # Setting max_keepalive_connections to 0 will resolve this, but is
+            # less efficient, as it will create a new connection
+            # (not client, so still better) for each request.
+            #
+            # Note: from my experiments with a simple python server, forcing the
+            # server to read/consume the request payload from the socket seems
+            # to resolve this, but I haven't figured out how to do that in our
+            # real actix-web backend server and whether it makes sense to do so.
+            #
+            # TODO: investigate if there are better ways to fix this rather than
+            # setting keepalive_expiry to 0. Other alternative: migrate to
+            # requests + aiohttp.
+            #
+            # limits=httpx.Limits(
+            #     max_keepalive_connections=0,
+            #     keepalive_expiry=0,
+            # ),
         )
 
         # Initialize resource objects
-        self.__agent = Agent(self.__client, self.__base_url, self.__project_api_key)
         self.__evals = Evals(self.__client, self.__base_url, self.__project_api_key)
+        self.__evaluators = Evaluators(
+            self.__client, self.__base_url, self.__project_api_key
+        )
         self.__browser_events = BrowserEvents(
             self.__client, self.__base_url, self.__project_api_key
         )
         self.__tags = Tags(self.__client, self.__base_url, self.__project_api_key)
-
-    @property
-    def agent(self) -> Agent:
-        """Get the Agent resource.
-
-        Returns:
-            Agent: The Agent resource instance.
-        """
-        return self.__agent
+        self.__datasets = Datasets(
+            self.__client, self.__base_url, self.__project_api_key
+        )
 
     @property
     def evals(self) -> Evals:
@@ -115,6 +131,24 @@ class LaminarClient:
         """
         return self.__tags
 
+    @property
+    def evaluators(self) -> Evaluators:
+        """Get the Evaluators resource.
+
+        Returns:
+            Evaluators: The Evaluators resource instance.
+        """
+        return self.__evaluators
+
+    @property
+    def datasets(self) -> Datasets:
+        """Get the Datasets resource.
+
+        Returns:
+            Datasets: The Datasets resource instance.
+        """
+        return self.__datasets
+
     def shutdown(self):
         """Shutdown the client by closing underlying connections."""
         self.__client.close()
@@ -155,5 +189,3 @@ class LaminarClient:
             "Content-Type": "application/json",
             "Accept": "application/json",
         }
-
-

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}]"
+        )