From b2bde00f23e57bf27118d7cd828566cfdf19894d Mon Sep 17 00:00:00 2001 From: fenilfaldu Date: Wed, 15 Oct 2025 21:19:39 +0530 Subject: [PATCH] auto-inject include_usage for streaming and capture usage from final chunk --- .../openinference/instrumentation/litellm/__init__.py | 10 ++++++++++ 1 file changed, 10 insertions(+) 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()) )