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

opik/message_processing/emulation/local_emulator_message_processor.py

@@ -0,0 +1,140 @@
+import datetime
+import logging
+from typing import List, Dict, Optional, Any, Union
+
+from opik.types import ErrorInfoDict, SpanType
+from . import models, emulator_message_processor
+from .. import messages
+from ...rest_api.types import span_write, trace_write
+
+
+LOGGER = logging.getLogger(__name__)
+
+
+class LocalEmulatorMessageProcessor(
+    emulator_message_processor.EmulatorMessageProcessor
+):
+    """This class serves as a replacement for the real backend and collects all logged messages
+    locally in memory to be used for evaluation.
+    """
+
+    def __init__(self, active: bool, merge_duplicates: bool = True) -> None:
+        super().__init__(active=active, merge_duplicates=merge_duplicates)
+
+    def process(
+        self,
+        message: Union[
+            messages.BaseMessage, span_write.SpanWrite, trace_write.TraceWrite
+        ],
+    ) -> None:
+        if not self.is_active():
+            return
+
+        if hasattr(message, "delivery_attempts") and message.delivery_attempts > 1:
+            # skip retries
+            LOGGER.debug("Skipping retry of the message: %s", message)
+            return
+
+        super().process(message)
+
+    def create_trace_model(
+        self,
+        trace_id: str,
+        start_time: datetime.datetime,
+        name: Optional[str],
+        project_name: str,
+        input: Any,
+        output: Any,
+        tags: Optional[List[str]],
+        metadata: Optional[Dict[str, Any]],
+        end_time: Optional[datetime.datetime],
+        spans: Optional[List[models.SpanModel]],
+        feedback_scores: Optional[List[models.FeedbackScoreModel]],
+        error_info: Optional[ErrorInfoDict],
+        thread_id: Optional[str],
+        last_updated_at: Optional[datetime.datetime] = None,
+    ) -> models.TraceModel:
+        if spans is None:
+            spans = []
+        if feedback_scores is None:
+            feedback_scores = []
+
+        return models.TraceModel(
+            id=trace_id,
+            start_time=start_time,
+            name=name,
+            project_name=project_name,
+            input=input,
+            output=output,
+            tags=tags,
+            metadata=metadata,
+            end_time=end_time,
+            spans=spans,
+            feedback_scores=feedback_scores,
+            error_info=error_info,
+            thread_id=thread_id,
+            last_updated_at=last_updated_at,
+        )
+
+    def create_span_model(
+        self,
+        span_id: str,
+        start_time: datetime.datetime,
+        name: Optional[str],
+        input: Any,
+        output: Any,
+        tags: Optional[List[str]],
+        metadata: Optional[Dict[str, Any]],
+        type: SpanType,
+        usage: Optional[Dict[str, Any]],
+        end_time: Optional[datetime.datetime],
+        project_name: str,
+        spans: Optional[List[models.SpanModel]],
+        feedback_scores: Optional[List[models.FeedbackScoreModel]],
+        model: Optional[str],
+        provider: Optional[str],
+        error_info: Optional[ErrorInfoDict],
+        total_cost: Optional[float],
+        last_updated_at: Optional[datetime.datetime],
+    ) -> models.SpanModel:
+        if spans is None:
+            spans = []
+        if feedback_scores is None:
+            feedback_scores = []
+
+        return models.SpanModel(
+            id=span_id,
+            start_time=start_time,
+            name=name,
+            input=input,
+            output=output,
+            tags=tags,
+            metadata=metadata,
+            type=type,
+            usage=usage,
+            end_time=end_time,
+            project_name=project_name,
+            spans=spans,
+            feedback_scores=feedback_scores,
+            model=model,
+            provider=provider,
+            error_info=error_info,
+            total_cost=total_cost,
+            last_updated_at=last_updated_at,
+        )
+
+    def create_feedback_score_model(
+        self,
+        score_id: str,
+        name: str,
+        value: float,
+        category_name: Optional[str],
+        reason: Optional[str],
+    ) -> models.FeedbackScoreModel:
+        return models.FeedbackScoreModel(
+            id=score_id,
+            name=name,
+            value=value,
+            category_name=category_name,
+            reason=reason,
+        )

opik/message_processing/emulation/models.py

