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

opik/integrations/litellm/litellm_completion_decorator.py

@@ -0,0 +1,242 @@
+import logging
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    List,
+    Optional,
+    Tuple,
+    Union,
+)
+from typing_extensions import override
+
+import opik.dict_utils as dict_utils
+import opik.llm_usage as llm_usage
+from opik.api_objects import span
+from opik.decorator import arguments_helpers, base_track_decorator
+from opik.types import LLMProvider
+
+import litellm
+import litellm.types.utils
+import litellm.litellm_core_utils.streaming_handler
+
+from . import stream_patchers, completion_chunks_aggregator
+
+LOGGER = logging.getLogger(__name__)
+
+KWARGS_KEYS_TO_LOG_AS_INPUTS: List[str] = [
+    "messages",
+    "functions",
+    "function_call",
+    "tools",
+    "tool_choice",
+    "response_format",
+    "stop",
+]
+RESPONSE_KEYS_TO_LOG_AS_OUTPUT: List[str] = ["choices"]
+
+# Sensitive parameters that should never be logged
+SENSITIVE_PARAMS_TO_EXCLUDE: List[str] = [
+    "api_key",
+    "aws_access_key_id",
+    "aws_secret_access_key",
+    "azure_ad_token",
+    "azure_ad_token_provider",
+    "azure_key",
+    "azure_password",
+    "azure_username",
+    "client_secret",
+    "vertex_credentials",
+    "vertex_project",
+    "vertex_location",
+    "anthropic_api_key",
+    "openai_api_key",
+    "cohere_api_key",
+    "replicate_api_key",
+    "huggingface_api_key",
+    "togetherai_api_key",
+    "baseten_api_key",
+    "openrouter_api_key",
+]
+
+PROVIDER_MAPPING: Dict[str, LLMProvider] = {
+    "openai": LLMProvider.OPENAI,
+    "vertex_ai": LLMProvider.GOOGLE_VERTEXAI,
+    "vertex_ai-language-models": LLMProvider.GOOGLE_VERTEXAI,
+    "gemini": LLMProvider.GOOGLE_AI,
+    "anthropic": LLMProvider.ANTHROPIC,
+    "vertex_ai-anthropic_models": LLMProvider.ANTHROPIC_VERTEXAI,
+    "bedrock": LLMProvider.BEDROCK,
+    "bedrock_converse": LLMProvider.BEDROCK,
+    "groq": LLMProvider.GROQ,
+}
+
+
+def _extract_provider_from_model(model_name: str) -> Optional[LLMProvider]:
+    try:
+        provider_info = litellm.get_llm_provider(model_name)
+        provider_name = provider_info[1] if len(provider_info) > 1 else None
+        if provider_name is None:
+            return None
+        return PROVIDER_MAPPING.get(provider_name, None)
+    except Exception:
+        return None
+
+
+def _convert_response_to_dict(
+    output: Union[
+        litellm.types.utils.ModelResponse,
+        Dict[str, Any],
+    ],
+) -> Dict[str, Any]:
+    if hasattr(output, "model_dump"):
+        return output.model_dump(mode="json")
+    elif isinstance(output, dict):
+        return output
+    else:
+        return dict(output)
+
+
+def _extract_usage_from_response(
+    response_dict: Dict[str, Any],
+) -> Optional[llm_usage.OpikUsage]:
+    usage_data = response_dict.get("usage")
+    if usage_data is None:
+        return None
+
+    opik_usage = llm_usage.try_build_opik_usage_or_log_error(
+        provider=LLMProvider.OPENAI,
+        usage=usage_data,
+        logger=LOGGER,
+        error_message="Failed to log token usage from litellm call",
+    )
+
+    if opik_usage is None:
+        opik_usage = llm_usage.build_opik_usage_from_unknown_provider(
+            usage=usage_data,
+        )
+
+    return opik_usage
+
+
+def _calculate_completion_cost(
+    output: Union[
+        litellm.types.utils.ModelResponse,
+        Dict[str, Any],
+    ],
+) -> Optional[float]:
+    try:
+        return litellm.completion_cost(completion_response=output)
+    except Exception as exception:
+        LOGGER.debug(
+            "Failed to calculate cost from litellm response: %s",
+            str(exception),
+            exc_info=True,
+        )
+        return None
+
+
+class LiteLLMCompletionTrackDecorator(base_track_decorator.BaseTrackDecorator):
+    @override
+    def _start_span_inputs_preprocessor(
+        self,
+        func: Callable,
+        track_options: arguments_helpers.TrackOptions,
+        args: Tuple,
+        kwargs: Dict[str, Any],
+    ) -> arguments_helpers.StartSpanParameters:
+        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 {}
+
+        # Filter out sensitive parameters before logging
+        filtered_kwargs = {
+            key: value
+            for key, value in kwargs.items()
+            if key not in SENSITIVE_PARAMS_TO_EXCLUDE
+        }
+
+        input_data, new_metadata = dict_utils.split_dict_by_keys(
+            filtered_kwargs, keys=KWARGS_KEYS_TO_LOG_AS_INPUTS
+        )
+        metadata = dict_utils.deepmerge(metadata, new_metadata)
+        metadata["created_from"] = "litellm"
+
+        model_name = kwargs.get("model", "")
+        provider = _extract_provider_from_model(model_name)
+
+        return arguments_helpers.StartSpanParameters(
+            name=name,
+            input=input_data,
+            type=track_options.type,
+            tags=["litellm"],
+            metadata=metadata,
+            project_name=track_options.project_name,
+            model=model_name,
+            provider=provider,
+        )
+
+    @override
+    def _end_span_inputs_preprocessor(
+        self,
+        output: Any,
+        capture_output: bool,
+        current_span_data: span.SpanData,
+    ) -> arguments_helpers.EndSpanParameters:
+        assert isinstance(
+            output,
+            (
+                litellm.types.utils.ModelResponse,
+                dict,
+            ),
+        ), f"Expected ModelResponse or dict, got {type(output)}"
+
+        response_dict = _convert_response_to_dict(output)
+        output_data, metadata = dict_utils.split_dict_by_keys(
+            response_dict, RESPONSE_KEYS_TO_LOG_AS_OUTPUT
+        )
+
+        model = response_dict.get("model")
+        provider = _extract_provider_from_model(model) if model else None
+        opik_usage = _extract_usage_from_response(response_dict)
+        total_cost = _calculate_completion_cost(output)
+
+        return arguments_helpers.EndSpanParameters(
+            output=output_data,
+            usage=opik_usage,
+            metadata=metadata,
+            model=model,
+            provider=provider.value if provider else None,
+            total_cost=total_cost,
+        )
+
+    @override
+    def _streams_handler(  # type: ignore
+        self,
+        output: Any,
+        capture_output: bool,
+        generations_aggregator: Optional[
+            Callable[
+                [List[litellm.types.utils.ModelResponse]],
+                Optional[litellm.types.utils.ModelResponse],
+            ]
+        ],
+    ) -> Optional[litellm.litellm_core_utils.streaming_handler.CustomStreamWrapper]:
+        assert (
+            generations_aggregator is not None
+        ), "LiteLLM decorator will always get aggregator function as input"
+
+        is_litellm_stream = isinstance(
+            output, litellm.litellm_core_utils.streaming_handler.CustomStreamWrapper
+        )
+        if not is_litellm_stream:
+            return None
+
+        span_to_end, trace_to_end = base_track_decorator.pop_end_candidates()
+
+        return stream_patchers.patch_stream(
+            stream=output,
+            span_to_end=span_to_end,
+            trace_to_end=trace_to_end,
+            generations_aggregator=completion_chunks_aggregator.aggregate,
+            finally_callback=self._after_call,
+        )

