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

opik/api_objects/attachment/attachment_context.py

@@ -0,0 +1,36 @@
+import dataclasses
+from typing import Literal
+
+from . import attachment
+
+
+@dataclasses.dataclass
+class AttachmentWithContext:
+    """
+    Represents an attachment along with its associated context.
+
+    This class is used to pair an attachment with additional contextual
+    information such as the entity type, entity ID, project name, and
+    context description. It is specifically useful when dealing with
+    attachments related to entities like spans or traces. The context
+    can help provide further insights or classification of the
+    attachment's purpose.
+
+    Attributes:
+        attachment_data: The actual attachment
+            object containing the associated data.
+        entity_type: The type of entity the
+            attachment is associated with. It must be either "span"
+            or "trace".
+        entity_id: The unique identifier of the related entity.
+        project_name: The name of the project to which the
+            attachment and its entity belong.
+        context: A brief context description for the attachment,
+            explaining its purpose or relevance.
+    """
+
+    attachment_data: attachment.Attachment
+    entity_type: Literal["span", "trace"]
+    entity_id: str
+    project_name: str
+    context: str

opik/api_objects/attachment/attachments_extractor.py

@@ -0,0 +1,153 @@
+import re
+from typing import Dict, Any, Literal, List, NamedTuple
+
+from . import attachment, attachment_context, decoder_base64
+
+
+class ExtractionResult(NamedTuple):
+    attachments: List[attachment.Attachment]
+    sanitized_data: Any
+
+
+class AttachmentsExtractor:
+    """
+    Extracts and processes attachments embedded as Base64 strings within data structures.
+
+    This class is designed to identify and decode Base64-encoded attachments located
+    within the provided data. It uses a regular expression pattern to search for
+    Base64 strings that meet a specified minimum length. Extracted attachments are
+    decoded and replaced with sanitized placeholders in the original data.
+    """
+
+    def __init__(self, min_attachment_size: int):
+        """
+        Initializes the class with a minimum attachment size and configures the base64
+        pattern for decoding attachments based on its length.
+
+        Args:
+            min_attachment_size: The minimum size of the attachment in characters
+                for it to be considered valid. This ensures that only large enough
+                base64 strings are matched to minimize false positives.
+        """
+        self._min_attachment_size = min_attachment_size
+        self.decoder = decoder_base64.Base64AttachmentDecoder()
+
+        # Pattern to match base64 strings (can be embedded in text)
+        # Requires at least min_attachment_size characters to reduce false positives
+        min_base64_groups = int(min_attachment_size / 4)
+        BASE64_PATTERN = (
+            r"(?:[A-Za-z0-9+/]{4}){"
+            + str(min_base64_groups)
+            + ",}(?:[A-Za-z0-9+/]{2}==|[A-Za-z0-9+/]{3}=)?"
+        )
+        self.pattern = re.compile(BASE64_PATTERN)
+
+    def extract_and_replace(
+        self,
+        data: Dict[str, Any],
+        entity_type: Literal["span", "trace"],
+        entity_id: str,
+        project_name: str,
+        context: Literal["input", "output", "metadata"],
+    ) -> List[attachment_context.AttachmentWithContext]:
+        # iterate over all items and extract attachments
+        attachments: List[attachment_context.AttachmentWithContext] = []
+        for key, value in data.items():
+            extraction_result = self._try_extract_attachments(value, context)
+            if extraction_result.attachments:
+                # replace the original value with the sanitized one and collect attachments
+                data[key] = extraction_result.sanitized_data
+                for extracted_attachment in extraction_result.attachments:
+                    attachments.append(
+                        attachment_context.AttachmentWithContext(
+                            attachment_data=extracted_attachment,
+                            entity_type=entity_type,
+                            entity_id=entity_id,
+                            project_name=project_name,
+                            context=context,
+                        )
+                    )
+
+        return attachments
+
+    def _try_extract_attachments(
+        self, data: Any, context: Literal["input", "output", "metadata"]
+    ) -> ExtractionResult:
+        """
+        Recursively extract attachments from data that can be a string, dict, list, or other type.
+
+        Args:
+            data: The data to process (can be str, dict, list, or other types)
+            context: The context where the data is located (input, output, or metadata)
+
+        Returns:
+            ExtractionResult with extracted attachments and sanitized data
+        """
+        # Handle string data - check for base64 attachments
+        if isinstance(data, str):
+            return self._extract_from_string(data, context)
+
+        # Handle dictionary data - recursively process each value
+        elif isinstance(data, dict):
+            return self._extract_from_dict(data, context)
+
+        # Handle list data - recursively process each element
+        elif isinstance(data, list):
+            return self._extract_from_list(data, context)
+
+        # For other types (int, bool, None, etc.), return as-is
+        else:
+            return ExtractionResult(attachments=[], sanitized_data=data)
+
+    def _extract_from_string(
+        self, data: str, context: Literal["input", "output", "metadata"]
+    ) -> ExtractionResult:
+        """Extract attachments from a string value."""
+        if len(data) < self._min_attachment_size:
+            # skip short strings
+            return ExtractionResult(attachments=[], sanitized_data=data)
+
+        attachments: List[attachment.Attachment] = []
+        sanitized_data = data
+        for match in self.pattern.finditer(data):
+            to_decode = match.group()
+            decoded_attachment = self.decoder.decode(to_decode, context)
+            if decoded_attachment is not None:
+                attachments.append(decoded_attachment)
+                sanitized_data = sanitized_data.replace(
+                    to_decode, f"[{decoded_attachment.file_name}]"
+                )
+
+        return ExtractionResult(attachments=attachments, sanitized_data=sanitized_data)
+
+    def _extract_from_dict(
+        self, data: Dict[str, Any], context: Literal["input", "output", "metadata"]
+    ) -> ExtractionResult:
+        """Recursively extract attachments from a dictionary."""
+        all_attachments: List[attachment.Attachment] = []
+        sanitized_dict = {}
+
+        for key, value in data.items():
+            result = self._try_extract_attachments(value, context)
+            sanitized_dict[key] = result.sanitized_data
+            all_attachments.extend(result.attachments)
+
+        return ExtractionResult(
+            attachments=all_attachments, sanitized_data=sanitized_dict
+        )
+
+    def _extract_from_list(
+        self, data: List[Any], context: Literal["input", "output", "metadata"]
+    ) -> ExtractionResult:
+        """Recursively extract attachments from a list."""
+        all_attachments: List[attachment.Attachment] = []
+        sanitized_list = []
+
+        for item in data:
+            result = self._try_extract_attachments(item, context)
+            sanitized_list.append(result.sanitized_data)
+            all_attachments.extend(result.attachments)
+
+        return ExtractionResult(
+            attachments=all_attachments, sanitized_data=sanitized_list
+        )

