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

(feat) ollama_chat acompletion without streaming

Browse source
This commit is contained in:
ishaan-jaff 2023-12-26 20:01:15 +05:30
parent 751d57379d
commit 3839213d28
1 changed files with 5 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

11
litellm/llms/ollama_chat.py
View file

@ -283,18 +283,19 @@ async def ollama_acompletion(url, data, model_response, encoding, logging_obj):
text = await resp.text() text = await resp.text()
raise OllamaError(status_code=resp.status, message=text) raise OllamaError(status_code=resp.status, message=text)
response_json = await resp.json()
## LOGGING ## LOGGING
logging_obj.post_call( logging_obj.post_call(
input=data["prompt"], input=data,
api_key="", api_key="",
original_response=resp.text, original_response=response_json,
additional_args={ additional_args={
"headers": None, "headers": None,
"api_base": url, "api_base": url,
}, },
) )
response_json = await resp.json()
## RESPONSE OBJECT ## RESPONSE OBJECT
model_response["choices"][0]["finish_reason"] = "stop" model_response["choices"][0]["finish_reason"] = "stop"
if data.get("format", "") == "json": if data.get("format", "") == "json":
@ -313,9 +314,7 @@ async def ollama_acompletion(url, data, model_response, encoding, logging_obj):
) )
model_response["choices"][0]["message"] = message model_response["choices"][0]["message"] = message
else: else:
model_response["choices"][0]["message"]["content"] = response_json[ model_response["choices"][0]["message"] = response_json["message"]
"response"
]
model_response["created"] = int(time.time()) model_response["created"] = int(time.time())
model_response["model"] = "ollama/" + data["model"] model_response["model"] = "ollama/" + data["model"]
prompt_tokens = response_json["prompt_eval_count"] # type: ignore prompt_tokens = response_json["prompt_eval_count"] # type: ignore