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

opik/integrations/llama_index/callback.py

@@ -67,6 +67,21 @@ class LlamaIndexCallbackHandler(base_handler.BaseCallbackHandler):
         self._map_event_id_to_span_data: Dict[str, span.SpanData] = {}
         self._map_event_id_to_output: Dict[str, Any] = {}
 
+        # For streaming: end_trace may be called before event_end, so we need to
+        # defer the trace output update until the event output is available
+        self._pending_root_output_updates: Dict[
+            str, Union[span.SpanData, trace.TraceData]
+        ] = {}
+
+    def _send_root_to_backend(
+        self, root: Union[span.SpanData, trace.TraceData]
+    ) -> None:
+        """Send root trace or span data to the backend."""
+        if isinstance(root, span.SpanData):
+            self._opik_client.span(**root.as_parameters)
+        elif isinstance(root, trace.TraceData):
+            self._opik_client.trace(**root.as_parameters)
+
     def start_trace(self, trace_id: Optional[str] = None) -> None:
         if (
             self._skip_index_construction_trace
@@ -111,17 +126,29 @@ class LlamaIndexCallbackHandler(base_handler.BaseCallbackHandler):
             return
 
         last_event = _get_last_event(trace_map)
-        last_event_output = self._map_event_id_to_output.get(last_event, None)
 
-        root.init_end_time().update(output=last_event_output)
+        # Check if the output for the last event is already available.
+        # For streaming calls, LlamaIndex calls end_trace() BEFORE event_end(),
+        # so the output won't be stored yet.
+        if last_event in self._map_event_id_to_output:
+            last_event_output = self._map_event_id_to_output.get(last_event)
+            root.init_end_time().update(output=last_event_output)
 
+            # Send the trace/span with output
+            self._send_root_to_backend(root)
+        else:
+            # Output not available yet (streaming scenario).
+            # Store the root so we can update it when event_end is called.
+            # Don't send the trace/span yet - it will be sent in on_event_end
+            # with the output and correct end_time to avoid race conditions.
+            # Note: We don't set end_time here because the actual end is when
+            # the last event ends, not when LlamaIndex calls end_trace().
+            self._pending_root_output_updates[last_event] = root
+
+        # Clean up context storage
         if isinstance(root, span.SpanData):
-            # We created a wrapper span (external trace existed)
-            self._opik_client.span(**root.as_parameters)
             self._opik_context_storage.pop_span_data(ensure_id=root.id)
         elif isinstance(root, trace.TraceData):
-            # We created a trace (no external trace existed)
-            self._opik_client.trace(**root.as_parameters)
             self._opik_context_storage.pop_trace_data(ensure_id=root.id)
 
         # Clean up
@@ -213,6 +240,16 @@ class LlamaIndexCallbackHandler(base_handler.BaseCallbackHandler):
         # Store output for end_trace
         self._map_event_id_to_output[event_id] = span_output
 
+        # Check if there's a pending root trace/span output update for this event.
+        # This happens when end_trace() was called before event_end() (streaming scenario).
+        if event_id in self._pending_root_output_updates:
+            root = self._pending_root_output_updates.pop(event_id)
+            # Set end_time now (the actual end) and update with output
+            root.init_end_time().update(output=span_output)
+
+            # Send the trace/span to the backend with correct end_time and output
+            self._send_root_to_backend(root)
+
         # Finalize span if it exists
         if event_id in self._map_event_id_to_span_data:
             span_data = self._map_event_id_to_span_data[event_id]

opik/integrations/openai/agents/opik_tracing_processor.py

@@ -82,24 +82,22 @@ class OpikTracingProcessor(tracing.TracingProcessor):
                     metadata=trace_metadata,
                     type="general",
                 )
