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

opik/message_processing/messages.py

@@ -1,15 +1,30 @@
 import dataclasses
 import datetime
-from typing import Optional, Any, Dict, List, Union
-from ..types import SpanType, ErrorInfoDict, LLMProvider
+from dataclasses import field
+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
 
 
 @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
-        # because it will try to deepcopy all object and will fail if there is non-serializable object
-        return {**self.__dict__}
+        # because it will try to deepcopy all objects and will fail if there is a non-serializable object
+        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
 
 
 @dataclasses.dataclass
@@ -25,12 +40,23 @@ class CreateTraceMessage(BaseMessage):
     tags: Optional[List[str]]
     error_info: Optional[ErrorInfoDict]
     thread_id: Optional[str]
+    last_updated_at: Optional[datetime.datetime]
+
+    def __post_init__(self) -> None:
+        if self.input is not None:
+            self.input = arguments_utils.recursive_shallow_copy(self.input)
+        if self.output is not None:
+            self.output = arguments_utils.recursive_shallow_copy(self.output)
 
     def as_payload_dict(self) -> Dict[str, Any]:
         data = super().as_payload_dict()
         data["id"] = data.pop("trace_id")
         return data
 
+    @staticmethod
+    def fields_to_anonymize() -> Set[str]:
+        return {"input", "output", "metadata"}
+
 
 @dataclasses.dataclass
 class UpdateTraceMessage(BaseMessage):
@@ -48,11 +74,21 @@ class UpdateTraceMessage(BaseMessage):
     error_info: Optional[ErrorInfoDict]
     thread_id: Optional[str]
 
+    def __post_init__(self) -> None:
+        if self.input is not None:
+            self.input = arguments_utils.recursive_shallow_copy(self.input)
+        if self.output is not None:
+            self.output = arguments_utils.recursive_shallow_copy(self.output)
+
     def as_payload_dict(self) -> Dict[str, Any]:
         data = super().as_payload_dict()
         data["id"] = data.pop("trace_id")
         return data
 
+    @staticmethod
+    def fields_to_anonymize() -> Set[str]:
+        return {"input", "output", "metadata"}
+
 
 @dataclasses.dataclass
 class CreateSpanMessage(BaseMessage):
@@ -73,6 +109,13 @@ class CreateSpanMessage(BaseMessage):
     provider: Optional[Union[LLMProvider, str]]
     error_info: Optional[ErrorInfoDict]
     total_cost: Optional[float]
+    last_updated_at: Optional[datetime.datetime]
+
+    def __post_init__(self) -> None:
+        if self.input is not None:
+            self.input = arguments_utils.recursive_shallow_copy(self.input)
+        if self.output is not None:
+            self.output = arguments_utils.recursive_shallow_copy(self.output)
 
     def as_payload_dict(self) -> Dict[str, Any]:
         data = super().as_payload_dict()
@@ -80,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):
@@ -100,17 +147,27 @@ class UpdateSpanMessage(BaseMessage):
     error_info: Optional[ErrorInfoDict]
     total_cost: Optional[float]
 
+    def __post_init__(self) -> None:
+        if self.input is not None:
+            self.input = arguments_utils.recursive_shallow_copy(self.input)
+        if self.output is not None:
+            self.output = arguments_utils.recursive_shallow_copy(self.output)
+
     def as_payload_dict(self) -> Dict[str, Any]:
         data = super().as_payload_dict()
         data["id"] = data.pop("span_id")
         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):
     """
-    There is no handler for that in message processor, it exists
+    There is no handler for that in the message processor, it exists
     only as an item of BatchMessage
     """
 
@@ -128,11 +185,6 @@ class AddFeedbackScoresBatchMessage(BaseMessage):
     batch: List[FeedbackScoreMessage]
     supports_batching: bool = True
 
-    def as_payload_dict(self) -> Dict[str, Any]:
-        data = super().as_payload_dict()
-        data.pop("supports_batching")
-        return data
-
 
 @dataclasses.dataclass
 class AddTraceFeedbackScoresBatchMessage(AddFeedbackScoresBatchMessage):
@@ -144,11 +196,101 @@ class AddSpanFeedbackScoresBatchMessage(AddFeedbackScoresBatchMessage):
     pass
 
 
+@dataclasses.dataclass
+class ThreadsFeedbackScoreMessage(FeedbackScoreMessage):
+    """
+    There is no handler for that in the message processor, it exists
+    only as an item of AddThreadsFeedbackScoresBatchMessage
+    """
+
+    def as_payload_dict(self) -> Dict[str, Any]:
+        data = super().as_payload_dict()
+        data["thread_id"] = data.pop("id")
+        return data
+
+
+@dataclasses.dataclass
+class AddThreadsFeedbackScoresBatchMessage(BaseMessage):
+    batch: List[ThreadsFeedbackScoreMessage]
+    supports_batching: bool = True
+
+
 @dataclasses.dataclass
 class CreateSpansBatchMessage(BaseMessage):
