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

opik/message_processing/streamer.py

@@ -4,10 +4,14 @@ import time
 from typing import List, Optional
 
 from . import messages, message_queue, queue_consumer
-from .. import synchronization
-from .batching import batch_manager
-from ..file_upload import base_upload_manager
 from .. import _logging
+from .. import synchronization
+from .preprocessing import (
+    attachments_preprocessor,
+    batching_preprocessor,
+    file_upload_preprocessor,
+)
+
 
 LOGGER = logging.getLogger(__name__)
 
@@ -17,52 +21,74 @@ class Streamer:
         self,
         queue: message_queue.MessageQueue[messages.BaseMessage],
         queue_consumers: List[queue_consumer.QueueConsumer],
-        batch_manager: Optional[batch_manager.BatchManager],
-        file_upload_manager: base_upload_manager.BaseFileUploadManager,
+        attachments_preprocessor: attachments_preprocessor.AttachmentsPreprocessor,
+        batch_preprocessor: batching_preprocessor.BatchingPreprocessor,
+        upload_preprocessor: file_upload_preprocessor.FileUploadPreprocessor,
     ) -> None:
         self._lock = threading.RLock()
         self._message_queue = queue
         self._queue_consumers = queue_consumers
-        self._batch_manager = batch_manager
-        self._file_upload_manager = file_upload_manager
+        self._attachments_preprocessor = attachments_preprocessor
+        self._batch_preprocessor = batch_preprocessor
+        self._upload_preprocessor = upload_preprocessor
 
         self._drain = False
 
-        self._start_queue_consumers()
+        self._idle = True
 
-        if self._batch_manager is not None:
-            self._batch_manager.start()
+        self._start_queue_consumers()
+        self._batch_preprocessor.start()
 
     def put(self, message: messages.BaseMessage) -> None:
         with self._lock:
             if self._drain:
                 return
 
-            if (
-                self._batch_manager is not None
-                and self._batch_manager.message_supports_batching(message)
-            ):
-                self._batch_manager.process_message(message)
-            elif base_upload_manager.message_supports_upload(message):
-                self._file_upload_manager.upload(message)
-            else:
-                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 new message has been added to the queue, and the oldest message has been discarded.",
-                        logger=LOGGER,
-                    )
-                self._message_queue.put(message)
+            self._idle = False
+            try:
+                # do embedded attachments pre-processing first (MUST ALWAYS BE DONE FIRST)
+                preprocessed_message = self._attachments_preprocessor.preprocess(
+                    message
+                )
+
+                # do file uploads pre-processing second
+                preprocessed_message = self._upload_preprocessor.preprocess(
+                    preprocessed_message
+                )
+
+                # do batching pre-processing third
+                preprocessed_message = self._batch_preprocessor.preprocess(
+                    preprocessed_message
+                )
+
+                # work with resulting message if not fully consumed by preprocessors
+                if preprocessed_message is not 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 new message has been added to the queue, and the oldest message has been discarded.",
+                            logger=LOGGER,
+                        )
+                    self._message_queue.put(preprocessed_message)
+            except Exception as ex:
+                LOGGER.error(
+                    "Failed to process message by streamer: %s", ex, exc_info=ex
+                )
+            self._idle = True
 
     def close(self, timeout: Optional[int]) -> bool:
         """
-        Stops data sending threads
+        Stops data processing threads
         """
         with self._lock:
+            synchronization.wait_for_done(
+                check_function=lambda: self._idle,
+                timeout=timeout,
+                sleep_time=0.1,
+            )
             self._drain = True
 
-        if self._batch_manager is not None:
-            self._batch_manager.stop()  # stopping causes adding remaining batch messages to the queue
+        self._batch_preprocessor.stop()  # stopping causes adding remaining batch messages to the queue
 
         self.flush(timeout)
         self._close_queue_consumers()
@@ -70,8 +96,17 @@ class Streamer:
         return self._message_queue.empty()
 
     def flush(self, timeout: Optional[float], upload_sleep_time: int = 5) -> bool:
-        if self._batch_manager is not None:
-            self._batch_manager.flush()
+        # wait for current pending messages processing to be completed
+        # this should be done before flushing batch preprocessor because some
+        # batch messages may be added to the queue during processing
+        with self._lock:
+            synchronization.wait_for_done(
+                check_function=lambda: self._idle,
+                timeout=timeout,
+                sleep_time=0.1,
+            )
+
+        self._batch_preprocessor.flush()
 
         start_time = time.time()
 
