diff --git a/python/instrumentation/openinference-instrumentation-litellm/src/openinference/instrumentation/litellm/__init__.py b/python/instrumentation/openinference-instrumentation-litellm/src/openinference/instrumentation/litellm/__init__.py index 9ac379199..d440b399b 100644 --- a/python/instrumentation/openinference-instrumentation-litellm/src/openinference/instrumentation/litellm/__init__.py +++ b/python/instrumentation/openinference-instrumentation-litellm/src/openinference/instrumentation/litellm/__init__.py @@ -638,6 +638,11 @@ def _completion_wrapper(self, *args: Any, **kwargs: Any) -> ModelResponse: return self.original_litellm_funcs["completion"](*args, **kwargs) # type:ignore if kwargs.get("stream", False): + if "stream_options" not in kwargs or kwargs["stream_options"] is None: + kwargs["stream_options"] = {"include_usage": True} + elif isinstance(kwargs["stream_options"], dict) and "include_usage" not in kwargs["stream_options"]: + kwargs["stream_options"]["include_usage"] = True + span = self._tracer.start_span( name="completion", attributes=dict(get_attributes_from_context()) ) @@ -668,6 +673,11 @@ async def _acompletion_wrapper( return await self.original_litellm_funcs["acompletion"](*args, **kwargs) # type:ignore if kwargs.get("stream", False): + if "stream_options" not in kwargs or kwargs["stream_options"] is None: + kwargs["stream_options"] = {"include_usage": True} + elif isinstance(kwargs["stream_options"], dict) and "include_usage" not in kwargs["stream_options"]: + kwargs["stream_options"]["include_usage"] = True + span = self._tracer.start_span( name="acompletion", attributes=dict(get_attributes_from_context()) )