-    batch: List[CreateSpanMessage]
+    batch: List[span_write.SpanWrite]
+
+    @staticmethod
+    def fields_to_anonymize() -> Set[str]:
+        return {"input", "output", "metadata"}
 
 
 @dataclasses.dataclass
 class CreateTraceBatchMessage(BaseMessage):
-    batch: List[CreateTraceMessage]
+    batch: List[trace_write.TraceWrite]
+
+    @staticmethod
+    def fields_to_anonymize() -> Set[str]:
+        return {"input", "output", "metadata"}
+
+
+@dataclasses.dataclass
+class GuardrailBatchItemMessage(BaseMessage):
+    """
+    There is no handler for that in the message processor, it exists
+    only as an item of BatchMessage
+    """
+
+    project_name: Optional[str]
+    entity_id: str
+    secondary_id: str
+    name: str
+    result: Union[Literal["passed", "failed"], Any]
+    config: Dict[str, Any]
+    details: Dict[str, Any]
+
+
+@dataclasses.dataclass
+class GuardrailBatchMessage(BaseMessage):
+    batch: List[GuardrailBatchItemMessage]
+    supports_batching: bool = True
+
+    def as_payload_dict(self) -> Dict[str, Any]:
+        data = super().as_payload_dict()
+        data.pop("supports_batching")
+        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
+    file_name: str
+    mime_type: Optional[str]
+    entity_type: AttachmentEntityType
+    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)

opik/message_processing/preprocessing/preprocessor.py

@@ -0,0 +1,36 @@
+import abc
+from typing import Optional
+
+from opik.message_processing import messages
+
+
+class MessagePreprocessor(abc.ABC):
+    """
+    Abstract base class for message preprocessing.
+
+    This class provides a common interface for pre-processing messages, allowing
+    derived classes to implement custom preprocessing logic tailored to specific
+    requirements. Instances of this class cannot be created directly; it must be
+    subclassed with the `preprocess` method implemented.
+    """
+
+    @abc.abstractmethod
+    def preprocess(
+        self, message: Optional[messages.BaseMessage]
+    ) -> Optional[messages.BaseMessage]:
+        """
+        Processes and preprocesses the given message to prepare it for further operations.
+
+        This is an abstract method and needs to be implemented in any concrete subclass. The
+        preprocessing step is typically used for transformations or checks on the given input
+        message before further processing.
+
+        Args:
+            message: The input message to be preprocessed. This can
+                optionally be None.
+
+        Returns:
+            The processed message after preprocessing. Returns None if the input message is None
+            or if a message was fully consumed here and no further processing is required.
+        """
+        pass

opik/message_processing/processors/attachments_extraction_processor.py

@@ -0,0 +1,146 @@
+import logging
+from typing import Optional, NamedTuple, List, Literal, cast
+
+from opik.api_objects.attachment import (
+    attachments_extractor,
+    attachment_context,
+    converters,
+)
+
+from . import message_processors
+from ..preprocessing import constants
+from .. import messages, streamer
+
+
+LOGGER = logging.getLogger(__name__)
+
+
+class EntityDetails(NamedTuple):
+    entity_type: Literal["span", "trace"]
+    entity_id: str
+    project_name: str
+
+
+class AttachmentsExtractionProcessor(message_processors.BaseMessageProcessor):
+    """
+    Class for processing message attachments through extraction and further handling.
+
+    The AttachmentsExtractionProcessor class is designed to handle attachments from incoming
+    messages. It checks the type of messages and processes them if they support
+    attachments. This includes extracting attachment data, replacing them with references,
+    and streaming processed or original messages through a pipeline. The class provides a
+    mechanism to toggle processing activity and ensures proper handling of messages with
+    embedded attachment information.
+    """
+
+    def __init__(
+        self,
+        min_attachment_size: int,
+        messages_streamer: streamer.Streamer,
+        url_override: str,
+        is_active: bool = True,
+    ):
+        """
+        Initializes an object with essential components for managing message streaming
+        and attachment extraction.
+
+        Args:
+            min_attachment_size: Minimum size for an attachment to be extracted.
+            messages_streamer: The streamer that is responsible for managing
+                messages broadcasts.
+            url_override: A custom URL to override default configurations if set.
+            is_active: Indicator of whether this instance is active. Default is True.
+        """
+        self._is_active = is_active
+        self.extractor = attachments_extractor.AttachmentsExtractor(min_attachment_size)
+        self.messages_streamer = messages_streamer
+        self._url_override = url_override
+
+        self.attachment_attributes = ["input", "output", "metadata"]
+
+    def is_active(self) -> bool:
+        return self._is_active
+
+    def process(self, message: messages.BaseMessage) -> None:
+        if not isinstance(message, messages.AttachmentSupportingMessage):
+            return
+
+        if self._is_active:
+            # do attachment processing only if the processor is active
+            try:
+                self._process_attachments_in_message(message.original_message)
+            except Exception as ex:
+                LOGGER.error(
+                    "Failed to process attachment support message: %s", ex, exc_info=ex
+                )
+
+        # put the original message into the streamer for further processing
+        original_message = message.original_message
+        setattr(original_message, constants.MARKER_ATTRIBUTE_NAME, True)
+        self.messages_streamer.put(original_message)
+
+    def _process_attachments_in_message(self, original: messages.BaseMessage) -> None:
+        entity_details = entity_type_from_attachment_message(original)
+        if entity_details is None:
+            LOGGER.error(
+                "Failed to extract entity details from message - %s. Skipping embedded attachments processing.",
+                original.__class__.__name__,
+            )
+            return
+
+        attachments = []
+
+        for attribute in self.attachment_attributes:
+            if getattr(original, attribute, None):
+                results = self.extractor.extract_and_replace(
+                    data=getattr(original, attribute),
+                    entity_type=entity_details.entity_type,
+                    entity_id=entity_details.entity_id,
+                    project_name=entity_details.project_name,
+                    context=cast(Literal["input", "output", "metadata"], attribute),
+                )
+                attachments.extend(results)
+
+        if len(attachments) > 0:
+            LOGGER.debug(
+                "Extracted %d attachments from %s (entity: %s/%s)",
+                len(attachments),
+                original.__class__.__name__,
+                entity_details.entity_type,
+                entity_details.entity_id,
+            )
+
+            self._process_attachments(attachments)
+        else:
+            LOGGER.debug(
+                "No attachments found in the message - %s.", original.__class__.__name__
+            )
+
+    def _process_attachments(
+        self, attachments: List[attachment_context.AttachmentWithContext]
+    ) -> None:
+        for attachment in attachments:
+            create_attachment_message = converters.attachment_to_message(
+                attachment_data=attachment.attachment_data,
+                entity_type=attachment.entity_type,
+                entity_id=attachment.entity_id,
+                project_name=attachment.project_name,
+                url_override=self._url_override,
+                delete_after_upload=True,  # make sure to delete attachments after upload to avoid leaking space and data
+            )
+            self.messages_streamer.put(create_attachment_message)
+
+
+def entity_type_from_attachment_message(
+    message: messages.BaseMessage,
+) -> Optional[EntityDetails]:
+    if isinstance(message, (messages.CreateSpanMessage, messages.UpdateSpanMessage)):
+        return EntityDetails("span", message.span_id, project_name=message.project_name)
+    elif isinstance(
+        message, (messages.CreateTraceMessage, messages.UpdateTraceMessage)
+    ):
+        return EntityDetails(
+            "trace", message.trace_id, project_name=message.project_name
+        )
+    else:
+        return None

