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

opik/message_processing/streamer.py

@@ -1,11 +1,17 @@
-import queue
 import threading
 import logging
-from typing import Any, List, Optional
+import time
+from typing import List, Optional
 
-from . import messages, queue_consumer
+from . import messages, message_queue, queue_consumer
+from .. import _logging
 from .. import synchronization
-from .batching import batch_manager
+from .preprocessing import (
+    attachments_preprocessor,
+    batching_preprocessor,
+    file_upload_preprocessor,
+)
+
 
 LOGGER = logging.getLogger(__name__)
 
@@ -13,66 +19,131 @@ LOGGER = logging.getLogger(__name__)
 class Streamer:
     def __init__(
         self,
-        message_queue: "queue.Queue[Any]",
+        queue: message_queue.MessageQueue[messages.BaseMessage],
         queue_consumers: List[queue_consumer.QueueConsumer],
-        batch_manager: Optional[batch_manager.BatchManager],
+        attachments_preprocessor: attachments_preprocessor.AttachmentsPreprocessor,
+        batch_preprocessor: batching_preprocessor.BatchingPreprocessor,
+        upload_preprocessor: file_upload_preprocessor.FileUploadPreprocessor,
     ) -> None:
         self._lock = threading.RLock()
-        self._message_queue = message_queue
+        self._message_queue = queue
         self._queue_consumers = queue_consumers
-        self._batch_manager = batch_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)
-            else:
-                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()
 
         return self._message_queue.empty()
 
-    def flush(self, timeout: Optional[int]) -> None:
-        if self._batch_manager is not None:
-            self._batch_manager.flush()
+    def flush(self, timeout: Optional[float], upload_sleep_time: int = 5) -> bool:
+        # 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()
 
         synchronization.wait_for_done(
-            check_function=lambda: (
-                self.workers_waiting()
-                and self._message_queue.empty()
-                and (self._batch_manager is None or self._batch_manager.is_empty())
-            ),
+            check_function=lambda: self._all_done(),
             timeout=timeout,
             sleep_time=0.1,
         )
 
-    def workers_waiting(self) -> bool:
-        return all([consumer.waiting for consumer in self._queue_consumers])
+        elapsed_time = time.time() - start_time
+        if timeout is not None:
+            timeout = timeout - elapsed_time
+            if timeout < 0.0:
+                timeout = 1.0
+
+        # flushing upload manager is blocking operation
+        upload_flushed = self._upload_preprocessor.flush(
+            timeout=timeout, sleep_time=upload_sleep_time
+        )
+
+        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_preprocessor.is_empty()
+        )
+
+    def workers_idling(self) -> bool:
+        return all([consumer.idling for consumer in self._queue_consumers])
+
+    def queue_size(self) -> int:
+        return self._message_queue.size()
 
     def _start_queue_consumers(self) -> None:
         for consumer in self._queue_consumers:

opik/message_processing/streamer_constructors.py

@@ -1,31 +1,82 @@
-import queue
-from typing import Any, List
+from typing import List, Optional
 