@@ -0,0 +1,162 @@
+from typing import List, Any, Optional, Dict
+
+from opik.config import OPIK_PROJECT_DEFAULT_NAME
+from opik.types import ErrorInfoDict
+
+import dataclasses
+import datetime
+
+
+@dataclasses.dataclass
+class FeedbackScoreModel:
+    """
+    Represents a model for a feedback score used to evaluate specific spans or traces.
+
+    This class stores and manages feedback scores linked to defined criteria, including
+    identifiers, names, values, categories, and explanations for each score.
+
+    Attributes:
+        id: Unique identifier for the feedback score.
+        name: Name associated with the feedback score.
+        value: The numerical value of the feedback score.
+        category_name: Category to which the feedback score belongs, if any.
+        reason: Reason or explanation for the feedback score, if available.
+    """
+
+    id: str
+    name: str
+    value: float
+    category_name: Optional[str] = None
+    reason: Optional[str] = None
+
+
+@dataclasses.dataclass
+class SpanModel:
+    """
+    Represents a span model used to describe specific points in a process, their metadata,
+    and associated data.
+
+    This class is used to store and manipulate structured data for events or
+    spans, including metadata, time markers, associated input/output, tags,
+    and additional properties. It serves as a representative structure for recording
+    and organizing event-specific information, often used in applications like
+    logging, distributed tracing, or data processing pipelines.
+
+    Attributes:
+        id: Unique identifier for the span.
+        start_time: Start time of the span.
+        name: Name of the span, if provided.
+        input: Input data associated with the span, if any.
+        output: Output data associated with the span, if any.
+        tags: List of tags linked to the span.
+        metadata: Additional metadata for the span.
+        type: Type of the span, defaulting to "general".
+        usage: Usage-related information for the span.
+        end_time: End time of the span, if available.
+        project_name: Name of the project the span is associated with,
+            defaulting to a predefined project name.
+        spans: List of nested spans related to this span.
+        feedback_scores: List of feedback scores associated
+            with the span.
+        model: Model identification used, if applicable.
+        provider: Provider of the span or associated services, if any.
+        error_info: Error information or diagnostics
+            for the span, if applicable.
+        total_cost: Total cost incurred associated with this span,
+            if relevant.
+        last_updated_at: Timestamp of when the span was
+            last updated, if available.
+    """
+
+    id: str
+    start_time: datetime.datetime
+    name: Optional[str] = None
+    input: Optional[Dict[str, Any]] = None
+    output: Optional[Dict[str, Any]] = None
+    tags: Optional[List[str]] = None
+    metadata: Optional[Dict[str, Any]] = None
+    type: str = "general"
+    usage: Optional[Dict[str, Any]] = None
+    end_time: Optional[datetime.datetime] = None
+    project_name: str = OPIK_PROJECT_DEFAULT_NAME
+    spans: List["SpanModel"] = dataclasses.field(default_factory=list)
+    feedback_scores: List[FeedbackScoreModel] = dataclasses.field(default_factory=list)
+    model: Optional[str] = None
+    provider: Optional[str] = None
+    error_info: Optional[ErrorInfoDict] = None
+    total_cost: Optional[float] = None
+    last_updated_at: Optional[datetime.datetime] = None
+
+
+@dataclasses.dataclass
+class ExperimentItemModel:
+    """
+    Represents an experiment item model that links a trace to a dataset item in an experiment.
+
+    Attributes:
+        id: Unique identifier for the experiment item.
+        experiment_id: The ID of the experiment this item belongs to.
+        trace_id: The ID of the trace associated with this experiment item.
+        dataset_item_id: The ID of the dataset item associated with this experiment item.
+    """
+
+    id: str
+    experiment_id: str
+    trace_id: str
+    dataset_item_id: str
+
+
+@dataclasses.dataclass
+class TraceModel:
+    """
+    Represents a trace model that encapsulates data about a trace, its related metadata,
+    and associated spans. It is used for tracking and analyzing data during execution
+    or processing tasks.
+
+    This class provides a structure to represent trace information, including the start
+    and end times, associated project details, input/output data, feedback scores, error
+    information, and thread association. It is designed to handle optional fields for
+    flexible use across various scenarios.
+
+    Attributes:
+        id: Unique identifier for the trace.
+        start_time: Timestamp representing the start of the trace.
+        name: Optional name for the trace, which can provide a descriptive
+            label.
+        project_name: Name of the project associated with the trace.
+        input: Optional dictionary containing the input data
+            associated with the trace.
+        output: Optional dictionary containing the output data
+            generated by the trace.
+        tags: Optional list of tags associated with the trace for
+            classification or filtering purposes.
+        metadata: Optional metadata providing additional
+            information about the trace.
+        end_time: Timestamp representing the end of the
+            trace.
+        spans: List of spans associated with the trace, representing
+            individual processing parts or segments within the trace.
+        feedback_scores: List of feedback scores associated
+            with the trace.
+        error_info: Optional dictionary containing information
+            about errors encountered during the trace.
+        thread_id: Optional identifier of the thread associated with the
+            trace.
+        last_updated_at: Timestamp for when the trace was
+            last updated.
+    """
+
+    id: str
+    start_time: datetime.datetime
+    name: Optional[str]
+    project_name: str
+    input: Optional[Dict[str, Any]] = None
+    output: Optional[Dict[str, Any]] = None
+    tags: Optional[List[str]] = None
+    metadata: Optional[Dict[str, Any]] = None
+    end_time: Optional[datetime.datetime] = None
+    spans: List[SpanModel] = dataclasses.field(default_factory=list)
+    feedback_scores: List[FeedbackScoreModel] = dataclasses.field(default_factory=list)
+    error_info: Optional[ErrorInfoDict] = None
+    thread_id: Optional[str] = None
+    last_updated_at: Optional[datetime.datetime] = None

