opik 1.9.41__py3-none-any.whl → 1.9.86__py3-none-any.whl

opik/api_objects/dataset/dataset.py

@@ -1,14 +1,28 @@
 import logging
 import functools
 import time
-from typing import Optional, Any, List, Dict, Sequence, Set, TYPE_CHECKING, Callable
+from typing import (
+    Optional,
+    Any,
+    List,
+    Dict,
+    Sequence,
+    Set,
+    TYPE_CHECKING,
+    Callable,
+    Iterator,
+)
 
 from opik.api_objects import rest_stream_parser
 from opik.rest_api import client as rest_api_client
-from opik.rest_api.types import dataset_item_write as rest_dataset_item
+from opik.rest_api.types import (
+    dataset_item_write as rest_dataset_item,
+    dataset_item as rest_dataset_item_read,
+)
 from opik.rest_api.core.api_error import ApiError
 from opik.message_processing.batching import sequence_splitter
 from opik.rate_limit import rate_limit
+from opik import id_helpers
 import opik.exceptions as exceptions
 import opik.config as config
 from opik.rest_client_configurator import retry_decorator
@@ -75,6 +89,7 @@ class Dataset:
         name: str,
         description: Optional[str],
         rest_client: rest_api_client.OpikApi,
+        dataset_items_count: Optional[int] = None,
     ) -> None:
         """
         A Dataset object. This object should not be created directly, instead use :meth:`opik.Opik.create_dataset` or :meth:`opik.Opik.get_dataset`.
@@ -82,6 +97,7 @@ class Dataset:
         self._name = name
         self._description = description
         self._rest_client = rest_client
+        self._dataset_items_count = dataset_items_count
 
         self._id_to_hash: Dict[str, str] = {}
         self._hashes: Set[str] = set()
@@ -103,13 +119,36 @@ class Dataset:
         """The description of the dataset."""
         return self._description
 
+    @property
+    def dataset_items_count(self) -> Optional[int]:
+        """
+        The total number of items in the dataset.
+
+        If the count is not cached locally, it will be fetched from the backend.
+        """
+        if self._dataset_items_count is None:
+            dataset_info = self._rest_client.datasets.get_dataset_by_identifier(
+                dataset_name=self._name
+            )
+            self._dataset_items_count = dataset_info.dataset_items_count
+        return self._dataset_items_count
+
     def _insert_batch_with_retry(
-        self, batch: List[rest_dataset_item.DatasetItemWrite]
+        self,
+        batch: List[rest_dataset_item.DatasetItemWrite],
+        batch_group_id: str,
     ) -> None:
-        """Insert a batch of dataset items with automatic retry on rate limit errors."""
+        """Insert a batch of dataset items with automatic retry on rate limit errors.
+
+        Args:
+            batch: List of dataset items to insert.
+            batch_group_id: UUIDv7 identifier that groups all batches from a single
+                user operation together. All batches sent as part of one insert/update
+                call share the same batch_group_id.
+        """
         _ensure_rest_api_call_respecting_rate_limit(
             lambda: self._rest_client.datasets.create_or_update_dataset_items(
-                dataset_name=self._name, items=batch
+                dataset_name=self._name, items=batch, batch_group_id=batch_group_id
             )
         )
         LOGGER.debug("Successfully sent dataset items batch of size %d", len(batch))
@@ -150,13 +189,15 @@ class Dataset:
             max_length=constants.DATASET_ITEMS_MAX_BATCH_SIZE,
         )
 
+        batch_group_id = id_helpers.generate_id()
+
         for batch in batches:
             LOGGER.debug("Sending dataset items batch of size %d", len(batch))
-            self._insert_batch_with_retry(batch)
+            self._insert_batch_with_retry(batch, batch_group_id=batch_group_id)
 
     def insert(self, items: Sequence[Dict[str, Any]]) -> None:
         """
-        Insert new items into the dataset.
+        Insert new items into the dataset. A new dataset version will be created.
 
         Args:
             items: List of dicts (which will be converted to dataset items)
