From 64d229caaa6e57f7dd169d82c5dc6a06d3aae5b1 Mon Sep 17 00:00:00 2001 From: ffreemt Date: Thu, 2 May 2024 19:30:01 +0800 Subject: [PATCH 1/2] Add return_exceptions to litellm.batch_completion for optionally returing exceptions and partial resuslt instead of throwing exceptions --- litellm/main.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/litellm/main.py b/litellm/main.py index 51ec954018d..11ab0a0b9d4 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -2303,6 +2303,7 @@ def batch_completion( user (str, optional): The user string for generating completions. Defaults to "". deployment_id (optional): The deployment ID for generating completions. Defaults to None. request_timeout (int, optional): The request timeout for generating completions. Defaults to None. + return_exceptions (bool): Whether to return exceptions and partial results when exceptions occur. Defaults to False. Returns: list: A list of completion results. @@ -2361,7 +2362,17 @@ def batch_completion( completions.append(future) # Retrieve the results from the futures - results = [future.result() for future in completions] + # 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: + results = [future.result() for future in completions] + return results From a7ec1772b1457594d3af48cdcb0a382279b841c7 Mon Sep 17 00:00:00 2001 From: ffreemt Date: Fri, 3 May 2024 11:28:38 +0800 Subject: [PATCH 2/2] Add litellm\tests\test_batch_completion_return_exceptions.py --- .gitignore | 2 ++ litellm/main.py | 3 +- ...test_batch_completion_return_exceptions.py | 29 +++++++++++++++++++ 3 files changed, 33 insertions(+), 1 deletion(-) create mode 100644 litellm/tests/test_batch_completion_return_exceptions.py diff --git a/.gitignore b/.gitignore index abc4ecb0ced..50085bd29d9 100644 --- a/.gitignore +++ b/.gitignore @@ -52,3 +52,5 @@ litellm/proxy/_new_secret_config.yaml litellm/proxy/_new_secret_config.yaml litellm/proxy/_super_secret_config.yaml litellm/proxy/_super_secret_config.yaml +.python-version +litellm/llms/tokenizers/9b5ad71b2ce5302211f9c61530b329a4922fc6a4 diff --git a/litellm/main.py b/litellm/main.py index 11ab0a0b9d4..8fc07b9bf6d 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -2143,7 +2143,7 @@ def completion( """ assume input to custom LLM api bases follow this format: resp = requests.post( - api_base, + api_base, json={ 'model': 'meta-llama/Llama-2-13b-hf', # model name 'params': { @@ -2280,6 +2280,7 @@ def batch_completion( deployment_id=None, request_timeout: Optional[int] = None, timeout: Optional[int] = 600, + return_exceptions: bool = False, # Optional liteLLM function params **kwargs, ): diff --git a/litellm/tests/test_batch_completion_return_exceptions.py b/litellm/tests/test_batch_completion_return_exceptions.py new file mode 100644 index 00000000000..b44146993f2 --- /dev/null +++ b/litellm/tests/test_batch_completion_return_exceptions.py @@ -0,0 +1,29 @@ +"""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)