phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

Make return-exceptions as default behavior in litellm.batch_completion

Browse source
This commit is contained in:
ffreemt 2024-05-24 11:09:11 +08:00
parent 2713272bba
commit 86d46308bf
2 changed files with 8 additions and 27 deletions
Download patch file Download diff file Expand all files Collapse all files

10
litellm/main.py
View file

@ -2332,11 +2332,6 @@ def batch_completion(
""" """
args = locals() 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 batch_messages = messages
completions = [] completions = []
model = model model = model
@ -2391,15 +2386,14 @@ def batch_completion(
# Retrieve the results from the futures # Retrieve the results from the futures
# results = [future.result() for future in completions] # results = [future.result() for future in completions]
if return_exceptions: # return exceptions if any
results = [] results = []
for future in completions: for future in completions:
try: try:
results.append(future.result()) results.append(future.result())
except Exception as exc: except Exception as exc:
results.append(exc) results.append(exc)
else: # original
results = [future.result() for future in completions]
return results return results

13
litellm/tests/test_batch_completion_return_exceptions.py
View file

@ -1,30 +1,17 @@
"""https://github.com/BerriAI/litellm/pull/3397/commits/a7ec1772b1457594d3af48cdcb0a382279b841c7#diff-44852387ceb00aade916d6b314dfd5d180499e54f35209ae9c07179febe08b4b.""" """https://github.com/BerriAI/litellm/pull/3397/commits/a7ec1772b1457594d3af48cdcb0a382279b841c7#diff-44852387ceb00aade916d6b314dfd5d180499e54f35209ae9c07179febe08b4b."""
"""Test batch_completion's return_exceptions.""" """Test batch_completion's return_exceptions."""
import pytest
import litellm import litellm
msg1 = [{"role": "user", "content": "hi 1"}] msg1 = [{"role": "user", "content": "hi 1"}]
msg2 = [{"role": "user", "content": "hi 2"}] 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(): def test_batch_completion_return_exceptions_true():
"""Test batch_completion's return_exceptions.""" """Test batch_completion's return_exceptions."""
res = litellm.batch_completion( res = litellm.batch_completion(
model="gpt-3.5-turbo", model="gpt-3.5-turbo",
messages=[msg1, msg2], messages=[msg1, msg2],
api_key="sk_xxx", # deliberately set invalid key api_key="sk_xxx", # deliberately set invalid key
return_exceptions=True,
) )
assert isinstance(res[0], litellm.exceptions.AuthenticationError) assert isinstance(res[0], litellm.exceptions.AuthenticationError)