From ed07de2729840a7a194e5b64d8f9458fa26f594d Mon Sep 17 00:00:00 2001 From: TheDiscoMole Date: Fri, 19 Jan 2024 23:36:24 +0100 Subject: [PATCH] changing ollama response parsing to expected behaviour --- litellm/llms/ollama.py | 9 ++++----- litellm/llms/ollama_chat.py | 12 ++++-------- 2 files changed, 8 insertions(+), 13 deletions(-) diff --git a/litellm/llms/ollama.py b/litellm/llms/ollama.py index 0f99622b3..89db4628c 100644 --- a/litellm/llms/ollama.py +++ b/litellm/llms/ollama.py @@ -199,12 +199,13 @@ def get_ollama_response( ## RESPONSE OBJECT model_response["choices"][0]["finish_reason"] = "stop" if optional_params.get("format", "") == "json": + function_call = json.loads(response_json["response"]) message = litellm.Message( content=None, tool_calls=[ { "id": f"call_{str(uuid.uuid4())}", - "function": {"arguments": response_json["response"], "name": ""}, + "function": {"name": function_call["name"], "arguments": json.dumps(function_call["arguments"])}, "type": "function", } ], @@ -295,15 +296,13 @@ async def ollama_acompletion(url, data, model_response, encoding, logging_obj): ## RESPONSE OBJECT model_response["choices"][0]["finish_reason"] = "stop" if data.get("format", "") == "json": + function_call = json.loads(response_json["response"]) message = litellm.Message( content=None, tool_calls=[ { "id": f"call_{str(uuid.uuid4())}", - "function": { - "arguments": response_json["response"], - "name": "", - }, + "function": {"name": function_call["name"], "arguments": json.dumps(function_call["arguments"])}, "type": "function", } ], diff --git a/litellm/llms/ollama_chat.py b/litellm/llms/ollama_chat.py index 31e3f0d16..e80f9d5d6 100644 --- a/litellm/llms/ollama_chat.py +++ b/litellm/llms/ollama_chat.py @@ -202,15 +202,13 @@ def get_ollama_response( ## RESPONSE OBJECT model_response["choices"][0]["finish_reason"] = "stop" if data.get("format", "") == "json": + function_call = json.loads(response_json["message"]["content"]) message = litellm.Message( content=None, tool_calls=[ { "id": f"call_{str(uuid.uuid4())}", - "function": { - "arguments": response_json["message"]["content"], - "name": "", - }, + "function": {"name": function_call["name"], "arguments": json.dumps(function_call["arguments"])}, "type": "function", } ], @@ -302,15 +300,13 @@ async def ollama_acompletion(url, data, model_response, encoding, logging_obj): ## RESPONSE OBJECT model_response["choices"][0]["finish_reason"] = "stop" if data.get("format", "") == "json": + function_call = json.loads(response_json["message"]["content"]) message = litellm.Message( content=None, tool_calls=[ { "id": f"call_{str(uuid.uuid4())}", - "function": { - "arguments": response_json["message"]["content"], - "name": "", - }, + "function": {"name": function_call["name"], "arguments": json.dumps(function_call["arguments"])}, "type": "function", } ],