phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

Merge pull request #3547 from BerriAI/litellm_support_stream_options_text_completion

Browse source 
[Feat] support `stream_options` on `litellm.text_completion`
This commit is contained in:
Ishaan Jaff 2024-05-09 18:05:58 -07:00 • committed by GitHub
commit 5eb12e30cc
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 59 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

19
litellm/utils.py
View file

@ -10062,16 +10062,19 @@ class CustomStreamWrapper:
text = ""
is_finished = False
finish_reason = None
usage = None
choices = getattr(chunk, "choices", [])
if len(choices) > 0:
text = choices[0].text
if choices[0].finish_reason is not None:
is_finished = True
finish_reason = choices[0].finish_reason
usage = getattr(chunk, "usage", None)
return {
"text": text,
"is_finished": is_finished,
"finish_reason": finish_reason,
"usage": usage,
}
except Exception as e:
@ -10601,6 +10604,11 @@ class CustomStreamWrapper:
print_verbose(f"completion obj content: {completion_obj['content']}")
if response_obj["is_finished"]:
self.received_finish_reason = response_obj["finish_reason"]
if (
self.stream_options
and self.stream_options.get("include_usage", False) == True
):
model_response.usage = response_obj["usage"]
elif self.custom_llm_provider == "azure_text":
response_obj = self.handle_azure_text_completion_chunk(chunk)
completion_obj["content"] = response_obj["text"]
@ -11130,9 +11138,10 @@ class CustomStreamWrapper:
class TextCompletionStreamWrapper:
def __init__(self, completion_stream, model):
def __init__(self, completion_stream, model, stream_options: Optional[dict] = None):
self.completion_stream = completion_stream
self.model = model
self.stream_options = stream_options
def __iter__(self):
return self
@ -11156,6 +11165,14 @@ class TextCompletionStreamWrapper:
text_choices["index"] = chunk["choices"][0]["index"]
text_choices["finish_reason"] = chunk["choices"][0]["finish_reason"]
response["choices"] = [text_choices]
# only pass usage when stream_options["include_usage"] is True
if (
self.stream_options
and self.stream_options.get("include_usage", False) == True
):
response["usage"] = chunk.get("usage", None)
return response
except Exception as e:
raise Exception(