-from . import queue_consumer, message_processors, streamer
-from ..rest_api import client as rest_api_client
+import httpx
+
+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,
-    n_consumers: int = 1,
+    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.MessageSender(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(message_processor, n_consumers, use_batching)
+    streamer = construct_streamer(
+        message_processor=message_processor,
+        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,
+    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: "queue.Queue[Any]" = queue.Queue()
+    message_queue_: message_queue.MessageQueue[messages.BaseMessage] = (
+        message_queue.MessageQueue(max_length=max_queue_size)
+    )
 
     queue_consumers: List[queue_consumer.QueueConsumer] = [
         queue_consumer.QueueConsumer(
-            message_queue=message_queue,
+            queue=message_queue_,
             message_processor=message_processor,
             name=f"QueueConsumerThread_{i}",
         )
@@ -33,15 +84,19 @@ def construct_streamer(
     ]
 
     batch_manager = (
-        batch_manager_constuctors.create_batch_manager(message_queue)
+        batch_manager_constuctors.create_batch_manager(message_queue_)
         if use_batching
         else None
     )
 
     streamer_ = streamer.Streamer(
-        message_queue=message_queue,
+        queue=message_queue_,
         queue_consumers=queue_consumers,
-        batch_manager=batch_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,15 +1,25 @@
-from typing import Any, Dict, List, Optional, Union
+import contextlib
+from typing import Any, Dict, List, Optional, Iterator, Union
 
 from opik import llm_usage
-from opik.api_objects import span, trace
-from opik.types import DistributedTraceHeadersDict, FeedbackScoreDict, LLMProvider
+from opik.api_objects import span, trace, opik_client, prompt
+from opik.api_objects.attachment import Attachment
+from opik.types import (
+    DistributedTraceHeadersDict,
+    FeedbackScoreDict,
+    LLMProvider,
+    ErrorInfoDict,
+)
+
+from opik import tracing_runtime_config
 
 from . import context_storage, exceptions
+from .decorator import error_info_collector
 
 
 def get_current_span_data() -> Optional[span.SpanData]:
     """
-    Returns current span created by track() decorator or None if no span was found.
+    Returns the current span created by track() decorator or None if no span was found.
     """
     span_data = context_storage.top_span_data()
     if span_data is None:
@@ -20,7 +30,7 @@ def get_current_span_data() -> Optional[span.SpanData]:
 
 def get_current_trace_data() -> Optional[trace.TraceData]:
     """
-    Returns current trace created by track() decorator or None if no trace was found.
+    Returns the current trace created by track() decorator or None if no trace was found.
     """
     trace_data = context_storage.get_trace_data()
     if trace_data is None:
@@ -31,13 +41,13 @@ def get_current_trace_data() -> Optional[trace.TraceData]:
 
 def get_distributed_trace_headers() -> DistributedTraceHeadersDict:
     """
-    Returns headers dictionary to be passed into tracked function on remote node.
+    Returns headers' dictionary to be passed into tracked function on remote node.
     Requires an existing span in the context, otherwise raises an error.
     """
     current_span_data = context_storage.top_span_data()
 
     if current_span_data is None:
-        raise Exception("There is no span in the context.")
+        raise exceptions.OpikException("There is no span in the context.")
 
     return DistributedTraceHeadersDict(
         opik_trace_id=current_span_data.trace_id,
@@ -56,6 +66,9 @@ def update_current_span(
     model: Optional[str] = None,
     provider: Optional[Union[str, LLMProvider]] = None,
     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.
@@ -66,17 +79,26 @@ def update_current_span(
         output: The output data of the span.
         metadata: The metadata of the span.
         tags: The tags of the span.
-        usage: Usage data for the span. In order for input, output and total tokens to be visible in the UI,
-            the usage must contain OpenAI-formatted keys (they can be passed additionaly to original usage on the top level of the dict):  prompt_tokens, completion_tokens and total_tokens.
+        usage: Usage data for the span. In order for input, output, and total tokens to be visible in the UI,
+            the usage must contain OpenAI-formatted keys (they can be passed additionally to the original usage on the top level of the dict): prompt_tokens, completion_tokens, and total_tokens.
             If OpenAI-formatted keys were not found, Opik will try to calculate them automatically if the usage
             format is recognized (you can see which provider's formats are recognized in opik.LLMProvider enum), but it is not guaranteed.
         feedback_scores: The feedback scores of the span.
         model: The name of LLM (in this case type parameter should be == llm)
         provider: The provider of LLM. You can find providers officially supported by Opik for cost tracking
-            in `opik.LLMProvider` enum. If your provider is not here, please open an issue in our github - https://github.com/comet-ml/opik.
-            If your provider not in the list, you can still specify it but the cost tracking will not be available
+            in `opik.LLMProvider` enum. If your provider is not here, please open an issue in our GitHub - https://github.com/comet-ml/opik.
+            If your provider is not in the list, you can still specify it, but the cost tracking will not be available
         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,
@@ -88,6 +110,9 @@ def update_current_span(
         "model": model,
         "provider": provider,
         "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:
@@ -104,6 +129,8 @@ def update_current_trace(
     tags: Optional[List[str]] = None,
     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.
@@ -117,7 +144,15 @@ def update_current_trace(
         feedback_scores: The feedback scores of the 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,
@@ -126,6 +161,8 @@ def update_current_trace(
         "tags": tags,
         "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:
@@ -134,10 +171,65 @@ def update_current_trace(
     current_trace_data.update(**new_params)
 
 
+@contextlib.contextmanager
+def trace_context(
+    trace_data: trace.TraceData,
+    client: opik_client.Opik,
+) -> Iterator[None]:
+    """
+    Provides a context manager to handle trace data within an execution context.
+
+    This function sets up trace data for the current context, ensuring it is
+    properly cleaned up and processed during the lifecycle of the context. It also
+    handles exceptions by collecting error information and associating it with
+    the trace data before raising the exception further. At the end of the context,
+    it finalizes the trace and logs it using the provided client.
+
+    Args:
+        trace_data: An instance of trace.TraceData containing information
+            about the current trace context, such as start time, end time, and
+            any relevant metadata for tracking execution.
+        client: An object of type opik_client.Opik used to report the trace
+            data, typically communicating with an external tracing or monitoring
+            system.
+
+    Yields:
+        None: The context manager yields control back to the caller, allowing
+            code execution within the defined trace context.
+
+    Raises:
+        Exception: The function raises any exceptions encountered within the
+            context after collecting error information and associating it with the
+            trace data.
+    """
+    if client.config.log_start_trace_span:
+        client.trace(**trace_data.as_start_parameters)
+
+    error_info: Optional[ErrorInfoDict] = None
+    try:
+        context_storage.set_trace_data(trace_data)
+        yield
+    except Exception as exception:
+        error_info = error_info_collector.collect(exception)
+        raise
+    finally:
+        trace_data = context_storage.pop_trace_data()  # type: ignore
+
+        assert trace_data is not None
+
+        if error_info is not None:
+            trace_data.error_info = error_info
+
+        trace_data.init_end_time()
+
+        client.trace(**trace_data.as_parameters)
+
+
 __all__ = [
     "get_current_span_data",
     "get_current_trace_data",
     "update_current_span",
     "update_current_trace",
     "get_distributed_trace_headers",
+    "trace_context",
 ]

opik/plugins/pytest/decorator.py

@@ -3,10 +3,10 @@ import os
 import functools
 import logging
 import opik
-from opik import opik_context
+import opik.opik_context as opik_context
 from . import test_runs_storage, test_run_content
 from opik.decorator import inspect_helpers
-from opik import config
+import opik.config as config
 
 LOGGER = logging.getLogger(__name__)
 

opik/plugins/pytest/experiment_runner.py

@@ -6,9 +6,10 @@ import json
 from opik.api_objects import opik_client
 from opik.api_objects.experiment import experiment_item
 from opik.api_objects.dataset import dataset_item
-from opik import id_helpers
+import opik.id_helpers as id_helpers
 
-from opik import datetime_helpers, dict_utils
+import opik.datetime_helpers as datetime_helpers
+import opik.dict_utils as dict_utils
 from . import test_runs_storage, test_run_content
 
 

opik/plugins/pytest/hooks.py

@@ -1,7 +1,7 @@
 import logging
-from opik 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()

opik/rate_limit/rate_limit.py

@@ -0,0 +1,25 @@
+import dataclasses
+from typing import Any, Dict, Optional
+
+
+@dataclasses.dataclass
+class CloudRateLimit:
+    remaining_limit: int
+    remaining_limit_reset_time: int
+
+    def retry_after(self) -> float:
+        return self.remaining_limit_reset_time
+
+
+def parse_rate_limit(rate_limit_response: Dict[str, Any]) -> Optional[CloudRateLimit]:
+    rate_limit_reset_time = int(rate_limit_response.get("RateLimit-Reset", 0))
+    if rate_limit_reset_time == 0:
+        rate_limit_reset_time = int(rate_limit_response.get("ratelimit-reset", 0))
+
+    if rate_limit_reset_time > 0:
+        return CloudRateLimit(
+            remaining_limit=0,
+            remaining_limit_reset_time=rate_limit_reset_time,
+        )
+
+    return None