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

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)

opik/message_processing/processors/message_processors_chain.py

@@ -0,0 +1,96 @@
+import logging
+from typing import Optional
+
+from opik.rest_api import client as rest_api_client
+
+from . import (
+    message_processors,
+    online_message_processor,
+)
+from ..emulation import local_emulator_message_processor
+
+
+LOGGER = logging.getLogger(__name__)
+
+
+def create_message_processors_chain(
+    rest_client: rest_api_client.OpikApi,
+) -> message_processors.ChainedMessageProcessor:
+    """
+    Creates a chain of message processors by combining an online processor and a
+    local emulator processor. The chain is primarily useful for processing messages
+    in a sequence where each processor in the chain contributes its functionality.
+
+    The online processor is initialized using the provided REST API client. The local
+    emulator processor is included but remains inactive by default. The constructed
+    chain ensures combined and streamlined processing, accommodating both online
+    and local simulation needs based on evaluation activation.
+
+    Args:
+        rest_client: REST API client instance used to configure the online message
+            processor.
+
+    Returns:
+        A chained message processor containing the online and local emulator processors.
+    """
+    online = online_message_processor.OpikMessageProcessor(rest_client=rest_client)
+    # is not active by default - will be activated during evaluation
+    local = local_emulator_message_processor.LocalEmulatorMessageProcessor(active=False)
+
+    return message_processors.ChainedMessageProcessor(processors=[online, local])
+
+
+def toggle_local_emulator_message_processor(
+    active: bool, chain: message_processors.ChainedMessageProcessor, reset: bool = True
+) -> None:
+    """
+    Toggles the state of the Local Emulator Message Processor within a given
+    ChainedMessageProcessor. This function either activates or deactivates the
+    processor based on the `active` parameter and resets its state if being
+    activated. Logs a warning if the Local Emulator Message Processor is not
+    found in the chain.
+
+    Args:
+        active: Determines whether to activate or deactivate the Local
+            Emulator Message Processor. If True, the processor is activated.
+        chain: The message processor
+            chain containing the Local Emulator Message Processor to be toggled.
+        reset: Determines whether to reset the Local Emulator Message Processor.
+            This can be used to clear the state of the Local Emulator before
+            evaluation. Also, it can be used to clean up the state of the Local Emulator
+            after evaluation to release system resources (memory).
+    """
+    local = chain.get_processor_by_type(
+        local_emulator_message_processor.LocalEmulatorMessageProcessor
+    )
+    if local is None:
+        LOGGER.warning("Local emulator message processor not found in the chain.")
+        return
+
+    if reset:
+        local.reset()
+
+    local.set_active(active=active)
+
+
+def get_local_emulator_message_processor(
+    chain: message_processors.ChainedMessageProcessor,
+) -> Optional[local_emulator_message_processor.LocalEmulatorMessageProcessor]:
+    """
+    Retrieves the local emulator message processor from a given chain of message processors.
+
+    This function searches through the provided chain and looks for a processor of type
+    LocalEmulatorMessageProcessor. If one is found, it is returned; otherwise, None is returned.
+
+    Args:
+        chain: A chain of message processors that may contain a
+            LocalEmulatorMessageProcessor.
+
+    Returns:
+        The LocalEmulatorMessageProcessor if found in the chain,
+        otherwise None.
+    """
+    local = chain.get_processor_by_type(
+        local_emulator_message_processor.LocalEmulatorMessageProcessor
+    )
+    return local

opik/message_processing/{message_processors.py → processors/online_message_processor.py}

@@ -1,22 +1,21 @@
-import abc
 import logging
-from typing import Any, Callable, Dict, Type
+from typing import Callable, Dict, Type, Any
+
 import pydantic
 import tenacity
 
