From a7ec1772b1457594d3af48cdcb0a382279b841c7 Mon Sep 17 00:00:00 2001
From: ffreemt <yucongo+fmt@gmail.com>
Date: Fri, 3 May 2024 11:28:38 +0800
Subject: [PATCH] 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 abc4ecb0c..50085bd29 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 11ab0a0b9..8fc07b9bf 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 000000000..b44146993
--- /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)