Added quota limit exception handling

Author: Andrej Simurka

docs/openapi.json

@@ -8,6 +8,7 @@
 from typing import Annotated, Any, Optional, cast
 
 from fastapi import APIRouter, Depends, HTTPException, Request, status
+from litellm.exceptions import RateLimitError
 from llama_stack_client import (
     APIConnectionError,
     AsyncLlamaStackClient,  # type: ignore
@@ -42,6 +43,7 @@
     ReferencedDocument,
     ToolCall,
     UnauthorizedResponse,
+    QuotaExceededResponse
 )
 from utils.endpoints import (
     check_configuration_loaded,
@@ -86,6 +88,10 @@
         "description": "User is not authorized",
         "model": ForbiddenResponse,
     },
+    429: {
+        "description": "The quota has been exceeded",
+        "model": QuotaExceededResponse,
+    },
     500: {
         "detail": {
             "response": "Unable to connect to Llama Stack",
@@ -421,6 +427,14 @@ async def query_endpoint_handler_base(  # pylint: disable=R0914
                 "cause": str(e),
             },
         ) from e
+    except RateLimitError as e:
+        raise HTTPException(
+            status_code=status.HTTP_429_TOO_MANY_REQUESTS,
+            detail={
+                "response": "Model quota exceeded",
+                "cause": f"The token quota for model {e.model} has been exceeded.",
+            },
+        ) from e
 
 
 @router.post("/query", responses=query_response)

src/app/endpoints/query.py

@@ -27,6 +27,7 @@
     QueryResponse,
     ReferencedDocument,
     UnauthorizedResponse,
+    QuotaExceededResponse
 )
 from utils.endpoints import (
     get_system_prompt,
@@ -59,6 +60,10 @@
         "description": "User is not authorized",
         "model": ForbiddenResponse,
     },