-import opik.logging_messages as logging_messages
-import opik.exceptions as exceptions
-from . import messages
-from ..jsonable_encoder import encode
-from .. import dict_utils
-from ..rate_limit import rate_limit
-from ..rest_api.types import (
+from opik import dict_utils, exceptions, logging_messages
+from opik.rate_limit import rate_limit
+from opik.rest_api import client as rest_api_client, core as rest_api_core
+from opik.rest_api.types import (
     feedback_score_batch_item,
-    guardrail,
     feedback_score_batch_item_thread,
+    guardrail,
+    experiment_item,
 )
-from ..rest_api import core as rest_api_core
-from ..rest_api import client as rest_api_client
+
+from . import message_processors
+from .. import encoder_helpers, messages
 
 LOGGER = logging.getLogger(__name__)
 
@@ -24,21 +23,16 @@ LOGGER = logging.getLogger(__name__)
 MessageProcessingHandler = Callable[[messages.BaseMessage], None]
 
 
-class BaseMessageProcessor(abc.ABC):
-    @abc.abstractmethod
-    def process(
-        self,
-        message: messages.BaseMessage,
-    ) -> None:
-        pass
-
-
-class OpikMessageProcessor(BaseMessageProcessor):
+class OpikMessageProcessor(message_processors.BaseMessageProcessor):
     def __init__(
-        self, rest_client: rest_api_client.OpikApi, batch_memory_limit_mb: int = 50
+        self,
+        rest_client: rest_api_client.OpikApi,
+        batch_memory_limit_mb: int = 50,
+        active: bool = True,
     ):
         self._rest_client = rest_client
         self._batch_memory_limit_mb = batch_memory_limit_mb
+        self._is_active = active
 
         self._handlers: Dict[Type, MessageProcessingHandler] = {
             messages.CreateSpanMessage: self._process_create_span_message,  # type: ignore
@@ -51,9 +45,17 @@ class OpikMessageProcessor(BaseMessageProcessor):
             messages.CreateSpansBatchMessage: self._process_create_spans_batch_message,  # type: ignore
             messages.CreateTraceBatchMessage: self._process_create_traces_batch_message,  # type: ignore
             messages.GuardrailBatchMessage: self._process_guardrail_batch_message,  # type: ignore
+            messages.CreateExperimentItemsBatchMessage: self._process_create_experiment_items_batch_message,  # type: ignore
+            messages.AttachmentSupportingMessage: self._noop_handler,  # type: ignore
         }
 
+    def is_active(self) -> bool:
+        return self._is_active
+
     def process(self, message: messages.BaseMessage) -> None:
+        if not self.is_active():
+            return
+
         message_type = type(message)
         handler = self._handlers.get(message_type)
         if handler is None:
@@ -92,6 +94,7 @@ class OpikMessageProcessor(BaseMessageProcessor):
                 f"{cause.__class__.__name__} - {cause}",
                 extra={"error_tracking_extra": error_tracking_extra},
             )
+            LOGGER.warning(logging_messages.MAKE_SURE_OPIK_IS_CONFIGURED_CORRECTLY)
         except pydantic.ValidationError as validation_error:
             error_tracking_extra = _generate_error_tracking_extra(
                 validation_error, message
@@ -112,6 +115,7 @@ class OpikMessageProcessor(BaseMessageProcessor):
                 exc_info=True,
                 extra={"error_tracking_extra": error_tracking_extra},
             )