opik/message_processing/encoder_helpers.py

@@ -0,0 +1,79 @@
+from typing import Dict, Any, Set, List, Literal
+
+import opik.hooks
+
+from .. import jsonable_encoder
+from ..anonymizer import anonymizer
+
+
+def encode_and_anonymize(
+    kwargs_dict: Dict[str, Any],
+    fields_to_anonymize: Set[str],
+    object_type: Literal["span", "trace"],
+) -> Dict[str, Any]:
+    """
+    Encodes and anonymizes the data in the given dictionary based on the specified
+    fields using registered anonymizers. If no anonymizers are registered, the
+    function simply encodes the dictionary without anonymization.
+
+    Args:
+        kwargs_dict: The dictionary containing the data to encode
+            and anonymize.
+        fields_to_anonymize: The set of fields within the dictionary to
+            anonymize.
+        object_type: A string indicating the type of object ('span' or 'trace')
+            that was used to create the kwargs_dict. This is passed to anonymizers
+            to provide context about the source object.
+
+    Returns:
+        A dictionary that has been encoded and, if applicable, anonymized.
+    """
+    # check if any anonymizer was registered
+    encoded_obj = jsonable_encoder.encode(kwargs_dict)
+    if not opik.hooks.has_anonymizers():
+        return encoded_obj
+
+    anonymizers = opik.hooks.get_anonymizers()
+    return anonymize_encoded_obj(
+        obj=encoded_obj,
+        fields_to_anonymize=fields_to_anonymize,
+        anonymizers=anonymizers,
+        object_type=object_type,
+    )
+
+
+def anonymize_encoded_obj(
+    obj: Dict[str, Any],
+    fields_to_anonymize: Set[str],
+    anonymizers: List[anonymizer.Anonymizer],
+    object_type: Literal["span", "trace"],
+) -> Dict[str, Any]:
+    """
+    Anonymizes specified fields in an encoded dictionary using the provided anonymizers.
+    This function iterates over the given set of field names and applies each anonymizer
+    to the corresponding field in the dictionary, if present. The anonymizers are expected
+    to implement an `anonymize` method that takes the field value, field name, and object type
+    as arguments. Only fields present in the dictionary and listed in `fields_to_anonymize`
+    are anonymized.
+
+    Args:
+        obj: The encoded dictionary whose fields are to be anonymized.
+        fields_to_anonymize: A set of field names within the dictionary to anonymize.
+        anonymizers: A list of anonymizer instances to apply to each field.
+        object_type: A string indicating the type of object ('span' or 'trace'),
+            providing context for anonymization.
+
+    Returns:
+        The dictionary with specified fields anonymized using the provided anonymizers.
+    """
+    if isinstance(obj, dict):
+        for field_name in fields_to_anonymize:
+            if field_name in obj:
+                for anonymizer_instance in anonymizers:
+                    obj[field_name] = anonymizer_instance.anonymize(
+                        obj[field_name],
+                        field_name=field_name,
+                        object_type=object_type,
+                    )
+
+    return obj

opik/message_processing/messages.py

@@ -1,9 +1,10 @@
 import dataclasses
 import datetime
 from dataclasses import field
-from typing import Optional, Any, Dict, List, Union, Literal
+from typing import Optional, Any, Dict, List, Union, Literal, Set
 
 from . import arguments_utils
+from .preprocessing import constants
 from ..rest_api.types import span_write, trace_write
 from ..types import SpanType, ErrorInfoDict, LLMProvider, AttachmentEntityType
 
@@ -11,6 +12,7 @@ from ..types import SpanType, ErrorInfoDict, LLMProvider, AttachmentEntityType
 @dataclasses.dataclass
 class BaseMessage:
     delivery_time: float = field(init=False, default=0.0)
