feat(router.py): allow user to call specific deployment via id

Allows easier health checks for specific deployments by just passing in model id
This commit is contained in:
Krrish Dholakia 2024-06-19 13:02:46 -07:00
parent 9cc104eb03
commit 16889b8478
3 changed files with 42 additions and 0 deletions

View file

@ -3005,6 +3005,10 @@ async def chat_completion(
llm_router is not None and data["model"] in router_model_names
): # model in router model list
tasks.append(llm_router.acompletion(**data))
elif (
llm_router is not None and data["model"] in llm_router.get_model_ids()
): # model in router model list
tasks.append(llm_router.acompletion(**data))
elif (
llm_router is not None
and llm_router.model_group_alias is not None
@ -3256,6 +3260,10 @@ async def completion(
llm_response = asyncio.create_task(
llm_router.atext_completion(**data, specific_deployment=True)
)
elif (
llm_router is not None and data["model"] in llm_router.get_model_ids()
): # model in router model list
llm_response = asyncio.create_task(llm_router.atext_completion(**data))
elif (
llm_router is not None
and data["model"] not in router_model_names
@ -3488,6 +3496,10 @@ async def embeddings(
llm_router is not None and data["model"] in llm_router.deployment_names
): # model in router deployments, calling a specific deployment on the router
response = await llm_router.aembedding(**data, specific_deployment=True)
elif (
llm_router is not None and data["model"] in llm_router.get_model_ids()
): # model in router deployments, calling a specific deployment on the router
response = await llm_router.aembedding(**data)
elif (
llm_router is not None
and data["model"] not in router_model_names