@@ -168,15 +209,17 @@ class Dataset:
         ]
         self.__internal_api__insert_items_as_dataclasses__(dataset_items)
 
+        # Invalidate the cached count so it will be fetched from backend on next access
+        self._dataset_items_count = None
+
     def __internal_api__sync_hashes__(self) -> None:
         """Updates all the hashes in the dataset"""
         LOGGER.debug("Start hash sync in dataset")
-        all_items = self.__internal_api__get_items_as_dataclasses__()
 
         self._id_to_hash = {}
         self._hashes = set()
 
-        for item in all_items:
+        for item in self.__internal_api__stream_items_as_dataclasses__():
             item_hash = item.content_hash()
             self._id_to_hash[item.id] = item_hash  # type: ignore
             self._hashes.add(item_hash)
@@ -201,9 +244,29 @@ class Dataset:
 
         self.insert(items)
 
+    def _delete_batch_with_retry(
+        self,
+        batch: List[str],
+        batch_group_id: str,
+    ) -> None:
+        """Delete a batch of dataset items with automatic retry on rate limit errors.
+
+        Args:
+            batch: List of item IDs to delete.
+            batch_group_id: UUIDv7 identifier that groups all batches from a single
+                user operation together. All batches sent as part of one delete
+                call share the same batch_group_id.
+        """
+        _ensure_rest_api_call_respecting_rate_limit(
+            lambda: self._rest_client.datasets.delete_dataset_items(
+                item_ids=batch, batch_group_id=batch_group_id
+            )
+        )
+        LOGGER.debug("Successfully deleted dataset items batch of size %d", len(batch))
+
     def delete(self, items_ids: List[str]) -> None:
         """
-        Delete items from the dataset.
+        Delete items from the dataset. A new dataset version will be created.
 
         Args:
             items_ids: List of item ids to delete.
@@ -212,9 +275,11 @@ class Dataset:
             items_ids, max_length=constants.DATASET_ITEMS_MAX_BATCH_SIZE
         )
 
+        batch_group_id = id_helpers.generate_id()
+
         for batch in batches:
             LOGGER.debug("Deleting dataset items batch: %s", batch)
-            self._rest_client.datasets.delete_dataset_items(item_ids=batch)
+            self._delete_batch_with_retry(batch, batch_group_id=batch_group_id)
 
             for item_id in batch:
                 if item_id in self._id_to_hash:
@@ -222,12 +287,18 @@ class Dataset:
                     self._hashes.discard(hash)
                     del self._id_to_hash[item_id]
 
+        # Invalidate the cached count so it will be fetched from backend on next access
+        self._dataset_items_count = None
+
     def clear(self) -> None:
         """
-        Delete all items from the given dataset.
+        Delete all items from the given dataset. A new dataset version will be created.
         """
-        all_items = self.__internal_api__get_items_as_dataclasses__()
-        item_ids = [item.id for item in all_items if item.id is not None]
+        item_ids = [
+            item.id
+            for item in self.__internal_api__stream_items_as_dataclasses__()
+            if item.id is not None
+        ]
 
         self.delete(item_ids)
 
@@ -240,7 +311,7 @@ class Dataset:
         Returns:
             A pandas DataFrame containing all items in the dataset.
         """
-        dataset_items = self.__internal_api__get_items_as_dataclasses__()
+        dataset_items = list(self.__internal_api__stream_items_as_dataclasses__())
 
         return converters.to_pandas(dataset_items, keys_mapping={})
 
@@ -251,7 +322,7 @@ class Dataset:
         Returns:
             A JSON string representation of all items in the dataset.
         """
-        dataset_items = self.__internal_api__get_items_as_dataclasses__()
+        dataset_items = list(self.__internal_api__stream_items_as_dataclasses__())
 
         return converters.to_json(dataset_items, keys_mapping={})
 
@@ -265,65 +336,88 @@ class Dataset:
         Returns:
             A list of dictionaries objects representing the samples.
         """
