(feat) ollama_chat - add streaming support

This commit is contained in:
ishaan-jaff 2023-12-25 23:38:01 +05:30
parent fa305054be
commit 916ba9a6b3

View file

@ -234,13 +234,13 @@ def ollama_completion_stream(url, data, logging_obj):
try: try:
if response.status_code != 200: if response.status_code != 200:
raise OllamaError( raise OllamaError(
status_code=response.status_code, message=response.text status_code=response.status_code, message=response.iter_lines()
) )
streamwrapper = litellm.CustomStreamWrapper( streamwrapper = litellm.CustomStreamWrapper(
completion_stream=response.iter_lines(), completion_stream=response.iter_lines(),
model=data["model"], model=data["model"],
custom_llm_provider="ollama", custom_llm_provider="ollama_chat",
logging_obj=logging_obj, logging_obj=logging_obj,
) )
for transformed_chunk in streamwrapper: for transformed_chunk in streamwrapper: