opik 1.6.4__py3-none-any.whl → 1.9.71__py3-none-any.whl

opik/api_objects/dataset/dataset.py

@@ -1,12 +1,16 @@
 import logging
-import json
 import functools
-from typing import Optional, Any, List, Dict, Sequence, Set, TYPE_CHECKING
+import time
+from typing import Optional, Any, List, Dict, Sequence, Set, TYPE_CHECKING, Callable
 
+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.core.api_error import ApiError
 from opik.message_processing.batching import sequence_splitter
-from opik import exceptions, config
+from opik.rate_limit import rate_limit
+import opik.exceptions as exceptions
+import opik.config as config
 from opik.rest_client_configurator import retry_decorator
 from .. import constants
 from . import dataset_item, converters
@@ -17,6 +21,54 @@ if TYPE_CHECKING:
 LOGGER = logging.getLogger(__name__)
 
 
+def _ensure_rest_api_call_respecting_rate_limit(
+    rest_callable: Callable[[], Any],
+) -> Any:
+    """
+    Execute a REST API call with automatic retry on rate limit (429) errors.
+
+    This function handles HTTP 429 rate limit errors by waiting for the duration
+    specified in the response headers and retrying the request. Regular retries
+    for other errors are handled by the underlying rest client.
+
+    Args:
+        rest_callable: A callable that performs the REST API call.
+
+    Returns:
+        The result of the successful REST API call.
+
+    Raises:
+        ApiError: If the error is not a 429 rate limit error.
+    """
+    while True:
+        try:
+            result = rest_callable()
+            return result
+        except ApiError as exception:
+            if exception.status_code == 429:
+                # Parse rate limit headers to get retry delay
+                if exception.headers is not None:
+                    rate_limiter = rate_limit.parse_rate_limit(exception.headers)
+                    if rate_limiter is not None:
+                        retry_after = rate_limiter.retry_after()
+                        LOGGER.info(
+                            "Rate limited (HTTP 429), retrying in %s seconds",
+                            retry_after,
+                        )
+                        time.sleep(retry_after)
+                        continue
+
+                # Fallback: wait 1 second if no header available
+                LOGGER.info(
+                    "Rate limited (HTTP 429) with no retry-after header, retrying in 1 second"
+                )
+                time.sleep(1)
+                continue
+
+            # Re-raise if not a 429 error
+            raise
+
+
 class Dataset:
     def __init__(
         self,
@@ -51,6 +103,17 @@ class Dataset:
         """The description of the dataset."""
         return self._description
 
+    def _insert_batch_with_retry(
+        self, batch: List[rest_dataset_item.DatasetItemWrite]
+    ) -> None:
+        """Insert a batch of dataset items with automatic retry on rate limit errors."""
+        _ensure_rest_api_call_respecting_rate_limit(
+            lambda: self._rest_client.datasets.create_or_update_dataset_items(
+                dataset_name=self._name, items=batch
+            )
+        )
+        LOGGER.debug("Successfully sent dataset items batch of size %d", len(batch))
+
     def __internal_api__insert_items_as_dataclasses__(
         self, items: List[dataset_item.DatasetItem]
     ) -> None:
@@ -89,9 +152,7 @@ class Dataset:
 
         for batch in batches:
             LOGGER.debug("Sending dataset items batch of size %d", len(batch))
-            self._rest_client.datasets.create_or_update_dataset_items(
-                dataset_name=self._name, items=batch
-            )
+            self._insert_batch_with_retry(batch)
 
     def insert(self, items: Sequence[Dict[str, Any]]) -> None:
         """
@@ -221,62 +282,63 @@ class Dataset:
         dataset_item_ids: Optional[List[str]] = None,
     ) -> List[dataset_item.DatasetItem]:
         results: List[dataset_item.DatasetItem] = []
-
-        if dataset_item_ids is not None:
-            dataset_item_ids = set(dataset_item_ids)  # type: ignore
-
-        while True:
-            stream = self._rest_client.datasets.stream_dataset_items(
-                dataset_name=self._name,
-                last_retrieved_id=results[-1].id if len(results) > 0 else None,
+        last_retrieved_id: Optional[str] = None
+        should_retrieve_more_items = True
+
+        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,
             )
 
-            previous_results_size = len(results)
-            if nb_samples is not None and len(results) == nb_samples:
-                break
+            if len(dataset_items) == 0:
+                should_retrieve_more_items = False
 
-            item_bytes = b"".join(stream)
-            for line in item_bytes.split(b"\n"):
-                if len(line) == 0:
-                    continue
+            for item in dataset_items:
+                dataset_item_id = item.id
+                last_retrieved_id = dataset_item_id
 
-                full_item_content: Dict[str, Any] = json.loads(
-                    line.decode("utf-8").strip()
-                )
-                data_item_content = (
-                    full_item_content["data"] if "data" in full_item_content else {}
-                )
-
-                dataset_item_id = full_item_content.get("id")
-
-                if dataset_item_ids is not None:
-                    if dataset_item_id not in dataset_item_ids:
+                if dataset_items_ids_left is not None:
+                    if dataset_item_id not in dataset_items_ids_left:
                         continue
                     else:
-                        dataset_item_ids.remove(dataset_item_id)
+                        dataset_items_ids_left.remove(dataset_item_id)
 
-                item = dataset_item.DatasetItem(
-                    id=full_item_content.get("id"),  # type: ignore
-                    trace_id=full_item_content.get("trace_id"),  # type: ignore
-                    span_id=full_item_content.get("span_id"),  # type: ignore
-                    source=full_item_content.get("source"),  # type: ignore
+                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,
                 )
 
-                results.append(item)
+                results.append(reconstructed_item)
 
-                # Break the loop if we have enough samples
+                # Stop retrieving if we have enough samples
                 if nb_samples is not None and len(results) == nb_samples:
+                    should_retrieve_more_items = False
                     break
 
-            # Break the loop if we have not received any new samples
-            if len(results) == previous_results_size:
-                break
+                # Stop retrieving if we found all filtered dataset items
+                if (
+                    dataset_items_ids_left is not None
+                    and len(dataset_items_ids_left) == 0
+                ):
+                    should_retrieve_more_items = False
+                    break
 
-        if dataset_item_ids and len(dataset_item_ids) > 0:
+        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_item_ids,
+                dataset_items_ids_left,
             )
 
         return results

opik/api_objects/dataset/rest_operations.py

@@ -1,8 +1,12 @@
+from __future__ import annotations
+
 from typing import List
 from opik.rest_api import OpikApi
-from opik import exceptions
+import opik.exceptions as exceptions
+from opik.message_processing import streamer
 from . import dataset
 from .. import experiment
+from ..experiment import experiments_client
 from ...rest_api.core.api_error import ApiError
 
 
@@ -55,7 +59,11 @@ def get_dataset_id(rest_client: OpikApi, dataset_name: str) -> str:
 
 
 def get_dataset_experiments(
-    rest_client: OpikApi, dataset_id: str, max_results: int = 1000
+    rest_client: OpikApi,
+    dataset_id: str,
+    max_results: int,
+    streamer: streamer.Streamer,
+    experiments_client: experiments_client.ExperimentsClient,
 ) -> List[experiment.Experiment]:
     page_size = 100
     experiments: List[experiment.Experiment] = []
@@ -78,7 +86,8 @@ def get_dataset_experiments(
                     name=experiment_.name,
                     dataset_name=experiment_.dataset_name,
                     rest_client=rest_client,
-                    # TODO: add prompt if exists
+                    streamer=streamer,
+                    experiments_client=experiments_client,
                 )
             )
 

opik/api_objects/experiment/experiment.py

@@ -1,13 +1,17 @@
 import functools
 import logging
-from typing import List, Optional
+from typing import List, Optional, TYPE_CHECKING
 
 from opik.message_processing.batching import sequence_splitter
+from opik.message_processing import messages, streamer
 from opik.rest_api import client as rest_api_client
-from opik.rest_api.types import experiment_item as rest_experiment_item
-from . import experiment_item
+from opik.rest_api import types as rest_api_types
+from . import experiment_item, experiments_client
 from .. import constants, helpers
-from ...api_objects.prompt import Prompt
+from ...api_objects.prompt import base_prompt
+
+if TYPE_CHECKING:
+    from opik.evaluation.metrics import score_result
 
 LOGGER = logging.getLogger(__name__)
 
@@ -19,34 +23,48 @@ class Experiment:
         name: Optional[str],
         dataset_name: str,
         rest_client: rest_api_client.OpikApi,
-        prompts: Optional[List[Prompt]] = None,
+        streamer: streamer.Streamer,
+        experiments_client: experiments_client.ExperimentsClient,
+        prompts: Optional[List[base_prompt.BasePrompt]] = None,
     ) -> None:
         self._id = id
         self._name = name
         self._dataset_name = dataset_name
         self._rest_client = rest_client
         self._prompts = prompts
+        self._streamer = streamer
+        self._experiments_client = experiments_client
 
     @property
     def id(self) -> str:
         return self._id
 
-    @functools.cached_property
-    def dataset_id(self) -> str:
-        return self._rest_client.datasets.get_dataset_by_identifier(
-            dataset_name=self._dataset_name
-        ).id
-
     @property
     def dataset_name(self) -> str:
         return self._dataset_name
 
-    @functools.cached_property
+    @property
     def name(self) -> str:
         if self._name is not None:
             return self._name
 
-        return self._rest_client.experiments.get_experiment_by_id(id=self.id).name
+        name = self._rest_client.experiments.get_experiment_by_id(id=self.id).name
+        self._name = name
+
+        return name
+
+    @functools.cached_property
+    def dataset_id(self) -> str:
+        return self._rest_client.datasets.get_dataset_by_identifier(
+            dataset_name=self._dataset_name
+        ).id
+
+    @property
+    def experiments_rest_client(self) -> rest_api_client.ExperimentsClient:
+        return self._rest_client.experiments
+
+    def get_experiment_data(self) -> rest_api_types.experiment_public.ExperimentPublic:
+        return self._rest_client.experiments.get_experiment_by_id(id=self.id)
 
     def insert(
         self,
@@ -62,8 +80,8 @@ class Experiment:
         Returns:
             None
         """
-        rest_experiment_items = [
-            rest_experiment_item.ExperimentItem(
+        experiment_item_messages = [
+            messages.ExperimentItemMessage(
                 id=helpers.generate_id(),
                 experiment_id=self.id,
                 dataset_item_id=item.dataset_item_id,
@@ -72,44 +90,62 @@ class Experiment:
             for item in experiment_items_references
         ]
 
+        # Split into batches for the streamer
         batches = sequence_splitter.split_into_batches(
-            rest_experiment_items, max_length=constants.EXPERIMENT_ITEMS_MAX_BATCH_SIZE
+            experiment_item_messages,
+            max_length=constants.FEEDBACK_SCORES_MAX_BATCH_SIZE,
         )
 
         for batch in batches:
-            LOGGER.debug("Sending experiment items batch: %s", batch)
-            self._rest_client.experiments.create_experiment_items(
-                experiment_items=batch,
+            create_experiment_items_batch_message = (
+                messages.CreateExperimentItemsBatchMessage(batch=batch)
             )
-            LOGGER.debug("Sent experiment items batch of size %d", len(batch))
+            self._streamer.put(create_experiment_items_batch_message)
 
-    def get_items(self) -> List[experiment_item.ExperimentItemContent]:
+    def get_items(
+        self,
+        max_results: Optional[int] = 10000,
+        truncate: bool = False,
+    ) -> List[experiment_item.ExperimentItemContent]:
         """
+        Retrieves and returns a list of experiment items for this experiment.
+
+        Args:
+            max_results: Maximum number of experiment items to retrieve. Defaults to 10000 if not specified.
+            truncate: Whether to truncate the items returned by the backend. Defaults to False.
+
         Returns:
-            List[ExperimentItemContent]: the list with contents of existing experiment items.
+            List of ExperimentItemContent objects for this experiment.
         """
-        result: List[experiment_item.ExperimentItemContent] = []
-
-        page = 0
-
-        while True:  # TODO: refactor this logic when backend implements a proper streaming endpoint
-            page += 1
-            dataset_items_page = (
-                self._rest_client.datasets.find_dataset_items_with_experiment_items(
-                    id=self.dataset_id,
-                    experiment_ids=f'["{self.id}"]',
-                    page=page,
-                    size=100,
-                )
-            )
-            if len(dataset_items_page.content) == 0:
-                break
+        if max_results is None:
+            max_results = 10000  # TODO: remove this once we have a proper way to get all experiment items
+
+        return self._experiments_client.find_experiment_items_for_dataset(
+            dataset_name=self.dataset_name,
+            experiment_ids=[self.id],
+            truncate=truncate,
+            max_results=max_results,
+        )
+
+    def log_experiment_scores(
+        self,
+        score_results: List["score_result.ScoreResult"],
+    ) -> None:
+        """Log experiment-level scores to the backend."""
+        experiment_scores: List[rest_api_types.ExperimentScore] = []
+
+        for score_result_ in score_results:
+            if score_result_.scoring_failed:
+                continue
 
-            for dataset_item in dataset_items_page.content:
-                rest_experiment_item_compare = dataset_item.experiment_items[0]
-                content = experiment_item.ExperimentItemContent.from_rest_experiment_item_compare(
-                    value=rest_experiment_item_compare
-                )
-                result.append(content)
+            experiment_score = rest_api_types.ExperimentScore(
+                name=score_result_.name,
+                value=score_result_.value,
+            )
+            experiment_scores.append(experiment_score)
 
-        return result
+        if experiment_scores:
+            self._rest_client.experiments.update_experiment(
+                id=self.id,
+                experiment_scores=experiment_scores,
+            )

opik/api_objects/experiment/experiment_item.py

@@ -24,6 +24,7 @@ class ExperimentItemContent:
     def from_rest_experiment_item_compare(
         cls,
         value: experiment_item_compare.ExperimentItemCompare,
+        dataset_item_data: Optional[Dict[str, Any]] = None,
     ) -> "ExperimentItemContent":
         if value.feedback_scores is None:
             feedback_scores: List[FeedbackScoreDict] = []
@@ -42,7 +43,7 @@ class ExperimentItemContent:
             id=value.id,
             trace_id=value.trace_id,
             dataset_item_id=value.dataset_item_id,
-            dataset_item_data=value.input,
+            dataset_item_data=dataset_item_data if dataset_item_data else value.input,
             evaluation_task_output=value.output,
             feedback_scores=feedback_scores,
         )

opik/api_objects/experiment/experiments_client.py

@@ -0,0 +1,64 @@
+import json
+from typing import List, Optional
+
+from . import rest_operations, experiment_item
+from .. import opik_query_language
+from ...rest_api import client as rest_api_client
+
+
+class ExperimentsClient:
+    """Client for managing and querying experiments in bulk."""
+
+    def __init__(self, rest_client: rest_api_client.OpikApi):
+        self._rest_client = rest_client
+
+    def find_experiment_items_for_dataset(
+        self,
+        dataset_name: str,
+        experiment_ids: List[str],
+        truncate: bool = True,
+        max_results: int = 1000,
+        filter_string: Optional[str] = None,
+    ) -> List[experiment_item.ExperimentItemContent]:
+        """
+        Find experiment items associated with a specific dataset among a list of experiments.
+
+        This method queries the dataset for experiment items matching the
+        criteria provided by the input parameters. It leverages the
+        ExperimentsClient to perform the underlying operation.
+
+        Args:
+            dataset_name: Name of the dataset to query for experiment items.
+            experiment_ids: List of experiment IDs to filter the results.
+            filter_string: Optional filter string to refine the
+                query based on additional criteria (dataset fields, feedback scores, etc.).
+            truncate: Whether to truncate image data stored in input, output,
+                or metadata. Defaults to True.
+            max_results: Maximum number of results to return. Defaults to 1000.
+
+        Returns:
+            A list of experiment item content objects that match the criteria.
+        """
+        # prepare filter expression
+        if filter_string is not None:
+            filter_expression = json.dumps(
+                opik_query_language.OpikQueryLanguage(
+                    filter_string
+                ).get_filter_expressions()
+            )
+        else:
+            filter_expression = None
+
+        # get dataset id
+        dataset_id = self._rest_client.datasets.get_dataset_by_identifier(
+            dataset_name=dataset_name
+        ).id
+
+        return rest_operations.find_experiment_items_for_dataset(
+            dataset_id=dataset_id,
+            experiment_ids=experiment_ids,
+            rest_client=self._rest_client,
+            max_results=max_results,
+            truncate=truncate,
+            filter_expression=filter_expression,
+        )

opik/api_objects/experiment/helpers.py

@@ -1,8 +1,12 @@
+import copy
 import logging
-from opik import jsonable_encoder
 from typing import Any, Dict, List, Mapping, Optional, Tuple
-import copy
-from .. import prompt
+
+from opik import id_helpers
+import opik.jsonable_encoder as jsonable_encoder
+
+from ..prompt import base_prompt
+
 
 LOGGER = logging.getLogger(__name__)
 
@@ -11,7 +15,7 @@ PromptVersion = Dict[str, str]
 
 def build_metadata_and_prompt_versions(
     experiment_config: Optional[Dict[str, Any]],
-    prompts: Optional[List[prompt.Prompt]],
+    prompts: Optional[List[base_prompt.BasePrompt]],
 ) -> Tuple[Optional[Dict[str, Any]], Optional[List[PromptVersion]]]:
     prompt_versions: Optional[List[PromptVersion]] = None
 
@@ -35,11 +39,22 @@ def build_metadata_and_prompt_versions(
 
     if prompts is not None and len(prompts) > 0:
         prompt_versions = []
-        experiment_config["prompts"] = []
+        experiment_config["prompts"] = {}
 
-        for prompt in prompts:
-            prompt_versions.append({"id": prompt.__internal_api__version_id__})
-            experiment_config["prompts"].append(prompt.prompt)
+        for prompt_obj in prompts:
+            prompt_versions.append({"id": prompt_obj.__internal_api__version_id__})
+            # Use __internal_api__to_info_dict__() to get the prompt content in a consistent way
+            prompt_info = prompt_obj.__internal_api__to_info_dict__()
+            # Extract the template/messages from the version dict
+            if "version" in prompt_info:
+                if "template" in prompt_info["version"]:
+                    experiment_config["prompts"][prompt_obj.name] = prompt_info[
+                        "version"
+                    ]["template"]
+                elif "messages" in prompt_info["version"]:
+                    experiment_config["prompts"][prompt_obj.name] = prompt_info[
+                        "version"
+                    ]["messages"]
 
     if experiment_config == {}:
         return None, None
@@ -50,9 +65,9 @@ def build_metadata_and_prompt_versions(
 
 
 def handle_prompt_args(
-    prompt: Optional[prompt.Prompt] = None,
-    prompts: Optional[List[prompt.Prompt]] = None,
-) -> Optional[List[prompt.Prompt]]:
+    prompt: Optional[base_prompt.BasePrompt] = None,
+    prompts: Optional[List[base_prompt.BasePrompt]] = None,
+) -> Optional[List[base_prompt.BasePrompt]]:
     if prompts is not None and len(prompts) > 0 and prompt is not None:
         LOGGER.warning(
             "Arguments `prompt` and `prompts` are mutually exclusive, `prompts` will be used`."
@@ -63,3 +78,12 @@ def handle_prompt_args(
         prompts = None
 
     return prompts
+
+
+def generate_unique_experiment_name(experiment_name_prefix: Optional[str]) -> str:
+    if experiment_name_prefix is None:
+        return id_helpers.generate_random_alphanumeric_string(12)
+
+    return (
+        f"{experiment_name_prefix}-{id_helpers.generate_random_alphanumeric_string(6)}"
+    )