mock responses for streaming

litellm/main.py

@@ -954,6 +954,30 @@ def batch_completion(
     return results
 
 
+def mock_completion(model: str, messages: List, stream: bool = False, mock_response: str = "This is a mock request"):
+    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
+    except:
+        raise Exception("Mock completion response failed")
 ### EMBEDDING ENDPOINTS ####################
 @client
 @timeout(  # type: ignore

litellm/tests/test_mock_request.py

@@ -13,9 +13,23 @@ def test_mock_request():
     try:
         model = "gpt-3.5-turbo"
         messages = [{"role": "user", "content": "Hey, I'm a mock request"}]
-        response = litellm.completion(model=model, messages=messages, mock_request=True)
+        response = litellm.mock_completion(model=model, messages=messages)
         print(response)
     except:
         traceback.print_exc()
 
-test_mock_request()
+def test_streaming_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, stream=True)
+        complete_response = "" 
+        for chunk in response: 
+            print(f"chunk: {chunk}")
+            complete_response += chunk["choices"][0]["delta"]["content"]
+        if complete_response == "": 
+            raise Exception("Empty response received")
+    except:
+        traceback.print_exc()
+
+test_streaming_mock_request()

pyproject.toml

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