fixes to mock completion

litellm/main.py

@@ -17,7 +17,8 @@ from litellm.utils import (
     CustomStreamWrapper,
     read_config_args,
     completion_with_fallbacks,
-    get_llm_provider
+    get_llm_provider,
+    mock_completion_streaming_obj
 )
 from .llms import anthropic
 from .llms import together_ai
@@ -72,30 +73,22 @@ async def acompletion(*args, **kwargs):
     else:
         return response
 
-## Use this in your testing pipeline, if you need to mock an LLM response
 def mock_completion(model: str, messages: List, stream: bool = False, mock_response: str = "This is a mock request", **kwargs):
     try:
-        model_response = ModelResponse()
-        if stream: # return a generator object, iterate through the text in chunks of 3 char / chunk
-            for i in range(0, len(mock_response), 3):
-                completion_obj = {"role": "assistant", "content": mock_response[i: i+3]}
-                yield {
-                "choices": 
-                    [
-                        {
-                            "delta": completion_obj,
-                            "finish_reason": None
-                        },
-                    ]
-                }
-        else:
-            ## RESPONSE OBJECT
-            completion_response = "This is a mock request"
-            model_response["choices"][0]["message"]["content"] = completion_response
-            model_response["created"] = time.time()
-            model_response["model"] = "MockResponse"
-            return model_response
+        model_response = ModelResponse(stream=stream)
+        if stream is True:
+            # don't try to access stream object,
+            response = mock_completion_streaming_obj(model_response, mock_response=mock_response, model=model)
+            return response
+        
+        completion_response = "This is a mock request"
+        model_response["choices"][0]["message"]["content"] = completion_response
+        model_response["created"] = time.time()
+        model_response["model"] = model
+        return model_response
+
     except:
+        traceback.print_exc()
         raise Exception("Mock completion response failed")
 
 @client

litellm/tests/test_mock_request.py

@@ -13,11 +13,13 @@ def test_mock_request():
     try:
         model = "gpt-3.5-turbo"
         messages = [{"role": "user", "content": "Hey, I'm a mock request"}]
-        response = litellm.mock_completion(model=model, messages=messages)
+        response = litellm.mock_completion(model=model, messages=messages, stream=False)
         print(response)
+        print(type(response))
     except:
         traceback.print_exc()
 
+# test_mock_request()
 def test_streaming_mock_request():
     try: 
         model = "gpt-3.5-turbo"

litellm/utils.py

@@ -2291,7 +2291,7 @@ class CustomStreamWrapper:
                 # Log the type of the received item
                 self.logging_obj.post_call(str(type(completion_stream)))
         if model in litellm.cohere_models:
-            # cohere does not return an iterator, so we need to wrap it in one
+            # these do not return an iterator, so we need to wrap it in one
             self.completion_stream = iter(completion_stream)
         else:
             self.completion_stream = completion_stream
@@ -2461,6 +2461,12 @@ class CustomStreamWrapper:
             raise StopAsyncIteration
 
 
+def mock_completion_streaming_obj(model_response, mock_response, model):
+    for i in range(0, len(mock_response), 3):
+        completion_obj = {"role": "assistant", "content": mock_response[i: i+3]}
+        model_response.choices[0].delta = completion_obj
+        yield model_response
+
 ########## Reading Config File ############################
 def read_config_args(config_path):
     try:

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "litellm"
-version = "0.1.624"
+version = "0.1.625"
 description = "Library to easily interface with LLM API providers"
 authors = ["BerriAI"]
 license = "MIT License"