-        dataset_items_as_dataclasses = self.__internal_api__get_items_as_dataclasses__(
-            nb_samples
-        )
         dataset_items_as_dicts = [
             {"id": item.id, **item.get_content()}
-            for item in dataset_items_as_dataclasses
+            for item in self.__internal_api__stream_items_as_dataclasses__(nb_samples)
         ]
 
         return dataset_items_as_dicts
 
-    @retry_decorator.opik_rest_retry
-    def __internal_api__get_items_as_dataclasses__(
+    def __internal_api__stream_items_as_dataclasses__(
         self,
         nb_samples: Optional[int] = None,
+        batch_size: Optional[int] = None,
         dataset_item_ids: Optional[List[str]] = None,
-    ) -> List[dataset_item.DatasetItem]:
-        results: List[dataset_item.DatasetItem] = []
+    ) -> Iterator[dataset_item.DatasetItem]:
+        """
+        Stream dataset items as a generator instead of loading all at once.
+
+        This method yields dataset items one at a time, enabling evaluation to start
+        processing items before the entire dataset is downloaded. This is particularly
+        useful for large datasets with heavy payloads (images, videos, audio).
+
+        Args:
+            nb_samples: Maximum number of items to retrieve. If None, all items are streamed.
+            batch_size: Maximum number of items to fetch per batch from the backend.
+                        If None, uses the default value from constants.DATASET_STREAM_BATCH_SIZE.
+            dataset_item_ids: Optional list of specific item IDs to retrieve. If provided,
+                            only items with matching IDs will be yielded.
+
+        Yields:
+            DatasetItem objects one at a time
+        """
+        if batch_size is None:
+            batch_size = constants.DATASET_STREAM_BATCH_SIZE
+
         last_retrieved_id: Optional[str] = None
         should_retrieve_more_items = True
-
+        items_yielded = 0
         dataset_items_ids_left = set(dataset_item_ids) if dataset_item_ids else None
 
         while should_retrieve_more_items:
-            dataset_items = rest_stream_parser.read_and_parse_stream(
-                stream=self._rest_client.datasets.stream_dataset_items(
-                    dataset_name=self._name,
-                    last_retrieved_id=last_retrieved_id,
-                ),
-                item_class=dataset_item.DatasetItem,
-                nb_samples=nb_samples,
-            )
+            # Wrap the streaming call in retry logic so we can resume from last_retrieved_id
+            @retry_decorator.opik_rest_retry
+            def _fetch_batch() -> List[rest_dataset_item_read.DatasetItem]:
+                return rest_stream_parser.read_and_parse_stream(
+                    stream=self._rest_client.datasets.stream_dataset_items(
+                        dataset_name=self._name,
+                        last_retrieved_id=last_retrieved_id,
+                        steam_limit=batch_size,
+                    ),
+                    item_class=rest_dataset_item_read.DatasetItem,
+                    nb_samples=nb_samples,
+                )
+
+            dataset_items = _fetch_batch()
 
             if len(dataset_items) == 0:
                 should_retrieve_more_items = False
+                break
 
             for item in dataset_items:
                 dataset_item_id = item.id
                 last_retrieved_id = dataset_item_id
 
+                # Filter by dataset_item_ids if provided
                 if dataset_items_ids_left is not None:
                     if dataset_item_id not in dataset_items_ids_left:
                         continue
                     else:
                         dataset_items_ids_left.remove(dataset_item_id)
 
-                data_item_content = item.get_content().get("data", {})
-
                 reconstructed_item = dataset_item.DatasetItem(
                     id=item.id,
                     trace_id=item.trace_id,
                     span_id=item.span_id,
                     source=item.source,
-                    **data_item_content,
+                    **item.data,
                 )
 
-                results.append(reconstructed_item)
+                yield reconstructed_item
+                items_yielded += 1
 
                 # Stop retrieving if we have enough samples
-                if nb_samples is not None and len(results) == nb_samples:
+                if nb_samples is not None and items_yielded >= nb_samples:
                     should_retrieve_more_items = False
                     break
 
@@ -335,14 +429,13 @@ class Dataset:
                     should_retrieve_more_items = False
                     break
 
