phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-25 10:44:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Preserving the Pydantic Message Object

Browse source 
Following statement replaces the Pydantic Message Object and initialize it with the dict
model_response["choices"][0]["message"] = response_json["message"]

We need to make sure message is always litellm.Message object

As a fix, based on the code of ollama.py file, i am updating just the content intead of entire object for both sync and async functions
This commit is contained in:
Rajan Paneru 2024-05-10 22:12:32 +09:30
parent 8eb842dcf5
commit 65b07bcb8c
1 changed files with 2 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

4
litellm/llms/ollama_chat.py
View file

@ -300,7 +300,7 @@ def get_ollama_response(
model_response["choices"][0]["message"] = message
model_response["choices"][0]["finish_reason"] = "tool_calls"
else:
model_response["choices"][0]["message"] = response_json["message"]
model_response["choices"][0]["message"]["content"] = response_json["message"]["content"]
model_response["created"] = int(time.time())
model_response["model"] = "ollama/" + model
prompt_tokens = response_json.get("prompt_eval_count", litellm.token_counter(messages=messages)) # type: ignore
@ -484,7 +484,7 @@ async def ollama_acompletion(
model_response["choices"][0]["message"] = message
model_response["choices"][0]["finish_reason"] = "tool_calls"
else:
model_response["choices"][0]["message"] = response_json["message"]
model_response["choices"][0]["message"]["content"] = response_json["message"]["content"]
model_response["created"] = int(time.time())
model_response["model"] = "ollama_chat/" + data["model"]