+    delivery_attempts: int = field(init=False, default=1)
 
     def as_payload_dict(self) -> Dict[str, Any]:
         # we are not using dataclasses.as_dict() here
@@ -18,6 +20,10 @@ class BaseMessage:
         data = {**self.__dict__}
         if "delivery_time" in data:
             data.pop("delivery_time")
+        if "delivery_attempts" in data:
+            data.pop("delivery_attempts")
+        if constants.MARKER_ATTRIBUTE_NAME in data:
+            data.pop(constants.MARKER_ATTRIBUTE_NAME)
         return data
 
 
@@ -47,6 +53,10 @@ class CreateTraceMessage(BaseMessage):
         data["id"] = data.pop("trace_id")
         return data
 
+    @staticmethod
+    def fields_to_anonymize() -> Set[str]:
+        return {"input", "output", "metadata"}
+
 
 @dataclasses.dataclass
 class UpdateTraceMessage(BaseMessage):
@@ -75,6 +85,10 @@ class UpdateTraceMessage(BaseMessage):
         data["id"] = data.pop("trace_id")
         return data
 
+    @staticmethod
+    def fields_to_anonymize() -> Set[str]:
+        return {"input", "output", "metadata"}
+
 
 @dataclasses.dataclass
 class CreateSpanMessage(BaseMessage):
@@ -109,6 +123,10 @@ class CreateSpanMessage(BaseMessage):
         data["total_estimated_cost"] = data.pop("total_cost")
         return data
 
+    @staticmethod
+    def fields_to_anonymize() -> Set[str]:
+        return {"input", "output", "metadata"}
+
 
 @dataclasses.dataclass
 class UpdateSpanMessage(BaseMessage):
@@ -141,6 +159,10 @@ class UpdateSpanMessage(BaseMessage):
         data["total_estimated_cost"] = data.pop("total_cost")
         return data
 
+    @staticmethod
+    def fields_to_anonymize() -> Set[str]:
+        return {"input", "output", "metadata"}
+
 
 @dataclasses.dataclass
 class FeedbackScoreMessage(BaseMessage):
@@ -197,11 +219,19 @@ class AddThreadsFeedbackScoresBatchMessage(BaseMessage):
 class CreateSpansBatchMessage(BaseMessage):
     batch: List[span_write.SpanWrite]
 
+    @staticmethod
+    def fields_to_anonymize() -> Set[str]:
+        return {"input", "output", "metadata"}
+
 
 @dataclasses.dataclass
 class CreateTraceBatchMessage(BaseMessage):
     batch: List[trace_write.TraceWrite]
 
+    @staticmethod
+    def fields_to_anonymize() -> Set[str]:
+        return {"input", "output", "metadata"}
+
 
 @dataclasses.dataclass
 class GuardrailBatchItemMessage(BaseMessage):
@@ -230,6 +260,25 @@ class GuardrailBatchMessage(BaseMessage):
         return data
 
 
+@dataclasses.dataclass
+class ExperimentItemMessage(BaseMessage):
+    """
+    There is no handler for that in the message processor, it exists
+    only as an item of CreateExperimentItemsBatchMessage
+    """
+
+    id: str
+    experiment_id: str
+    trace_id: str
+    dataset_item_id: str
+
+
+@dataclasses.dataclass
+class CreateExperimentItemsBatchMessage(BaseMessage):
+    batch: List[ExperimentItemMessage]
+    supports_batching: bool = True
+
+
 @dataclasses.dataclass
 class CreateAttachmentMessage(BaseMessage):
     file_path: str
@@ -239,3 +288,9 @@ class CreateAttachmentMessage(BaseMessage):
     entity_id: str
     project_name: str
     encoded_url_override: str
+    delete_after_upload: bool = False
+
+
+@dataclasses.dataclass
+class AttachmentSupportingMessage(BaseMessage):
+    original_message: BaseMessage

opik/message_processing/preprocessing/attachments_preprocessor.py