+        # Warn if some requested items were not found
         if dataset_items_ids_left and len(dataset_items_ids_left) > 0:
             LOGGER.warning(
                 "The following dataset items were not found in the dataset: %s",
                 dataset_items_ids_left,
             )
 
-        return results
-
     def insert_from_json(
         self,
         json_array: str,

opik/api_objects/dataset/rest_operations.py

@@ -31,6 +31,7 @@ def get_datasets(
                 name=dataset_fern.name,
                 description=dataset_fern.description,
                 rest_client=rest_client,
+                dataset_items_count=dataset_fern.dataset_items_count,
             )
 
             if sync_items:
@@ -88,6 +89,7 @@ def get_dataset_experiments(
                     rest_client=rest_client,
                     streamer=streamer,
                     experiments_client=experiments_client,
+                    tags=experiment_.tags,
                 )
             )
 

opik/api_objects/experiment/experiment.py

@@ -26,6 +26,7 @@ class Experiment:
         streamer: streamer.Streamer,
         experiments_client: experiments_client.ExperimentsClient,
         prompts: Optional[List[base_prompt.BasePrompt]] = None,
+        tags: Optional[List[str]] = None,
     ) -> None:
         self._id = id
         self._name = name
@@ -34,6 +35,7 @@ class Experiment:
         self._prompts = prompts
         self._streamer = streamer
         self._experiments_client = experiments_client
+        self._tags = tags
 
     @property
     def id(self) -> str:
@@ -53,6 +55,10 @@ class Experiment:
 
         return name
 
