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

opik/__init__.py

@@ -6,22 +6,31 @@ from .api_objects.experiment.experiment_item import (
     ExperimentItemReferences,
 )
 from .api_objects.opik_client import Opik
-from .api_objects.prompt import Prompt
+from .api_objects.prompt import Prompt, ChatPrompt
 from .api_objects.prompt.types import PromptType
 from .api_objects.span import Span
 from .api_objects.trace import Trace
 from .configurator.configure import configure
 from .decorator.tracker import flush_tracker, track
-from .evaluation import evaluate, evaluate_experiment, evaluate_prompt
+from .evaluation import (
+    evaluate,
+    evaluate_experiment,
+    evaluate_on_dict_items,
+    evaluate_prompt,
+)
 from .integrations.sagemaker import auth as sagemaker_auth
 from .plugins.pytest.decorator import llm_unit
 from .types import LLMProvider
 from . import opik_context
-from .decorator.tracing_runtime_config import (
+from .tracing_runtime_config import (
     is_tracing_active,
     reset_tracing_to_config_default,
     set_tracing_active,
 )
+from .decorator.context_manager.span_context_manager import start_as_current_span
+from .decorator.context_manager.trace_context_manager import start_as_current_trace
+from .simulation import SimulatedUser, run_simulation
+from .api_objects.local_recording import record_traces_locally
 
 
 _logging.setup()
@@ -33,6 +42,7 @@ __all__ = [
     "evaluate",
     "evaluate_prompt",
     "evaluate_experiment",
+    "evaluate_on_dict_items",
     "ExperimentItemContent",
     "ExperimentItemReferences",
     "track",
@@ -45,11 +55,17 @@ __all__ = [
     "llm_unit",
     "configure",
     "Prompt",
+    "ChatPrompt",
     "PromptType",
     "LLMProvider",
     "reset_tracing_to_config_default",
     "set_tracing_active",
     "is_tracing_active",
+    "start_as_current_span",
+    "start_as_current_trace",
+    "SimulatedUser",
+    "run_simulation",
+    "record_traces_locally",
 ]
 
 sagemaker_auth.setup_aws_sagemaker_session_hook()

opik/anonymizer/__init__.py

@@ -0,0 +1,5 @@
+from .anonymizer import Anonymizer
+from .factory import create_anonymizer
+from .recursive_anonymizer import RecursiveAnonymizer
+
+__all__ = ["Anonymizer", "create_anonymizer", "RecursiveAnonymizer"]

opik/anonymizer/anonymizer.py

@@ -0,0 +1,12 @@
+import abc
+from typing import Dict, Any, Union, List
+
+AnonymizerDataType = Union[Dict[str, Any], str, List[Any]]
+
+
+class Anonymizer(abc.ABC):
+    """Abstract base class for anonymizing sensitive data in various data structures."""
+
+    @abc.abstractmethod
+    def anonymize(self, data: AnonymizerDataType, **kwargs: Any) -> AnonymizerDataType:
+        pass

opik/anonymizer/factory.py

@@ -0,0 +1,80 @@
+from typing import Union, List, Dict, Callable, Tuple
+
+from . import anonymizer, rules_anonymizer, rules
+
+RulesType = Union[
+    List[Dict[str, str]],
+    List[Tuple[str, str]],
+    List[Callable[[str], str]],
+    List[Union[Dict[str, str], Tuple[str, str], Callable[[str], str]]],
+    Dict[str, str],
+    Tuple[str, str],
+    Callable[[str], str],
+]
+
+
+def create_anonymizer(
+    anonymizer_rules: RulesType, max_depth: int = 10
+) -> anonymizer.Anonymizer:
+    """Create an anonymizer with the specified rules.
+
+    Args:
+        anonymizer_rules: Anonymizer rules specification in various formats:
+            - Dict with "regex" and "replace" keys for a single regex rule
+            - Tuple with (regex, replacement) for a single regex rule
+            - Callable that takes a string and returns anonymized string
+            - List of any of the above for multiple rules
+        max_depth: Maximum recursion depth for nested data structures.
+
+    Returns:
+        An Anonymizer instance configured with the specified rules.
+
+    Raises:
+        ValueError: If a rule format is invalid.
+    """
+    rule_objects: List[rules.Rule] = []
+
+    if callable(anonymizer_rules):
+        # Single function rule
+        rule_objects.append(rules.FunctionRule(anonymizer_rules))
+    elif isinstance(anonymizer_rules, dict):
+        # Single dictionary rule
+        _check_dictionary_rule(anonymizer_rules)
+        rule_objects.append(
+            rules.RegexRule(anonymizer_rules["regex"], anonymizer_rules["replace"])
+        )
+    elif isinstance(anonymizer_rules, tuple):
+        # Single tuple rule
+        _check_tuple_rule(anonymizer_rules)
+        regex_pattern, replacement = anonymizer_rules
+        rule_objects.append(rules.RegexRule(regex_pattern, replacement))
+    elif isinstance(anonymizer_rules, list):
+        # List of rules
+        for rule in anonymizer_rules:
+            if callable(rule) and not isinstance(rule, (dict, tuple)):
+                rule_objects.append(rules.FunctionRule(rule))
+            elif isinstance(rule, dict):
+                _check_dictionary_rule(rule)
+                rule_objects.append(rules.RegexRule(rule["regex"], rule["replace"]))
+            elif isinstance(rule, tuple):
+                _check_tuple_rule(rule)
+                regex_pattern, replacement = rule
+                rule_objects.append(rules.RegexRule(regex_pattern, replacement))
+            else:
+                raise ValueError(f"Unsupported rule type in list: {type(rule)}")
+    else:
+        raise ValueError(f"Unsupported rules type: {type(anonymizer_rules)}")
+
+    return rules_anonymizer.RulesAnonymizer(rule_objects, max_depth=max_depth)
+
+
+def _check_dictionary_rule(rule: Dict[str, str]) -> None:
+    if "regex" not in rule or "replace" not in rule:
+        raise ValueError("Dictionary rule must have 'regex' and 'replace' keys")
+
+
+def _check_tuple_rule(rule: Tuple[str, str]) -> None:
+    if len(rule) != 2:
+        raise ValueError(
+            "Tuple rule must have exactly 2 elements: (regex, replacement)"
+        )

opik/anonymizer/recursive_anonymizer.py

@@ -0,0 +1,64 @@
+import abc
+from typing import Any, Optional
+
+from . import anonymizer
+
+
+class RecursiveAnonymizer(anonymizer.Anonymizer):
+    """Abstract base class for anonymizing sensitive data in various data structures.
+
+    This class provides a framework for recursively anonymizing text content within
+    nested data structures such as dictionaries, lists, and strings. Subclasses must
+    implement the anonymize_text() method to define the specific anonymization logic.
+    """
+
+    def __init__(self, max_depth: int = 10):
+        """Initialize the Anonymizer with depth limiting.
+
+        Args:
+            max_depth: Maximum recursion depth to prevent infinite loops when
+                      processing deeply nested or circular data structures.
+                      Defaults to 10.
+        """
+        self.max_depth = max_depth
+
+    def anonymize(
+        self, data: anonymizer.AnonymizerDataType, **kwargs: Any
+    ) -> anonymizer.AnonymizerDataType:
+        return self._recursive_anonymize(data, **kwargs)
+
+    @abc.abstractmethod
+    def anonymize_text(self, data: str, **kwargs: Any) -> str:
+        pass
+
+    def _recursive_anonymize(
+        self,
+        data: anonymizer.AnonymizerDataType,
+        depth: int = 0,
+        field_name: Optional[str] = None,
+        **kwargs: Any,
+    ) -> anonymizer.AnonymizerDataType:
+        if depth >= self.max_depth:
+            return data
+
+        if field_name is None:
+            field_name = ""
+
+        if isinstance(data, str):
+            return self.anonymize_text(data, field_name=field_name, **kwargs)
+        elif isinstance(data, dict):
+            return {
+                key: self._recursive_anonymize(
+                    value, depth + 1, field_name=f"{field_name}.{key}", **kwargs
+                )
+                for key, value in data.items()
+            }
+        elif isinstance(data, list):
+            return [
+                self._recursive_anonymize(
+                    item, depth + 1, field_name=f"{field_name}.{i}", **kwargs
+                )
+                for i, item in enumerate(data)
+            ]
+        else:
+            return data

opik/anonymizer/rules.py

@@ -0,0 +1,56 @@
+import abc
+import re
+from typing import Callable
+
+
+class Rule(abc.ABC):
+    """Abstract base class for text anonymization rules.
+
+    Rules define specific patterns or conditions for anonymizing sensitive
+    information in text. Subclasses must implement the apply() method to
+    define the anonymization logic.
+    """
+
+    @abc.abstractmethod
+    def apply(self, text: str) -> str:
+        pass
+
+
+class RegexRule(Rule):
+    """A rule that uses regular expressions to find and replace patterns in text.
+
+    This rule compiles a regular expression pattern and applies it to input text,
+    replacing all matches with a specified replacement string.
+    """
+
+    def __init__(self, regex: str, replacement: str):
+        """Initialize the regex rule with a pattern and replacement.
+
+        Args:
+            regex: Regular expression pattern to match sensitive data.
+            replacement: String to replace matched patterns with.
+        """
+        self.pattern = re.compile(regex)
+        self.replacement = replacement
+
+    def apply(self, text: str) -> str:
+        return self.pattern.sub(self.replacement, text)
+
+
+class FunctionRule(Rule):
+    """A rule that applies a custom function to anonymize text.
+
+    This rule allows for flexible anonymization by accepting any callable
+    that takes a string as input and returns an anonymized string.
+    """
+
+    def __init__(self, func: Callable[[str], str]):
+        """Initialize the function rule with a custom anonymization function.
+
+        Args:
+            func: A callable that takes a string and returns an anonymized version.
+        """
+        self.func = func
+
+    def apply(self, text: str) -> str:
+        return self.func(text)

opik/anonymizer/rules_anonymizer.py

@@ -0,0 +1,35 @@
+from typing import List, Any
+
+from . import recursive_anonymizer, rules
+
+
+class RulesAnonymizer(recursive_anonymizer.RecursiveAnonymizer):
+    """An anonymizer that applies a list of rules sequentially to text data.
+
+    This class takes a list of Rule objects and applies them to
+    anonymize sensitive information in text.
+    """
+
+    def __init__(self, anonymizer_rules: List[rules.Rule], max_depth: int = 10):
+        """Initialize the RulesAnonymizer with a list of rules.
+
+        Args:
+            anonymizer_rules: List of Rule objects to apply for anonymization.
+            max_depth: Maximum recursion depth for nested data structures.
+        """
+        super().__init__(max_depth)
+        self.rules = anonymizer_rules
+
+    def anonymize_text(self, data: str, **kwargs: Any) -> str:
+        """Apply all rules sequentially to the input text.
+
+        Args:
+            data: The text to anonymize.
+
+        Returns:
+            The anonymized text after applying all rules.
+        """
+        result = data
+        for rule in self.rules:
+            result = rule.apply(result)
+        return result

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