Merge pull request #5098 from BerriAI/litellm_provider_wildcard_routing

[Feat-Router + Proxy] Add provider wildcard routing
This commit is contained in:
Ishaan Jaff 2024-08-07 14:51:42 -07:00 • committed by GitHub
commit 04b201efed
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 204 additions and 57 deletions

View file

@ -284,52 +284,58 @@ curl --location 'http://0.0.0.0:4000/v1/model/info' \
--data ''
```
## Wildcard Model Name (Add ALL MODELS from env)
## Provider specific wildcard routing
**Proxy all models from a provider**
Dynamically call any model from any given provider without the need to predefine it in the config YAML file. As long as the relevant keys are in the environment (see [providers list](../providers/)), LiteLLM will make the call correctly.
Use this if you want to **proxy all models from a specific provider without defining them on the config.yaml**
1. Setup config.yaml
```
**Step 1** - define provider specific routing on config.yaml
```yaml
model_list:
- model_name: "*" # all requests where model not in your config go to this deployment
# provider specific wildcard routing
- model_name: "anthropic/*"
litellm_params:
model: "*" # passes our validation check that a real provider is given
model: "anthropic/*"
api_key: os.environ/ANTHROPIC_API_KEY
- model_name: "groq/*"
litellm_params:
model: "groq/*"
api_key: os.environ/GROQ_API_KEY
```
2. Start LiteLLM proxy
Step 2 - Run litellm proxy
```
litellm --config /path/to/config.yaml
```shell
$ litellm --config /path/to/config.yaml
```
3. Try claude 3-5 sonnet from anthropic
Step 3 Test it
```bash
curl -X POST 'http://0.0.0.0:4000/chat/completions' \
-H 'Content-Type: application/json' \
-H 'Authorization: Bearer sk-1234' \
-D '{
"model": "claude-3-5-sonnet-20240620",
"messages": [
{"role": "user", "content": "Hey, how'\''s it going?"},
{
"role": "assistant",
"content": "I'\''m doing well. Would like to hear the rest of the story?"
},
{"role": "user", "content": "Na"},
{
"role": "assistant",
"content": "No problem, is there anything else i can help you with today?"
},
{
"role": "user",
"content": "I think you'\''re getting cut off sometimes"
}
Test with `anthropic/` - all models with `anthropic/` prefix will get routed to `anthropic/*`
```shell
curl http://localhost:4000/v1/chat/completions \
-H "Content-Type: application/json" \
-H "Authorization: Bearer sk-1234" \
-d '{
"model": "anthropic/claude-3-sonnet-20240229",
"messages": [
{"role": "user", "content": "Hello, Claude!"}
]
}
'
}'
```
Test with `groq/` - all models with `groq/` prefix will get routed to `groq/*`
```shell
curl http://localhost:4000/v1/chat/completions \
-H "Content-Type: application/json" \
-H "Authorization: Bearer sk-1234" \
-d '{
"model": "groq/llama3-8b-8192",
"messages": [
{"role": "user", "content": "Hello, Claude!"}
]
}'
```
## Load Balancing

View file

@ -8,9 +8,15 @@ model_list:
litellm_params:
model: fireworks_ai/accounts/fireworks/models/llama-v3-70b-instruct
api_key: "os.environ/FIREWORKS"
- model_name: "*"
# provider specific wildcard routing
- model_name: "anthropic/*"
litellm_params:
model: "*"
model: "anthropic/*"
api_key: os.environ/ANTHROPIC_API_KEY
- model_name: "groq/*"
litellm_params:
model: "groq/*"
api_key: os.environ/GROQ_API_KEY
- model_name: "*"
litellm_params:
model: openai/*

View file

@ -3007,7 +3007,10 @@ async def chat_completion(
elif (
llm_router is not None
and data["model"] not in router_model_names
and llm_router.default_deployment is not None
and (
llm_router.default_deployment is not None
or len(llm_router.provider_default_deployments) > 0
)
): # model in router deployments, calling a specific deployment on the router
tasks.append(llm_router.acompletion(**data))
elif user_model is not None: # `litellm --model <your-model-name>`
@ -3275,7 +3278,10 @@ async def completion(
elif (
llm_router is not None
and data["model"] not in router_model_names
and llm_router.default_deployment is not None
and (
llm_router.default_deployment is not None
or len(llm_router.provider_default_deployments) > 0
)
): # model in router deployments, calling a specific deployment on the router
llm_response = asyncio.create_task(llm_router.atext_completion(**data))
elif user_model is not None: # `litellm --model <your-model-name>`
@ -3541,7 +3547,10 @@ async def embeddings(
elif (
llm_router is not None
and data["model"] not in router_model_names
and llm_router.default_deployment is not None
and (
llm_router.default_deployment is not None
or len(llm_router.provider_default_deployments) > 0
)
): # model in router deployments, calling a specific deployment on the router
tasks.append(llm_router.aembedding(**data))
elif user_model is not None: # `litellm --model <your-model-name>`
@ -3708,7 +3717,10 @@ async def image_generation(
elif (
llm_router is not None
and data["model"] not in router_model_names
and llm_router.default_deployment is not None
and (
llm_router.default_deployment is not None
or len(llm_router.provider_default_deployments) > 0
)
): # model in router deployments, calling a specific deployment on the router
response = await llm_router.aimage_generation(**data)
elif user_model is not None: # `litellm --model <your-model-name>`
@ -3850,7 +3862,10 @@ async def audio_speech(
elif (
llm_router is not None
and data["model"] not in router_model_names
and llm_router.default_deployment is not None
and (
llm_router.default_deployment is not None
or len(llm_router.provider_default_deployments) > 0
)
): # model in router deployments, calling a specific deployment on the router
response = await llm_router.aspeech(**data)
elif user_model is not None: # `litellm --model <your-model-name>`
@ -4020,7 +4035,10 @@ async def audio_transcriptions(
elif (
llm_router is not None
and data["model"] not in router_model_names
and llm_router.default_deployment is not None
and (
llm_router.default_deployment is not None
or len(llm_router.provider_default_deployments) > 0
)
): # model in router deployments, calling a specific deployment on the router
response = await llm_router.atranscription(**data)
elif user_model is not None: # `litellm --model <your-model-name>`
@ -5270,7 +5288,10 @@ async def moderations(
elif (
llm_router is not None
and data.get("model") not in router_model_names
and llm_router.default_deployment is not None
and (
llm_router.default_deployment is not None
or len(llm_router.provider_default_deployments) > 0
)
): # model in router deployments, calling a specific deployment on the router
response = await llm_router.amoderation(**data)
elif user_model is not None: # `litellm --model <your-model-name>`
@ -5421,7 +5442,10 @@ async def anthropic_response(
elif (
llm_router is not None
and data["model"] not in router_model_names
and llm_router.default_deployment is not None
and (
llm_router.default_deployment is not None
or len(llm_router.provider_default_deployments) > 0
)
): # model in router deployments, calling a specific deployment on the router
llm_response = asyncio.create_task(llm_router.aadapter_completion(**data))
elif user_model is not None: # `litellm --model <your-model-name>`

View file

@ -17,6 +17,7 @@ import inspect
import json
import logging
import random
import re
import threading
import time
import traceback
@ -310,6 +311,7 @@ class Router:
)
self.default_deployment = None # use this to track the users default deployment, when they want to use model = *
self.default_max_parallel_requests = default_max_parallel_requests
self.provider_default_deployments: Dict[str, List] = {}
if model_list is not None:
model_list = copy.deepcopy(model_list)
@ -3607,6 +3609,10 @@ class Router:
),
)
provider_specific_deployment = re.match(
f"{custom_llm_provider}/*", deployment.model_name
)
# Check if user is trying to use model_name == "*"
# this is a catch all model for their specific api key
if deployment.model_name == "*":
@ -3615,6 +3621,17 @@ class Router:
self.router_general_settings.pass_through_all_models = True
else:
self.default_deployment = deployment.to_json(exclude_none=True)
# Check if user is using provider specific wildcard routing
# example model_name = "databricks/*" or model_name = "anthropic/*"
elif provider_specific_deployment:
if custom_llm_provider in self.provider_default_deployments:
self.provider_default_deployments[custom_llm_provider].append(
deployment.to_json(exclude_none=True)
)
else:
self.provider_default_deployments[custom_llm_provider] = [
deployment.to_json(exclude_none=True)
]
# Azure GPT-Vision Enhancements, users can pass os.environ/
data_sources = deployment.litellm_params.get("dataSources", []) or []
@ -4452,12 +4469,37 @@ class Router:
)
model = self.model_group_alias[model]
if model not in self.model_names and self.default_deployment is not None:
updated_deployment = copy.deepcopy(
self.default_deployment
) # self.default_deployment
updated_deployment["litellm_params"]["model"] = model
return model, updated_deployment
if model not in self.model_names:
# check if provider/ specific wildcard routing
try:
(
_,
custom_llm_provider,
_,
_,
) = litellm.get_llm_provider(model=model)
# check if custom_llm_provider
if custom_llm_provider in self.provider_default_deployments:
_provider_deployments = self.provider_default_deployments[
custom_llm_provider
]
provider_deployments = []
for deployment in _provider_deployments:
dep = copy.deepcopy(deployment)
dep["litellm_params"]["model"] = model
provider_deployments.append(dep)
return model, provider_deployments
except:
# get_llm_provider raises exception when provider is unknown
pass
# check if default deployment is set
if self.default_deployment is not None:
updated_deployment = copy.deepcopy(
self.default_deployment
) # self.default_deployment
updated_deployment["litellm_params"]["model"] = model
return model, updated_deployment
## get healthy deployments
### get all deployments

View file

@ -60,6 +60,63 @@ def test_router_multi_org_list():
assert len(router.get_model_list()) == 3
@pytest.mark.asyncio()
async def test_router_provider_wildcard_routing():
"""
Pass list of orgs in 1 model definition,
expect a unique deployment for each to be created
"""
router = litellm.Router(
model_list=[
{
"model_name": "openai/*",
"litellm_params": {
"model": "openai/*",
"api_key": os.environ["OPENAI_API_KEY"],
"api_base": "https://api.openai.com/v1",
},
},
{
"model_name": "anthropic/*",
"litellm_params": {
"model": "anthropic/*",
"api_key": os.environ["ANTHROPIC_API_KEY"],
},
},
{
"model_name": "groq/*",
"litellm_params": {
"model": "groq/*",
"api_key": os.environ["GROQ_API_KEY"],
},
},
]
)
print("router model list = ", router.get_model_list())
response1 = await router.acompletion(
model="anthropic/claude-3-sonnet-20240229",
messages=[{"role": "user", "content": "hello"}],
)
print("response 1 = ", response1)
response2 = await router.acompletion(
model="openai/gpt-3.5-turbo",
messages=[{"role": "user", "content": "hello"}],
)
print("response 2 = ", response2)
response3 = await router.acompletion(
model="groq/llama3-8b-8192",
messages=[{"role": "user", "content": "hello"}],
)
print("response 3 = ", response3)
def test_router_specific_model_via_id():
"""
Call a specific deployment by it's id

View file

@ -86,12 +86,16 @@ model_list:
model: openai/*
api_key: os.environ/OPENAI_API_KEY
# Pass through all llm requests to litellm.completion/litellm.embedding
# if user passes model="anthropic/claude-3-opus-20240229" proxy will make requests to anthropic claude-3-opus-20240229 using ANTHROPIC_API_KEY
- model_name: "*"
# provider specific wildcard routing
- model_name: "anthropic/*"
litellm_params:
model: "*"
model: "anthropic/*"
api_key: os.environ/ANTHROPIC_API_KEY
- model_name: "groq/*"
litellm_params:
model: "groq/*"
api_key: os.environ/GROQ_API_KEY
- model_name: mistral-embed
litellm_params:
model: mistral/mistral-embed

View file

@ -119,7 +119,9 @@ async def chat_completion(session, key, model: Union[str, List] = "gpt-4"):
print()
if status != 200:
raise Exception(f"Request did not return a 200 status code: {status}")
raise Exception(
f"Request did not return a 200 status code: {status}, response text={response_text}"
)
response_header_check(
response
@ -485,6 +487,12 @@ async def test_proxy_all_models():
session=session, key=LITELLM_MASTER_KEY, model="groq/llama3-8b-8192"
)
await chat_completion(
session=session,
key=LITELLM_MASTER_KEY,
model="anthropic/claude-3-sonnet-20240229",
)
@pytest.mark.asyncio
async def test_batch_chat_completions():