mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
fix(route_llm_request.py): map team model from list in route llm request
remove unnecessary proxy model table lookup for model alias Fixes issue where aliases weren't being consistently written to model table
This commit is contained in:
parent
058503ea11
commit
323c54d117
2 changed files with 50 additions and 1 deletions
|
|
@ -33,6 +33,25 @@ class ProxyModelNotFoundError(HTTPException):
|
|||
super().__init__(status_code=status.HTTP_400_BAD_REQUEST, detail=detail)
|
||||
|
||||
|
||||
def get_team_id_from_data(data: dict) -> Optional[str]:
|
||||
"""
|
||||
Get the team id from the data's metadata or litellm_metadata params.
|
||||
"""
|
||||
if (
|
||||
"metadata" in data
|
||||
and data["metadata"] is not None
|
||||
and "user_api_key_team_id" in data["metadata"]
|
||||
):
|
||||
return data["metadata"].get("user_api_key_team_id")
|
||||
elif (
|
||||
"litellm_metadata" in data
|
||||
and data["litellm_metadata"] is not None
|
||||
and "user_api_key_team_id" in data["litellm_metadata"]
|
||||
):
|
||||
return data["litellm_metadata"].get("user_api_key_team_id")
|
||||
return None
|
||||
|
||||
|
||||
async def route_request(
|
||||
data: dict,
|
||||
llm_router: Optional[LitellmRouter],
|
||||
|
|
@ -55,6 +74,7 @@ async def route_request(
|
|||
"""
|
||||
Common helper to route the request
|
||||
"""
|
||||
team_id = get_team_id_from_data(data)
|
||||
router_model_names = llm_router.model_names if llm_router is not None else []
|
||||
if "api_key" in data or "api_base" in data:
|
||||
return getattr(llm_router, f"{route_type}")(**data)
|
||||
|
|
@ -78,7 +98,16 @@ async def route_request(
|
|||
models = [model.strip() for model in data.pop("model").split(",")]
|
||||
return llm_router.abatch_completion(models=models, **data)
|
||||
elif llm_router is not None:
|
||||
if (
|
||||
team_model_name = (
|
||||
llm_router.map_team_model(data["model"], team_id)
|
||||
if team_id is not None
|
||||
else None
|
||||
)
|
||||
if team_model_name is not None:
|
||||
data["model"] = team_model_name
|
||||
return getattr(llm_router, f"{route_type}")(**data)
|
||||
|
||||
elif (
|
||||
data["model"] in router_model_names
|
||||
or data["model"] in llm_router.get_model_ids()
|
||||
):
|
||||
|
|
|
|||
|
|
@ -5407,6 +5407,26 @@ class Router:
|
|||
ids.append(id)
|
||||
return ids
|
||||
|
||||
def map_team_model(self, team_model_name: str, team_id: str) -> Optional[str]:
|
||||
"""
|
||||
Map a team model name to a team-specific model name.
|
||||
|
||||
Returns:
|
||||
- team_model_name: str - the team-specific model name
|
||||
- None: if no team-specific model name is found
|
||||
"""
|
||||
for model in self.model_list:
|
||||
model_team_id = model["model_info"].get("team_id")
|
||||
model_team_public_model_name = model["model_info"].get(
|
||||
"team_public_model_name"
|
||||
)
|
||||
if (
|
||||
model_team_id == team_id
|
||||
and model_team_public_model_name == team_model_name
|
||||
):
|
||||
return model["model_name"]
|
||||
return None
|
||||
|
||||
def _get_all_deployments(
|
||||
self, model_name: str, model_alias: Optional[str] = None
|
||||
) -> List[DeploymentTypedDict]:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue