diff --git a/litellm/__pycache__/main.cpython-311.pyc b/litellm/__pycache__/main.cpython-311.pyc index 62fec350c1..4745a44d98 100644 Binary files a/litellm/__pycache__/main.cpython-311.pyc and b/litellm/__pycache__/main.cpython-311.pyc differ diff --git a/litellm/__pycache__/utils.cpython-311.pyc b/litellm/__pycache__/utils.cpython-311.pyc index 58e8830f85..4977702a25 100644 Binary files a/litellm/__pycache__/utils.cpython-311.pyc and b/litellm/__pycache__/utils.cpython-311.pyc differ diff --git a/litellm/main.py b/litellm/main.py index 3d63f7e36b..a4803b209f 100644 --- a/litellm/main.py +++ b/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 diff --git a/litellm/tests/test_mock_request.py b/litellm/tests/test_mock_request.py index 637aed3e7b..33a1ac8ba0 100644 --- a/litellm/tests/test_mock_request.py +++ b/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() \ No newline at end of file +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() \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 13f83e81e4..d1c37d7bf0 100644 --- a/pyproject.toml +++ b/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"