tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-24 12:26:58 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(ollama.py): add support for async streaming

Browse Source
This commit is contained in:
Krrish Dholakia
2023-12-12 16:44:20 -08:00
parent dfdb17ae36
commit 8e7116635f
1 changed files with 43 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/llms/ollama.py
+43 -2
View File
@@ -141,8 +141,12 @@ def get_ollama_response_stream(
additional_args={"api_base": url, "complete_input_dict": data, "headers": {}, "acompletion": acompletion,},
)
if acompletion is True:
response = ollama_acompletion(url=url, data=data, model_response=model_response, encoding=encoding, logging_obj=logging_obj)
if optional_params.get("stream", False):
response = ollama_async_streaming(url=url, data=data, model_response=model_response, encoding=encoding, logging_obj=logging_obj)
else:
response = ollama_acompletion(url=url, data=data, model_response=model_response, encoding=encoding, logging_obj=logging_obj)
return response
else:
return ollama_completion_stream(url=url, data=data)
@@ -178,8 +182,45 @@ def ollama_completion_stream(url, data):
traceback.print_exc()
session.close()
async def ollama_acompletion(url, data, model_response, encoding, logging_obj):
async def ollama_async_streaming(url, data, model_response, encoding, logging_obj):
try:
timeout = aiohttp.ClientTimeout(total=600) # 10 minutes
async with aiohttp.ClientSession(timeout=timeout) as session:
resp = await session.post(url, json=data)
if resp.status != 200:
text = await resp.text()
raise OllamaError(status_code=resp.status, message=text)
async for line in resp.content.iter_any():
if line:
try:
json_chunk = line.decode("utf-8")
chunks = json_chunk.split("\n")
completion_string = ""
for chunk in chunks:
if chunk.strip() != "":
j = json.loads(chunk)
if "error" in j:
completion_obj = {
"role": "assistant",
"content": "",
"error": j
}
yield completion_obj
if "response" in j:
completion_obj = {
"role": "assistant",
"content": j["response"],
}
yield completion_obj
except Exception as e:
traceback.print_exc()
except Exception as e:
traceback.print_exc()
async def ollama_acompletion(url, data, model_response, encoding, logging_obj):
try:
timeout = aiohttp.ClientTimeout(total=600) # 10 minutes
async with aiohttp.ClientSession(timeout=timeout) as session: