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

opik/message_processing/batching/batch_manager.py

@@ -1,3 +1,4 @@
+import threading
 from typing import Type, Dict
 from .. import messages
 from . import base_batcher
@@ -15,30 +16,41 @@ class BatchManager:
         self._flushing_thread = flushing_thread.FlushingThread(
             batchers=list(self._message_to_batcher_mapping.values())
         )
+        self._lock = threading.RLock()
 
     def start(self) -> None:
         self._flushing_thread.start()
 
     def stop(self) -> None:
-        self._flushing_thread.close()
+        with self._lock:
+            # stop the flushing thread
+            self._flushing_thread.close()
+            # force flush all pending messages
+            self.flush()
 
     def message_supports_batching(self, message: messages.BaseMessage) -> bool:
+        if message is None:
+            return False
+
         if hasattr(message, "supports_batching"):
             return message.supports_batching
 
         return message.__class__ in self._message_to_batcher_mapping
 
     def process_message(self, message: messages.BaseMessage) -> None:
-        self._message_to_batcher_mapping[type(message)].add(message)
+        with self._lock:
+            self._message_to_batcher_mapping[type(message)].add(message)
 
     def is_empty(self) -> bool:
-        return all(
-            [
-                batcher.is_empty()
-                for batcher in self._message_to_batcher_mapping.values()
-            ]
-        )
+        with self._lock:
+            return all(
+                [
+                    batcher.is_empty()
+                    for batcher in self._message_to_batcher_mapping.values()
+                ]
+            )
 
     def flush(self) -> None:
-        for batcher in self._message_to_batcher_mapping.values():
-            batcher.flush()
+        with self._lock:
+            for batcher in self._message_to_batcher_mapping.values():
+                batcher.flush()

opik/message_processing/batching/batchers.py

@@ -37,8 +37,9 @@ class CreateSpanMessageBatcher(base_batcher.BaseBatcher):
         return batches
 
     def add(self, message: messages.CreateSpanMessage) -> None:  # type: ignore
-        # remove any duplicate spans from the batch that was already added
-        self._remove_matching_messages(lambda x: x.span_id == message.span_id)  # type: ignore
+        # remove any duplicate start span message from the batch that was already added
+        if message.end_time is not None:
+            self._remove_matching_messages(lambda x: x.span_id == message.span_id)  # type: ignore
 
         return super().add(message)
 
@@ -73,8 +74,9 @@ class CreateTraceMessageBatcher(base_batcher.BaseBatcher):
         return batches
 
     def add(self, message: messages.CreateTraceMessage) -> None:  # type: ignore
-        # remove any duplicate traces from the batch that was already added
-        self._remove_matching_messages(lambda x: x.trace_id == message.trace_id)  # type: ignore
+        # remove any duplicate start trace message from the batch that was already added
+        if message.end_time is not None:
+            self._remove_matching_messages(lambda x: x.trace_id == message.trace_id)  # type: ignore
 
         return super().add(message)
 
@@ -99,26 +101,21 @@ class BaseAddFeedbackScoresBatchMessageBatcher(base_batcher.BaseBatcher):
             messages.AddThreadsFeedbackScoresBatchMessage,
         ],
     ) -> None:
-        with self._lock:
-            new_messages = message.batch
-            n_new_messages = len(new_messages)
-            n_accumulated_messages = len(self._accumulated_messages)
+        new_messages = message.batch
+        n_new_messages = len(new_messages)
+        n_accumulated_messages = len(self._accumulated_messages)
 
-            if n_new_messages + n_accumulated_messages >= self._max_batch_size:
-                free_space_in_accumulator = (
-                    self._max_batch_size - n_accumulated_messages
-                )
+        if n_new_messages + n_accumulated_messages >= self._max_batch_size:
+            free_space_in_accumulator = self._max_batch_size - n_accumulated_messages
 
