mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-26 19:24:27 +00:00
(feat) support timeout on bedrock
This commit is contained in:
parent
43232012d7
commit
cd9005d6e6
2 changed files with 57 additions and 2 deletions
|
@ -389,7 +389,7 @@ class Router:
|
||||||
model_name = data["model"]
|
model_name = data["model"]
|
||||||
for k, v in self.default_litellm_params.items():
|
for k, v in self.default_litellm_params.items():
|
||||||
if (
|
if (
|
||||||
k not in kwargs
|
k not in kwargs and v is not None
|
||||||
): # prioritize model-specific params > default router params
|
): # prioritize model-specific params > default router params
|
||||||
kwargs[k] = v
|
kwargs[k] = v
|
||||||
elif k == "metadata":
|
elif k == "metadata":
|
||||||
|
@ -409,13 +409,24 @@ class Router:
|
||||||
else:
|
else:
|
||||||
model_client = potential_model_client
|
model_client = potential_model_client
|
||||||
self.total_calls[model_name] += 1
|
self.total_calls[model_name] += 1
|
||||||
|
|
||||||
|
timeout = (
|
||||||
|
data.get(
|
||||||
|
"timeout", None
|
||||||
|
) # timeout set on litellm_params for this deployment
|
||||||
|
or self.timeout # timeout set on router
|
||||||
|
or kwargs.get(
|
||||||
|
"timeout", None
|
||||||
|
) # this uses default_litellm_params when nothing is set
|
||||||
|
)
|
||||||
|
|
||||||
response = await litellm.acompletion(
|
response = await litellm.acompletion(
|
||||||
**{
|
**{
|
||||||
**data,
|
**data,
|
||||||
"messages": messages,
|
"messages": messages,
|
||||||
"caching": self.cache_responses,
|
"caching": self.cache_responses,
|
||||||
"client": model_client,
|
"client": model_client,
|
||||||
"timeout": self.timeout,
|
"timeout": timeout,
|
||||||
**kwargs,
|
**kwargs,
|
||||||
}
|
}
|
||||||
)
|
)
|
||||||
|
|
|
@ -85,3 +85,47 @@ def test_router_timeouts():
|
||||||
|
|
||||||
print("Response:", response)
|
print("Response:", response)
|
||||||
print("********** TOKENS USED SO FAR = ", total_tokens_used)
|
print("********** TOKENS USED SO FAR = ", total_tokens_used)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_router_timeouts_bedrock():
|
||||||
|
import openai
|
||||||
|
|
||||||
|
# Model list for OpenAI and Anthropic models
|
||||||
|
model_list = [
|
||||||
|
{
|
||||||
|
"model_name": "bedrock",
|
||||||
|
"litellm_params": {
|
||||||
|
"model": "bedrock/anthropic.claude-instant-v1",
|
||||||
|
"timeout": 0.001,
|
||||||
|
},
|
||||||
|
"tpm": 80000,
|
||||||
|
},
|
||||||
|
]
|
||||||
|
|
||||||
|
# Configure router
|
||||||
|
router = Router(
|
||||||
|
model_list=model_list,
|
||||||
|
routing_strategy="usage-based-routing",
|
||||||
|
debug_level="DEBUG",
|
||||||
|
set_verbose=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
litellm.set_verbose = True
|
||||||
|
try:
|
||||||
|
response = await router.acompletion(
|
||||||
|
model="bedrock",
|
||||||
|
messages=[{"role": "user", "content": "hello, who are u"}],
|
||||||
|
)
|
||||||
|
print(response)
|
||||||
|
pytest.fail("Did not raise error `openai.APITimeoutError`")
|
||||||
|
except openai.APITimeoutError as e:
|
||||||
|
print(
|
||||||
|
"Passed: Raised correct exception. Got openai.APITimeoutError\nGood Job", e
|
||||||
|
)
|
||||||
|
print(type(e))
|
||||||
|
pass
|
||||||
|
except Exception as e:
|
||||||
|
pytest.fail(
|
||||||
|
f"Did not raise error `openai.APITimeoutError`. Instead raised error type: {type(e)}, Error: {e}"
|
||||||
|
)
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue