Quota helpers

Author: tisnik

src/utils/quota.py

@@ -0,0 +1,67 @@
+"""Quota handling helper functions."""
+
+import psycopg2
+
+from fastapi import HTTPException, status
+
+from quota.quota_limiter import QuotaLimiter
+
+from log import get_logger
+
+logger = get_logger(__name__)
+
+
+def consume_tokens(
+    quota_limiters: list[QuotaLimiter],
+    user_id: str,
+    input_tokens: int,
+    output_tokens: int,
+) -> None:
+    """Consume tokens from cluster and/or user quotas."""
+    # check if any quota limiter is configured
+    for quota_limiter in quota_limiters:
+        quota_limiter.consume_tokens(
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            subject_id=user_id,
+        )
+
+
+def check_tokens_available(quota_limiters: list[QuotaLimiter], user_id: str) -> None:
+    """Check if tokens are available for user."""
+    try:
+        for quota_limiter in quota_limiters:
+            quota_limiter.ensure_available_quota(subject_id=user_id)
+    except psycopg2.Error as pg_error:
+        message = "Error communicating with quota database backend"
+        logger.error(message)
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail={
+                "response": message,
+                "cause": str(pg_error),
+            },
+        ) from pg_error
+    except Exception as quota_exceed_error:
+        message = "The quota has been exceeded"
+        logger.error(message)
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail={
+                "response": message,
+                "cause": str(quota_exceed_error),
+            },
+        ) from quota_exceed_error
+
+
+def get_available_quotas(
+    quota_limiters: list[QuotaLimiter],
+    user_id: str,
+) -> dict[str, int]:
+    """Get quota available from all quota limiters."""
+    available_quotas: dict[str, int] = {}
+    for quota_limiter in quota_limiters:
+        name = quota_limiter.__class__.__name__
+        available_quota = quota_limiter.available_quota(user_id)
+        available_quotas[name] = available_quota
+    return available_quotas