-                messages_that_fit_in_batch = new_messages[:free_space_in_accumulator]
-                messages_that_dont_fit_in_batch = new_messages[
-                    free_space_in_accumulator:
-                ]
+            messages_that_fit_in_batch = new_messages[:free_space_in_accumulator]
+            messages_that_dont_fit_in_batch = new_messages[free_space_in_accumulator:]
 
-                self._accumulated_messages += messages_that_fit_in_batch
-                new_messages = messages_that_dont_fit_in_batch
-                self.flush()
+            self._accumulated_messages += messages_that_fit_in_batch
+            new_messages = messages_that_dont_fit_in_batch
+            self.flush()
 
-            self._accumulated_messages += new_messages
+        self._accumulated_messages += new_messages
 
 
 class AddSpanFeedbackScoresBatchMessageBatcher(
@@ -195,23 +192,18 @@ class CreateExperimentItemsBatchMessageBatcher(base_batcher.BaseBatcher):
     def add(  # type: ignore
         self, message: messages.CreateExperimentItemsBatchMessage
     ) -> None:
-        with self._lock:
-            new_messages = message.batch
-            n_new_messages = len(new_messages)
-            n_accumulated_messages = len(self._accumulated_messages)
-
-            if n_new_messages + n_accumulated_messages >= self._max_batch_size:
-                free_space_in_accumulator = (
-                    self._max_batch_size - n_accumulated_messages
-                )
-
-                messages_that_fit_in_batch = new_messages[:free_space_in_accumulator]
-                messages_that_dont_fit_in_batch = new_messages[
-                    free_space_in_accumulator:
-                ]
-
-                self._accumulated_messages += messages_that_fit_in_batch
-                new_messages = messages_that_dont_fit_in_batch
-                self.flush()
-
-            self._accumulated_messages += new_messages
+        new_messages = message.batch
+        n_new_messages = len(new_messages)
+        n_accumulated_messages = len(self._accumulated_messages)
+
+        if n_new_messages + n_accumulated_messages >= self._max_batch_size:
+            free_space_in_accumulator = self._max_batch_size - n_accumulated_messages
+
+            messages_that_fit_in_batch = new_messages[:free_space_in_accumulator]
+            messages_that_dont_fit_in_batch = new_messages[free_space_in_accumulator:]
+
+            self._accumulated_messages += messages_that_fit_in_batch
+            new_messages = messages_that_dont_fit_in_batch
+            self.flush()
+
+        self._accumulated_messages += new_messages

opik/message_processing/batching/flushing_thread.py

@@ -17,9 +17,6 @@ class FlushingThread(threading.Thread):
         self._closed = False
 
     def close(self) -> None:
-        for batcher in self._batchers:
-            batcher.flush()
-
         self._closed = True
 
     def run(self) -> None:

opik/message_processing/emulation/emulator_message_processor.py

@@ -9,7 +9,9 @@ from opik import dict_utils
 from opik.rest_api.types import span_write, trace_write
 from opik.types import ErrorInfoDict, SpanType
 from . import models
-from .. import message_processors, messages
+from .. import messages
+from ..processors import message_processors
+
 
 LOGGER = logging.getLogger(__name__)
 
@@ -77,6 +79,12 @@ class EmulatorMessageProcessor(message_processors.BaseMessageProcessor, abc.ABC)
             self._span_to_feedback_scores: Dict[
                 str, List[models.FeedbackScoreModel]
             ] = collections.defaultdict(list)
+            self._trace_to_attachments: Dict[str, List[models.AttachmentModel]] = (
+                collections.defaultdict(list)
+            )
+            self._span_to_attachments: Dict[str, List[models.AttachmentModel]] = (
+                collections.defaultdict(list)
+            )
             self._experiment_items: List[models.ExperimentItemModel] = []
 
     def is_active(self) -> bool:
@@ -111,6 +119,7 @@ class EmulatorMessageProcessor(message_processors.BaseMessageProcessor, abc.ABC)
 
             for trace in self._trace_trees:
                 trace.feedback_scores = self._trace_to_feedback_scores[trace.id]
+                trace.attachments = self._trace_to_attachments[trace.id] or None
 
             self._trace_trees.sort(key=lambda x: x.start_time)
             return self._trace_trees
@@ -176,6 +185,7 @@ class EmulatorMessageProcessor(message_processors.BaseMessageProcessor, abc.ABC)
             for span_id in all_span_ids:
                 span = self._span_observations[span_id]
                 span.feedback_scores = self._span_to_feedback_scores[span_id]
+                span.attachments = self._span_to_attachments[span_id] or None
 
             self._span_trees.sort(key=lambda x: x.start_time)
 
@@ -353,6 +363,8 @@ class EmulatorMessageProcessor(message_processors.BaseMessageProcessor, abc.ABC)
             messages.CreateSpansBatchMessage: self._handle_create_spans_batch_message,  # type: ignore
             messages.CreateTraceBatchMessage: self._handle_create_traces_batch_message,  # type: ignore
             messages.CreateExperimentItemsBatchMessage: self._handle_create_experiment_items_batch_message,  # type: ignore
+            messages.AttachmentSupportingMessage: self._noop_handler,  # type: ignore
+            messages.CreateAttachmentMessage: self._handle_create_attachment_message,  # type: ignore
         }
 
     def _handle_create_trace_message(
@@ -553,6 +565,29 @@ class EmulatorMessageProcessor(message_processors.BaseMessageProcessor, abc.ABC)
             )
             self._experiment_items.append(experiment_item)
 
