add kwargs to mock_completion

Allows replacing `completion` with `mock_completion` without encountering errors.
This commit is contained in:
Toni Engelhardt 2023-09-14 10:27:19 +01:00 committed by GitHub
parent da9546bb99
commit 630b5d2c1f
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23

View file

@ -952,7 +952,7 @@ def batch_completion(
return results
## 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"):
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