diff --git a/litellm/main.py b/litellm/main.py index 186b87060b..60e71a412e 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -3195,6 +3195,7 @@ def text_completion( Union[str, List[str]] ] = None, # Optional: Sequences where the API will stop generating further tokens. stream: Optional[bool] = None, # Optional: Whether to stream back partial progress. + stream_options: Optional[dict] = None, suffix: Optional[ str ] = None, # Optional: The suffix that comes after a completion of inserted text. @@ -3272,6 +3273,8 @@ def text_completion( optional_params["stop"] = stop if stream is not None: optional_params["stream"] = stream + if stream_options is not None: + optional_params["stream_options"] = stream_options if suffix is not None: optional_params["suffix"] = suffix if temperature is not None: @@ -3382,7 +3385,9 @@ def text_completion( if kwargs.get("acompletion", False) == True: return response if stream == True or kwargs.get("stream", False) == True: - response = TextCompletionStreamWrapper(completion_stream=response, model=model) + response = TextCompletionStreamWrapper( + completion_stream=response, model=model, stream_options=stream_options + ) return response transformed_logprobs = None # only supported for TGI models