opik/integrations/litellm/opik_tracker.py

@@ -0,0 +1,43 @@
+from typing import Callable, Optional, TypeVar
+
+from . import litellm_completion_decorator
+from . import completion_chunks_aggregator
+
+F = TypeVar("F", bound=Callable)
+
+
+def track_completion(
+    project_name: Optional[str] = None,
+) -> Callable[[F], F]:
+    """Decorator for tracking LiteLLM function calls with Opik.
+
+    Can be used within other Opik-tracked functions to create proper span hierarchy.
+
+    Supported (streaming and non-streaming modes):
+    * `litellm.completion`
+    * `litellm.acompletion`
+
+    Example:
+        ```python
+        import litellm
+        from opik.integrations.litellm import track_completion
+
+        tracked_completion = track_completion(project_name="my-project")(litellm.completion)
+        response = tracked_completion(model="gpt-3.5-turbo", messages=[...])
+        ```
+
+    Args:
+        project_name: The name of the project to log data.
+
+    Returns:
+        Decorator function that wraps the completion function with Opik tracking.
+    """
+
+    decorator_factory = litellm_completion_decorator.LiteLLMCompletionTrackDecorator()
+
+    return decorator_factory.track(  # type: ignore
+        type="llm",
+        name=None,  # Use the function's name (completion or acompletion)
+        project_name=project_name,
+        generations_aggregator=completion_chunks_aggregator.aggregate,
+    )

opik/integrations/litellm/stream_patchers.py

