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

opik/__init__.py

@@ -1,40 +1,71 @@
 from . import _logging, environment, error_tracking, package_version
+from .api_objects.attachment import Attachment
 from .api_objects.dataset import Dataset
 from .api_objects.experiment.experiment_item import (
     ExperimentItemContent,
     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 .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()
 
 __version__ = package_version.VERSION
 __all__ = [
     "__version__",
+    "Attachment",
     "evaluate",
     "evaluate_prompt",
     "evaluate_experiment",
+    "evaluate_on_dict_items",
     "ExperimentItemContent",
     "ExperimentItemReferences",
     "track",
     "flush_tracker",
     "Opik",
+    "opik_context",
     "Trace",
     "Span",
     "Dataset",
     "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/__init__.py

@@ -0,0 +1,5 @@
+from .attachment import Attachment
+from .client import AttachmentClient
+
+
+__all__ = ["Attachment", "AttachmentClient"]

opik/api_objects/attachment/attachment.py

@@ -0,0 +1,20 @@
+from typing import Optional
+
+import pydantic
+
+
+class Attachment(pydantic.BaseModel):
+    """
+    Represents an Attachment to be added to the Trace or Span.
+
+    Args:
+        data: The data to be added to the Attachment as a path to the file.
+        file_name: The custom filename to assign to the data in the attachment.
+            If not provided, the original filename of the data will be used.
+        content_type: The MIME type of the data to be added to the attachment.
+            If not specified, it will be inferred from the data file.
+    """
+
+    data: str
+    file_name: Optional[str] = None
+    content_type: Optional[str] = None

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
+        )