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

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/processors/online_message_processor.py

@@ -0,0 +1,324 @@
+import logging
+from typing import Callable, Dict, Type, Any
+
+import pydantic
+import tenacity
+
+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,
+    feedback_score_batch_item_thread,
+    guardrail,
+    experiment_item,
+)
+
+from . import message_processors
+from .. import encoder_helpers, messages
+
+LOGGER = logging.getLogger(__name__)
+
+
+MessageProcessingHandler = Callable[[messages.BaseMessage], None]
+
+
+class OpikMessageProcessor(message_processors.BaseMessageProcessor):
+    def __init__(
+        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
+            messages.CreateTraceMessage: self._process_create_trace_message,  # type: ignore
+            messages.UpdateSpanMessage: self._process_update_span_message,  # type: ignore
+            messages.UpdateTraceMessage: self._process_update_trace_message,  # type: ignore
+            messages.AddTraceFeedbackScoresBatchMessage: self._process_add_trace_feedback_scores_batch_message,  # type: ignore
+            messages.AddSpanFeedbackScoresBatchMessage: self._process_add_span_feedback_scores_batch_message,  # type: ignore
+            messages.AddThreadsFeedbackScoresBatchMessage: self._process_add_threads_feedback_scores_batch_message,  # type: ignore
+            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:
+            LOGGER.debug("Unknown type of message - %s", message_type.__name__)
+            return
+
+        try:
+            handler(message)
+        except rest_api_core.ApiError as exception:
+            if exception.status_code == 409:
+                # sometimes a retry mechanism works in a way that it sends the same request 2 times.
+                # if the backend rejects the second request, we don't want users to see an error.
+                return
+            elif exception.status_code == 429:
+                if exception.headers is not None:
+                    rate_limiter = rate_limit.parse_rate_limit(exception.headers)
+                    if rate_limiter is not None:
+                        raise exceptions.OpikCloudRequestsRateLimited(
+                            headers=exception.headers,
+                            retry_after=rate_limiter.retry_after(),
+                        )
+
+            error_tracking_extra = _generate_error_tracking_extra(exception, message)
+            LOGGER.error(
+                logging_messages.FAILED_TO_PROCESS_MESSAGE_IN_BACKGROUND_STREAMER,
+                message_type.__name__,
+                str(exception),
+                extra={"error_tracking_extra": error_tracking_extra},
+            )
+        except tenacity.RetryError as retry_error:
+            cause = retry_error.last_attempt.exception()
+            error_tracking_extra = _generate_error_tracking_extra(cause, message)
+            LOGGER.error(
+                logging_messages.FAILED_TO_PROCESS_MESSAGE_IN_BACKGROUND_STREAMER,
+                message_type.__name__,
+                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
+            )
+            LOGGER.error(
+                "Failed to process message: '%s' due to input data validation error:\n%s\n",
+                message_type.__name__,
+                validation_error,
+                exc_info=True,
+                extra={"error_tracking_extra": error_tracking_extra},
+            )
+        except Exception as exception:
+            error_tracking_extra = _generate_error_tracking_extra(exception, message)
+            LOGGER.error(
+                logging_messages.FAILED_TO_PROCESS_MESSAGE_IN_BACKGROUND_STREAMER,
+                message_type.__name__,
+                str(exception),
+                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,
+        message: messages.CreateSpanMessage,
+    ) -> None:
+        create_span_kwargs = message.as_payload_dict()
+        cleaned_create_span_kwargs = dict_utils.remove_none_from_dict(
+            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)
+
+    def _process_create_trace_message(
+        self,
+        message: messages.CreateTraceMessage,
+    ) -> None:
+        create_trace_kwargs = message.as_payload_dict()
+        cleaned_create_trace_kwargs = dict_utils.remove_none_from_dict(
+            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)
+
+    def _process_update_span_message(
+        self,
+        message: messages.UpdateSpanMessage,
+    ) -> None:
+        update_span_kwargs = message.as_payload_dict()
+
+        cleaned_update_span_kwargs = dict_utils.remove_none_from_dict(
+            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)
+
+    def _process_update_trace_message(
+        self,
+        message: messages.UpdateTraceMessage,
+    ) -> None:
+        update_trace_kwargs = message.as_payload_dict()
+
+        cleaned_update_trace_kwargs = dict_utils.remove_none_from_dict(
+            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)
+
+    def _process_add_span_feedback_scores_batch_message(
+        self,
+        message: messages.AddSpanFeedbackScoresBatchMessage,
+    ) -> None:
+        scores = [
+            feedback_score_batch_item.FeedbackScoreBatchItem(**score_message.__dict__)
+            for score_message in message.batch
+        ]
+
+        LOGGER.debug("Add spans feedbacks scores request of size: %d", len(scores))
+
+        self._rest_client.spans.score_batch_of_spans(
+            scores=scores,
+        )
+        LOGGER.debug("Sent batch of spans feedback scores %d", len(scores))
+
+    def _process_add_trace_feedback_scores_batch_message(
+        self,
+        message: messages.AddTraceFeedbackScoresBatchMessage,
+    ) -> None:
+        scores = [
+            feedback_score_batch_item.FeedbackScoreBatchItem(**score_message.__dict__)
+            for score_message in message.batch
+        ]
+
+        LOGGER.debug("Add traces feedbacks scores request: %d", len(scores))
+
+        self._rest_client.traces.score_batch_of_traces(
+            scores=scores,
+        )
+        LOGGER.debug("Sent batch of traces feedbacks scores of size %d", len(scores))
+
+    def _process_add_threads_feedback_scores_batch_message(
+        self,
+        message: messages.AddThreadsFeedbackScoresBatchMessage,
+    ) -> None:
+        scores = [
+            feedback_score_batch_item_thread.FeedbackScoreBatchItemThread(
+                **score_message.as_payload_dict()
+            )
+            for score_message in message.batch
+        ]
+
+        try:
+            LOGGER.debug("Add threads feedbacks scores request of size %d", len(scores))
+            self._rest_client.traces.score_batch_of_threads(
+                scores=scores,
+            )
+            LOGGER.debug(
+                "Sent batch of threads feedbacks scores of size %d", len(scores)
+            )
+        except rest_api_core.ApiError as exception:
+            # In the case of AddThreadsFeedbackScoresBatchMessage, the backend will reject the request
+            # if thread is not closed which can happen if the user is unaware of this fact.
+            # Thus, we display the warning message.
+            if exception.status_code == 409:
+                LOGGER.warning(
+                    "Threads feedbacks scores batch was rejected by the backend, reason: '%s'",
+                    exception.body,
+                )
+            # propagate further to be handled in a unified error handler
+            raise exception
+
+    def _process_create_spans_batch_message(
+        self, message: messages.CreateSpansBatchMessage
+    ) -> None:
+        LOGGER.debug("Create spans batch request of size %d", len(message.batch))
+        self._rest_client.spans.create_spans(spans=message.batch)
+        LOGGER.debug("Sent spans batch of size %d", len(message.batch))
+
+    def _process_create_traces_batch_message(
+        self, message: messages.CreateTraceBatchMessage
+    ) -> None:
+        LOGGER.debug("Create trace batch request of size %d", len(message.batch))
+        self._rest_client.traces.create_traces(traces=message.batch)
+        LOGGER.debug("Sent trace batch of size %d", len(message.batch))
+
+    def _process_guardrail_batch_message(
+        self,
+        message: messages.GuardrailBatchMessage,
+    ) -> None:
+        batch = []
+
+        for message_item in message.batch:
+            guardrail_batch_item_message = guardrail.Guardrail(**message_item.__dict__)
+            batch.append(guardrail_batch_item_message)
+
+        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
+) -> Dict[str, Any]:
+    result: Dict[str, Any] = {"exception": exception}
+
+    if isinstance(exception, rest_api_core.ApiError):
+        fingerprint = [
+            type(message).__name__,
+            type(exception).__name__,
+            str(exception.status_code),
+        ]
+        result["fingerprint"] = fingerprint
+        result["status_code"] = exception.status_code
+
+    return result