@@ -0,0 +1,70 @@
+from typing import Optional, Union
+
+from opik.message_processing import messages
+
+from . import constants, preprocessor
+
+
+class AttachmentsPreprocessor(preprocessor.MessagePreprocessor):
+    def __init__(self, enabled: bool = True) -> None:
+        self._enabled = enabled
+
+    def preprocess(
+        self, message: Optional[messages.BaseMessage]
+    ) -> Optional[messages.BaseMessage]:
+        """
+        Processes a given message and ensures that it is converted into a specialized
+        message type if applicable. If the message is already pre-processed, it
+        returns the original message to avoid infinite recursion.
+
+        Args:
+            message: The message object to be processed.
+
+        Returns:
+            The processed message, either in its original form
+            or converted into a message type supporting embedded attachments.
+        """
+        if not self._enabled:
+            return message
+
+        if message is None:
+            # possibly already pre-processed by other preprocessors
+            return None
+
+        if hasattr(message, constants.MARKER_ATTRIBUTE_NAME):
+            # already pre-processed - just return the original message to avoid infinite recursion
+            return message
+
+        if _has_potential_content_with_attachments(message):
+            return messages.AttachmentSupportingMessage(message)
+        else:
+            return message
+
+
+def _has_potential_content_with_attachments(message: messages.BaseMessage) -> bool:
+    # Check if it's an Update message - always process these
+    if isinstance(message, (messages.UpdateSpanMessage, messages.UpdateTraceMessage)):
+        return _message_has_field_of_interest_set(message)
+
+    # Check if it's a Create message with end_time set - only process these
+    if isinstance(message, (messages.CreateSpanMessage, messages.CreateTraceMessage)):
+        if message.end_time is not None:
+            return _message_has_field_of_interest_set(message)
+
+    # All other message types should not be wrapped
+    return False
+
+
+def _message_has_field_of_interest_set(
+    message: Union[
+        messages.UpdateSpanMessage,
+        messages.UpdateTraceMessage,
+        messages.CreateSpanMessage,
+        messages.CreateTraceMessage,
+    ],
+) -> bool:
+    return (
+        message.input is not None
+        or message.output is not None
+        or message.metadata is not None
+    )

opik/message_processing/preprocessing/batching_preprocessor.py

@@ -0,0 +1,53 @@
+from typing import Optional
+
+from . import preprocessor
+from .. import messages
+from ..batching import batch_manager
+
+
+class BatchingPreprocessor(preprocessor.MessagePreprocessor):
+    """
+    Handles message batching during preprocessing.
+
+    The BatchingPreprocessor class processes messages, enabling efficient message
+    batching if a batching manager is provided. It supports starting, stopping,
+    flushing, and checking the state of the batching manager, ensuring that
+    messages are processed or delegated based on their batching capabilities.
+    """
+
+    def __init__(self, batching_manager: Optional[batch_manager.BatchManager]) -> None:
+        self._batch_manager = batching_manager
+
+    def preprocess(
+        self, message: Optional[messages.BaseMessage]
+    ) -> Optional[messages.BaseMessage]:
+        if message is None:
+            # possibly already processed
+            return None
+
+        if (
+            self._batch_manager is not None
+            and self._batch_manager.message_supports_batching(message)
+        ):
+            self._batch_manager.process_message(message)
+            return None
+
+        return message
+
+    def start(self) -> None:
+        if self._batch_manager is not None:
+            self._batch_manager.start()
+
+    def stop(self) -> None:
+        if self._batch_manager is not None:
+            self._batch_manager.stop()
+
+    def flush(self) -> None:
+        if self._batch_manager is not None:
+            self._batch_manager.flush()
+
+    def is_empty(self) -> bool:
+        if self._batch_manager is not None:
+            return self._batch_manager.is_empty()
+
+        return True

opik/message_processing/preprocessing/constants.py

@@ -0,0 +1 @@
+MARKER_ATTRIBUTE_NAME = "_preprocessed_for_attachments"

opik/message_processing/preprocessing/file_upload_preprocessor.py

@@ -0,0 +1,38 @@
+from typing import Optional
+
+from opik.file_upload import base_upload_manager
+
+from . import preprocessor
+from .. import messages
+
+
+class FileUploadPreprocessor(preprocessor.MessagePreprocessor):
+    """
+    Preprocesses messages to handle file uploads.
+
+    This class is responsible for processing messages to determine if they support
+    file uploads and delegating the upload task to a file upload manager. It also
+    provides functionality to flush pending uploads with configurable timeout and
+    sleep intervals.
+    """
+
+    def __init__(
+        self, file_upload_manager: base_upload_manager.BaseFileUploadManager
+    ) -> None:
+        self._file_upload_manager = file_upload_manager
+
+    def preprocess(
+        self, message: Optional[messages.BaseMessage]
+    ) -> Optional[messages.BaseMessage]:
+        if message is None:
+            # possibly already processed
+            return None
+
+        if base_upload_manager.message_supports_upload(message):
+            self._file_upload_manager.upload(message)
+            return None
+
+        return message
+
+    def flush(self, timeout: Optional[float], sleep_time: int) -> bool:
+        return self._file_upload_manager.flush(timeout=timeout, sleep_time=sleep_time)