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

fix(proxy_server.py): restrict proxy to just listed models

Browse source
This commit is contained in:
Krrish Dholakia 2024-02-12 09:30:59 -08:00
parent 5d59fa166c
commit 8f989235ea
1 changed files with 18 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

21
litellm/proxy/proxy_server.py
View file

@ -2372,8 +2372,13 @@ async def chat_completion(
llm_router is not None and data["model"] in llm_router.deployment_names llm_router is not None and data["model"] in llm_router.deployment_names
): # model in router deployments, calling a specific deployment on the router ): # model in router deployments, calling a specific deployment on the router
response = await llm_router.acompletion(**data, specific_deployment=True) response = await llm_router.acompletion(**data, specific_deployment=True)
else: # router is not set elif user_model is not None: # `litellm --model <your-model-name>`
response = await litellm.acompletion(**data) response = await litellm.acompletion(**data)
else:
raise HTTPException(
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
detail={"message": "Invalid model name passed in"},
)
# Post Call Processing # Post Call Processing
data["litellm_status"] = "success" # used for alerting data["litellm_status"] = "success" # used for alerting
@ -2584,8 +2589,13 @@ async def embeddings(
llm_router is not None and data["model"] in llm_router.deployment_names llm_router is not None and data["model"] in llm_router.deployment_names
): # model in router deployments, calling a specific deployment on the router ): # model in router deployments, calling a specific deployment on the router
response = await llm_router.aembedding(**data, specific_deployment=True) response = await llm_router.aembedding(**data, specific_deployment=True)
else: elif user_model is not None: # `litellm --model <your-model-name>`
response = await litellm.aembedding(**data) response = await litellm.aembedding(**data)
else:
raise HTTPException(
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
detail={"message": "Invalid model name passed in"},
)
### ALERTING ### ### ALERTING ###
data["litellm_status"] = "success" # used for alerting data["litellm_status"] = "success" # used for alerting
@ -2719,8 +2729,13 @@ async def image_generation(
response = await llm_router.aimage_generation( response = await llm_router.aimage_generation(
**data **data
) # ensure this goes the llm_router, router will do the correct alias mapping ) # ensure this goes the llm_router, router will do the correct alias mapping
else: elif user_model is not None: # `litellm --model <your-model-name>`
response = await litellm.aimage_generation(**data) response = await litellm.aimage_generation(**data)
else:
raise HTTPException(
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
detail={"message": "Invalid model name passed in"},
)
### ALERTING ### ### ALERTING ###
data["litellm_status"] = "success" # used for alerting data["litellm_status"] = "success" # used for alerting