From c50d8afc865ea941b7eea3f1ebad53bafcecf311 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Mon, 15 Jan 2024 17:17:43 -0800
Subject: [PATCH] (test) test_post_call_rule

---
 litellm/tests/test_rules.py | 77 +++++++++++++++++++++----------------
 1 file changed, 43 insertions(+), 34 deletions(-)

diff --git a/litellm/tests/test_rules.py b/litellm/tests/test_rules.py
index 664b9db08d..92ec20fe42 100644
--- a/litellm/tests/test_rules.py
+++ b/litellm/tests/test_rules.py
@@ -19,15 +19,6 @@ def my_pre_call_rule(input: str):
     return True
 
 
-def my_post_call_rule(input: str):
-    input = input.lower()
-    print(f"input: {input}")
-    print(f"INSIDE MY POST CALL RULE, len(input) - {len(input)}")
-    if "sorry" in input:
-        return False
-    return True
-
-
 ## Test 1: Pre-call rule
 def test_pre_call_rule():
     try:
@@ -55,32 +46,50 @@ def test_pre_call_rule():
     litellm.pre_call_rules = []
 
 
+def my_post_call_rule(input: str):
+    input = input.lower()
+    print(f"input: {input}")
+    print(f"INSIDE MY POST CALL RULE, len(input) - {len(input)}")
+    if len(input) < 200:
+        return {
+            "decision": False,
+            "message": "This violates LiteLLM Proxy Rules. Response too short",
+        }
+    return {"decision": True}
+
+
 # test_pre_call_rule()
-## Test 2: Post-call rule
+# Test 2: Post-call rule
 # commenting out of ci/cd since llm's have variable output which was causing our pipeline to fail erratically.
-# def test_post_call_rule():
-#     try:
-#         litellm.pre_call_rules = []
-#         litellm.post_call_rules = [my_post_call_rule]
-#         ### completion
-#         response = completion(model="gpt-3.5-turbo",
-#                       messages=[{"role": "user", "content": "say sorry"}],
-#                       fallbacks=["deepinfra/Gryphe/MythoMax-L2-13b"])
-#         pytest.fail(f"Completion call should have been failed. ")
-#     except:
-#         pass
-#     print(f"MAKING ACOMPLETION CALL")
-#     # litellm.set_verbose = True
-#     ### async completion
-#     async def test_async_response():
-#         messages=[{"role": "user", "content": "say sorry"}]
-#         try:
-#             response = await acompletion(model="gpt-3.5-turbo", messages=messages)
-#             pytest.fail(f"acompletion call should have been failed.")
-#         except Exception as e:
-#             pass
-#     asyncio.run(test_async_response())
-#     litellm.pre_call_rules = []
-#     litellm.post_call_rules = []
+def test_post_call_rule():
+    try:
+        litellm.pre_call_rules = []
+        litellm.post_call_rules = [my_post_call_rule]
+        ### completion
+        response = completion(
+            model="gpt-3.5-turbo",
+            messages=[{"role": "user", "content": "say sorry"}],
+            max_tokens=2,
+        )
+        pytest.fail(f"Completion call should have been failed. ")
+    except Exception as e:
+        print("Got exception", e)
+        print(type(e))
+        print(vars(e))
+        pass
+    # print(f"MAKING ACOMPLETION CALL")
+    # litellm.set_verbose = True
+    ### async completion
+    # async def test_async_response():
+    #     messages=[{"role": "user", "content": "say sorry"}]
+    #     try:
+    #         response = await acompletion(model="gpt-3.5-turbo", messages=messages)
+    #         pytest.fail(f"acompletion call should have been failed.")
+    #     except Exception as e:
+    #         pass
+    # asyncio.run(test_async_response())
+    litellm.pre_call_rules = []
+    litellm.post_call_rules = []
+
 
 # test_post_call_rule()