@@ -0,0 +1,151 @@
+import logging
+import functools
+from typing import Any, List, Optional, Callable, TypeVar
+
+from opik.api_objects import trace, span
+from opik.decorator import generator_wrappers, error_info_collector
+import litellm.litellm_core_utils.streaming_handler
+
+
+LOGGER = logging.getLogger(__name__)
+
+StreamItem = TypeVar("StreamItem")
+AggregatedResult = TypeVar("AggregatedResult")
+
+_original_next = (
+    litellm.litellm_core_utils.streaming_handler.CustomStreamWrapper.__next__
+)
+_original_anext = (
+    litellm.litellm_core_utils.streaming_handler.CustomStreamWrapper.__anext__
+)
+
+
+def _create_sync_next_wrapper(
+    original_next: Callable,
+    generations_aggregator: Callable,
+    finally_callback: generator_wrappers.FinishGeneratorCallback,
+) -> Callable:
+    @functools.wraps(original_next)
+    def wrapper(
+        self: litellm.litellm_core_utils.streaming_handler.CustomStreamWrapper,
+    ) -> Any:
+        if not hasattr(self, "_opik_accumulated_items"):
+            if hasattr(self, "opik_tracked_instance"):
+                self._opik_accumulated_items = []
+                self._opik_error_info = None
+
+        try:
+            item = original_next(self)
+            if hasattr(self, "_opik_accumulated_items"):
+                self._opik_accumulated_items.append(item)
+            return item
+        except StopIteration:
+            if hasattr(self, "_opik_accumulated_items"):
+                try:
+                    output = generations_aggregator(self._opik_accumulated_items)
+                    finally_callback(
+                        output=output,
+                        error_info=self._opik_error_info,
+                        capture_output=True,
+                        generators_span_to_end=self.span_to_end,
+                        generators_trace_to_end=self.trace_to_end,
+                    )
+                finally:
+                    if hasattr(self, "_opik_accumulated_items"):
+                        delattr(self, "_opik_accumulated_items")
+                    if hasattr(self, "_opik_error_info"):
+                        delattr(self, "_opik_error_info")
+                    if hasattr(self, "opik_tracked_instance"):
+                        delattr(self, "opik_tracked_instance")
+            raise
+        except Exception as exception:
+            if hasattr(self, "_opik_accumulated_items"):
+                self._opik_error_info = error_info_collector.collect(exception)
+                LOGGER.debug(
+                    "Exception raised from LiteLLM stream: %s",
+                    str(exception),
+                    exc_info=True,
+                )
+            raise
+
+    return wrapper
+
+
+def _create_async_next_wrapper(
+    original_anext: Callable,
+    generations_aggregator: Callable,
+    finally_callback: generator_wrappers.FinishGeneratorCallback,
+) -> Callable:
+    @functools.wraps(original_anext)
+    async def wrapper(
+        self: litellm.litellm_core_utils.streaming_handler.CustomStreamWrapper,
+    ) -> Any:
+        if not hasattr(self, "_opik_accumulated_items_async"):
+            if hasattr(self, "opik_tracked_instance_async"):
+                self._opik_accumulated_items_async = []
+                self._opik_error_info_async = None
+
+        try:
+            item = await original_anext(self)
+            if hasattr(self, "_opik_accumulated_items_async"):
+                self._opik_accumulated_items_async.append(item)
+            return item
+        except StopAsyncIteration:
+            if hasattr(self, "_opik_accumulated_items_async"):
+                try:
+                    output = generations_aggregator(self._opik_accumulated_items_async)
+                    finally_callback(
+                        output=output,
+                        error_info=self._opik_error_info_async,
+                        capture_output=True,
+                        generators_span_to_end=self.span_to_end_async,
+                        generators_trace_to_end=self.trace_to_end_async,
+                    )
+                finally:
+                    if hasattr(self, "_opik_accumulated_items_async"):
+                        delattr(self, "_opik_accumulated_items_async")
+                    if hasattr(self, "_opik_error_info_async"):
+                        delattr(self, "_opik_error_info_async")
+                    if hasattr(self, "opik_tracked_instance_async"):
+                        delattr(self, "opik_tracked_instance_async")
+            raise
+        except Exception as exception:
+            if hasattr(self, "_opik_accumulated_items_async"):
+                self._opik_error_info_async = error_info_collector.collect(exception)
+                LOGGER.debug(
+                    "Exception raised from LiteLLM async stream: %s",
+                    str(exception),
+                    exc_info=True,
+                )
+            raise
+
+    return wrapper
+
+
+def patch_stream(
+    stream: litellm.litellm_core_utils.streaming_handler.CustomStreamWrapper,
+    span_to_end: span.SpanData,
+    trace_to_end: Optional[trace.TraceData],
+    generations_aggregator: Callable[[List[StreamItem]], Optional[AggregatedResult]],
+    finally_callback: generator_wrappers.FinishGeneratorCallback,
+) -> litellm.litellm_core_utils.streaming_handler.CustomStreamWrapper:
+    litellm.litellm_core_utils.streaming_handler.CustomStreamWrapper.__next__ = (
+        _create_sync_next_wrapper(
+            _original_next, generations_aggregator, finally_callback
+        )
+    )
+    litellm.litellm_core_utils.streaming_handler.CustomStreamWrapper.__anext__ = (
+        _create_async_next_wrapper(
+            _original_anext, generations_aggregator, finally_callback
+        )
+    )
+
+    stream.opik_tracked_instance = True
+    stream.span_to_end = span_to_end
+    stream.trace_to_end = trace_to_end
+
+    stream.opik_tracked_instance_async = True
+    stream.span_to_end_async = span_to_end
+    stream.trace_to_end_async = trace_to_end
+
+    return stream