+    def _handle_create_attachment_message(
+        self, message: messages.CreateAttachmentMessage
+    ) -> None:
+        """Handle attachment messages by adding them to the appropriate span or trace.
+
+        Attachments are stored in temporary dictionaries and will be connected to their
+        spans/traces when the trace trees are built, similar to how feedback scores work.
+        """
+        attachment_model = models.AttachmentModel(
+            file_path=message.file_path,
+            file_name=message.file_name,
+            content_type=message.mime_type,
+        )
+
+        if message.entity_type == "span":
+            self._span_to_attachments[message.entity_id].append(attachment_model)
+        elif message.entity_type == "trace":
+            self._trace_to_attachments[message.entity_id].append(attachment_model)
+
+    def _noop_handler(self, message: messages.BaseMessage) -> None:
+        # just ignore the message
+        pass
+
     @property
     def experiment_items(self) -> List[models.ExperimentItemModel]:
         """Returns the list of experiment items collected."""

opik/message_processing/emulation/models.py

@@ -30,6 +30,25 @@ class FeedbackScoreModel:
     reason: Optional[str] = None
 
 
+@dataclasses.dataclass
+class AttachmentModel:
+    """
+    Represents a model for an attachment associated with a span or trace.
+
+    This class stores metadata about files or data attached to spans or traces,
+    including the file path, name, and content type.
+
+    Attributes:
+        file_path: Path to the attached file.
+        file_name: Name of the attached file.
+        content_type: MIME type of the attached file.
+    """
+
+    file_path: str
+    file_name: str
+    content_type: Optional[str] = None
+
+
 @dataclasses.dataclass
 class SpanModel:
     """
@@ -86,6 +105,7 @@ class SpanModel:
     error_info: Optional[ErrorInfoDict] = None
     total_cost: Optional[float] = None
     last_updated_at: Optional[datetime.datetime] = None
+    attachments: Optional[List[AttachmentModel]] = None
 
 
 @dataclasses.dataclass
@@ -160,3 +180,4 @@ class TraceModel:
     error_info: Optional[ErrorInfoDict] = None
     thread_id: Optional[str] = None
     last_updated_at: Optional[datetime.datetime] = None
+    attachments: Optional[List[AttachmentModel]] = None

opik/message_processing/messages.py

@@ -4,6 +4,7 @@ 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
 
@@ -21,6 +22,8 @@ class BaseMessage:
             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
 
 
@@ -285,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)

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