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

Merge pull request #3537 from BerriAI/litellm_support_stream_options_param

Browse source 
[Feat] support `stream_options` param for OpenAI
This commit is contained in:
Ishaan Jaff 2024-05-09 08:34:08 -07:00 • committed by GitHub
commit 0b1885ca99
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 101 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

39
litellm/utils.py
View file

@ -612,6 +612,7 @@ class ModelResponse(OpenAIObject):
system_fingerprint=None,
usage=None,
stream=None,
stream_options=None,
response_ms=None,
hidden_params=None,
**params,
@ -658,6 +659,12 @@ class ModelResponse(OpenAIObject):
usage = usage
elif stream is None or stream == False:
usage = Usage()
elif (
stream == True
and stream_options is not None
and stream_options.get("include_usage") == True
):
usage = Usage()
if hidden_params:
self._hidden_params = hidden_params
@ -4839,6 +4846,7 @@ def get_optional_params(
top_p=None,
n=None,
stream=False,
stream_options=None,
stop=None,
max_tokens=None,
presence_penalty=None,
@ -4908,6 +4916,7 @@ def get_optional_params(
"top_p": None,
"n": None,
"stream": None,
"stream_options": None,
"stop": None,
"max_tokens": None,
"presence_penalty": None,
@ -5779,6 +5788,8 @@ def get_optional_params(
optional_params["n"] = n
if stream is not None:
optional_params["stream"] = stream
if stream_options is not None:
optional_params["stream_options"] = stream_options
if stop is not None:
optional_params["stop"] = stop
if max_tokens is not None:
@ -6085,6 +6096,7 @@ def get_supported_openai_params(model: str, custom_llm_provider: str):
"top_p",
"n",
"stream",
"stream_options",
"stop",
"max_tokens",
"presence_penalty",
@ -9502,7 +9514,12 @@ def get_secret(
# replicate/anthropic/cohere
class CustomStreamWrapper:
def __init__(
self, completion_stream, model, custom_llm_provider=None, logging_obj=None
self,
completion_stream,
model,
custom_llm_provider=None,
logging_obj=None,
stream_options=None,
):
self.model = model
self.custom_llm_provider = custom_llm_provider
@ -9528,6 +9545,7 @@ class CustomStreamWrapper:
self.response_id = None
self.logging_loop = None
self.rules = Rules()
self.stream_options = stream_options
def __iter__(self):
return self
@ -9968,6 +9986,7 @@ class CustomStreamWrapper:
is_finished = False
finish_reason = None
logprobs = None
usage = None
original_chunk = None # this is used for function/tool calling
if len(str_line.choices) > 0:
if (
@ -10002,12 +10021,15 @@ class CustomStreamWrapper:
else:
logprobs = None
usage = getattr(str_line, "usage", None)
return {
"text": text,
"is_finished": is_finished,
"finish_reason": finish_reason,
"logprobs": logprobs,
"original_chunk": str_line,
"usage": usage,
}
except Exception as e:
traceback.print_exc()
@ -10310,7 +10332,9 @@ class CustomStreamWrapper:
raise e
def model_response_creator(self):
model_response = ModelResponse(stream=True, model=self.model)
model_response = ModelResponse(
stream=True, model=self.model, stream_options=self.stream_options
)
if self.response_id is not None:
model_response.id = self.response_id
else:
@ -10630,6 +10654,12 @@ class CustomStreamWrapper:
if response_obj["logprobs"] is not None:
model_response.choices[0].logprobs = response_obj["logprobs"]
if (
self.stream_options is not None
and self.stream_options["include_usage"] == True
):
model_response.usage = response_obj["usage"]
model_response.model = self.model
print_verbose(
f"model_response finish reason 3: {self.received_finish_reason}; response_obj={response_obj}"
@ -10717,6 +10747,11 @@ class CustomStreamWrapper:
except Exception as e:
model_response.choices[0].delta = Delta()
else:
if (
self.stream_options is not None
and self.stream_options["include_usage"] == True
):
return model_response
return
print_verbose(
f"model_response.choices[0].delta: {model_response.choices[0].delta}; completion_obj: {completion_obj}"