@@ -88,17 +123,20 @@ class Streamer:
                 timeout = 1.0
 
         # flushing upload manager is blocking operation
-        upload_flushed = self._file_upload_manager.flush(
+        upload_flushed = self._upload_preprocessor.flush(
             timeout=timeout, sleep_time=upload_sleep_time
         )
 
-        return upload_flushed and self._all_done()
+        flushed = upload_flushed and self._all_done()
+        LOGGER.debug(f"Streamer flushed completely: {flushed}")
+
+        return flushed
 
     def _all_done(self) -> bool:
         return (
             self.workers_idling()
             and self._message_queue.empty()
-            and (self._batch_manager is None or self._batch_manager.is_empty())
+            and self._batch_preprocessor.is_empty()
         )
 
     def workers_idling(self) -> bool:

opik/message_processing/streamer_constructors.py

@@ -2,42 +2,72 @@ from typing import List, Optional
 
 import httpx
 
-from . import queue_consumer, messages, message_processors, message_queue, streamer
-from ..file_upload import upload_manager, base_upload_manager
-from ..rest_api import client as rest_api_client
+from . import (
+    queue_consumer,
+    messages,
+    message_queue,
+    streamer,
+)
 from .batching import batch_manager_constuctors
+from .preprocessing import (
+    attachments_preprocessor,
+    batching_preprocessor,
+    file_upload_preprocessor,
+)
+from .processors import attachments_extraction_processor, message_processors
+from ..file_upload import upload_manager
+from ..rest_api import client as rest_api_client
 
 
 def construct_online_streamer(
     rest_client: rest_api_client.OpikApi,
     httpx_client: httpx.Client,
     use_batching: bool,
+    use_attachment_extraction: bool,
+    min_base64_embedded_attachment_size: int,
     file_upload_worker_count: int,
     n_consumers: int,
     max_queue_size: int,
+    message_processor: message_processors.ChainedMessageProcessor,
+    url_override: str,
 ) -> streamer.Streamer:
-    message_processor = message_processors.OpikMessageProcessor(rest_client=rest_client)
-
     file_uploader = upload_manager.FileUploadManager(
         rest_client=rest_client,
         httpx_client=httpx_client,
         worker_count=file_upload_worker_count,
     )
 
-    return construct_streamer(
+    streamer = construct_streamer(
         message_processor=message_processor,
-        file_upload_manager=file_uploader,
+        upload_preprocessor=file_upload_preprocessor.FileUploadPreprocessor(
+            file_uploader
+        ),
         n_consumers=n_consumers,
         use_batching=use_batching,
+        use_attachment_extraction=use_attachment_extraction,
         max_queue_size=max_queue_size,
     )
 
+    # add attachment extraction processor to the beginning of the processing chain
+    attachment_extraction = (
+        attachments_extraction_processor.AttachmentsExtractionProcessor(
+            messages_streamer=streamer,
+            min_attachment_size=min_base64_embedded_attachment_size,
+            url_override=url_override,
+            is_active=use_attachment_extraction,
+        )
+    )
+    message_processor.add_first(attachment_extraction)
+
+    return streamer
+
 
 def construct_streamer(
     message_processor: message_processors.BaseMessageProcessor,
-    file_upload_manager: base_upload_manager.BaseFileUploadManager,
+    upload_preprocessor: file_upload_preprocessor.FileUploadPreprocessor,
     n_consumers: int,
     use_batching: bool,
+    use_attachment_extraction: bool,
     max_queue_size: Optional[int],
 ) -> streamer.Streamer:
     message_queue_: message_queue.MessageQueue[messages.BaseMessage] = (
@@ -62,8 +92,11 @@ def construct_streamer(
     streamer_ = streamer.Streamer(
         queue=message_queue_,
         queue_consumers=queue_consumers,
-        batch_manager=batch_manager,
-        file_upload_manager=file_upload_manager,
+        batch_preprocessor=batching_preprocessor.BatchingPreprocessor(batch_manager),
+        upload_preprocessor=upload_preprocessor,
+        attachments_preprocessor=attachments_preprocessor.AttachmentsPreprocessor(
+            use_attachment_extraction
+        ),
     )
 
     return streamer_

opik/opik_context.py

@@ -1,8 +1,8 @@
 import contextlib
-from typing import Any, Dict, List, Optional, Union, Iterator
+from typing import Any, Dict, List, Optional, Iterator, Union
 
-import opik.llm_usage as llm_usage
-from opik.api_objects import span, trace, opik_client
+from opik import llm_usage
+from opik.api_objects import span, trace, opik_client, prompt
 from opik.api_objects.attachment import Attachment
 from opik.types import (
     DistributedTraceHeadersDict,
@@ -11,7 +11,7 @@ from opik.types import (
     ErrorInfoDict,
 )
 
-import opik.decorator.tracing_runtime_config as tracing_runtime_config
+from opik import tracing_runtime_config
 
 from . import context_storage, exceptions
 from .decorator import error_info_collector
@@ -68,6 +68,7 @@ def update_current_span(
     total_cost: Optional[float] = None,
     attachments: Optional[List[Attachment]] = None,
     error_info: Optional[ErrorInfoDict] = None,
+    prompts: Optional[List[prompt.BasePrompt]] = None,
 ) -> None:
     """
     Update the current span with the provided parameters. This method is usually called within a tracked function.
@@ -90,10 +91,14 @@ def update_current_span(
         total_cost: The cost of the span in USD. This value takes priority over the cost calculated by Opik from the usage.
         attachments: The list of attachments to be uploaded to the span.
         error_info: The error information of the span.
+        prompts: The list of prompts used in the span.
     """
     if not tracing_runtime_config.is_tracing_active():
         return
 
+    if prompts is not None:
+        prompts = [p.__internal_api__to_info_dict__() for p in prompts]
+
     new_params = {
         "name": name,
         "input": input,
@@ -107,6 +112,7 @@ def update_current_span(
         "total_cost": total_cost,
         "attachments": attachments,
         "error_info": error_info,
+        "prompts": prompts,
     }
     current_span_data = context_storage.top_span_data()
     if current_span_data is None:
@@ -124,6 +130,7 @@ def update_current_trace(
     feedback_scores: Optional[List[FeedbackScoreDict]] = None,
     thread_id: Optional[str] = None,
     attachments: Optional[List[Attachment]] = None,
+    prompts: Optional[List[prompt.BasePrompt]] = None,
 ) -> None:
     """
     Update the current trace with the provided parameters. This method is usually called within a tracked function.
@@ -138,10 +145,14 @@ def update_current_trace(
         thread_id: Used to group multiple traces into a thread.
             The identifier is user-defined and has to be unique per project.
         attachments: The list of attachments to be uploaded to the trace.
+        prompts: The list of prompts used in the trace.
     """
     if not tracing_runtime_config.is_tracing_active():
         return
 
+    if prompts is not None:
+        prompts = [p.__internal_api__to_info_dict__() for p in prompts]
+
     new_params = {
         "name": name,
         "input": input,
@@ -151,6 +162,7 @@ def update_current_trace(
         "feedback_scores": feedback_scores,
         "thread_id": thread_id,
         "attachments": attachments,
+        "prompts": prompts,
     }
     current_trace_data = context_storage.get_trace_data()
     if current_trace_data is None:

opik/plugins/pytest/hooks.py

@@ -1,7 +1,7 @@
 import logging
 import opik._logging as _logging
 from typing import List, Any, Generator, TYPE_CHECKING
-from opik.types import FeedbackScoreDict
+from opik.types import BatchFeedbackScoreDict
 
 from opik.api_objects import opik_client
 from . import test_runs_storage, experiment_runner, summary
@@ -48,13 +48,15 @@ def pytest_sessionfinish(session: "pytest.Session", exitstatus: Any) -> None:
         return
 
     try:
-        traces_feedback_scores: List[FeedbackScoreDict] = []
+        traces_feedback_scores: List[BatchFeedbackScoreDict] = []
 
         for item in llm_test_items:
             report: "pytest.TestReport" = item.report
             trace_id = test_runs_storage.TEST_RUNS_TO_TRACE_DATA[item.nodeid].id
             traces_feedback_scores.append(
-                {"id": trace_id, "name": "Passed", "value": report.passed}
+                BatchFeedbackScoreDict(
+                    id=trace_id, name="Passed", value=float(report.passed)
+                )
             )
 
         client = opik_client.get_client_cached()