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

opik/integrations/openai/opik_tracker.py

@@ -2,7 +2,11 @@ from typing import Optional, TypeVar
 
 import openai
 
-from . import chat_completion_chunks_aggregator, openai_decorator
+from . import (
+    chat_completion_chunks_aggregator,
+    openai_chat_completions_decorator,
+)
+import opik.semantic_version as semantic_version
 
 OpenAIClient = TypeVar("OpenAIClient", openai.OpenAI, openai.AsyncOpenAI)
 
@@ -11,12 +15,18 @@ def track_openai(
     openai_client: OpenAIClient,
     project_name: Optional[str] = None,
 ) -> OpenAIClient:
-    """Adds Opik tracking to an OpenAI client.
+    """Adds Opik tracking wrappers to an OpenAI client.
+
+    The client is always patched; however every wrapped call checks
+    `opik.is_tracing_active()` before emitting
+    any telemetry. If tracing is disabled at call time, the wrapped function
+    executes normally but no span/trace is sent.
 
     Tracks calls to:
     * `openai_client.chat.completions.create()`, including support for stream=True mode.
     * `openai_client.beta.chat.completions.parse()`
     * `openai_client.beta.chat.completions.stream()`
+    * `openai_client.responses.create()`
 
     Can be used within other Opik-tracked functions.
 