opik/api_objects/attachment/client.py

@@ -206,6 +206,7 @@ class AttachmentClient:
             entity_id=entity_id,
             project_name=project_name,
             encoded_url_override=encoded_url_override,
+            delete_after_upload=False,
         )
 
         file_uploader.upload_attachment(

opik/api_objects/attachment/converters.py

@@ -13,6 +13,7 @@ def attachment_to_message(
     entity_id: str,
     project_name: str,
     url_override: str,
+    delete_after_upload: bool = False,
 ) -> messages.CreateAttachmentMessage:
     if attachment_data.data is None:
         raise ValueError("Attachment data cannot be None")
@@ -32,6 +33,7 @@ def attachment_to_message(
         entity_id=entity_id,
         project_name=project_name,
         encoded_url_override=base_url_path,
+        delete_after_upload=delete_after_upload,
     )
 
 

opik/api_objects/attachment/decoder.py

@@ -0,0 +1,18 @@
+import abc
+from typing import Any, Optional
+
+from . import attachment
+
+
+class AttachmentDecoder(abc.ABC):
+    """
+    Abstract base class for decoding file attachments.
+
+    This class serves as an interface for decoding raw attachment data into
+    an `Attachment` object. Implementing classes should define the specific
+    logic to handle various attachment decoding formats.
+    """
+
+    @abc.abstractmethod
+    def decode(self, raw_data: str, **kwargs: Any) -> Optional[attachment.Attachment]:
+        pass

opik/api_objects/attachment/decoder_base64.py

@@ -0,0 +1,83 @@
+import base64
+import binascii
+import logging
+import tempfile
+from typing import Any, Optional, Literal
+
+from . import attachment, decoder, decoder_helpers
+
+LOGGER = logging.getLogger(__name__)
+
+
+class Base64AttachmentDecoder(decoder.AttachmentDecoder):
+    """Decodes base64 encoded attachment data.
+
+    This decoder decodes base64 strings, detects MIME types from content, and creates Attachment objects.
+    """
+
+    def decode(
+        self,
+        raw_data: str,
+        context: Literal["input", "output", "metadata"] = "input",
+        **kwargs: Any,
+    ) -> Optional[attachment.Attachment]:
+        """Decode base64 encoded data into an Attachment object.
+
+        Args:
+            raw_data: Base64 encoded string data
+            context: Context string for filename generation.
+
+        Returns:
+            Attachment object with decoded data, or None if decoding fails or type is not recognizable
+        """
+        if not isinstance(raw_data, str):
+            LOGGER.warning("Attachment data is not a string, skipping.")
+            return None
+
+        try:
+            # Decode base64 string to bytes
+            decoded_bytes = base64.b64decode(raw_data, validate=True)
+
+            # Detect MIME type from content
+            mime_type = decoder_helpers.detect_mime_type(decoded_bytes)
+
+            # Skip if not a recognizable file type
+            if not mime_type or mime_type in ("application/octet-stream", "text/plain"):
+                LOGGER.debug("Attachment type is not recognized, skipping.")
+                return None
+
+            # Get file extension from the MIME type
+            extension = decoder_helpers.get_file_extension(mime_type)
+
+            # Generate filename
+            file_name = decoder_helpers.create_attachment_filename(
+                context, extension=extension
+            )
+
+            # Save decoded bytes to a temporary file
+            temp_file = tempfile.NamedTemporaryFile(
+                mode="wb", delete=False, suffix=extension
+            )
+            temp_file.write(decoded_bytes)
+            temp_file.flush()
+            temp_file.close()
+
+            # Return Attachment object with a file path
+            return attachment.Attachment(
+                data=temp_file.name, file_name=file_name, content_type=mime_type
+            )
+
+        except (ValueError, binascii.Error) as e:
+            LOGGER.debug(
+                "Failed to decode attachment data, reason: invalid base64. Reason: %s",
+                e,
+                exc_info=True,
+            )
+            # Not valid base64, return None
+            return None
+        except Exception as ex:
+            LOGGER.warning(
+                "Failed to decode attachment data, reason: %s", ex, exc_info=True
+            )
+            # Unexpected error, return None to avoid crashing the pipeline
+            return None

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/constants.py

@@ -5,3 +5,5 @@ FEEDBACK_SCORES_MAX_BATCH_SIZE = 1000
 EXPERIMENT_ITEMS_MAX_BATCH_SIZE = 1000
 DATASET_ITEMS_MAX_BATCH_SIZE = 1000
 DELETE_TRACE_BATCH_SIZE = 1000
+
+DATASET_STREAM_BATCH_SIZE = 2000