From 86d46308bf11394a9b493c5e9ff2ea360b9be247 Mon Sep 17 00:00:00 2001 From: ffreemt Date: Fri, 24 May 2024 11:09:11 +0800 Subject: [PATCH] Make return-exceptions as default behavior in litellm.batch_completion --- litellm/main.py | 22 +++++++------------ ...test_batch_completion_return_exceptions.py | 13 ----------- 2 files changed, 8 insertions(+), 27 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index 826eb5c12..a4f1cde0c 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -2332,11 +2332,6 @@ def batch_completion( """ args = locals() - # extra kw for dealing with exceptions - return_exceptions = args.get("kwargs").get("return_exceptions", False) - if "return_exceptions" in args.get("kwargs"): - args.get("kwargs").pop("return_exceptions") - batch_messages = messages completions = [] model = model @@ -2391,15 +2386,14 @@ def batch_completion( # Retrieve the results from the futures # results = [future.result() for future in completions] - if return_exceptions: - results = [] - for future in completions: - try: - results.append(future.result()) - except Exception as exc: - results.append(exc) - else: # original - results = [future.result() for future in completions] + # return exceptions if any + results = [] + for future in completions: + try: + results.append(future.result()) + except Exception as exc: + results.append(exc) + return results diff --git a/litellm/tests/test_batch_completion_return_exceptions.py b/litellm/tests/test_batch_completion_return_exceptions.py index df330f65d..4e285771f 100644 --- a/litellm/tests/test_batch_completion_return_exceptions.py +++ b/litellm/tests/test_batch_completion_return_exceptions.py @@ -1,30 +1,17 @@ """https://github.com/BerriAI/litellm/pull/3397/commits/a7ec1772b1457594d3af48cdcb0a382279b841c7#diff-44852387ceb00aade916d6b314dfd5d180499e54f35209ae9c07179febe08b4b.""" """Test batch_completion's return_exceptions.""" -import pytest import litellm msg1 = [{"role": "user", "content": "hi 1"}] msg2 = [{"role": "user", "content": "hi 2"}] -def test_batch_completion_return_exceptions_default(): - """Test batch_completion's return_exceptions.""" - with pytest.raises(Exception): - _ = litellm.batch_completion( - model="gpt-3.5-turbo", - messages=[msg1, msg2], - api_key="sk_xxx", # deliberately set invalid key - # return_exceptions=False, - ) - - def test_batch_completion_return_exceptions_true(): """Test batch_completion's return_exceptions.""" res = litellm.batch_completion( model="gpt-3.5-turbo", messages=[msg1, msg2], api_key="sk_xxx", # deliberately set invalid key - return_exceptions=True, ) assert isinstance(res[0], litellm.exceptions.AuthenticationError) \ No newline at end of file