@@ -32,33 +42,104 @@ def track_openai(
 
     openai_client.opik_tracked = True
 
-    decorator_factory = openai_decorator.OpenaiTrackDecorator()
+    _patch_openai_chat_completions(openai_client, project_name)
+
+    if hasattr(openai_client, "responses"):
+        _patch_openai_responses(openai_client, project_name)
+
+    return openai_client
+
 
+def _patch_openai_chat_completions(
+    openai_client: OpenAIClient,
+    project_name: Optional[str] = None,
+) -> None:
+    chat_completions_decorator_factory = (
+        openai_chat_completions_decorator.OpenaiChatCompletionsTrackDecorator()
+    )
     if openai_client.base_url.host != "api.openai.com":
-        decorator_factory.provider = openai_client.base_url.host
+        chat_completions_decorator_factory.provider = openai_client.base_url.host
 
-    completions_create_decorator = decorator_factory.track(
+    completions_create_decorator = chat_completions_decorator_factory.track(
         type="llm",
         name="chat_completion_create",
         generations_aggregator=chat_completion_chunks_aggregator.aggregate,
         project_name=project_name,
     )
-    completions_parse_decorator = decorator_factory.track(
+    completions_parse_decorator = chat_completions_decorator_factory.track(
         type="llm",
         name="chat_completion_parse",
         generations_aggregator=chat_completion_chunks_aggregator.aggregate,
         project_name=project_name,
     )
+    completions_stream_decorator = chat_completions_decorator_factory.track(
+        type="llm",
+        name="chat_completion_stream",
+        generations_aggregator=chat_completion_chunks_aggregator.aggregate,
+        project_name=project_name,
+    )
 
-    # OpenAI implemented beta.chat.completions.stream() in a way that it
-    # calls chat.completions.create(stream=True) under the hood.
-    # So decorating `create` will automatically work for tracking `stream`.
     openai_client.chat.completions.create = completions_create_decorator(
         openai_client.chat.completions.create
     )
-
-    openai_client.beta.chat.completions.parse = completions_parse_decorator(
-        openai_client.beta.chat.completions.parse
+    if semantic_version.SemanticVersion.parse(openai.__version__) < "1.92.0":  # type: ignore
+        # beta.chat.completions.stream() calls chat.completions.create(stream=True)
+        # under the hood.
+        # So decorating `create` will automatically work for tracking `stream`.
+        openai_client.beta.chat.completions.parse = completions_parse_decorator(
+            openai_client.beta.chat.completions.parse
+        )
+    else:
+        # OpenAI reworked beta API.
+        # * chat.completion.stream calls chat.completion.create under the hood, so
+        #   it doesn't need to be decorated again.
+        # * But beta.chat.completion.stream does not call chat.completion.create, so
+        #   it needs to be decorated!
+        openai_client.beta.chat.completions.stream = completions_stream_decorator(
+            openai_client.beta.chat.completions.stream
+        )
+
+        openai_client.chat.completions.parse = completions_parse_decorator(
+            openai_client.chat.completions.parse
+        )
+        openai_client.beta.chat.completions.parse = completions_parse_decorator(
+            openai_client.beta.chat.completions.parse
+        )
+
+
+def _patch_openai_responses(
+    openai_client: OpenAIClient,
+    project_name: Optional[str] = None,
+) -> None:
+    from . import (
+        response_events_aggregator,
+        openai_responses_decorator,
     )
 
-    return openai_client
+    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
+
+    if hasattr(openai_client.responses, "create"):
+        responses_create_decorator = responses_decorator_factory.track(
+            type="llm",
+            name="responses_create",
+            generations_aggregator=response_events_aggregator.aggregate,
+            project_name=project_name,
+        )
+        openai_client.responses.create = responses_create_decorator(
+            openai_client.responses.create
+        )
+
+    if hasattr(openai_client.responses, "parse"):
+        responses_parse_decorator = responses_decorator_factory.track(
+            type="llm",
+            name="responses_parse",
+            generations_aggregator=response_events_aggregator.aggregate,
+            project_name=project_name,
+        )
+        openai_client.responses.parse = responses_parse_decorator(
+            openai_client.responses.parse
+        )

opik/integrations/openai/response_events_aggregator.py

@@ -0,0 +1,36 @@
+import logging
+from typing import List, Optional
+
+from openai.types import responses as openai_responses
+
+import opik.logging_messages as logging_messages
+
+LOGGER = logging.getLogger(__name__)
+
+
+def aggregate(
+    items: List[openai_responses.ResponseStreamEvent],
+) -> Optional[openai_responses.Response]:
+    try:
+        completed_event = [
+            event
+            for event in items
+            if isinstance(
+                event,
+                (
+                    openai_responses.ResponseCompletedEvent,
+                    openai_responses.ResponseErrorEvent,
+                    openai_responses.ResponseIncompleteEvent,
+                ),
+            )
+        ]
+
+        response = completed_event[0].response
+        return response
+    except Exception as exception:
+        LOGGER.error(
+            logging_messages.FAILED_TO_PARSE_OPENAI_STREAM_CONTENT,
+            str(exception),
+            exc_info=True,
+        )
+        return None

opik/integrations/openai/stream_patchers.py

@@ -1,27 +1,37 @@
 import logging
-from typing import Iterator, AsyncIterator, Any, List, Optional, Callable
-from opik.types import ErrorInfoDict
+from typing import Iterator, AsyncIterator, List, Optional, Callable, TypeVar
 
+from opik.types import ErrorInfoDict
 from opik.api_objects import trace, span
 from opik.decorator import generator_wrappers, error_info_collector
-from openai.types.chat import chat_completion_chunk, chat_completion
 import functools
 import openai
+import openai.lib.streaming.chat
+
 
 LOGGER = logging.getLogger(__name__)
 
+# Raw low-level stream methods
 original_stream_iter_method = openai.Stream.__iter__
 original_async_stream_aiter_method = openai.AsyncStream.__aiter__
 
+# Stream manager (factory object) methods
+original_chat_completion_stream_manager_enter_method = (
+    openai.lib.streaming.chat.ChatCompletionStreamManager.__enter__
+)
+original_async_chat_completion_stream_manager_aenter_method = (
+    openai.lib.streaming.chat.AsyncChatCompletionStreamManager.__aenter__
+)
+
+StreamItem = TypeVar("StreamItem")
+AggregatedResult = TypeVar("AggregatedResult")
+
 
 def patch_sync_stream(
     stream: openai.Stream,
     span_to_end: span.SpanData,
     trace_to_end: Optional[trace.TraceData],
-    generations_aggregator: Callable[
-        [List[chat_completion_chunk.ChatCompletionChunk]],
-        chat_completion.ChatCompletion,
-    ],
+    generations_aggregator: Callable[[List[StreamItem]], Optional[AggregatedResult]],
     finally_callback: generator_wrappers.FinishGeneratorCallback,
 ) -> openai.Stream:
     """
@@ -31,15 +41,16 @@ def patch_sync_stream(
     for event in stream:
         print(event)
     ```
+
     """
 
     def Stream__iter__decorator(dunder_iter_func: Callable) -> Callable:
         @functools.wraps(dunder_iter_func)
         def wrapper(
             self: openai.Stream,
-        ) -> Iterator[Any]:
+        ) -> Iterator[StreamItem]:
             try:
-                accumulated_items: List[chat_completion_chunk.ChatCompletionChunk] = []
+                accumulated_items: List[StreamItem] = []
                 error_info: Optional[ErrorInfoDict] = None
                 for item in dunder_iter_func(self):
                     accumulated_items.append(item)
@@ -85,10 +96,7 @@ def patch_async_stream(
     stream: openai.AsyncStream,
     span_to_end: span.SpanData,
     trace_to_end: Optional[trace.TraceData],
-    generations_aggregator: Callable[
-        [List[chat_completion_chunk.ChatCompletionChunk]],
-        chat_completion.ChatCompletion,
-    ],
+    generations_aggregator: Callable[[List[StreamItem]], Optional[AggregatedResult]],
     finally_callback: generator_wrappers.FinishGeneratorCallback,
 ) -> openai.Stream:
     """
@@ -104,9 +112,9 @@ def patch_async_stream(
         @functools.wraps(dunder_aiter_func)
         async def wrapper(
             self: openai.AsyncStream,
-        ) -> AsyncIterator[Any]:
+        ) -> AsyncIterator[StreamItem]:
             try:
-                accumulated_items: List[chat_completion_chunk.ChatCompletionChunk] = []
+                accumulated_items: List[StreamItem] = []
                 error_info: Optional[ErrorInfoDict] = None
 
                 async for item in dunder_aiter_func(self):
@@ -149,3 +157,105 @@ def patch_async_stream(
     stream.trace_to_end = trace_to_end
 
     return stream
+
+
+def patch_sync_chat_completion_stream_manager(
+    chat_completion_stream_manager: openai.lib.streaming.chat.ChatCompletionStreamManager,
+    span_to_end: span.SpanData,
+    trace_to_end: Optional[trace.TraceData],
+    generations_aggregator: Callable[[List[StreamItem]], Optional[AggregatedResult]],
+    finally_callback: generator_wrappers.FinishGeneratorCallback,
+) -> openai.lib.streaming.chat.ChatCompletionStreamManager:
+    """
+    User flow that caused this non-trivial patching:
+
+    ```
+    stream_manager = openai_client.chat.completions.create(stream=True)
+
+    with stream_manager as stream:
+        for event in stream:
+            print(event)
+    ```
+
+    `create` method with stream=True returns a stream_manager object that creates a Stream object in the context
+    manager. We need to patch the __enter__ method to return our patched stream.
+    """
+
+    def ChatCompletionStreamManager__enter__decorator(enter_func: Callable) -> Callable:
+        @functools.wraps(enter_func)
+        def wrapper(
+            self: openai.lib.streaming.chat.ChatCompletionStreamManager,
+        ) -> openai.lib.streaming.chat.ChatCompletionStream:
+            chat_completion_stream = enter_func(self)
+
+            chat_completion_stream._raw_stream = patch_sync_stream(
+                stream=chat_completion_stream._raw_stream,
+                span_to_end=span_to_end,
+                trace_to_end=trace_to_end,
+                generations_aggregator=generations_aggregator,
+                finally_callback=finally_callback,
+            )
+
+            return chat_completion_stream
+
+        return wrapper
+
+    # Replace the original __enter__ method with our decorated version
+    openai.lib.streaming.chat.ChatCompletionStreamManager.__enter__ = (
+        ChatCompletionStreamManager__enter__decorator(
+            original_chat_completion_stream_manager_enter_method
+        )
+    )
+
+    return chat_completion_stream_manager
+
+
+def patch_async_chat_completion_stream_manager(
+    async_chat_completion_stream_manager: openai.lib.streaming.chat.AsyncChatCompletionStreamManager,
+    span_to_end: span.SpanData,
+    trace_to_end: Optional[trace.TraceData],
+    generations_aggregator: Callable[[List[StreamItem]], Optional[AggregatedResult]],
+    finally_callback: generator_wrappers.FinishGeneratorCallback,
+) -> openai.lib.streaming.chat.AsyncChatCompletionStreamManager:
+    """
+    User flow that caused this non-trivial patching:
+
+    ```
+    async_stream_manager = openai_client.chat.completions.create(stream=True)
+
+    async with async_stream_manager as async_stream:
+        async for event in async_stream:
+            print(event)
+    ```
+
+    For more details see patch_sync_message_stream_manager docstring
+    """
+
+    def AsyncChatCompletionStreamManager__aenter__decorator(
+        aenter_func: Callable,
+    ) -> Callable:
+        @functools.wraps(aenter_func)
+        async def wrapper(
+            self: openai.lib.streaming.chat.AsyncChatCompletionStreamManager,
+        ) -> openai.lib.streaming.chat.AsyncChatCompletionStream:
+            async_chat_completion_stream = await aenter_func(self)
+
+            async_chat_completion_stream._raw_stream = patch_async_stream(
+                stream=async_chat_completion_stream._raw_stream,
+                span_to_end=span_to_end,
+                trace_to_end=trace_to_end,
+                generations_aggregator=generations_aggregator,
+                finally_callback=finally_callback,
+            )
+
+            return async_chat_completion_stream
+
+        return wrapper
+
+    openai.lib.streaming.chat.AsyncChatCompletionStreamManager.__aenter__ = (
+        AsyncChatCompletionStreamManager__aenter__decorator(
+            original_async_chat_completion_stream_manager_aenter_method
+        )
+    )
+
+    return async_chat_completion_stream_manager

opik/integrations/sagemaker/auth.py

@@ -43,4 +43,8 @@ def setup_aws_sagemaker_session_hook() -> None:
 
         client.auth = sagemaker_auth
 
-    opik.hooks.register_httpx_client_hook(sagemaker_auth_client_hook)
+    opik.hooks.add_httpx_client_hook(
+        opik.hooks.HttpxClientHook(
+            client_modifier=sagemaker_auth_client_hook, client_init_arguments=None
+        )
+    )

opik/jsonable_encoder.py

@@ -1,3 +1,4 @@
+import base64
 import dataclasses
 import datetime as dt
 import logging
@@ -45,10 +46,19 @@ def encode(obj: Any, seen: Optional[Set[int]] = None) -> Any:
             if isinstance(obj, type_):
                 return encode(encoder(obj), seen)
 
-        if dataclasses.is_dataclass(obj) or isinstance(obj, pydantic.BaseModel):
+        if dataclasses.is_dataclass(obj):
             obj_dict = obj.__dict__
             return encode(obj_dict, seen)
 
+        if isinstance(obj, pydantic.BaseModel):
+            obj_dict = {**obj.__dict__}
+            obj_dict.update(
+                obj.__pydantic_extra__
+                if isinstance(obj.__pydantic_extra__, dict)
+                else {}
+            )
+            return encode(obj_dict, seen)
+
         if isinstance(obj, Enum):
             return encode(obj.value, seen)
         if isinstance(obj, PurePath):
@@ -59,6 +69,8 @@ def encode(obj: Any, seen: Optional[Set[int]] = None) -> Any:
             return datetime_utils.serialize_datetime(obj)
         if isinstance(obj, dt.date):
             return str(obj)
+        if isinstance(obj, bytes):
+            return base64.b64encode(obj).decode("utf-8")
         if isinstance(obj, dict):
             encoded_dict = {}
             allowed_keys = set(obj.keys())
@@ -68,6 +80,7 @@ def encode(obj: Any, seen: Optional[Set[int]] = None) -> Any:
                     encoded_value = encode(value, seen)
                     encoded_dict[encoded_key] = encoded_value
             return encoded_dict
+
         if isinstance(obj, (list, set, frozenset, GeneratorType, tuple)):
             encoded_list = []
             for item in obj:
@@ -77,6 +90,9 @@ def encode(obj: Any, seen: Optional[Set[int]] = None) -> Any:
         if np is not None and isinstance(obj, np.ndarray):
             return encode(obj.tolist(), seen)
 
+        if _is_pydantic_iterator_validator(obj):
+            return "<Pydantic ValidatorIterator serialization is not supported>"
+
     except Exception:
         LOGGER.debug("Failed to serialize object.", exc_info=True)
 
@@ -90,3 +106,15 @@ def encode(obj: Any, seen: Optional[Set[int]] = None) -> Any:
     data = str(obj)
 
     return data
+
+
+def _is_pydantic_iterator_validator(obj: Any) -> bool:
+    if "ValidatorIterator" == obj.__class__.__name__ and "pydantic" in obj.__module__:
+        # ValidatorIterator is not defined in python code and is added to the pydantic-core
+        # namespace during the runtime.
+        # This class fully replaces the original generator object, so it is not possible
+        # to extract any extra information from the object.
+
+        return True
+
+    return False

opik/llm_usage/base_original_provider_usage.py

@@ -2,7 +2,7 @@ import abc
 from typing import Dict, Any
 
 import pydantic
-from opik import dict_utils
+import opik.dict_utils as dict_utils
 
 
 class BaseOriginalProviderUsage(pydantic.BaseModel, abc.ABC):
@@ -16,6 +16,20 @@ class BaseOriginalProviderUsage(pydantic.BaseModel, abc.ABC):
     def to_backend_compatible_flat_dict(self, parent_key_prefix: str) -> Dict[str, int]:
         result = {**self.__dict__}
 
+        return self.flatten_result_and_add_model_extra(
+            result=result, parent_key_prefix=parent_key_prefix
+        )
+
+    @classmethod
+    @abc.abstractmethod
+    def from_original_usage_dict(
+        cls, usage: Dict[str, Any]
+    ) -> "BaseOriginalProviderUsage":
+        pass
+
+    def flatten_result_and_add_model_extra(
+        self, result: Dict[str, Any], parent_key_prefix: str
+    ) -> Dict[str, int]:
         result = dict_utils.flatten_dict(
             d=result, delim=".", parent_key=parent_key_prefix
         )
@@ -29,10 +43,3 @@ class BaseOriginalProviderUsage(pydantic.BaseModel, abc.ABC):
         result = dict_utils.keep_only_values_of_type(d=result, value_type=int)
 
         return result
-
-    @classmethod
-    @abc.abstractmethod
-    def from_original_usage_dict(
-        cls, usage: Dict[str, Any]
-    ) -> "BaseOriginalProviderUsage":
-        pass

opik/llm_usage/bedrock_usage.py

@@ -1,9 +1,9 @@
-from typing import Dict, Any
+from typing import Dict, Any, Optional
 from . import base_original_provider_usage
 
 
 class BedrockUsage(base_original_provider_usage.BaseOriginalProviderUsage):
-    """Anthropic calls token usage data. Updated 11.03.2025"""
+    """Bedrock calls token usage data. Updated 04.08.2025"""
 
     inputTokens: int
     """The number of input tokens which were used."""
@@ -11,6 +11,12 @@ class BedrockUsage(base_original_provider_usage.BaseOriginalProviderUsage):
     outputTokens: int
     """The number of output tokens which were used."""
 
+    cacheReadInputTokens: Optional[int] = None
+    """The number of input tokens which were read from the cache."""
+
+    cacheWriteInputTokens: Optional[int] = None
+    """The number of input tokens which were written to the cache."""
+
     def to_backend_compatible_flat_dict(self, parent_key_prefix: str) -> Dict[str, int]:
         return super().to_backend_compatible_flat_dict(parent_key_prefix)
 

opik/llm_usage/google_usage.py

@@ -1,11 +1,13 @@
 from typing import Optional, Dict, Any
+
+
 from . import base_original_provider_usage
 
 
 class GoogleGeminiUsage(base_original_provider_usage.BaseOriginalProviderUsage):
     """Google AI / VertexAI calls token usage data. Updated 11.03.2025"""
 
-    candidates_token_count: int
+    candidates_token_count: Optional[int]
     """Number of tokens in the response(s)."""
 
     prompt_token_count: int
@@ -17,6 +19,9 @@ class GoogleGeminiUsage(base_original_provider_usage.BaseOriginalProviderUsage):
     cached_content_token_count: Optional[int] = None
     """Output only. Number of tokens in the cached part in the input (the cached content)."""
 
+    thoughts_token_count: Optional[int] = None
+    """Number of tokens spent for reasoning. Only available for Gemini models with reasoning enabled. (Gemini-2.5 and above)"""
+
     def to_backend_compatible_flat_dict(self, parent_key_prefix: str) -> Dict[str, int]:
         return super().to_backend_compatible_flat_dict(parent_key_prefix)
 

opik/llm_usage/llm_usage_info.py

@@ -10,4 +10,10 @@ from . import opik_usage
 class LLMUsageInfo(pydantic.BaseModel):
     provider: Optional[Union[LLMProvider, str]] = None
     model: Optional[str] = None
+    """
+    The model name (or model id) that can be used to identify the model for cost calculation.
+    """
     usage: Optional[Union[Dict[str, Any], opik_usage.OpikUsage]] = None
+    """
+    The dictionary with token usage in the original format of the provider.
+    """

opik/llm_usage/{openai_usage.py → openai_chat_completions_usage.py}

@@ -1,7 +1,6 @@
 import pydantic
 from typing import Optional, Dict, Any
 from . import base_original_provider_usage
-from opik import dict_utils
 
 
 class OpenAICompletionsUsage(base_original_provider_usage.BaseOriginalProviderUsage):
@@ -33,19 +32,10 @@ class OpenAICompletionsUsage(base_original_provider_usage.BaseOriginalProviderUs
         if self.prompt_tokens_details is not None:
             result["prompt_tokens_details"] = self.prompt_tokens_details.model_dump()
 
-        result = dict_utils.flatten_dict(
-            d=result, delim=".", parent_key=parent_key_prefix
+        return self.flatten_result_and_add_model_extra(
+            result=result, parent_key_prefix=parent_key_prefix
         )
 
-        if self.model_extra is not None:
-            model_extra = dict_utils.flatten_dict(
-                d=self.model_extra, delim=".", parent_key=parent_key_prefix
-            )
-            result.update(model_extra)
-
-        result = dict_utils.keep_only_values_of_type(d=result, value_type=int)
-        return result
-
     @classmethod
     def from_original_usage_dict(
         cls, usage_dict: Dict[str, Any]

opik/llm_usage/{openai_agent_usage.py → openai_responses_usage.py}

@@ -1,11 +1,10 @@
 import pydantic
 from typing import Optional, Dict, Any
 from . import base_original_provider_usage
-from opik import dict_utils
 
 
-class OpenAIAgentUsage(base_original_provider_usage.BaseOriginalProviderUsage):
-    """OpenAI Agentic calls token usage data. Updated 11.03.2025"""
+class OpenAIResponsesUsage(base_original_provider_usage.BaseOriginalProviderUsage):
+    """OpenAI Responses API calls token usage data. Updated 11.03.2025"""
 
     output_tokens: int
     """Number of tokens in the generated response."""
@@ -31,21 +30,14 @@ class OpenAIAgentUsage(base_original_provider_usage.BaseOriginalProviderUsage):
         if self.input_tokens_details is not None:
             result["input_tokens_details"] = self.input_tokens_details.model_dump()
 
-        result = dict_utils.flatten_dict(
-            d=result, delim=".", parent_key=parent_key_prefix
+        return self.flatten_result_and_add_model_extra(
+            result=result, parent_key_prefix=parent_key_prefix
         )
 
-        if self.model_extra is not None:
-            model_extra = dict_utils.flatten_dict(
-                d=self.model_extra, delim=".", parent_key=parent_key_prefix
-            )
-            result.update(model_extra)
-
-        result = dict_utils.keep_only_values_of_type(d=result, value_type=int)
-        return result
-
     @classmethod
-    def from_original_usage_dict(cls, usage_dict: Dict[str, Any]) -> "OpenAIAgentUsage":
+    def from_original_usage_dict(
+        cls, usage_dict: Dict[str, Any]
+    ) -> "OpenAIResponsesUsage":
         usage_dict = {**usage_dict}
         output_tokens_details_raw = usage_dict.pop("output_tokens_details", None)
         input_tokens_details_raw = usage_dict.pop("input_tokens_details", None)