(feat) proxy: add mode in model info

This commit is contained in:
ishaan-jaff 2023-12-06 11:29:59 -08:00
parent 7c77cc3cfa
commit fd86876164
2 changed files with 6 additions and 3 deletions

View file

@ -68,11 +68,12 @@ async def _perform_health_check(model_list: list):
tasks = [] tasks = []
for model in model_list: for model in model_list:
litellm_params = model["litellm_params"] litellm_params = model["litellm_params"]
model_info = model.get("model_info", {})
litellm_params["model"] = litellm.utils.remove_model_id(litellm_params["model"]) litellm_params["model"] = litellm.utils.remove_model_id(litellm_params["model"])
litellm_params["messages"] = _get_random_llm_message() litellm_params["messages"] = _get_random_llm_message()
prepped_params.append(litellm_params) prepped_params.append(litellm_params)
if model.get("mode", None) == "embedding": if model_info.get("mode", None) == "embedding":
# this is an embedding model # this is an embedding model
tasks.append(_check_embedding_model(litellm_params)) tasks.append(_check_embedding_model(litellm_params))
else: else:

View file

@ -11,6 +11,7 @@ model_list:
api_base: os.environ/AZURE_API_BASE api_base: os.environ/AZURE_API_BASE
api_key: os.environ/AZURE_API_KEY api_key: os.environ/AZURE_API_KEY
api_version: "2023-07-01-preview" api_version: "2023-07-01-preview"
model_info:
mode: embedding mode: embedding
- model_name: openai-gpt-3.5 - model_name: openai-gpt-3.5
litellm_params: litellm_params:
@ -20,6 +21,7 @@ model_list:
litellm_params: litellm_params:
model: text-embedding-ada-002 model: text-embedding-ada-002
api_key: os.environ/OPENAI_API_KEY api_key: os.environ/OPENAI_API_KEY
model_info:
mode: embedding mode: embedding