opik/message_processing/queue_consumer.py

@@ -1,9 +1,15 @@
-import time
-import queue
+import logging
 import threading
-from typing import Any, Optional
+import time
+from queue import Empty
+from typing import Optional
+
+from . import message_queue, messages
+from .processors import message_processors
+from .. import exceptions, _logging
+
+LOGGER = logging.getLogger(__name__)
 
-from . import message_processors
 
 SLEEP_BETWEEN_LOOP_ITERATIONS = 0.1
 
@@ -11,15 +17,16 @@ SLEEP_BETWEEN_LOOP_ITERATIONS = 0.1
 class QueueConsumer(threading.Thread):
     def __init__(
         self,
-        message_queue: "queue.Queue[Any]",
+        queue: message_queue.MessageQueue[messages.BaseMessage],
         message_processor: message_processors.BaseMessageProcessor,
         name: Optional[str] = None,
     ):
         super().__init__(daemon=True, name=name)
-        self._message_queue = message_queue
+        self._message_queue = queue
         self._message_processor = message_processor
         self._processing_stopped = False
-        self.waiting = True
+        self.idling = True
+        self.next_message_time = time.monotonic()
 
     def run(self) -> None:
         while self._processing_stopped is False:
@@ -28,21 +35,62 @@ class QueueConsumer(threading.Thread):
         return
 
     def _loop(self) -> None:
+        now = time.monotonic()
+        if now < self.next_message_time:
+            # mark as not idling because we still have work to do
+            self.idling = False
+            time.sleep(SLEEP_BETWEEN_LOOP_ITERATIONS)
+            return
+
+        message = None
         try:
-            self.waiting = True
+            self.idling = True
             message = self._message_queue.get(timeout=SLEEP_BETWEEN_LOOP_ITERATIONS)
-            self.waiting = False
+            self.idling = False
 
             if message is None:
                 return
 
-            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)
 
-        except queue.Empty:
+        except Empty:
             time.sleep(SLEEP_BETWEEN_LOOP_ITERATIONS)
-        except Exception:
-            # TODO
+        except exceptions.OpikCloudRequestsRateLimited as limit_exception:
+            LOGGER.info(
+                "Ingestion rate limited, retrying in %s seconds, remaining queue size: %d, details: %s",
+                limit_exception.retry_after,
+                len(self._message_queue)
+                + 1,  # add 1 to account for the current message
+                limit_exception.headers,
+            )
+            # set the next iteration time to avoid rate limiting
+            self.next_message_time = now + limit_exception.retry_after
+            if message is not None:
+                message.delivery_time = self.next_message_time
+                # put a message back to keep an order in the queue
+                self._push_message_back(message)
+        except Exception as ex:
+            LOGGER.error(
+                "Failed to process message, unexpected error: %s", ex, exc_info=ex
+            )
             pass
 
     def close(self) -> None:
         self._processing_stopped = True
+
+    def _push_message_back(self, message: messages.BaseMessage) -> None:
+        if self._message_queue.accept_put_without_discarding() is False:
+            _logging.log_once_at_level(
+                logging.WARNING,
+                "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)