-                _, opik_span_data = (
-                    span_creation_handler.create_span_respecting_context(
-                        start_span_arguments=start_span_arguments,
-                        distributed_trace_headers=None,
-                        opik_context_storage=self._opik_context_storage,
-                    )
+                result = span_creation_handler.create_span_respecting_context(
+                    start_span_arguments=start_span_arguments,
+                    distributed_trace_headers=None,
+                    opik_context_storage=self._opik_context_storage,
                 )
-                self._opik_spans_data_map[trace.trace_id] = opik_span_data
-                self._opik_context_storage.add_span_data(opik_span_data)
+                self._opik_spans_data_map[trace.trace_id] = result.span_data
+                self._opik_context_storage.add_span_data(result.span_data)
                 self._openai_trace_id_to_external_opik_parent_span_id[
                     trace.trace_id
-                ] = opik_span_data.parent_span_id
+                ] = result.span_data.parent_span_id
 
                 if (
                     self._opik_client.config.log_start_trace_span
                     and tracing_runtime_config.is_tracing_active()
                 ):
-                    self._opik_client.span(**opik_span_data.as_start_parameters)
+                    self._opik_client.span(**result.span_data.as_start_parameters)
 
         except Exception:
             LOGGER.error("on_trace_start failed", exc_info=True)

opik/integrations/openai/opik_tracker.py

@@ -1,6 +1,7 @@
 from typing import Optional, TypeVar
 
 import openai
