phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-27 11:43:54 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Merge pull request #371 from promptmetheus/simplify-mock-logic

Browse source 
Simplify mock logic
This commit is contained in:
Krish Dholakia 2023-09-14 09:23:45 -07:00 • committed by GitHub
commit 1e9097d997
1 changed files with 31 additions and 25 deletions
Download patch file Download diff file Expand all files Collapse all files

56
litellm/main.py
View file

@ -72,6 +72,32 @@ async def acompletion(*args, **kwargs):
else: else:
return response 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
except:
raise Exception("Mock completion response failed")
@client @client
@timeout( # type: ignore @timeout( # type: ignore
600 600
@ -96,6 +122,7 @@ def completion(
# Optional liteLLM function params # Optional liteLLM function params
*, *,
return_async=False, return_async=False,
mock_response: Optional[str] = None,
api_key: Optional[str] = None, api_key: Optional[str] = None,
api_version: Optional[str] = None, api_version: Optional[str] = None,
api_base: Optional[str] = None, api_base: Optional[str] = None,
@ -118,6 +145,10 @@ def completion(
caching = False, caching = False,
cache_params = {}, # optional to specify metadata for caching cache_params = {}, # optional to specify metadata for caching
) -> ModelResponse: ) -> ModelResponse:
# If `mock_response` is set, execute the `mock_completion` method instead.
if mock_response:
return mock_completion(model, messages, stream=stream, mock_response=mock_response)
args = locals() args = locals()
try: try:
logging = litellm_logging_obj logging = litellm_logging_obj
@ -978,31 +1009,6 @@ def batch_completion(
results = [future.result() for future in completions] results = [future.result() for future in completions]
return results 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", **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
except:
raise Exception("Mock completion response failed")
### EMBEDDING ENDPOINTS #################### ### EMBEDDING ENDPOINTS ####################
@client @client
@timeout( # type: ignore @timeout( # type: ignore