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

opik/api_objects/attachment/decoder_helpers.py

@@ -0,0 +1,137 @@
+import mimetypes
+import random
+import time
+from typing import Optional
+
+
+# The attachment file name regex
+ATTACHMENT_FILE_NAME_REGEX = r"(?:input|output|metadata)-attachment-\d+-\d+-sdk\.\w+"
+ATTACHMENT_FILE_NAME_PLACEHOLDER_REGEX = (
+    r"\[((?:input|output|metadata)-attachment-\d+-\d+-sdk\.\w+)\]"
+)
+
+
+def get_file_extension(mime_type: str) -> str:
+    """Convert MIME type to file extension.
+
+    Mirrors the Java getFileExtension() method in AttachmentStripperService.
+
+    Args:
+        mime_type: The MIME type (e.g., "image/png", "application/pdf")
+
+    Returns:
+        File extension without a leading dot (e.g., "png", "pdf")
+    """
+    if not mime_type:
+        return "bin"
+
+    # Try to get extension from mimetypes module
+    extension = mimetypes.guess_extension(mime_type, strict=False)
+
+    if extension:
+        # Remove the leading dot
+        extension = extension.lstrip(".")
+        # Handle special cases where mimetypes returns less common extensions
+        if mime_type == "image/jpeg" and extension == "jpe":
+            return "jpg"
+        return extension
+
+    # Fallback: extract from the MIME type (e.g., "image/png" -> "png")
+    if "/" in mime_type:
+        subtype = mime_type.split("/")[1]
+        # Handle special cases like "svg+xml" -> "svg"
+        if "+" in subtype:
+            subtype = subtype.split("+")[0]
+        # Remove any parameters (e.g., "jpeg; charset=utf-8" -> "jpeg")
+        subtype = subtype.split(";")[0].strip()
+        return subtype
+
+    return "bin"
+
+
+def detect_mime_type(data: bytes) -> Optional[str]:
+    """Detect MIME type from byte content using magic bytes.
+
+    This provides basic MIME type detection similar to Apache Tika in the Java implementation.
+    It checks common file format magic bytes.
+
+    Args:
+        data: The byte data to analyze
+
+    Returns:
+        Detected MIME type string, or "application/octet-stream" if unknown
+    """
+    if len(data) < 4:
+        return "application/octet-stream"
+
+    # Check common file format magic bytes
+    # PNG
+    if data[:8] == b"\x89PNG\r\n\x1a\n":
+        return "image/png"
+
+    # JPEG
+    if data[:2] == b"\xff\xd8" and data[-2:] == b"\xff\xd9":
+        return "image/jpeg"
+
+    # GIF
+    if data[:6] in (b"GIF87a", b"GIF89a"):
+        return "image/gif"
+
+    # PDF
+    if data[:4] == b"%PDF":
+        return "application/pdf"
+
+    # WebP
+    if data[:4] == b"RIFF" and data[8:12] == b"WEBP":
+        return "image/webp"
+
+    # SVG (XML-based, check for SVG tag)
+    try:
+        text = data[:1024].decode("utf-8", errors="ignore")
+        if "<svg" in text.lower():
+            return "image/svg+xml"
+    except Exception:
+        pass
+
+    # MP4
+    if len(data) >= 12 and data[4:8] == b"ftyp":
+        return "video/mp4"
+
+    # JSON
+    try:
+        text = data[:100].decode("utf-8", errors="strict").strip()
+        if text.startswith("{") or text.startswith("["):
+            return "application/json"
+    except Exception:
+        pass
+
+    # Default to octet-stream for unknown types
+    return "application/octet-stream"
+
+
+def create_attachment_filename(context: str, extension: str) -> str:
+    """
+    Generates a unique attachment filename based on the provided context and file extension.
+
+    This function creates a filename by combining the given context, a randomly generated
+    prefix to ensure uniqueness, the current timestamp in milliseconds, and the provided
+    file extension. The generated filename aligns with the backend convention for naming
+    attachments, which includes specific formatting and structure.
+
+    Args:
+        context: The context to use as the base for the filename (e.g., "input",
+            "output", or "metadata").
+        extension: The file extension to use for the filename (e.g., "png",
+            "jpg", "txt").
+
+    Returns:
+        A generated filename string in the format
+            "{context}-attachment-{random_prefix}-{timestamp}.{extension}".
+    """
+    # The backend has the following naming convention: r"\\[((?:input|output|metadata)-attachment-\\d+-\\d+\\.\\w+)\\]"
+    # Example: [input-attachment-1-1704067200000.png]
+
+    timestamp = int(round(time.time() * 1000))
+    # we need to generate a large enough random prefix to avoid collisions
+    random_prefix = random.randint(1, 99999999)
+    return f"{context}-attachment-{random_prefix}-{timestamp}-sdk.{extension}"

opik/api_objects/data_helpers.py