+    429: {
+        "description": "The quota has been exceeded",
+        "model": QuotaExceededResponse,
+    },
     500: {
         "detail": {
             "response": "Unable to connect to Llama Stack",

src/app/endpoints/query_v2.py

@@ -8,6 +8,7 @@
 from datetime import UTC, datetime
 from typing import Annotated, Any, AsyncGenerator, AsyncIterator, Iterator, cast
 
+from litellm.exceptions import RateLimitError
 from fastapi import APIRouter, Depends, HTTPException, Request, status
 from fastapi.responses import StreamingResponse
 from llama_stack_client import (
@@ -48,7 +49,7 @@
 from models.config import Action
 from models.database.conversations import UserConversation
 from models.requests import QueryRequest
-from models.responses import ForbiddenResponse, UnauthorizedResponse
+from models.responses import ForbiddenResponse, UnauthorizedResponse, QuotaExceededResponse
 from utils.endpoints import (
     check_configuration_loaded,
     create_referenced_documents_with_metadata,
@@ -104,6 +105,10 @@
         "description": "User is not authorized",
         "model": ForbiddenResponse,
     },
+    429: {
+        "description": "The quota has been exceeded",
+        "model": QuotaExceededResponse,
+    },
     500: {
         "detail": {
             "response": "Unable to connect to Llama Stack",
@@ -922,6 +927,14 @@ async def response_generator(
                 "cause": str(e),
             },
         ) from e
+    except RateLimitError as e:
+        raise HTTPException(
+            status_code=status.HTTP_429_TOO_MANY_REQUESTS,
+            detail={
+                "response": "Model quota exceeded",
+                "cause": f"The token quota for model {e.model} has been exceeded.",
+            },
+        ) from e
     except Exception as e:  # pylint: disable=broad-except
         # Handle other errors with OLS-compatible error response
         # This broad exception catch is intentional to ensure all errors

src/app/endpoints/streaming_query.py

@@ -1142,6 +1142,71 @@ def __init__(self, user_id: str, resource: str, resource_id: str):
     }
 
 
+class QuotaExceededResponse(AbstractErrorResponse):
+    """429 Too Many Requests - LLM quota exceeded."""
+
+    def __init__(self, user_id: str, model_name: str, limit: int, period: str):
+        """Initialize a QuotaExceededResponse."""
+        super().__init__(
+            detail=DetailModel(
+                response="Quota exceeded",
+                cause=(
+                    f"User {user_id} has exceeded the {limit}-request {period} quota "
+                    f"for {model_name}. Please upgrade your plan or wait until the next cycle."
+                ),
+            )
+        )
+    
+    model_config = {
+        "json_schema_extra": {
+            "examples": [
+                {
+                    "detail": {
+                        "response": "The quota has been exceeded",
+                        "cause": "User 123 has no available tokens."
+                    }
+                },
+                {
+                    "detail": {
+                        "response": "The quota has been exceeded",
+                        "cause": "Cluster has no available tokens."
+                    }
+                },
+                {
+                    "detail": {
+                        "response": "The quota has been exceeded",
+                        "cause": "Unknown subject 999 has no available tokens."
+                    }
+                },
+                {
+                    "detail": {
+                        "response": "The quota has been exceeded",
+                        "cause": "User 123 has 5 tokens, but 10 tokens are needed."
+                    }
+                },
+                {
+                    "detail": {
+                        "response": "The quota has been exceeded",
+                        "cause": "Cluster has 500 tokens, but 900 tokens are needed."
+                    }
+                },
+                {
+                    "detail": {
+                        "response": "The quota has been exceeded",
+                        "cause": "Unknown subject 999 has 3 tokens, but 6 tokens are needed."
+                    }
+                },
+                {
+                    "detail": {
+                        "response": "The model quota has been exceeded",
+                        "cause": "The token quota for model gpt-4-turbo has been exceeded."
+                    }
+                }
+            ]
+        }
+    }
+
+
 class InvalidFeedbackStoragePathResponse(AbstractErrorResponse):
     """500 Internal Error - Invalid feedback storage path."""
 

src/models/responses.py

@@ -10,6 +10,7 @@
 import pytest
 from pytest_mock import MockerFixture
 from fastapi import HTTPException, Request, status
+from litellm.exceptions import RateLimitError
 
 from llama_stack_client import APIConnectionError
 from llama_stack_client.types import UserMessage  # type: ignore
@@ -2261,3 +2262,47 @@ async def test_get_topic_summary_create_turn_parameters(mocker: MockerFixture) -
         stream=False,
         toolgroups=None,
     )
+
+
+@pytest.mark.asyncio
+async def test_query_endpoint_quota_exceeded(
+    mocker: MockerFixture, dummy_request: Request) -> None:
+    """Test that query endpoint raises HTTP 429 when model quota is exceeded."""
+    query_request = QueryRequest(
+        query="What is OpenStack?",
+        provider="openai",
+        model="gpt-4-turbo",
+    )
+    mock_client = mocker.AsyncMock()
+    mock_agent = mocker.AsyncMock()
+    mock_agent.create_turn.side_effect = RateLimitError(
+        model="gpt-4-turbo", llm_provider="openai", message=""
+    )
+    mocker.patch(
+        "app.endpoints.query.get_agent",
+        return_value=(mock_agent, "conv-123", "sess-123")
+    )
+    mocker.patch(
+        "app.endpoints.query.select_model_and_provider_id",
+        return_value=("openai/gpt-4-turbo", "gpt-4-turbo", "openai")
+    )
+    mocker.patch("app.endpoints.query.validate_model_provider_override")
+    mocker.patch(
+        "app.endpoints.query.AsyncLlamaStackClientHolder.get_client",
+        return_value=mock_client
+    )
+    mocker.patch(
+        "app.endpoints.query.handle_mcp_headers_with_toolgroups",
+        return_value={}
+    )
+    
+    with pytest.raises(HTTPException) as exc_info:
+        await query_endpoint_handler(
+            dummy_request, 
+            query_request=query_request, 
+            auth=MOCK_AUTH
+        )
+    assert exc_info.value.status_code == status.HTTP_429_TOO_MANY_REQUESTS
+    detail = exc_info.value.detail
+    assert detail["response"] == "Model quota exceeded"
+    assert "gpt-4-turbo" in detail["cause"]

tests/unit/app/endpoints/test_query.py

@@ -1,10 +1,12 @@
 # pylint: disable=redefined-outer-name, import-error
 """Unit tests for the /query (v2) REST API endpoint using Responses API."""
 
+from litellm.exceptions import RateLimitError
 import pytest
 from fastapi import HTTPException, status, Request
 
 from llama_stack_client import APIConnectionError
+from pytest_mock import MockerFixture
 
 from models.requests import QueryRequest, Attachment
 from models.config import ModelContextProtocolServer
@@ -16,6 +18,14 @@
     query_endpoint_handler_v2,
 )
 
+# User ID must be proper UUID
+MOCK_AUTH = (
+    "00000001-0001-0001-0001-000000000001",
+    "mock_username",
+    False,
+    "mock_token",
+)
+
 
 @pytest.fixture
 def dummy_request() -> Request:
@@ -421,3 +431,43 @@ def _raise(*_args, **_kwargs):
     assert exc.value.status_code == status.HTTP_500_INTERNAL_SERVER_ERROR
     assert "Unable to connect to Llama Stack" in str(exc.value.detail)
     fail_metric.inc.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_query_endpoint_quota_exceeded(
+    mocker: MockerFixture, dummy_request: Request) -> None:
+    """Test that query endpoint raises HTTP 429 when model quota is exceeded."""
+    query_request = QueryRequest(
+        query="What is OpenStack?",
+        provider="openai",
+        model="gpt-4-turbo",
+        attachments=[],
+    )
+    mock_client = mocker.AsyncMock()
+    mock_client.responses.create.side_effect = RateLimitError(
+        model="gpt-4-turbo", llm_provider="openai", message=""
+    )
+    mocker.patch(
+        "app.endpoints.query.select_model_and_provider_id",
+        return_value=("openai/gpt-4-turbo", "gpt-4-turbo", "openai")
+    )
+    mocker.patch("app.endpoints.query.validate_model_provider_override")
+    mocker.patch(
+        "app.endpoints.query.AsyncLlamaStackClientHolder.get_client",
+        return_value=mock_client
+    )
+    mocker.patch(
+        "app.endpoints.query.handle_mcp_headers_with_toolgroups",
+        return_value={}
+    )
+    
+    with pytest.raises(HTTPException) as exc_info:
+        await query_endpoint_handler_v2(
+            dummy_request, 
+            query_request=query_request, 
+            auth=MOCK_AUTH
+        )
+    assert exc_info.value.status_code == status.HTTP_429_TOO_MANY_REQUESTS
+    detail = exc_info.value.detail
+    assert detail["response"] == "Model quota exceeded"
+    assert "gpt-4-turbo" in detail["cause"]

tests/unit/app/endpoints/test_query_v2.py

@@ -6,6 +6,7 @@
 
 import json
 
+from litellm.exceptions import RateLimitError
 import pytest
 from pytest_mock import MockerFixture
 
@@ -1789,6 +1790,51 @@ async def test_streaming_query_handles_none_event(mocker: MockerFixture) -> None
     assert isinstance(response, StreamingResponse)
 
 
+@pytest.mark.asyncio
+async def test_query_endpoint_quota_exceeded(
+    mocker: MockerFixture) -> None:
+    """Test that streaming query endpoint raises HTTP 429 when model quota is exceeded."""
+    query_request = QueryRequest(
+        query="What is OpenStack?",
+        provider="openai",
+        model="gpt-4-turbo",
+    )
+    request = Request(scope={"type": "http"})
+    mock_client = mocker.AsyncMock()
+    mock_agent = mocker.AsyncMock()
+    mock_agent.create_turn.side_effect = RateLimitError(
+        model="gpt-4-turbo", llm_provider="openai", message=""
+    )
+    mocker.patch(
+        "app.endpoints.streaming_query.get_agent",
+        return_value=(mock_agent, "conv-123", "sess-123")
+    )
+    mocker.patch(
+        "app.endpoints.streaming_query.select_model_and_provider_id",
+        return_value=("openai/gpt-4-turbo", "gpt-4-turbo", "openai")
+    )
+    mocker.patch("app.endpoints.streaming_query.validate_model_provider_override")
+    mocker.patch(
+        "app.endpoints.streaming_query.AsyncLlamaStackClientHolder.get_client",
+        return_value=mock_client
+    )
+    mocker.patch(
+        "app.endpoints.streaming_query.handle_mcp_headers_with_toolgroups",
+        return_value={}
+    )
+    
+    with pytest.raises(HTTPException) as exc_info:
+        await streaming_query_endpoint_handler(
+            request, 
+            query_request=query_request, 
+            auth=MOCK_AUTH
+        )
+    assert exc_info.value.status_code == status.HTTP_429_TOO_MANY_REQUESTS
+    detail = exc_info.value.detail
+    assert detail["response"] == "Model quota exceeded"
+    assert "gpt-4-turbo" in detail["cause"]
+
+
 # ============================================================================
 # OLS Compatibility Tests
 # ============================================================================