From 60e6aa93ac247e49cf550d8252f7ff4b4353a219 Mon Sep 17 00:00:00 2001 From: Nikita Pawar Date: Mon, 6 Jul 2026 22:34:16 +0530 Subject: [PATCH] feat(ai-backend): add per-user AI quotas and provider usage budgets --- services/ai/openapi.golden.json | 280 ++++++++++++++++++ services/ai/src/vidyasetu_ai/api/router.py | 5 + services/ai/src/vidyasetu_ai/api/routes/ai.py | 78 +++++ services/ai/src/vidyasetu_ai/core/quota.py | 267 +++++++++++++++++ services/ai/src/vidyasetu_ai/db/database.py | 49 +-- .../002_add_quota_and_budget_tables.py | 131 ++++++++ services/ai/src/vidyasetu_ai/db/models.py | 88 ++++++ services/ai/tests/test_quota.py | 221 ++++++++++++++ 8 files changed, 1099 insertions(+), 20 deletions(-) create mode 100644 services/ai/src/vidyasetu_ai/api/routes/ai.py create mode 100644 services/ai/src/vidyasetu_ai/core/quota.py create mode 100644 services/ai/src/vidyasetu_ai/db/migrations/versions/002_add_quota_and_budget_tables.py create mode 100644 services/ai/tests/test_quota.py diff --git a/services/ai/openapi.golden.json b/services/ai/openapi.golden.json index 5a928a53..6d70973e 100644 --- a/services/ai/openapi.golden.json +++ b/services/ai/openapi.golden.json @@ -73,10 +73,154 @@ } ] } + }, + "/api/v1/ai/request": { + "post": { + "tags": [ + "ai" + ], + "summary": "Check Quota Only", + "description": "Check request budget and concurrency, increment requests count and active concurrency.", + "operationId": "check_quota_only_api_v1_ai_request_post", + "responses": { + "200": { + "description": "Successful Response", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/StandardResponse" + } + } + } + }, + "422": { + "description": "Validation Error", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + } + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ] + } + }, + "/api/v1/ai/job/start": { + "post": { + "tags": [ + "ai" + ], + "summary": "Start Job", + "description": "Start a long-running job by checking quota limits and reserving estimated tokens.", + "operationId": "start_job_api_v1_ai_job_start_post", + "requestBody": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/JobStartInput" + } + } + }, + "required": true + }, + "responses": { + "200": { + "description": "Successful Response", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/JobStartResponse" + } + } + } + }, + "422": { + "description": "Validation Error", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + } + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ] + } + }, + "/api/v1/ai/job/reconcile": { + "post": { + "tags": [ + "ai" + ], + "summary": "Reconcile Job", + "description": "Reconcile a completed/failed job and adjust usage tracker tokens and active concurrency.", + "operationId": "reconcile_job_api_v1_ai_job_reconcile_post", + "requestBody": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/JobReconcileInput" + } + } + }, + "required": true + }, + "responses": { + "200": { + "description": "Successful Response", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/StandardResponse" + } + } + } + }, + "422": { + "description": "Validation Error", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + } + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ] + } } }, "components": { "schemas": { + "HTTPValidationError": { + "properties": { + "detail": { + "items": { + "$ref": "#/components/schemas/ValidationError" + }, + "type": "array", + "title": "Detail" + } + }, + "type": "object", + "title": "HTTPValidationError" + }, "HealthResponse": { "properties": { "status": { @@ -107,6 +251,71 @@ ], "title": "HealthResponse" }, + "JobReconcileInput": { + "properties": { + "reservation_id": { + "type": "string", + "format": "uuid", + "title": "Reservation Id" + }, + "actual_tokens": { + "anyOf": [ + { + "type": "integer", + "minimum": 0.0 + }, + { + "type": "null" + } + ], + "title": "Actual Tokens" + }, + "status": { + "type": "string", + "pattern": "^(COMPLETED|FAILED|CANCELLED)$", + "title": "Status", + "default": "COMPLETED" + } + }, + "type": "object", + "required": [ + "reservation_id" + ], + "title": "JobReconcileInput" + }, + "JobStartInput": { + "properties": { + "estimated_tokens": { + "type": "integer", + "minimum": 1.0, + "title": "Estimated Tokens" + } + }, + "type": "object", + "required": [ + "estimated_tokens" + ], + "title": "JobStartInput" + }, + "JobStartResponse": { + "properties": { + "status": { + "type": "string", + "title": "Status" + }, + "reservation_id": { + "type": "string", + "format": "uuid", + "title": "Reservation Id" + } + }, + "type": "object", + "required": [ + "status", + "reservation_id" + ], + "title": "JobStartResponse" + }, "ReadinessResponse": { "properties": { "status": { @@ -120,6 +329,24 @@ }, "type": "object", "title": "Checks" + }, + "loaded_models": { + "items": { + "type": "string" + }, + "type": "array", + "title": "Loaded Models" + }, + "load_error": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Load Error" } }, "type": "object", @@ -127,6 +354,59 @@ "status" ], "title": "ReadinessResponse" + }, + "StandardResponse": { + "properties": { + "status": { + "type": "string", + "title": "Status" + } + }, + "type": "object", + "required": [ + "status" + ], + "title": "StandardResponse" + }, + "ValidationError": { + "properties": { + "loc": { + "items": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "integer" + } + ] + }, + "type": "array", + "title": "Location" + }, + "msg": { + "type": "string", + "title": "Message" + }, + "type": { + "type": "string", + "title": "Error Type" + }, + "input": { + "title": "Input" + }, + "ctx": { + "type": "object", + "title": "Context" + } + }, + "type": "object", + "required": [ + "loc", + "msg", + "type" + ], + "title": "ValidationError" } }, "securitySchemes": { diff --git a/services/ai/src/vidyasetu_ai/api/router.py b/services/ai/src/vidyasetu_ai/api/router.py index 7b9ed8c5..7a3515af 100644 --- a/services/ai/src/vidyasetu_ai/api/router.py +++ b/services/ai/src/vidyasetu_ai/api/router.py @@ -1,5 +1,6 @@ from fastapi import APIRouter, Depends +from vidyasetu_ai.api.routes.ai import router as ai_router from vidyasetu_ai.core.security import require_internal_api_key from vidyasetu_ai.schemas.health import HealthResponse @@ -14,3 +15,7 @@ ) def internal_ping() -> HealthResponse: return HealthResponse(status="ok", service="vidyasetu-ai") + + +api_router.include_router(ai_router) + diff --git a/services/ai/src/vidyasetu_ai/api/routes/ai.py b/services/ai/src/vidyasetu_ai/api/routes/ai.py new file mode 100644 index 00000000..dda64e39 --- /dev/null +++ b/services/ai/src/vidyasetu_ai/api/routes/ai.py @@ -0,0 +1,78 @@ +import uuid +from typing import Annotated + +from fastapi import APIRouter, Depends +from pydantic import BaseModel, Field +from sqlalchemy.ext.asyncio import AsyncSession + +from vidyasetu_ai.core.quota import QuotaManager, enforce_quota, get_user_context +from vidyasetu_ai.core.security import require_internal_api_key +from vidyasetu_ai.db.database import get_db + +router = APIRouter(prefix="/ai", tags=["ai"]) + + +class JobStartInput(BaseModel): + estimated_tokens: int = Field(..., ge=1) + + +class JobReconcileInput(BaseModel): + reservation_id: uuid.UUID + actual_tokens: int | None = Field(None, ge=0) + status: str = Field("COMPLETED", pattern="^(COMPLETED|FAILED|CANCELLED)$") + + +class StandardResponse(BaseModel): + status: str + + +class JobStartResponse(BaseModel): + status: str + reservation_id: uuid.UUID + + +@router.post( + "/request", + response_model=StandardResponse, + dependencies=[Depends(enforce_quota)], +) +async def check_quota_only() -> StandardResponse: + """Check request budget and concurrency, increment requests count and active concurrency.""" + return StandardResponse(status="allowed") + + +@router.post( + "/job/start", + response_model=JobStartResponse, + dependencies=[Depends(require_internal_api_key)], +) +async def start_job( + body: JobStartInput, + user_context: Annotated[tuple[str, bool], Depends(get_user_context)], + db: Annotated[AsyncSession, Depends(get_db)], +) -> JobStartResponse: + """Start a long-running job by checking quota limits and reserving estimated tokens.""" + user_id, is_byok = user_context + reservation_id = await QuotaManager.check_and_update_usage( + db, user_id, is_byok, estimated_tokens=body.estimated_tokens + ) + return JobStartResponse(status="started", reservation_id=reservation_id) + + +@router.post( + "/job/reconcile", + response_model=StandardResponse, + dependencies=[Depends(require_internal_api_key)], +) +async def reconcile_job( + body: JobReconcileInput, + db: Annotated[AsyncSession, Depends(get_db)], +) -> StandardResponse: + """Reconcile a completed/failed job and adjust usage tracker tokens and active concurrency.""" + await QuotaManager.reconcile_job( + db, + reservation_id=body.reservation_id, + actual_tokens=body.actual_tokens, + status=body.status, + ) + return StandardResponse(status="reconciled") diff --git a/services/ai/src/vidyasetu_ai/core/quota.py b/services/ai/src/vidyasetu_ai/core/quota.py new file mode 100644 index 00000000..2567b68b --- /dev/null +++ b/services/ai/src/vidyasetu_ai/core/quota.py @@ -0,0 +1,267 @@ +import logging +import uuid +from datetime import datetime, timezone +from typing import Annotated + +from fastapi import Depends, Header, HTTPException, status +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from vidyasetu_ai.core.security import require_internal_api_key +from vidyasetu_ai.db.database import get_db +from vidyasetu_ai.db.models import ( + AiJobReservation, + UserQuotaAssignment, + UserQuotaPolicy, + UserUsageTracker, +) + +logger = logging.getLogger("vidyasetu_ai.quota") + + +class QuotaLimitExceeded(HTTPException): + def __init__(self, code: str, detail: str, retry_after: int | None = None): + headers = {} + if retry_after is not None: + headers["Retry-After"] = str(max(1, int(retry_after))) + super().__init__( + status_code=429, + detail={"code": code, "message": detail, "retry_after": retry_after}, + headers=headers, + ) + + +class QuotaManager: + @staticmethod + async def get_user_policy(db: AsyncSession, user_id: str, is_byok: bool) -> UserQuotaPolicy: + # Check custom assignment + stmt = select(UserQuotaAssignment).where(UserQuotaAssignment.user_id == user_id) + result = await db.execute(stmt) + assignment = result.scalar_one_or_none() + + if assignment: + policy_id = assignment.byok_policy_id if is_byok else assignment.platform_policy_id + stmt = select(UserQuotaPolicy).where(UserQuotaPolicy.id == policy_id) + result = await db.execute(stmt) + policy = result.scalar_one_or_none() + if policy: + return policy + + # Fallback to defaults + default_name = "default_byok" if is_byok else "default_platform" + stmt = select(UserQuotaPolicy).where(UserQuotaPolicy.name == default_name) + result = await db.execute(stmt) + policy = result.scalar_one_or_none() + if not policy: + raise HTTPException( + status_code=500, + detail=f"Default quota policy '{default_name}' not seeded", + ) + return policy + + @classmethod + async def check_and_update_usage( + cls, + db: AsyncSession, + user_id: str, + is_byok: bool, + estimated_tokens: int = 0, + ) -> uuid.UUID | None: + now = datetime.now(timezone.utc) + policy = await cls.get_user_policy(db, user_id, is_byok) + + # Lock the row for this user & request type to prevent concurrency race conditions + stmt = ( + select(UserUsageTracker) + .where(UserUsageTracker.user_id == user_id, UserUsageTracker.is_byok == is_byok) + .with_for_update() + ) + result = await db.execute(stmt) + tracker = result.scalar_one_or_none() + + if not tracker: + tracker = UserUsageTracker( + user_id=user_id, + is_byok=is_byok, + window_start=now, + requests_used=0, + tokens_used=0, + active_concurrency=0, + ) + db.add(tracker) + await db.flush() + + # Check window reset + now_naive = now.replace(tzinfo=None) + window_start_naive = tracker.window_start.replace(tzinfo=None) if tracker.window_start.tzinfo else tracker.window_start + elapsed = (now_naive - window_start_naive).total_seconds() + if elapsed >= policy.window_seconds: + tracker.window_start = now + tracker.requests_used = 0 + tracker.tokens_used = 0 + + # Calculate remaining time in current window for Retry-After + retry_after = int(max(0.0, policy.window_seconds - elapsed)) + + + # 1. Concurrency limits + if tracker.active_concurrency >= policy.max_concurrency: + logger.info( + f"Quota Decision: REJECTED | User: {user_id} | BYOK: {is_byok} | " + f"Reason: CONCURRENCY_LIMIT_EXCEEDED | Current: {tracker.active_concurrency} | " + f"Max: {policy.max_concurrency}" + ) + raise QuotaLimitExceeded( + code="CONCURRENCY_LIMIT_EXCEEDED", + detail="Too many concurrent requests. Please wait for active tasks to finish.", + retry_after=5, + ) + + # 2. Request limits + if tracker.requests_used + 1 > policy.max_requests: + logger.info( + f"Quota Decision: REJECTED | User: {user_id} | BYOK: {is_byok} | " + f"Reason: REQUEST_LIMIT_EXCEEDED | Current: {tracker.requests_used} | " + f"Max: {policy.max_requests}" + ) + raise QuotaLimitExceeded( + code="REQUEST_LIMIT_EXCEEDED", + detail="Request budget exceeded for the current window.", + retry_after=retry_after, + ) + + # 3. Token limits + if policy.max_tokens is not None: + if tracker.tokens_used + estimated_tokens > policy.max_tokens: + logger.info( + f"Quota Decision: REJECTED | User: {user_id} | BYOK: {is_byok} | " + f"Reason: TOKEN_BUDGET_EXCEEDED | Current Tokens: {tracker.tokens_used} | " + f"Req: {estimated_tokens} | Max: {policy.max_tokens}" + ) + raise QuotaLimitExceeded( + code="TOKEN_BUDGET_EXCEEDED", + detail="Token budget exceeded for the current window.", + retry_after=retry_after, + ) + + # Log decision + logger.info( + f"Quota Decision: ALLOWED | User: {user_id} | BYOK: {is_byok} | " + f"Requests: {tracker.requests_used + 1}/{policy.max_requests} | " + f"Tokens: {tracker.tokens_used + estimated_tokens}/{policy.max_tokens or 'unlimited'} | " + f"Concurrency: {tracker.active_concurrency + 1}/{policy.max_concurrency}" + ) + + # Commit updates + tracker.requests_used += 1 + tracker.active_concurrency += 1 + tracker.tokens_used += estimated_tokens + tracker.updated_at = now + + reservation_id = None + if estimated_tokens > 0: + reservation_id = uuid.uuid4() + reservation = AiJobReservation( + id=reservation_id, + user_id=user_id, + is_byok=is_byok, + reserved_tokens=estimated_tokens, + status="PENDING", + ) + db.add(reservation) + + await db.commit() + return reservation_id + + @classmethod + async def reconcile_job( + cls, + db: AsyncSession, + reservation_id: uuid.UUID, + actual_tokens: int | None = None, + status: str = "COMPLETED", + ) -> None: + stmt = ( + select(AiJobReservation) + .where(AiJobReservation.id == reservation_id) + .with_for_update() + ) + result = await db.execute(stmt) + reservation = result.scalar_one_or_none() + + if not reservation or reservation.status != "PENDING": + return + + stmt = ( + select(UserUsageTracker) + .where( + UserUsageTracker.user_id == reservation.user_id, + UserUsageTracker.is_byok == reservation.is_byok, + ) + .with_for_update() + ) + result = await db.execute(stmt) + tracker = result.scalar_one_or_none() + + if tracker: + tracker.active_concurrency = max(0, tracker.active_concurrency - 1) + if status == "COMPLETED" and actual_tokens is not None: + diff = actual_tokens - reservation.reserved_tokens + tracker.tokens_used = max(0, tracker.tokens_used + diff) + else: + # Reclaim all reserved tokens on failure or cancel + tracker.tokens_used = max(0, tracker.tokens_used - reservation.reserved_tokens) + tracker.updated_at = datetime.now(timezone.utc) + + reservation.status = status + await db.commit() + + @classmethod + async def release_concurrency( + cls, + db: AsyncSession, + user_id: str, + is_byok: bool, + ) -> None: + stmt = ( + select(UserUsageTracker) + .where( + UserUsageTracker.user_id == user_id, + UserUsageTracker.is_byok == is_byok, + ) + .with_for_update() + ) + result = await db.execute(stmt) + tracker = result.scalar_one_or_none() + + if tracker: + tracker.active_concurrency = max(0, tracker.active_concurrency - 1) + tracker.updated_at = datetime.now(timezone.utc) + await db.commit() + + +# FastAPI Dependencies +def get_user_context( + x_user_id: Annotated[str | None, Header(include_in_schema=False)] = None, + x_user_is_byok: Annotated[bool | None, Header(include_in_schema=False)] = None, +) -> tuple[str, bool]: + if not x_user_id: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="X-User-Id header is required for AI endpoints", + ) + return x_user_id, bool(x_user_is_byok) + + +async def enforce_quota( + user_context: Annotated[tuple[str, bool], Depends(get_user_context)], + db: Annotated[AsyncSession, Depends(get_db)], + internal_api_key: Annotated[None, Depends(require_internal_api_key)] = None, +): + user_id, is_byok = user_context + await QuotaManager.check_and_update_usage(db, user_id, is_byok) + try: + yield user_id, is_byok + finally: + await QuotaManager.release_concurrency(db, user_id, is_byok) + diff --git a/services/ai/src/vidyasetu_ai/db/database.py b/services/ai/src/vidyasetu_ai/db/database.py index 659933b0..e4d1a6e0 100644 --- a/services/ai/src/vidyasetu_ai/db/database.py +++ b/services/ai/src/vidyasetu_ai/db/database.py @@ -1,20 +1,29 @@ -from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine -from sqlalchemy.orm import DeclarativeBase - -from vidyasetu_ai.core.config import get_settings - - -class Base(DeclarativeBase): - pass - - -def build_engine(database_url: str | None = None): - url = database_url or get_settings().database_url - # Replace postgresql:// with postgresql+asyncpg:// for async driver - if url.startswith("postgresql://"): - url = url.replace("postgresql://", "postgresql+asyncpg://", 1) - return create_async_engine(url, echo=False, pool_pre_ping=True) - - -def build_session_factory(engine): - return async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False) \ No newline at end of file +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine +from sqlalchemy.orm import DeclarativeBase + +from vidyasetu_ai.core.config import get_settings + + +class Base(DeclarativeBase): + pass + + +def build_engine(database_url: str | None = None): + url = database_url or get_settings().database_url + # Replace postgresql:// with postgresql+asyncpg:// for async driver + if url.startswith("postgresql://"): + url = url.replace("postgresql://", "postgresql+asyncpg://", 1) + return create_async_engine(url, echo=False, pool_pre_ping=True) + + +def build_session_factory(engine): + return async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False) + + +_engine = build_engine() +_session_factory = build_session_factory(_engine) + + +async def get_db(): + async with _session_factory() as session: + yield session \ No newline at end of file diff --git a/services/ai/src/vidyasetu_ai/db/migrations/versions/002_add_quota_and_budget_tables.py b/services/ai/src/vidyasetu_ai/db/migrations/versions/002_add_quota_and_budget_tables.py new file mode 100644 index 00000000..e72524bf --- /dev/null +++ b/services/ai/src/vidyasetu_ai/db/migrations/versions/002_add_quota_and_budget_tables.py @@ -0,0 +1,131 @@ +"""Add quota and budget tables. + +Revision ID: 002 +Revises: 001 +Create Date: 2026-07-06 +""" + +import sqlalchemy as sa +from alembic import op +from sqlalchemy.dialects.postgresql import UUID + +revision = "002" +down_revision = "001" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + # user_quota_policies + op.create_table( + "user_quota_policies", + sa.Column("id", UUID(as_uuid=True), primary_key=True), + sa.Column("name", sa.Text(), nullable=False, unique=True), + sa.Column("max_requests", sa.Integer(), nullable=False), + sa.Column("max_tokens", sa.BigInteger(), nullable=True), + sa.Column("max_concurrency", sa.Integer(), nullable=False), + sa.Column("window_seconds", sa.Integer(), nullable=False, server_default="86400"), + sa.Column( + "created_at", + sa.DateTime(timezone=True), + server_default=sa.func.now(), + nullable=False, + ), + sa.Column( + "updated_at", + sa.DateTime(timezone=True), + server_default=sa.func.now(), + nullable=False, + ), + ) + + # user_quota_assignments + op.create_table( + "user_quota_assignments", + sa.Column("user_id", sa.Text(), primary_key=True), + sa.Column( + "platform_policy_id", + UUID(as_uuid=True), + sa.ForeignKey("user_quota_policies.id", ondelete="RESTRICT"), + nullable=False, + ), + sa.Column( + "byok_policy_id", + UUID(as_uuid=True), + sa.ForeignKey("user_quota_policies.id", ondelete="RESTRICT"), + nullable=False, + ), + sa.Column( + "created_at", + sa.DateTime(timezone=True), + server_default=sa.func.now(), + nullable=False, + ), + sa.Column( + "updated_at", + sa.DateTime(timezone=True), + server_default=sa.func.now(), + nullable=False, + ), + ) + + # user_usage_trackers + op.create_table( + "user_usage_trackers", + sa.Column("user_id", sa.Text(), primary_key=True), + sa.Column("is_byok", sa.Boolean(), primary_key=True, server_default="false"), + sa.Column( + "window_start", + sa.DateTime(timezone=True), + server_default=sa.func.now(), + nullable=False, + ), + sa.Column("requests_used", sa.Integer(), nullable=False, server_default="0"), + sa.Column("tokens_used", sa.BigInteger(), nullable=False, server_default="0"), + sa.Column("active_concurrency", sa.Integer(), nullable=False, server_default="0"), + sa.Column( + "updated_at", + sa.DateTime(timezone=True), + server_default=sa.func.now(), + nullable=False, + ), + ) + + # ai_job_reservations + op.create_table( + "ai_job_reservations", + sa.Column("id", UUID(as_uuid=True), primary_key=True), + sa.Column("user_id", sa.Text(), nullable=False), + sa.Column("is_byok", sa.Boolean(), nullable=False, server_default="false"), + sa.Column("reserved_tokens", sa.Integer(), nullable=False), + sa.Column("status", sa.Text(), nullable=False, server_default="PENDING"), + sa.Column( + "created_at", + sa.DateTime(timezone=True), + server_default=sa.func.now(), + nullable=False, + ), + sa.Column( + "updated_at", + sa.DateTime(timezone=True), + server_default=sa.func.now(), + nullable=False, + ), + ) + + # Seed default policies + op.execute( + """ + INSERT INTO user_quota_policies (id, name, max_requests, max_tokens, max_concurrency, window_seconds, created_at, updated_at) + VALUES + ('a0eebc99-9c0b-4ef8-bb6d-6bb9bd380a11', 'default_platform', 100, 50000, 2, 86400, NOW(), NOW()), + ('a0eebc99-9c0b-4ef8-bb6d-6bb9bd380a22', 'default_byok', 500, NULL, 5, 86400, NOW(), NOW()) + """ + ) + + +def downgrade() -> None: + op.drop_table("ai_job_reservations") + op.drop_table("user_usage_trackers") + op.drop_table("user_quota_assignments") + op.drop_table("user_quota_policies") diff --git a/services/ai/src/vidyasetu_ai/db/models.py b/services/ai/src/vidyasetu_ai/db/models.py index 36b71b94..9853b248 100644 --- a/services/ai/src/vidyasetu_ai/db/models.py +++ b/services/ai/src/vidyasetu_ai/db/models.py @@ -81,3 +81,91 @@ class AiChunk(Base): Index("ix_ai_chunks_document_id", "document_id"), Index("ix_ai_chunks_embedding_model", "embedding_model"), ) + + +class UserQuotaPolicy(Base): + __tablename__ = "user_quota_policies" + + id: Mapped[uuid.UUID] = mapped_column( + UUID(as_uuid=True), primary_key=True, default=uuid.uuid4 + ) + name: Mapped[str] = mapped_column(Text, nullable=False, unique=True) + max_requests: Mapped[int] = mapped_column(Integer, nullable=False) + max_tokens: Mapped[int | None] = mapped_column(BigInteger, nullable=True) + max_concurrency: Mapped[int] = mapped_column(Integer, nullable=False) + window_seconds: Mapped[int] = mapped_column(Integer, nullable=False, default=86400) + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), server_default=func.now(), nullable=False + ) + updated_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), + server_default=func.now(), + onupdate=func.now(), + nullable=False, + ) + + +class UserQuotaAssignment(Base): + __tablename__ = "user_quota_assignments" + + user_id: Mapped[str] = mapped_column(Text, primary_key=True) + platform_policy_id: Mapped[uuid.UUID] = mapped_column( + UUID(as_uuid=True), + ForeignKey("user_quota_policies.id", ondelete="RESTRICT"), + nullable=False, + ) + byok_policy_id: Mapped[uuid.UUID] = mapped_column( + UUID(as_uuid=True), + ForeignKey("user_quota_policies.id", ondelete="RESTRICT"), + nullable=False, + ) + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), server_default=func.now(), nullable=False + ) + updated_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), + server_default=func.now(), + onupdate=func.now(), + nullable=False, + ) + + +class UserUsageTracker(Base): + __tablename__ = "user_usage_trackers" + + user_id: Mapped[str] = mapped_column(Text, primary_key=True) + is_byok: Mapped[bool] = mapped_column(primary_key=True, default=False) + window_start: Mapped[datetime] = mapped_column( + DateTime(timezone=True), server_default=func.now(), nullable=False + ) + requests_used: Mapped[int] = mapped_column(Integer, nullable=False, default=0) + tokens_used: Mapped[int] = mapped_column(BigInteger, nullable=False, default=0) + active_concurrency: Mapped[int] = mapped_column(Integer, nullable=False, default=0) + updated_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), + server_default=func.now(), + onupdate=func.now(), + nullable=False, + ) + + +class AiJobReservation(Base): + __tablename__ = "ai_job_reservations" + + id: Mapped[uuid.UUID] = mapped_column( + UUID(as_uuid=True), primary_key=True, default=uuid.uuid4 + ) + user_id: Mapped[str] = mapped_column(Text, nullable=False) + is_byok: Mapped[bool] = mapped_column(nullable=False, default=False) + reserved_tokens: Mapped[int] = mapped_column(Integer, nullable=False) + status: Mapped[str] = mapped_column(Text, nullable=False, default="PENDING") + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), server_default=func.now(), nullable=False + ) + updated_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), + server_default=func.now(), + onupdate=func.now(), + nullable=False, + ) + diff --git a/services/ai/tests/test_quota.py b/services/ai/tests/test_quota.py new file mode 100644 index 00000000..a42902ea --- /dev/null +++ b/services/ai/tests/test_quota.py @@ -0,0 +1,221 @@ +import uuid +import pytest +import pytest_asyncio +from fastapi.testclient import TestClient +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine + +from vidyasetu_ai.core.quota import QuotaManager +from vidyasetu_ai.db.database import Base, get_db +from vidyasetu_ai.db.models import ( + AiJobReservation, + UserQuotaAssignment, + UserQuotaPolicy, + UserUsageTracker, +) +from vidyasetu_ai.main import app + +VALID_KEY = "test-internal-api-key-at-least-32-characters" + +# Create in-memory async SQLite database +test_engine = create_async_engine("sqlite+aiosqlite:///:memory:", echo=False) +test_session_factory = async_sessionmaker( + test_engine, class_=AsyncSession, expire_on_commit=False +) + + +async def override_get_db(): + async with test_session_factory() as session: + yield session + + +app.dependency_overrides[get_db] = override_get_db + + +@pytest_asyncio.fixture(autouse=True) +async def setup_db(): + async with test_engine.begin() as conn: + await conn.run_sync( + lambda connection: Base.metadata.create_all( + connection, + tables=[ + Base.metadata.tables["user_quota_policies"], + Base.metadata.tables["user_quota_assignments"], + Base.metadata.tables["user_usage_trackers"], + Base.metadata.tables["ai_job_reservations"], + ], + ) + ) + + # Seed default policies + async with test_session_factory() as session: + p1 = UserQuotaPolicy( + id=uuid.UUID("a0eebc99-9c0b-4ef8-bb6d-6bb9bd380a11"), + name="default_platform", + max_requests=3, + max_tokens=1000, + max_concurrency=2, + window_seconds=10, + ) + p2 = UserQuotaPolicy( + id=uuid.UUID("a0eebc99-9c0b-4ef8-bb6d-6bb9bd380a22"), + name="default_byok", + max_requests=5, + max_tokens=None, + max_concurrency=3, + window_seconds=10, + ) + session.add_all([p1, p2]) + await session.commit() + + yield + + async with test_engine.begin() as conn: + await conn.run_sync( + lambda connection: Base.metadata.drop_all( + connection, + tables=[ + Base.metadata.tables["ai_job_reservations"], + Base.metadata.tables["user_usage_trackers"], + Base.metadata.tables["user_quota_assignments"], + Base.metadata.tables["user_quota_policies"], + ], + ) + ) + + +@pytest.mark.asyncio +async def test_get_user_policy_defaults() -> None: + async with test_session_factory() as session: + p_platform = await QuotaManager.get_user_policy(session, "user1", is_byok=False) + assert p_platform.name == "default_platform" + assert p_platform.max_requests == 3 + + p_byok = await QuotaManager.get_user_policy(session, "user1", is_byok=True) + assert p_byok.name == "default_byok" + assert p_byok.max_tokens is None + + +@pytest.mark.asyncio +async def test_get_user_policy_custom() -> None: + async with test_session_factory() as session: + # Create a custom policy + custom_p = UserQuotaPolicy( + id=uuid.UUID("a0eebc99-9c0b-4ef8-bb6d-6bb9bd380a99"), + name="custom_premium", + max_requests=100, + max_tokens=100000, + max_concurrency=10, + window_seconds=60, + ) + assignment = UserQuotaAssignment( + user_id="vip_user", + platform_policy_id=custom_p.id, + byok_policy_id=custom_p.id, + ) + session.add_all([custom_p, assignment]) + await session.commit() + + p = await QuotaManager.get_user_policy(session, "vip_user", is_byok=False) + assert p.name == "custom_premium" + assert p.max_requests == 100 + + +def test_api_quota_rate_limits(client: TestClient) -> None: + headers = { + "X-Internal-API-Key": VALID_KEY, + "X-User-Id": "user123", + "X-User-Is-BYOK": "false", + } + + # First request: Allowed + response = client.post("/api/v1/ai/request", headers=headers) + assert response.status_code == 200 + assert response.json()["status"] == "allowed" + + # Second request: Allowed + response = client.post("/api/v1/ai/request", headers=headers) + assert response.status_code == 200 + + # Third request: Allowed + response = client.post("/api/v1/ai/request", headers=headers) + assert response.status_code == 200 + + # Fourth request: 429 Limit Exceeded + response = client.post("/api/v1/ai/request", headers=headers) + assert response.status_code == 429 + res_body = response.json()["detail"] + assert res_body["code"] == "REQUEST_LIMIT_EXCEEDED" + assert "Retry-After" in response.headers + + +def test_api_byok_bypasses_tokens_limit(client: TestClient) -> None: + headers = { + "X-Internal-API-Key": VALID_KEY, + "X-User-Id": "byok_user", + "X-User-Is-BYOK": "true", + } + + # Start a job requesting huge amount of tokens + response = client.post( + "/api/v1/ai/job/start", + headers=headers, + json={"estimated_tokens": 999999}, + ) + assert response.status_code == 200 + assert response.json()["status"] == "started" + + +def test_concurrency_and_reconciliation(client: TestClient) -> None: + headers = { + "X-Internal-API-Key": VALID_KEY, + "X-User-Id": "job_user", + "X-User-Is-BYOK": "false", + } + + # Start job 1 (concurrency 1) + r1 = client.post( + "/api/v1/ai/job/start", + headers=headers, + json={"estimated_tokens": 100}, + ) + assert r1.status_code == 200 + res_id_1 = r1.json()["reservation_id"] + + # Start job 2 (concurrency 2) + r2 = client.post( + "/api/v1/ai/job/start", + headers=headers, + json={"estimated_tokens": 100}, + ) + assert r2.status_code == 200 + res_id_2 = r2.json()["reservation_id"] + + # Start job 3 (concurrency 3 -> exceeds max_concurrency of 2) + r3 = client.post( + "/api/v1/ai/job/start", + headers=headers, + json={"estimated_tokens": 100}, + ) + assert r3.status_code == 429 + assert r3.json()["detail"]["code"] == "CONCURRENCY_LIMIT_EXCEEDED" + + # Reconcile/complete job 1 (concurrency becomes 1) + reconcile_res = client.post( + "/api/v1/ai/job/reconcile", + headers=headers, + json={ + "reservation_id": res_id_1, + "actual_tokens": 50, + "status": "COMPLETED", + }, + ) + assert reconcile_res.status_code == 200 + + # Start job 3 again (should succeed now) + r3 = client.post( + "/api/v1/ai/job/start", + headers=headers, + json={"estimated_tokens": 100}, + ) + assert r3.status_code == 200