+            LOGGER.warning(logging_messages.MAKE_SURE_OPIK_IS_CONFIGURED_CORRECTLY)
 
     def _process_create_span_message(
         self,
@@ -121,7 +125,12 @@ class OpikMessageProcessor(BaseMessageProcessor):
         cleaned_create_span_kwargs = dict_utils.remove_none_from_dict(
             create_span_kwargs
         )
-        cleaned_create_span_kwargs = encode(cleaned_create_span_kwargs)
+        cleaned_create_span_kwargs = encoder_helpers.encode_and_anonymize(
+            cleaned_create_span_kwargs,
+            fields_to_anonymize=message.fields_to_anonymize(),
+            object_type="span",
+        )
+
         LOGGER.debug("Create span request: %s", cleaned_create_span_kwargs)
         self._rest_client.spans.create_span(**cleaned_create_span_kwargs)
 
@@ -133,7 +142,12 @@ class OpikMessageProcessor(BaseMessageProcessor):
         cleaned_create_trace_kwargs = dict_utils.remove_none_from_dict(
             create_trace_kwargs
         )
-        cleaned_create_trace_kwargs = encode(cleaned_create_trace_kwargs)
+        cleaned_create_trace_kwargs = encoder_helpers.encode_and_anonymize(
+            cleaned_create_trace_kwargs,
+            fields_to_anonymize=message.fields_to_anonymize(),
+            object_type="trace",
+        )
+
         LOGGER.debug("Create trace request: %s", cleaned_create_trace_kwargs)
         self._rest_client.traces.create_trace(**cleaned_create_trace_kwargs)
 
@@ -146,7 +160,12 @@ class OpikMessageProcessor(BaseMessageProcessor):
         cleaned_update_span_kwargs = dict_utils.remove_none_from_dict(
             update_span_kwargs
         )
-        cleaned_update_span_kwargs = encode(cleaned_update_span_kwargs)
+        cleaned_update_span_kwargs = encoder_helpers.encode_and_anonymize(
+            cleaned_update_span_kwargs,
+            fields_to_anonymize=message.fields_to_anonymize(),
+            object_type="span",
+        )
+
         LOGGER.debug("Update span request: %s", cleaned_update_span_kwargs)
         self._rest_client.spans.update_span(**cleaned_update_span_kwargs)
 
@@ -159,7 +178,12 @@ class OpikMessageProcessor(BaseMessageProcessor):
         cleaned_update_trace_kwargs = dict_utils.remove_none_from_dict(
             update_trace_kwargs
         )
-        cleaned_update_trace_kwargs = encode(cleaned_update_trace_kwargs)
+        cleaned_update_trace_kwargs = encoder_helpers.encode_and_anonymize(
+            cleaned_update_trace_kwargs,
+            fields_to_anonymize=message.fields_to_anonymize(),
+            object_type="trace",
+        )
+
         LOGGER.debug("Update trace request: %s", cleaned_update_trace_kwargs)
         self._rest_client.traces.update_trace(**cleaned_update_trace_kwargs)
         LOGGER.debug("Sent trace %s", message.trace_id)
@@ -253,6 +277,35 @@ class OpikMessageProcessor(BaseMessageProcessor):
 
         self._rest_client.guardrails.create_guardrails(guardrails=batch)
 
+    def _process_create_experiment_items_batch_message(
+        self,
+        message: messages.CreateExperimentItemsBatchMessage,
+    ) -> None:
+        experiment_items_batch = [
+            experiment_item.ExperimentItem(
+                id=item.id,
+                experiment_id=item.experiment_id,
+                dataset_item_id=item.dataset_item_id,
+                trace_id=item.trace_id,
+            )
+            for item in message.batch
+        ]
+
+        LOGGER.debug(
+            "Create experiment items batch request of size %d",
+            len(experiment_items_batch),
+        )
+        self._rest_client.experiments.create_experiment_items(
+            experiment_items=experiment_items_batch
+        )
+        LOGGER.debug(
+            "Sent experiment items batch of size %d", len(experiment_items_batch)
+        )
+
+    def _noop_handler(self, message: messages.BaseMessage) -> None:
+        # just ignore the message
+        pass
+
 
 def _generate_error_tracking_extra(
     exception: Exception, message: messages.BaseMessage
@@ -260,8 +313,11 @@ def _generate_error_tracking_extra(
     result: Dict[str, Any] = {"exception": exception}
 
     if isinstance(exception, rest_api_core.ApiError):
-        fingerprint = [type(message).__name__, type(exception).__name__]
-        fingerprint.append(str(exception.status_code))
+        fingerprint = [
+            type(message).__name__,
+            type(exception).__name__,
+            str(exception.status_code),
+        ]
         result["fingerprint"] = fingerprint
         result["status_code"] = exception.status_code
 

opik/message_processing/queue_consumer.py

@@ -4,7 +4,8 @@ import time
 from queue import Empty
 from typing import Optional
 
-from . import message_processors, message_queue, messages
+from . import message_queue, messages
+from .processors import message_processors
 from .. import exceptions, _logging
 
 LOGGER = logging.getLogger(__name__)
@@ -49,8 +50,9 @@ class QueueConsumer(threading.Thread):
 
             if message is None:
                 return
-            elif message.delivery_time <= now:
-                self._message_processor.process(message)
+
+            if message.delivery_time <= now:
+                self._process_message(message)
             else:
                 # put a message back to keep an order in the queue
                 self._push_message_back(message)
@@ -87,4 +89,8 @@ class QueueConsumer(threading.Thread):
                 "The message queue size limit has been reached. The current message has been returned to the queue, and the newest message has been discarded.",
                 logger=LOGGER,
             )
+        message.delivery_attempts += 1
         self._message_queue.put_back(message)
+
+    def _process_message(self, message: messages.BaseMessage) -> None:
+        self._message_processor.process(message)