fix: Handle None response and update InvalidRequestError test

This commit is contained in:
Paul Gauthier (aider) 2024-11-08 10:02:48 -08:00
parent b79c09cf58
commit 14d02bc843
2 changed files with 5 additions and 1 deletions

View file

@ -71,6 +71,8 @@ def simple_send_with_retries(model_name, messages, extra_params=None):
} }
_hash, response = send_completion(**kwargs) _hash, response = send_completion(**kwargs)
if not response or not hasattr(response, 'choices') or not response.choices:
return None
return response.choices[0].message.content return response.choices[0].message.content
except litellm_ex.exceptions_tuple() as err: except litellm_ex.exceptions_tuple() as err:
ex_info = litellm_ex.get_ex_info(err) ex_info = litellm_ex.get_ex_info(err)

View file

@ -85,7 +85,9 @@ class TestSendChat(unittest.TestCase):
mock.status_code = 400 mock.status_code = 400
mock_completion.side_effect = litellm.InvalidRequestError( mock_completion.side_effect = litellm.InvalidRequestError(
"Invalid request", response=mock, llm_provider="test_provider", model="test_model" message="Invalid request",
llm_provider="test_provider",
model="test_model"
) )
result = simple_send_with_retries(self.mock_model, self.mock_messages) result = simple_send_with_retries(self.mock_model, self.mock_messages)