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 c0e3a47650
commit 90df061ffb

View file

@ -952,7 +952,7 @@ def batch_completion(
return results return results
## Use this in your testing pipeline, if you need to mock an LLM 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"): def mock_completion(model: str, messages: List, stream: bool = False, mock_response: str = "This is a mock request", **kwargs):
try: try:
model_response = ModelResponse() model_response = ModelResponse()
if stream: # return a generator object, iterate through the text in chunks of 3 char / chunk if stream: # return a generator object, iterate through the text in chunks of 3 char / chunk