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

(fix) return usage in mock_completion

Browse source
This commit is contained in:
ishaan-jaff 2024-01-19 11:25:47 -08:00
parent 6a695477ba
commit cb40f58cd3
1 changed files with 9 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

13
litellm/main.py
View file

@ -343,6 +343,11 @@ def mock_completion(
model_response["choices"][0]["message"]["content"] = mock_response model_response["choices"][0]["message"]["content"] = mock_response
model_response["created"] = int(time.time()) model_response["created"] = int(time.time())
model_response["model"] = model model_response["model"] = model
model_response.usage = Usage(
prompt_tokens=10, completion_tokens=20, total_tokens=30
)
return model_response return model_response
except: except:
@ -534,10 +539,6 @@ def completion(
non_default_params = { non_default_params = {
k: v for k, v in kwargs.items() if k not in default_params k: v for k, v in kwargs.items() if k not in default_params
} # model-specific params - pass them straight to the model/provider } # model-specific params - pass them straight to the model/provider
if mock_response:
return mock_completion(
model, messages, stream=stream, mock_response=mock_response
)
if timeout is None: if timeout is None:
timeout = ( timeout = (
kwargs.get("request_timeout", None) or 600 kwargs.get("request_timeout", None) or 600
@ -674,6 +675,10 @@ def completion(
optional_params=optional_params, optional_params=optional_params,
litellm_params=litellm_params, litellm_params=litellm_params,
) )
if mock_response:
return mock_completion(
model, messages, stream=stream, mock_response=mock_response
)
if custom_llm_provider == "azure": if custom_llm_provider == "azure":
# azure configs # azure configs
api_type = get_secret("AZURE_API_TYPE") or "azure" api_type = get_secret("AZURE_API_TYPE") or "azure"