Mirrowel commited on
Commit
3083d86
·
1 Parent(s): b76364c

feat(client): set default include_usage for streaming responses

Browse files

This change ensures that when streaming is enabled, the `stream_options` dictionary is initialized if absent, and `include_usage` is set to `True` by default. This provides consistent behavior for including usage data in streaming completions.

Files changed (1) hide show
  1. src/rotator_library/client.py +4 -0
src/rotator_library/client.py CHANGED
@@ -737,6 +737,10 @@ class RotatingClient:
737
  The completion response object, or an async generator for streaming responses, or None if all retries fail.
738
  """
739
  if kwargs.get("stream"):
 
 
 
 
740
  return self._streaming_acompletion_with_retry(request=request, pre_request_callback=pre_request_callback, **kwargs)
741
  else:
742
  return self._execute_with_retry(litellm.acompletion, request=request, pre_request_callback=pre_request_callback, **kwargs)
 
737
  The completion response object, or an async generator for streaming responses, or None if all retries fail.
738
  """
739
  if kwargs.get("stream"):
740
+ if "stream_options" not in kwargs:
741
+ kwargs["stream_options"] = {}
742
+ if "include_usage" not in kwargs["stream_options"]:
743
+ kwargs["stream_options"]["include_usage"] = True
744
  return self._streaming_acompletion_with_retry(request=request, pre_request_callback=pre_request_callback, **kwargs)
745
  else:
746
  return self._execute_with_retry(litellm.acompletion, request=request, pre_request_callback=pre_request_callback, **kwargs)