scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-10 22:41:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 83

fix(proxy): share the Bedrock count-tokens predicate between auth and the passthrough handler

Browse source
This commit is contained in:
mateo-berri 2026-09-03 16:41:31 -07:00
parent 3b43b5a3d6
commit bd9593b74d
3 changed files with 19 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
litellm/proxy/auth/auth_utils.py
View file

@ -1989,13 +1989,14 @@ def get_model_from_request(
def _model_from_bedrock_route(route: str) -> str | None:
bedrock_endpoint: Final = re.sub(r"^/bedrock/", "", route, flags=re.IGNORECASE)
if "count_tokens" in bedrock_endpoint.lower() or "count-tokens" in bedrock_endpoint.lower():
return None
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
_extract_model_from_bedrock_endpoint,
is_bedrock_count_tokens_endpoint,
)
bedrock_endpoint: Final = re.sub(r"^/bedrock/", "", route, flags=re.IGNORECASE)
if is_bedrock_count_tokens_endpoint(bedrock_endpoint):
return None
try:
return _extract_model_from_bedrock_endpoint(bedrock_endpoint)
except ValueError:

7
litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
View file

@ -703,6 +703,10 @@ BEDROCK_ENDPOINT_ACTIONS: Final = {
BEDROCK_STREAMING_ACTIONS: Final = {"invoke-with-response-stream", "converse-stream"}
def is_bedrock_count_tokens_endpoint(endpoint: str) -> bool:
return "count_tokens" in endpoint or "count-tokens" in endpoint
def _extract_model_from_bedrock_endpoint(endpoint: str) -> str:
"""
Extract model name from Bedrock endpoint path.
@ -977,8 +981,7 @@ async def bedrock_llm_proxy_route(
request_body: Final = await _read_request_body(request=request)
# Special handling for count_tokens endpoints
if "count_tokens" in endpoint or "count-tokens" in endpoint:
if is_bedrock_count_tokens_endpoint(endpoint):
return await handle_bedrock_count_tokens(
endpoint=endpoint,
request=request,

10
tests/test_litellm/proxy/auth/test_auth_utils.py
View file

@ -508,6 +508,16 @@ def test_get_model_from_request_bedrock_count_tokens_uses_body_model():
)
def test_get_model_from_request_bedrock_uppercase_count_tokens_segment_is_not_count_tokens():
assert (
get_model_from_request(
request_data={"model": "us.anthropic.claude-haiku-4-5-20251001-v1:0"},
route="/bedrock/model/us.anthropic.claude-sonnet-4-6/invoke/COUNT_TOKENS",
)
== "us.anthropic.claude-sonnet-4-6"
)
def test_get_model_from_request_bedrock_unparseable_endpoint_keeps_body_model():
assert (
get_model_from_request(