(feat) completion ollama raise exception when ollama resp != 200

This commit is contained in:
ishaan-jaff 2023-11-10 08:54:03 -08:00
parent 8f10d1a3e3
commit 2f07460333

View file

@ -10,6 +10,14 @@ try:
except ImportError:
async_generator_imported = False # this should not throw an error, it will impact the 'import litellm' statement
class OllamaError(Exception):
def __init__(self, status_code, message):
self.status_code = status_code
self.message = message
super().__init__(
self.message
) # Call the base class constructor with the parameters it needs
class OllamaConfig():
"""
Reference: https://github.com/jmorganca/ollama/blob/main/docs/api.md#parameters
@ -123,6 +131,8 @@ def get_ollama_response_stream(
session = requests.Session()
with session.post(url, json=data, stream=True) as resp:
if resp.status_code != 200:
raise OllamaError(status_code=resp.status_code, message=resp.text)
for line in resp.iter_lines():
if line:
try:
@ -174,6 +184,8 @@ if async_generator_imported:
session = requests.Session()
with session.post(url, json=data, stream=True) as resp:
if resp.status_code != 200:
raise OllamaError(status_code=resp.status_code, message=resp.text)
for line in resp.iter_lines():
if line:
try: