mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
fix(proxy): fall back to local tokenizer on provider token count errors
This commit is contained in:
parent
dd70a19dbc
commit
5f24c5bf4b
2 changed files with 82 additions and 2 deletions
|
|
@ -10781,9 +10781,13 @@ async def _try_provider_token_count(
|
|||
tools=tools,
|
||||
system=system,
|
||||
)
|
||||
except httpx.HTTPStatusError as e:
|
||||
except Exception as e:
|
||||
error_message = getattr(e, "message", None) or str(e)
|
||||
status_code = getattr(e, "status_code", None) or e.response.status_code
|
||||
status_code = getattr(e, "status_code", None)
|
||||
if status_code is None and hasattr(e, "response") and hasattr(e.response, "status_code"):
|
||||
status_code = e.response.status_code
|
||||
status_code = status_code or 500
|
||||
|
||||
if litellm.disable_token_counter is True:
|
||||
raise ProxyException(
|
||||
message=error_message,
|
||||
|
|
|
|||
|
|
@ -1446,3 +1446,79 @@ async def test_try_provider_token_count_http_error_disabled():
|
|||
assert exc_info.value.code == "500"
|
||||
finally:
|
||||
litellm.disable_token_counter = original_disable
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_try_provider_token_count_generic_error_fallback():
|
||||
"""
|
||||
Test that _try_provider_token_count catches a generic Exception
|
||||
and returns None (falling back to local tokenizer) when litellm.disable_token_counter != True.
|
||||
"""
|
||||
from litellm.proxy.proxy_server import _try_provider_token_count
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
import litellm
|
||||
|
||||
original_disable = getattr(litellm, "disable_token_counter", False)
|
||||
litellm.disable_token_counter = False
|
||||
|
||||
try:
|
||||
mock_provider_counter = MagicMock()
|
||||
mock_provider_counter.should_use_token_counting_api.return_value = True
|
||||
|
||||
# Make count_tokens raise a generic Exception
|
||||
error = ValueError("Some weird provider error")
|
||||
mock_provider_counter.count_tokens = AsyncMock(side_effect=error)
|
||||
|
||||
result = await _try_provider_token_count(
|
||||
provider_counter=mock_provider_counter,
|
||||
custom_llm_provider="openai",
|
||||
model_to_use="gpt-4",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
contents=None,
|
||||
deployment=None,
|
||||
request_model="gpt-4"
|
||||
)
|
||||
|
||||
# It should return None on generic Exception instead of raising Exception
|
||||
assert result is None
|
||||
finally:
|
||||
litellm.disable_token_counter = original_disable
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_try_provider_token_count_generic_error_disabled():
|
||||
"""
|
||||
Test that _try_provider_token_count raises ProxyException on generic Exception
|
||||
when litellm.disable_token_counter == True.
|
||||
"""
|
||||
from litellm.proxy.proxy_server import _try_provider_token_count
|
||||
from litellm.proxy._types import ProxyException
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
import litellm
|
||||
|
||||
original_disable = getattr(litellm, "disable_token_counter", False)
|
||||
litellm.disable_token_counter = True
|
||||
|
||||
try:
|
||||
mock_provider_counter = MagicMock()
|
||||
mock_provider_counter.should_use_token_counting_api.return_value = True
|
||||
|
||||
# Make count_tokens raise a generic Exception
|
||||
error = ValueError("Some weird provider error")
|
||||
mock_provider_counter.count_tokens = AsyncMock(side_effect=error)
|
||||
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await _try_provider_token_count(
|
||||
provider_counter=mock_provider_counter,
|
||||
custom_llm_provider="openai",
|
||||
model_to_use="gpt-4",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
contents=None,
|
||||
deployment=None,
|
||||
request_model="gpt-4"
|
||||
)
|
||||
|
||||
# Default status code for generic exceptions is 500
|
||||
assert exc_info.value.code == "500"
|
||||
finally:
|
||||
litellm.disable_token_counter = original_disable
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue