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

(feat) ollama_chat - streaming

Browse source
This commit is contained in:
ishaan-jaff 2023-12-25 23:38:47 +05:30
parent 916ba9a6b3
commit f27a93a4dc
1 changed files with 43 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

43
litellm/utils.py
View file

@ -6781,6 +6781,41 @@ class CustomStreamWrapper:
except Exception as e:
raise e
def handle_ollama_chat_stream(self, chunk):
# for ollama_chat/ provider
try:
if isinstance(chunk, dict):
json_chunk = chunk
else:
json_chunk = json.loads(chunk)
if "error" in json_chunk:
raise Exception(f"Ollama Error - {json_chunk}")
text = ""
is_finished = False
finish_reason = None
if json_chunk["done"] == True:
text = ""
is_finished = True
finish_reason = "stop"
return {
"text": text,
"is_finished": is_finished,
"finish_reason": finish_reason,
}
elif "message" in json_chunk:
print_verbose(f"delta content: {json_chunk}")
text = json_chunk["message"]["content"]
return {
"text": text,
"is_finished": is_finished,
"finish_reason": finish_reason,
}
else:
raise Exception(f"Ollama Error - {json_chunk}")
except Exception as e:
raise e
def handle_bedrock_stream(self, chunk):
if hasattr(chunk, "get"):
chunk = chunk.get("chunk")
@ -6993,6 +7028,14 @@ class CustomStreamWrapper:
model_response.choices[0].finish_reason = response_obj[
"finish_reason"
]
elif self.custom_llm_provider == "ollama_chat":
response_obj = self.handle_ollama_chat_stream(chunk)
completion_obj["content"] = response_obj["text"]
print_verbose(f"completion obj content: {completion_obj['content']}")
if response_obj["is_finished"]:
model_response.choices[0].finish_reason = response_obj[
"finish_reason"
]
elif self.custom_llm_provider == "text-completion-openai":
response_obj = self.handle_openai_text_completion_chunk(chunk)
completion_obj["content"] = response_obj["text"]