fix(ollama): map session usage limit and rate limit errors to RateLimitError

Ollama's "session usage limit" error was falling through the exception
mapper as a generic APIConnectionError (status 500). APIConnectionErrors
are explicitly excluded from cooldown logic, so the failed deployment
was never cooled down and retries kept hitting the same order=1
deployment instead of failing over to order=2.

Map "session usage limit" and generic "rate limit" strings in the Ollama
exception handler to RateLimitError (429), which triggers cooldown on
the deployment and allows the router to select the next-priority
deployment on retry.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
tombii 2026-03-03 15:08:08 +01:00
parent 974c02fae4
commit 5b8fef098d
2 changed files with 45 additions and 0 deletions

View file

@ -2040,6 +2040,14 @@ def exception_type( # type: ignore # noqa: PLR0915
llm_provider="ollama",
model=model,
)
elif "session usage limit" in error_str or "rate limit" in error_str.lower():
exception_mapping_worked = True
raise RateLimitError(
message=f"OllamaException: {original_exception}",
llm_provider="ollama",
model=model,
response=getattr(original_exception, "response", None),
)
elif custom_llm_provider == "vllm":
if hasattr(original_exception, "status_code"):
if original_exception.status_code == 0:

View file

@ -281,6 +281,43 @@ def test_vertex_ai_rate_limit_error_mapping(error_message, should_raise_rate_lim
)
ollama_rate_limit_test_cases = [
# Positive cases — session/rate-limit messages that should map to RateLimitError
('{"error": "you (tom9876) have reached your session usage limit, please wait or upgrade to continue"}', True),
("session usage limit reached", True),
("rate limit exceeded", True),
("Rate Limit: too many requests", True),
# Negative case — generic error, not a rate limit
("no such file or directory", False),
]
@pytest.mark.parametrize(
"error_message, should_raise_rate_limit", ollama_rate_limit_test_cases
)
def test_ollama_rate_limit_error_mapping(error_message, should_raise_rate_limit):
"""
Tests that the exception_type function correctly maps Ollama's
session usage limit and rate limit errors to litellm.RateLimitError.
"""
for provider in ("ollama", "ollama_chat"):
original_exception = Exception(error_message)
if should_raise_rate_limit:
with pytest.raises(litellm.RateLimitError):
exception_type(
model="llama3",
original_exception=original_exception,
custom_llm_provider=provider,
)
else:
with pytest.raises((litellm.BadRequestError, litellm.APIConnectionError)):
exception_type(
model="llama3",
original_exception=original_exception,
custom_llm_provider=provider,
)
class TestExtractAndRaiseLitellmException:
"""Tests for extract_and_raise_litellm_exception function"""