From 4761345311f98ada5bdf6d7bcf0e89b5489429b5 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Tue, 30 Apr 2024 19:30:18 -0700
Subject: [PATCH] fix(main.py): fix mock completion response

---
 litellm/main.py | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/litellm/main.py b/litellm/main.py
index 569418eca03..cdea40d1199 100644
--- a/litellm/main.py
+++ b/litellm/main.py
@@ -360,7 +360,7 @@ def mock_completion(
     model: str,
     messages: List,
     stream: Optional[bool] = False,
-    mock_response: str = "This is a mock request",
+    mock_response: Union[str, Exception] = "This is a mock request",
     logging=None,
     **kwargs,
 ):
@@ -388,10 +388,11 @@ def mock_completion(
     """
     try:
         ## LOGGING
-        logging.pre_call(
-            input=messages,
-            api_key="mock-key",
-        )
+        if logging is not None:
+            logging.pre_call(
+                input=messages,
+                api_key="mock-key",
+            )
         if isinstance(mock_response, Exception):
             raise litellm.APIError(
                 status_code=500,  # type: ignore