+import opik
 
 from . import (
     chat_completion_chunks_aggregator,
@@ -11,6 +12,13 @@ import opik.semantic_version as semantic_version
 OpenAIClient = TypeVar("OpenAIClient", openai.OpenAI, openai.AsyncOpenAI)
 
 
+def _get_provider(openai_client: OpenAIClient) -> str:
+    """Get the provider name from the OpenAI client's base URL."""
+    if openai_client.base_url.host != "api.openai.com":
+        return openai_client.base_url.host
+    return "openai"
+
+
 def track_openai(
     openai_client: OpenAIClient,
     project_name: Optional[str] = None,
@@ -27,6 +35,9 @@ def track_openai(
     * `openai_client.beta.chat.completions.parse()`
     * `openai_client.beta.chat.completions.stream()`
     * `openai_client.responses.create()`
+    * `openai_client.videos.create()`, `videos.create_and_poll()`, `videos.poll()`,
+      `videos.list()`, `videos.delete()`, `videos.remix()`, `videos.download_content()`,
+      and `write_to_file()` on downloaded content
 
     Can be used within other Opik-tracked functions.
 
@@ -47,6 +58,9 @@ def track_openai(
     if hasattr(openai_client, "responses"):
         _patch_openai_responses(openai_client, project_name)
 
+    if hasattr(openai_client, "videos"):
+        _patch_openai_videos(openai_client, project_name)
+
     return openai_client
 
 
@@ -57,8 +71,7 @@ def _patch_openai_chat_completions(
     chat_completions_decorator_factory = (
         openai_chat_completions_decorator.OpenaiChatCompletionsTrackDecorator()
     )
-    if openai_client.base_url.host != "api.openai.com":
-        chat_completions_decorator_factory.provider = openai_client.base_url.host
+    chat_completions_decorator_factory.provider = _get_provider(openai_client)
 
     completions_create_decorator = chat_completions_decorator_factory.track(
         type="llm",
@@ -119,8 +132,7 @@ def _patch_openai_responses(
     responses_decorator_factory = (
         openai_responses_decorator.OpenaiResponsesTrackDecorator()
     )
-    if openai_client.base_url.host != "api.openai.com":
-        responses_decorator_factory.provider = openai_client.base_url.host
+    responses_decorator_factory.provider = _get_provider(openai_client)
 
     if hasattr(openai_client.responses, "create"):
         responses_create_decorator = responses_decorator_factory.track(
@@ -143,3 +155,86 @@ def _patch_openai_responses(
         openai_client.responses.parse = responses_parse_decorator(
             openai_client.responses.parse
         )
+
+
+def _patch_openai_videos(
+    openai_client: OpenAIClient,
+    project_name: Optional[str] = None,
+) -> None:
+    from . import videos
+
+    provider = _get_provider(openai_client)
+    create_decorator_factory = videos.VideosCreateTrackDecorator(provider=provider)
+    download_decorator_factory = videos.VideosDownloadTrackDecorator()
+
+    video_metadata = {"created_from": "openai", "type": "openai_videos"}
+    video_tags = ["openai"]
+
+    if hasattr(openai_client.videos, "create"):
+        decorator = create_decorator_factory.track(
+            type="llm",
+            name="videos.create",
+            project_name=project_name,
+        )
+        openai_client.videos.create = decorator(openai_client.videos.create)
+
+    if hasattr(openai_client.videos, "create_and_poll"):
+        decorator = opik.track(
+            name="videos.create_and_poll",
+            tags=video_tags,
+            metadata=video_metadata,
+            project_name=project_name,
+        )
+        openai_client.videos.create_and_poll = decorator(
+            openai_client.videos.create_and_poll
+        )
+
+    if hasattr(openai_client.videos, "remix"):
+        decorator = create_decorator_factory.track(
+            type="llm",
+            name="videos.remix",
+            project_name=project_name,
+        )
+        openai_client.videos.remix = decorator(openai_client.videos.remix)
+
+    # Note: videos.retrieve is intentionally NOT patched to avoid too many spans
+    # since it's called frequently during polling operations.
+
+    if hasattr(openai_client.videos, "poll"):
+        decorator = opik.track(
+            name="videos.poll",
+            tags=video_tags,
+            metadata=video_metadata,
+            project_name=project_name,
+        )
+        openai_client.videos.poll = decorator(openai_client.videos.poll)
+
+    if hasattr(openai_client.videos, "delete"):
+        decorator = opik.track(
+            name="videos.delete",
+            tags=video_tags,
+            metadata=video_metadata,
+            project_name=project_name,
+        )
+        openai_client.videos.delete = decorator(openai_client.videos.delete)
+
+    # Patch download_content - also patches write_to_file on returned instances
+    # download_content returns a lazy response object, write_to_file does the actual download
+    if hasattr(openai_client.videos, "download_content"):
+        decorator = download_decorator_factory.track(
+            type="general",
+            name="videos.download_content",
+            project_name=project_name,
+        )
+        openai_client.videos.download_content = decorator(
+            openai_client.videos.download_content
+        )
+
+    if hasattr(openai_client.videos, "list"):
+        decorator = opik.track(
+            name="videos.list",
+            tags=video_tags,
+            metadata=video_metadata,
+            project_name=project_name,
+        )
+        openai_client.videos.list = decorator(openai_client.videos.list)

opik/integrations/openai/videos/__init__.py

@@ -0,0 +1,9 @@
+from .videos_create_decorator import VideosCreateTrackDecorator
+from .videos_download_decorator import VideosDownloadTrackDecorator
+from . import binary_response_write_to_file_decorator
+
+__all__ = [
+    "VideosCreateTrackDecorator",
+    "VideosDownloadTrackDecorator",
+    "binary_response_write_to_file_decorator",
+]

opik/integrations/openai/videos/binary_response_write_to_file_decorator.py

@@ -0,0 +1,88 @@
+"""
+Decorator for HttpxBinaryResponseContent.write_to_file method.
+"""
+
+import functools
+import logging
+import os
+import shutil
+import tempfile
+from typing import Callable, Optional, Any
+
+import opik
+from opik.api_objects import attachment, span
+
+LOGGER = logging.getLogger(__name__)
+
+
+def create_write_to_file_decorator(
+    project_name: Optional[str] = None,
+) -> Callable[[Callable[[str], None]], Callable[[str], None]]:
+    """
+    Create a decorator that tracks write_to_file calls.
+
+    Uses functools.wraps with opik context manager to avoid limitations
+    of BaseTrackDecorator when the wrapped function returns None.
+    """
+
+    def decorator(func: Callable[[str], None]) -> Callable[[str], None]:
+        @functools.wraps(func)
+        def wrapper(file: str) -> Any:
+            with opik.start_as_current_span(
+                name="videos.write_to_file",
+                input={"file": str(file)},
+                metadata={"created_from": "openai", "type": "openai_videos"},
+                tags=["openai"],
+                type="general",
+                project_name=project_name,
+            ) as span_data:
+                result = func(file)
+                _attach_video_file(file, span_data)
+
+                # The result is None but we still return it in case the return value
+                # is changed in the future.
+                return result
+
+        return wrapper
+
+    return decorator
+
+
+def _attach_video_file(file_path: str, span_data: span.SpanData) -> None:
+    """
+    Create a temporary copy of the video file and attach it to the span.
+
+    A copy is made to ensure the file remains available for upload even if
+    the user deletes the original file. The temp file is created with
+    delete=False so it persists until the upload manager processes it.
+    """
+    try:
+        file_str = str(file_path)
+        file_name = os.path.basename(file_str)
+        _, extension = os.path.splitext(file_name)
+
+        # Create a temporary copy of the video file
+        temp_file = tempfile.NamedTemporaryFile(
+            mode="wb", delete=False, suffix=extension
+        )
+        shutil.copyfileobj(open(file_str, "rb"), temp_file)
+        temp_file.flush()
+        temp_file.close()
+
+        video_attachment = attachment.Attachment(
+            data=temp_file.name,
+            file_name=file_name,
+            content_type="video/mp4",
+        )
+        span_data.update(attachments=[video_attachment])
+
+        LOGGER.debug(
+            "Video attachment created from temporary copy: %s -> %s",
+            file_str,
+            temp_file.name,
+        )
+    except Exception as e:
+        LOGGER.error(
+            "Failed to attach video to span: %s",
+            str(e),
+        )

opik/integrations/openai/videos/videos_create_decorator.py

@@ -0,0 +1,159 @@
+"""
+Decorator for OpenAI video creation methods (create, remix).
+
+Output type: openai.types.video.Video
+
+This decorator is used only for LLM spans that generate videos.
+"""
+
+import logging
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    List,
+    Optional,
+    Tuple,
+    TYPE_CHECKING,
+)
+
+from typing_extensions import override
+
+import opik.dict_utils as dict_utils
+from opik.api_objects import span
+from opik.decorator import arguments_helpers, base_track_decorator
+
+if TYPE_CHECKING:
+    from openai.types.video import Video
+
+LOGGER = logging.getLogger(__name__)
+
+# Input parameters to log for video generation
+VIDEO_CREATE_KWARGS_KEYS_TO_LOG_AS_INPUTS = [
+    "prompt",
+    "seconds",
+    "size",
+]
+
+# Input parameters to log for video remix
+VIDEO_REMIX_KWARGS_KEYS_TO_LOG_AS_INPUTS = [
+    "video_id",
+    "prompt",
+]
+
+# Response keys to log as output for Video object
+VIDEO_RESPONSE_KEYS_TO_LOG_AS_OUTPUT = [
+    "id",
+    "status",
+    "prompt",
+    "seconds",
+    "size",
+    "progress",
+    "error",
+]
+
+
+class VideosCreateTrackDecorator(base_track_decorator.BaseTrackDecorator):
+    """
+    Decorator for tracking OpenAI video creation methods (LLM spans).
+
+    Handles: videos.create, videos.remix
+    """
+
+    def __init__(self, provider: str) -> None:
+        super().__init__()
+        self.provider = provider
+
+    @override
+    def _start_span_inputs_preprocessor(
+        self,
+        func: Callable,
+        track_options: arguments_helpers.TrackOptions,
+        args: Tuple,
+        kwargs: Dict[str, Any],
+    ) -> arguments_helpers.StartSpanParameters:
+        assert kwargs is not None, "Expected kwargs to be not None in videos API calls"
+
+        name = track_options.name if track_options.name is not None else func.__name__
+
+        metadata = track_options.metadata if track_options.metadata is not None else {}
+
+        # Determine which keys to log based on the method name
+        func_name = func.__name__
+        if func_name == "remix":
+            keys_to_log = VIDEO_REMIX_KWARGS_KEYS_TO_LOG_AS_INPUTS
+        else:
+            keys_to_log = VIDEO_CREATE_KWARGS_KEYS_TO_LOG_AS_INPUTS
+
+        input_data, new_metadata = dict_utils.split_dict_by_keys(
+            kwargs, keys=keys_to_log
+        )
+        metadata = dict_utils.deepmerge(metadata, new_metadata)
+        metadata.update(
+            {
+                "created_from": "openai",
+                "type": "openai_videos",
+            }
+        )
+
+        tags = ["openai"]
+        model = kwargs.get("model", None)
+
+        result = arguments_helpers.StartSpanParameters(
+            name=name,
+            input=input_data,
+            type=track_options.type,
+            tags=tags,
+            metadata=metadata,
+            project_name=track_options.project_name,
+            model=model,
+            provider=self.provider,
+        )
+
+        return result
+
+    @override
+    def _end_span_inputs_preprocessor(
+        self,
+        output: Optional["Video"],
+        capture_output: bool,
+        current_span_data: span.SpanData,
+    ) -> arguments_helpers.EndSpanParameters:
+        if output is None:
+            return arguments_helpers.EndSpanParameters(
+                output=None,
+                usage=None,
+                metadata={},
+                model=None,
+                provider=self.provider,
+            )
+
+        result_dict: Dict[str, Any] = output.model_dump(mode="json")
+
+        output_data, metadata = dict_utils.split_dict_by_keys(
+            result_dict, VIDEO_RESPONSE_KEYS_TO_LOG_AS_OUTPUT
+        )
+
+        # Add video generation info to metadata for cost calculation
+        if result_dict.get("seconds") is not None:
+            metadata["video_seconds"] = int(result_dict["seconds"])
+        if result_dict.get("size") is not None:
+            metadata["video_size"] = result_dict["size"]
+
+        return arguments_helpers.EndSpanParameters(
+            output=output_data,
+            usage=None,
+            metadata=metadata,
+            model=result_dict.get("model"),
+            provider=self.provider,
+        )
+
+    @override
+    def _streams_handler(
+        self,
+        output: Any,
+        capture_output: bool,
+        generations_aggregator: Optional[Callable[[List[Any]], Any]],
+    ) -> Optional[Any]:
+        NOT_A_STREAM = None
+        return NOT_A_STREAM

opik/integrations/openai/videos/videos_download_decorator.py

@@ -0,0 +1,110 @@
+"""
+Decorator for OpenAI video download_content method.
+"""
+
+import logging
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    List,
+    Optional,
+    Tuple,
+)
+
+from openai._legacy_response import HttpxBinaryResponseContent
+from typing_extensions import override
+
+from opik.api_objects import span
+from opik.decorator import arguments_helpers, base_track_decorator
+from . import binary_response_write_to_file_decorator
+
+LOGGER = logging.getLogger(__name__)
+
+
+class VideosDownloadTrackDecorator(base_track_decorator.BaseTrackDecorator):
+    """
+    Decorator for tracking OpenAI videos.download_content method.
+
+    Also patches the returned HttpxBinaryResponseContent instance's write_to_file
+    method to create a tracked span when the video is actually downloaded.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.provider = "openai"
+
+    @override
+    def _start_span_inputs_preprocessor(
+        self,
+        func: Callable,
+        track_options: arguments_helpers.TrackOptions,
+        args: Tuple,
+        kwargs: Dict[str, Any],
+    ) -> arguments_helpers.StartSpanParameters:
+        assert kwargs is not None, "Expected kwargs to be not None in videos API calls"
+
+        name = track_options.name if track_options.name is not None else func.__name__
+
+        metadata = track_options.metadata if track_options.metadata is not None else {}
+
+        metadata.update(
+            {
+                "created_from": "openai",
+                "type": "openai_videos",
+            }
+        )
+
+        tags = ["openai"]
+
+        result = arguments_helpers.StartSpanParameters(
+            name=name,
+            input=kwargs,
+            type=track_options.type,
+            tags=tags,
+            metadata=metadata,
+            project_name=track_options.project_name,
+        )
+
+        return result
+
+    @override
+    def _end_span_inputs_preprocessor(
+        self,
+        output: Any,
+        capture_output: bool,
+        current_span_data: span.SpanData,
+    ) -> arguments_helpers.EndSpanParameters:
+        # Patch write_to_file on the returned instance
+        if output is not None:
+            _track_instance_write_to_file(output, current_span_data.project_name)
+
+        result = arguments_helpers.EndSpanParameters(
+            output={"output": output} if not isinstance(output, dict) else output,
+            usage=None,
+            metadata={},
+            model=None,
+        )
+
+        return result
+
+    @override
+    def _streams_handler(
+        self,
+        output: Any,
+        capture_output: bool,
+        generations_aggregator: Optional[Callable[[List[Any]], Any]],
+    ) -> Optional[Any]:
+        NOT_A_STREAM = None
+        return NOT_A_STREAM
+
+
+def _track_instance_write_to_file(
+    instance: HttpxBinaryResponseContent,
+    project_name: Optional[str],
+) -> None:
+    """Patch write_to_file on this specific instance to track the download."""
+    decorator = binary_response_write_to_file_decorator.create_write_to_file_decorator(
+        project_name=project_name,
+    )
+    instance.write_to_file = decorator(instance.write_to_file)  # type: ignore[method-assign]

opik/message_processing/batching/base_batcher.py

@@ -1,4 +1,3 @@
-import threading
 import time
 import abc
 
@@ -24,25 +23,22 @@ class BaseBatcher(abc.ABC):
         self._batch_memory_limit_mb: int = batch_memory_limit_mb
 
         self._last_time_flush_callback_called: float = time.time()
-        self._lock = threading.RLock()
 
     def flush(self) -> None:
-        with self._lock:
-            if len(self._accumulated_messages) > 0:
-                batch_messages = self._create_batches_from_accumulated_messages()
-                self._accumulated_messages = []
+        if len(self._accumulated_messages) > 0:
+            batch_messages = self._create_batches_from_accumulated_messages()
+            self._accumulated_messages = []
 
-                for batch_message in batch_messages:
-                    self._flush_callback(batch_message)
-            self._last_time_flush_callback_called = time.time()
+            for batch_message in batch_messages:
+                self._flush_callback(batch_message)
+        self._last_time_flush_callback_called = time.time()
 
     def is_ready_to_flush(self) -> bool:
         elapsed = time.time() - self._last_time_flush_callback_called
         return elapsed >= self._flush_interval_seconds
 
     def is_empty(self) -> bool:
-        with self._lock:
-            return len(self._accumulated_messages) == 0
+        return len(self._accumulated_messages) == 0
 
     @abc.abstractmethod
     def _create_batches_from_accumulated_messages(
@@ -51,10 +47,9 @@ class BaseBatcher(abc.ABC):
 
     @abc.abstractmethod
     def add(self, message: messages.BaseMessage) -> None:
-        with self._lock:
-            self._accumulated_messages.append(message)
-            if len(self._accumulated_messages) >= self._max_batch_size:
-                self.flush()
+        self._accumulated_messages.append(message)
+        if len(self._accumulated_messages) >= self._max_batch_size:
+            self.flush()
 
     def _remove_matching_messages(
         self, filter_func: Callable[[messages.BaseMessage], bool]
@@ -65,10 +60,9 @@ class BaseBatcher(abc.ABC):
         Args:
             filter_func: A function that takes a BaseMessage and returns True if the message should be removed
         """
-        with self._lock:
-            self._accumulated_messages = list(
-                filter(lambda x: not filter_func(x), self._accumulated_messages)
-            )
+        self._accumulated_messages = list(
+            filter(lambda x: not filter_func(x), self._accumulated_messages)
+        )
 
     def size(self) -> int:
         """
@@ -81,5 +75,4 @@ class BaseBatcher(abc.ABC):
         Returns:
             int: The total number of accumulated messages.
         """
-        with self._lock:
-            return len(self._accumulated_messages)
+        return len(self._accumulated_messages)