From 6561e0838ea36695847d38f648b81c46d225ce08 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Sat, 11 May 2024 13:09:17 -0700
Subject: [PATCH] test - router.batch_acompletion

---
 litellm/tests/test_router_batch_completion.py | 60 +++++++++++++++++++
 1 file changed, 60 insertions(+)
 create mode 100644 litellm/tests/test_router_batch_completion.py

diff --git a/litellm/tests/test_router_batch_completion.py b/litellm/tests/test_router_batch_completion.py
new file mode 100644
index 000000000..f2873b18d
--- /dev/null
+++ b/litellm/tests/test_router_batch_completion.py
@@ -0,0 +1,60 @@
+#### What this tests ####
+# This tests litellm router with batch completion
+
+import sys, os, time, openai
+import traceback, asyncio
+import pytest
+
+sys.path.insert(
+    0, os.path.abspath("../..")
+)  # Adds the parent directory to the system path
+import litellm
+from litellm import Router
+from litellm.router import Deployment, LiteLLM_Params, ModelInfo
+from concurrent.futures import ThreadPoolExecutor
+from collections import defaultdict
+from dotenv import load_dotenv
+import os, httpx
+
+load_dotenv()
+
+
+@pytest.mark.asyncio
+async def test_batch_completion_multiple_models():
+    litellm.set_verbose = True
+
+    router = litellm.Router(
+        model_list=[
+            {
+                "model_name": "gpt-3.5-turbo",
+                "litellm_params": {
+                    "model": "gpt-3.5-turbo",
+                },
+            },
+            {
+                "model_name": "groq-llama",
+                "litellm_params": {
+                    "model": "groq/llama3-8b-8192",
+                },
+            },
+        ]
+    )
+
+    response = await router.abatch_completion(
+        models=["gpt-3.5-turbo", "groq-llama"],
+        messages=[
+            {"role": "user", "content": "is litellm becoming a better product ?"}
+        ],
+        max_tokens=15,
+    )
+
+    print(response)
+    assert len(response) == 2
+
+    models_in_responses = []
+    for individual_response in response:
+        _model = individual_response["model"]
+        models_in_responses.append(_model)
+
+    # assert both models are different
+    assert models_in_responses[0] != models_in_responses[1]