test(batch_completion): add tests for batch_completion and rate limit handler

litellm/tests/test_batch_completions.py

@@ -12,17 +12,24 @@ import litellm
 from litellm import batch_completion, batch_completion_models, completion, batch_completion_models_all_responses
 # litellm.set_verbose=True
 
-# def test_batch_completions():
-#     messages = [[{"role": "user", "content": "Hey, how's it going"}] for _ in range(5)]
-#     model = "gpt-3.5-turbo"
-#     try:
-#         result = batch_completion(model=model, messages=messages)
-#         print(result)
-#         print(len(result))
-#     except Timeout as e:
-#         pass
-#     except Exception as e:
-#         pytest.fail(f"An error occurred: {e}")
+def test_batch_completions():
+    messages = [[{"role": "user", "content": "write a short poem"}] for _ in range(3)]
+    model = "gpt-3.5-turbo"
+    try:
+        result = batch_completion(
+            model=model, 
+            messages=messages,
+            max_tokens=10,
+            temperature=0.2
+        )
+        print(result)
+        print(len(result))
+        assert(len(result)==3)
+    except Timeout as e:
+        pass
+    except Exception as e:
+        pytest.fail(f"An error occurred: {e}")
+# test_batch_completions()
 
 def test_batch_completions_models():
     try:
@@ -58,27 +65,28 @@ def test_batch_completion_models_all_responses():
 
 
 
-import asyncio
-##### USAGE ################
+def test_rate_limit_handler():
+    import asyncio
+    ##### USAGE ################
 
-jobs = [
-    {"model": "gpt-3.5-turbo-16k", "messages": [{"content": "Please provide a summary of the latest scientific discoveries.", "role": "user"}]},
-    {"model": "gpt-3.5-turbo-16k", "messages": [{"content": "Please provide a summary of the latest scientific discoveries.", "role": "user"}]},
-]
+    jobs = [
+        {"model": "gpt-3.5-turbo-16k", "messages": [{"content": "Please provide a summary of the latest scientific discoveries.", "role": "user"}]},
+        {"model": "gpt-3.5-turbo-16k", "messages": [{"content": "Please provide a summary of the latest scientific discoveries.", "role": "user"}]},
+    ]
 
-from litellm import RateLimitManager
+    from litellm import RateLimitManager
 
-handler = RateLimitManager(
-    max_requests_per_minute = 60,
-    max_tokens_per_minute = 20000
-)
-
-try:
-    asyncio.run(
-     handler.batch_completion(
-        jobs = jobs,
-        api_key=os.environ['OPENAI_API_KEY'],
+    handler = RateLimitManager(
+        max_requests_per_minute = 60,
+        max_tokens_per_minute = 20000
     )
-)
-except Exception as e:
-    print(e)
+
+    try:
+        asyncio.run(
+        handler.batch_completion(
+            jobs = jobs,
+            api_key=os.environ['OPENAI_API_KEY'],
+        )
+    )
+    except Exception as e:
+        print(e)