fix(utils.py): fix streaming chunks to not return role, unless set

This commit is contained in:
Krrish Dholakia 2024-02-01 09:53:33 -08:00
parent 76f053f85b
commit d46df34ff5
2 changed files with 7 additions and 6 deletions

View file

@ -145,8 +145,8 @@ def get_ollama_response(
): # completion(top_k=3) > cohere_config(top_k=3) <- allows for dynamic variables to be passed in ): # completion(top_k=3) > cohere_config(top_k=3) <- allows for dynamic variables to be passed in
optional_params[k] = v optional_params[k] = v
optional_params["stream"] = optional_params.get("stream", False) stream = optional_params.pop("stream", False)
data = {"model": model, "prompt": prompt, **optional_params} data = {"model": model, "prompt": prompt, "options": optional_params}
## LOGGING ## LOGGING
logging_obj.pre_call( logging_obj.pre_call(
input=None, input=None,
@ -159,7 +159,7 @@ def get_ollama_response(
}, },
) )
if acompletion is True: if acompletion is True:
if optional_params.get("stream", False) == True: if stream == True:
response = ollama_async_streaming( response = ollama_async_streaming(
url=url, url=url,
data=data, data=data,
@ -176,7 +176,7 @@ def get_ollama_response(
logging_obj=logging_obj, logging_obj=logging_obj,
) )
return response return response
elif optional_params.get("stream", False) == True: elif stream == True:
return ollama_completion_stream(url=url, data=data, logging_obj=logging_obj) return ollama_completion_stream(url=url, data=data, logging_obj=logging_obj)
response = requests.post(url=f"{url}", json=data, timeout=litellm.request_timeout) response = requests.post(url=f"{url}", json=data, timeout=litellm.request_timeout)
@ -254,7 +254,7 @@ async def ollama_async_streaming(url, data, model_response, encoding, logging_ob
) as response: ) as response:
if response.status_code != 200: if response.status_code != 200:
raise OllamaError( raise OllamaError(
status_code=response.status_code, message=response.text status_code=response.status_code, message=await response.aread()
) )
streamwrapper = litellm.CustomStreamWrapper( streamwrapper = litellm.CustomStreamWrapper(

View file

@ -233,7 +233,8 @@ class Delta(OpenAIObject):
def __init__(self, content=None, role=None, **params): def __init__(self, content=None, role=None, **params):
super(Delta, self).__init__(**params) super(Delta, self).__init__(**params)
self.content = content self.content = content
self.role = role if role is not None:
self.role = role
def __contains__(self, key): def __contains__(self, key):
# Define custom behavior for the 'in' operator # Define custom behavior for the 'in' operator