+    @property
+    def tags(self) -> Optional[List[str]]:
+        return self._tags
+
     @functools.cached_property
     def dataset_id(self) -> str:
         return self._rest_client.datasets.get_dataset_by_identifier(

opik/api_objects/helpers.py

@@ -13,7 +13,7 @@ from ..rest_api.types import (
     trace_filter_public,
     trace_thread_filter,
 )
-from ..types import FeedbackScoreDict
+from ..types import BatchFeedbackScoreDict
 
 LOGGER = logging.getLogger(__name__)
 
@@ -139,7 +139,7 @@ def parse_search_expressions(
 
 
 def parse_feedback_score_messages(
-    scores: List[FeedbackScoreDict],
+    scores: List[BatchFeedbackScoreDict],
     project_name: str,
     parsed_item_class: Type[ScoreMessageT],
     logger: logging.Logger,
@@ -155,9 +155,13 @@ def parse_feedback_score_messages(
 
     score_messages = [
         parsed_item_class(
+            id=score_dict["id"],
+            name=score_dict["name"],
+            value=score_dict["value"],
             source=constants.FEEDBACK_SCORE_SOURCE_SDK,
-            project_name=project_name,
-            **score_dict,
+            project_name=score_dict.get("project_name") or project_name,
+            reason=score_dict.get("reason"),
+            category_name=score_dict.get("category_name"),
         )
         for score_dict in valid_scores
     ]

opik/api_objects/local_recording.py

@@ -1,9 +1,10 @@
 import contextlib
 from typing import Iterator, List
 from typing import Optional
+
 from . import opik_client
-from ..message_processing import message_processors_chain
 from ..message_processing.emulation import local_emulator_message_processor, models
+from ..message_processing.processors import message_processors_chain
 
 
 class _LocalRecordingHandle:
@@ -53,7 +54,7 @@ def record_traces_locally(
 
     # Disallow nested/local concurrent recordings in the same process
     existing_local = message_processors_chain.get_local_emulator_message_processor(
-        chain=client._message_processor
+        chain=client.__internal_api__message_processor__
     )
     if existing_local is not None and existing_local.is_active():
         raise RuntimeError(
@@ -61,10 +62,10 @@ def record_traces_locally(
         )
 
     message_processors_chain.toggle_local_emulator_message_processor(
-        active=True, chain=client._message_processor, reset=True
+        active=True, chain=client.__internal_api__message_processor__, reset=True
     )
     local = message_processors_chain.get_local_emulator_message_processor(
-        chain=client._message_processor
+        chain=client.__internal_api__message_processor__
     )
     if local is None:
         # Should not happen given the default chain, but guard just in case
@@ -76,5 +77,5 @@ def record_traces_locally(
     finally:
         client.flush()
         message_processors_chain.toggle_local_emulator_message_processor(
-            active=False, chain=client._message_processor, reset=True
+            active=False, chain=client.__internal_api__message_processor__, reset=True
         )

opik/api_objects/observation_data.py

@@ -0,0 +1,101 @@
+import dataclasses
+import datetime
+import logging
+from typing import Any, Dict, List, Optional, TypeVar
+
+import opik.api_objects.attachment as attachment
+import opik.datetime_helpers as datetime_helpers
+from opik.types import ErrorInfoDict, FeedbackScoreDict
+from . import data_helpers
+
+LOGGER = logging.getLogger(__name__)
+
+ObservationDataT = TypeVar("ObservationDataT", bound="ObservationData")
+
+
+@dataclasses.dataclass(kw_only=True)
+class ObservationData:
+    """
+    Base class for TraceData and SpanData containing common attributes and methods.
+
+    This class uses Python 3.10's kw_only=True feature to allow optional parameters
+    to be defined in the parent class while child classes can have required parameters.
+    """
+
+    name: Optional[str] = None
+    start_time: Optional[datetime.datetime] = dataclasses.field(
+        default_factory=datetime_helpers.local_timestamp
+    )
+    end_time: Optional[datetime.datetime] = None
+    metadata: Optional[Dict[str, Any]] = None
+    input: Optional[Dict[str, Any]] = None
+    output: Optional[Dict[str, Any]] = None
+    tags: Optional[List[str]] = None
+    feedback_scores: Optional[List[FeedbackScoreDict]] = None
+    project_name: Optional[str] = None
+    error_info: Optional[ErrorInfoDict] = None
+    attachments: Optional[List[attachment.Attachment]] = None
+
+    def update(self: ObservationDataT, **new_data: Any) -> ObservationDataT:
+        """
+        Updates the attributes of the object with the provided key-value pairs. This method checks if
+        an attribute exists before updating it and merges the data appropriately for specific
+        keywords like metadata, output, input, attachments, and tags. If a key doesn't correspond
+        to an attribute of the object or the provided value is None, the update is skipped.
+
+        Args:
+            **new_data: Key-value pairs of attributes to update. Keys should match existing
+                attributes on the object, and values that are None will not update.
+
+        Returns:
+            The updated object instance (preserves the actual subclass type).
+        """
+        for key, value in new_data.items():
+            if value is None:
+                continue
+
+            if key not in self.__dict__ and key != "prompts":
+                LOGGER.debug(
+                    "An attempt to update observation with parameter name it doesn't have: %s",
+                    key,
+                )
+                continue
+
+            if key == "metadata":
+                self.metadata = data_helpers.merge_metadata(
+                    self.metadata, new_metadata=value
+                )
+                continue
+            elif key == "output":
+                self.output = data_helpers.merge_outputs(self.output, new_outputs=value)
+                continue
+            elif key == "input":
+                self.input = data_helpers.merge_inputs(self.input, new_inputs=value)
+                continue
+            elif key == "attachments":
+                self._update_attachments(value)
+                continue
+            elif key == "tags":
+                self.tags = data_helpers.merge_tags(self.tags, new_tags=value)
+                continue
+            elif key == "prompts":
+                self.metadata = data_helpers.merge_metadata(
+                    self.metadata, new_metadata=new_data.get("metadata"), prompts=value
+                )
+                continue
+
+            self.__dict__[key] = value
+
+        return self
+
+    def init_end_time(self: ObservationDataT) -> ObservationDataT:
+        """Initialize the end_time to the current timestamp."""
+        self.end_time = datetime_helpers.local_timestamp()
+        return self
+
+    def _update_attachments(self, attachments: List[attachment.Attachment]) -> None:
+        """Merge new attachments with existing ones."""
+        if self.attachments is None:
+            self.attachments = attachments
+        else:
+            self.attachments.extend(attachments)