opik/message_processing/processors/message_processors.py

@@ -0,0 +1,92 @@
+import abc
+import logging
+from typing import List, Optional, TypeVar, Type
+
+from .. import messages
+
+import opik.exceptions
+
+
+LOGGER = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+
+class BaseMessageProcessor(abc.ABC):
+    @abc.abstractmethod
+    def process(self, message: messages.BaseMessage) -> None:
+        pass
+
+    @abc.abstractmethod
+    def is_active(self) -> bool:
+        return False
+
+
+class ChainedMessageProcessor(BaseMessageProcessor):
+    """
+    Processes messages through a chain of message processors.
+
+    This class allows for the sequential processing of a message by a list of
+    `BaseMessageProcessor` instances. Each processor in the chain is invoked in the
+    order provided. If an exception occurs during the processing by a specific
+    processor, it is logged, and the process continues with the next processor in
+    the chain.
+    """
+
+    def __init__(self, processors: List[BaseMessageProcessor]) -> None:
+        self._processors = processors
+
+    def is_active(self) -> bool:
+        return True
+
+    def process(self, message: messages.BaseMessage) -> None:
+        rate_limit_error: Optional[opik.exceptions.OpikCloudRequestsRateLimited] = None
+
+        for processor in self._processors:
+            try:
+                processor.process(message)
+            except opik.exceptions.OpikCloudRequestsRateLimited as ex:
+                rate_limit_error = ex
+            except Exception as ex:
+                LOGGER.error(
+                    "Unexpected error while processing message: %s with message processor: %s",
+                    ex,
+                    type(processor),
+                    exc_info=True,
+                )
+
+        # Rate limit error is a special case that is handled by the caller.
+        if rate_limit_error is not None:
+            raise rate_limit_error
+
+    def get_processor_by_type(self, processor_type: Type[T]) -> Optional[T]:
+        """
+        Retrieves a processor from the available processors that matches the specified type.
+
+        This method iterates through the list of processors and checks if any of them is
+        an instance of the given class type. If a match is found, it returns the processor.
+
+        Args:
+            processor_type: Concrete processor class to search for.
+
+        Returns:
+            The processor matching the specified type if found, else None.
+        """
+        for processor in self._processors:
+            if isinstance(processor, processor_type):
+                return processor
+        return None
+
+    def add_first(self, processor: BaseMessageProcessor) -> None:
+        """
+        Inserts a processor at the first position in the list of processors.
+
+        This method allows prioritizing a given processor by placing it
+        at the beginning of the internal processor list. As a result, the
+        provided processor will be executed before others.
+
+        Args:
+            processor: The message processor to be
+                added to the front of the processor list.
+        """
+        self._processors.insert(0, processor)