@@ -0,0 +1,79 @@
+from typing import Optional, Any, Dict, List, Union
+
+import pydantic
+
+from opik import dict_utils
+
+
+def merge_tags(
+    existing_tags: Optional[List[str]], new_tags: Optional[List[str]]
+) -> Optional[List[str]]:
+    """Merge tag lists, preserving existing tags and adding new ones.
+
+    If both existing_tags and new_tags are None or empty, return None."""
+    if existing_tags is None and new_tags is None:
+        return None
+
+    result = list(existing_tags or [])
+    if new_tags:
+        for tag in new_tags:
+            if tag not in result:
+                result.append(tag)
+
+    return result if result else None
+
+
+def merge_metadata(
+    existing_metadata: Optional[Dict[str, Any]],
+    new_metadata: Optional[Union[Dict[str, Any], pydantic.BaseModel]],
+    prompts: Optional[List[Dict[str, Any]]] = None,
+) -> Optional[Dict[str, Any]]:
+    """Merge the existing metadata dictionary with new data, with new values taking precedence.
+
+    If both existing_metadata and new_metadata are None or empty, return None.
+    """
+    if prompts is not None:
+        new_metadata = new_metadata or {}
+        new_metadata["opik_prompts"] = prompts
+
+    return _merge_dictionary_with_data(existing_metadata, new_data=new_metadata)
+
+
+def merge_inputs(
+    existing_inputs: Optional[Dict[str, Any]],
+    new_inputs: Optional[Union[Dict[str, Any], pydantic.BaseModel]],
+) -> Optional[Dict[str, Any]]:
+    """Merge the existing input dictionary with new data, with new values taking precedence.
+
+    If both existing_inputs and new_inputs are None or empty, return None."""
+    return _merge_dictionary_with_data(existing_inputs, new_data=new_inputs)
+
+
+def merge_outputs(
+    existing_outputs: Optional[Dict[str, Any]],
+    new_outputs: Optional[Union[Dict[str, Any], pydantic.BaseModel]],
+) -> Optional[Dict[str, Any]]:
+    """Merge the existing output dictionary with new data, with new values taking precedence.
+
+    If both existing_outputs and new_outputs are None or empty, return None."""
+    return _merge_dictionary_with_data(existing_outputs, new_data=new_outputs)
+
+
+def _merge_dictionary_with_data(
+    existing_dict: Optional[Dict[str, Any]],
+    new_data: Optional[Union[Dict[str, Any], pydantic.BaseModel]],
+) -> Optional[Dict[str, Any]]:
+    """Merge the dictionary with new data, with new values taking precedence.
+
+    If both existing_dict and new_data are None or empty, return None."""
+    if existing_dict is None and new_data is None:
+        return None
+
+    if isinstance(new_data, pydantic.BaseModel):
+        new_data = new_data.model_dump()
+
+    result = dict(existing_dict or {})
+    if new_data:
+        result = dict_utils.deepmerge(result, new_data)
+
+    return result if result else None

opik/api_objects/dataset/dataset.py

@@ -1,11 +1,14 @@
 import logging
 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.rate_limit import rate_limit
 import opik.exceptions as exceptions
 import opik.config as config
 from opik.rest_client_configurator import retry_decorator
@@ -18,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,
@@ -52,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:
@@ -90,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:
         """

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
 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,15 +1,17 @@
 import functools
 import logging
-from typing import List, Optional
+from typing import List, Optional, TYPE_CHECKING
 
-import opik.rest_api
 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 opik.rest_api.types import experiment_public
-from . import experiment_item
-from .. import constants, helpers, rest_stream_parser
-from ...api_objects.prompt import Prompt
+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 base_prompt
+
+if TYPE_CHECKING:
+    from opik.evaluation.metrics import score_result
 
 LOGGER = logging.getLogger(__name__)
 
@@ -21,13 +23,17 @@ 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:
@@ -57,7 +63,7 @@ class Experiment:
     def experiments_rest_client(self) -> rest_api_client.ExperimentsClient:
         return self._rest_client.experiments
 
-    def get_experiment_data(self) -> experiment_public.ExperimentPublic:
+    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(
@@ -74,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,
@@ -84,68 +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,
-        max_results: Optional[int] = None,
+        max_results: Optional[int] = 10000,
         truncate: bool = False,
     ) -> List[experiment_item.ExperimentItemContent]:
         """
-        Retrieves and returns a list of experiment items by streaming from the backend in batches, with an option to
-        truncate the results for each batch.
-
-        This method streams experiment items from a backend service in chunks up to the specified `max_results`
-        or until the available items are exhausted. It handles batch-wise retrieval and parsing, ensuring the client
-        receives a list of `ExperimentItemContent` objects, while respecting the constraints on maximum retrieval size
-        from the backend. If truncation is enabled, the backend may return truncated details for each item.
+        Retrieves and returns a list of experiment items for this experiment.
 
         Args:
-            max_results: Maximum number of experiment items to retrieve.
-            truncate: Whether to truncate the items returned by the backend.
+            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 of ExperimentItemContent objects for this experiment.
         """
-        result: List[experiment_item.ExperimentItemContent] = []
-        max_endpoint_batch_size = rest_stream_parser.MAX_ENDPOINT_BATCH_SIZE
-
-        while True:
-            if max_results is None:
-                current_batch_size = max_endpoint_batch_size
-            else:
-                current_batch_size = min(
-                    max_results - len(result), max_endpoint_batch_size
-                )
-
-            items_stream = self._rest_client.experiments.stream_experiment_items(
-                experiment_name=self.name,
-                limit=current_batch_size,
-                last_retrieved_id=result[-1].id if len(result) > 0 else None,
-                truncate=truncate,
-            )
+        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,
+        )
 
-            experiment_item_compare_current_batch = (
-                rest_stream_parser.read_and_parse_stream(
-                    stream=items_stream,
-                    item_class=opik.rest_api.ExperimentItemCompare,
-                )
-            )
+    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 item in experiment_item_compare_current_batch:
-                converted_item = experiment_item.ExperimentItemContent.from_rest_experiment_item_compare(
-                    value=item
-                )
-                result.append(converted_item)
+        for score_result_ in score_results:
+            if score_result_.scoring_failed:
+                continue
 
-            if current_batch_size > len(experiment_item_compare_current_batch):
-                break
+            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
-import opik.jsonable_encoder as 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)}"
+    )