From 6d8a0d436d4de20bd49eae28f7221c68443b8e30 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Sun, 4 May 2025 09:23:23 +0000 Subject: [PATCH 01/24] feat-72: ai usage quota - adds new db table to track - checks for both ai requests --- backend/.env.example | 2 ++ backend/src/core/config.py | 2 ++ backend/src/flashcards/api.py | 4 ++++ backend/src/flashcards/models.py | 7 ++++++ backend/src/flashcards/services.py | 38 ++++++++++++++++++++++++++++-- 5 files changed, 51 insertions(+), 2 deletions(-) diff --git a/backend/.env.example b/backend/.env.example index 49dab74..ae93e51 100644 --- a/backend/.env.example +++ b/backend/.env.example @@ -21,6 +21,8 @@ POSTGRES_PASSWORD=changethis # AI AI_MODEL= AI_API_KEY= +AI_MAX_USAGE_QUOTA=30 +AI_QUOTA_TIME_RANGE_DAYS=30 # time in days COLLECTION_GENERATION_PROMPT="I want to generate flashcards on a specific topic for efficient studying. Please create a set of flashcards covering key concepts, definitions, important details, and examples, with a focus on progressively building understanding of the topic. The flashcards should aim to provide a helpful learning experience by using structured explanations, real-world examples and formatting. Each flashcard should follow this format: Front (Question/Prompt): A clear and concise question or term to test recall, starting with introductory concepts and moving toward more complex details. Back (Answer): If the front is a concept or topic, provide a detailed explanation, broken down into clear paragraphs with easy-to-understand language. If possible, include a real-world example, analogy or illustrative diagrams to make the concept more memorable and relatable. If the front is a vocabulary word (for language learning), provide a direct translation in the target language. Optional Hint: A short clue to aid recall, especially for more complex concepts. Important: Use valid Markdown format for the back of the flashcard." CARD_GENERATION_PROMPT="I want to generate a flashcard on a specific topic. The contents of the flashcard should provide helpful information that aim to help the learner retain the concepts given. The flashcard must follow this format: Front (Question/Prompt): A clear and concise question or term to test recall. Back (Answer): If the front is a concept or topic, provide a detailed explanation, broken down into clear paragraphs with easy-to-understand language. If possible, include a real-world example, analogy or illustrative diagrams to make the concept more memorable and relatable. If the front is a vocabulary word (for language learning), provide a direct translation in the target language. Important: Use valid Markdown format for the back of the flashcard." \ No newline at end of file diff --git a/backend/src/core/config.py b/backend/src/core/config.py index bd66a54..e5a2a4a 100644 --- a/backend/src/core/config.py +++ b/backend/src/core/config.py @@ -64,6 +64,8 @@ def SQLALCHEMY_DATABASE_URI(self) -> PostgresDsn: AI_API_KEY: str | None = None AI_MODEL: str | None = None + AI_MAX_USAGE_QUOTA: int | None = None + AI_QUOTA_TIME_RANGE_DAYS: int | None = None COLLECTION_GENERATION_PROMPT: str | None = None CARD_GENERATION_PROMPT: str | None = None diff --git a/backend/src/flashcards/api.py b/backend/src/flashcards/api.py index f657d87..ea0ff2c 100644 --- a/backend/src/flashcards/api.py +++ b/backend/src/flashcards/api.py @@ -52,6 +52,8 @@ async def create_collection( if collection_in.prompt: try: + if not services.is_within_ai_usage_quota(session, current_user.id): + raise HTTPException(status_code=429, detail="Quota for AI usage is reached.") flashcard_collection = await services.generate_ai_collection( provider, collection_in.prompt ) @@ -145,6 +147,8 @@ async def create_card( if not access_checked: raise HTTPException(status_code=404, detail="Collection not found") if card_in.prompt: + if not services.is_within_ai_usage_quota(session, current_user.id): + raise HTTPException(status_code=429, detail="Quota for AI usage is reached.") card_base = await services.generate_ai_flashcard(card_in.prompt, provider) card_in.front = card_base.front card_in.back = card_base.back diff --git a/backend/src/flashcards/models.py b/backend/src/flashcards/models.py index 3f2c95b..73846df 100644 --- a/backend/src/flashcards/models.py +++ b/backend/src/flashcards/models.py @@ -63,3 +63,10 @@ class PracticeCard(SQLModel, table=True): updated_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) session: PracticeSession = Relationship(back_populates="practice_cards") card: Card = Relationship(back_populates="practice_cards") + +class AIUsageQuota(SQLModel, table=True): + id: uuid.UUID | None = Field(default_factory=uuid.uuid4, primary_key=True) + user_id: uuid.UUID = Field(foreign_key="user.id", index=True, ondelete="CASCADE") + user: "User" = Relationship(back_populates="collections") + usage_count: int = Field(default=0) + last_reset_time: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) \ No newline at end of file diff --git a/backend/src/flashcards/services.py b/backend/src/flashcards/services.py index da32518..2330958 100644 --- a/backend/src/flashcards/services.py +++ b/backend/src/flashcards/services.py @@ -1,7 +1,7 @@ import json import random import uuid -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from typing import Literal from google import genai @@ -9,10 +9,17 @@ from sqlmodel import Session, func, select from src.ai_models.gemini.exceptions import AIGenerationError +from src.core.config import settings from .ai_config import get_card_config, get_flashcard_config from .exceptions import EmptyCollectionError -from .models import Card, Collection, PracticeCard, PracticeSession +from .models import ( + AIUsageQuota, + Card, + Collection, + PracticeCard, + PracticeSession +) from .schemas import ( AIFlashcardCollection, CardBase, @@ -408,3 +415,30 @@ async def generate_ai_flashcard(prompt: str, provider) -> CardBase: raise AIGenerationError(f"Invalid AI response format: {str(e)}") except Exception as e: raise AIGenerationError(f"Error processing AI response: {str(e)}") + + +def is_within_ai_usage_quota(session: Session, user_id: uuid.UUID) -> bool: + if _has_exceeded_usage_quota(session, user_id): + return False + _increase_usage_quota(session, user_id) + return True + +def _has_exceeded_usage_quota(session: Session, user_id: uuid.UUID) -> bool: + statement = select(AIUsageQuota).filter_by(user_id=user_id) + quota = session.exec(statement).first() + if not quota: + return False + return quota.usage_count >= settings.AI_MAX_USAGE_QUOTA + + +def _increase_usage_quota(session: Session, user_id: uuid.UUID): + statement = select(AIUsageQuota).filter_by(user_id=user_id) + quota = session.exec(statement).first() + if not quota: + quota = AIUsageQuota(user_id=user_id) + session.add(quota) + if datetime.now() - quota.last_reset_time > timedelta(days=settings.AI_QUOTA_TIME_RANGE_DAYS): + quota.usage_count = 0 + quota.last_reset_time = datetime.now() + quota.usage_count += 1 + session.commit() From 0b58fd810db697a5b0fd12f663d2feecad315ed2 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Sun, 4 May 2025 09:30:24 +0000 Subject: [PATCH 02/24] feat-72: remove unneeded reference --- backend/src/flashcards/models.py | 1 - 1 file changed, 1 deletion(-) diff --git a/backend/src/flashcards/models.py b/backend/src/flashcards/models.py index 73846df..3d1225f 100644 --- a/backend/src/flashcards/models.py +++ b/backend/src/flashcards/models.py @@ -67,6 +67,5 @@ class PracticeCard(SQLModel, table=True): class AIUsageQuota(SQLModel, table=True): id: uuid.UUID | None = Field(default_factory=uuid.uuid4, primary_key=True) user_id: uuid.UUID = Field(foreign_key="user.id", index=True, ondelete="CASCADE") - user: "User" = Relationship(back_populates="collections") usage_count: int = Field(default=0) last_reset_time: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) \ No newline at end of file From 531eacc5f62386e88a63572c68b7f37827d8d1b7 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Sun, 4 May 2025 10:10:12 +0000 Subject: [PATCH 03/24] feat-72: include new schema version for new table --- .../d1ea38d75310_add_ai_usage_quota_tables.py | 31 +++++++++++++++++++ backend/src/flashcards/services.py | 2 +- 2 files changed, 32 insertions(+), 1 deletion(-) create mode 100644 backend/src/alembic/versions/d1ea38d75310_add_ai_usage_quota_tables.py diff --git a/backend/src/alembic/versions/d1ea38d75310_add_ai_usage_quota_tables.py b/backend/src/alembic/versions/d1ea38d75310_add_ai_usage_quota_tables.py new file mode 100644 index 0000000..c0bddee --- /dev/null +++ b/backend/src/alembic/versions/d1ea38d75310_add_ai_usage_quota_tables.py @@ -0,0 +1,31 @@ +"""Add AI Usage Quota tables + +Revision ID: d1ea38d75310 +Revises: cb16ae472c1e +Create Date: 2025-05-04 09:59:20.325131 + +""" +from alembic import op +import sqlalchemy as sa +import sqlmodel.sql.sqltypes + + +# revision identifiers, used by Alembic. +revision = 'd1ea38d75310' +down_revision = 'cb16ae472c1e' +branch_labels = None +depends_on = None + + +def upgrade(): + op.create_table( + 'aiusagequota', + sa.Column('id', sa.UUID(), primary_key=True, nullable=False), + sa.Column('user_id', sa.UUID(), sa.ForeignKey('user.id', ondelete='CASCADE'), index=True, nullable=False), + sa.Column('usage_count', sa.Integer, default=0, nullable=False), + sa.Column('last_reset_time', sa.DateTime(timezone=True), nullable=False), + ) + + +def downgrade(): + op.drop_table('aiusagequota') diff --git a/backend/src/flashcards/services.py b/backend/src/flashcards/services.py index 2330958..c6774cf 100644 --- a/backend/src/flashcards/services.py +++ b/backend/src/flashcards/services.py @@ -437,7 +437,7 @@ def _increase_usage_quota(session: Session, user_id: uuid.UUID): if not quota: quota = AIUsageQuota(user_id=user_id) session.add(quota) - if datetime.now() - quota.last_reset_time > timedelta(days=settings.AI_QUOTA_TIME_RANGE_DAYS): + if datetime.now(timezone.utc) - quota.last_reset_time => timedelta(days=settings.AI_QUOTA_TIME_RANGE_DAYS): quota.usage_count = 0 quota.last_reset_time = datetime.now() quota.usage_count += 1 From 9a3f24c6b5e3a72c80de4f06fc107cde689b386f Mon Sep 17 00:00:00 2001 From: Sebastian Date: Sun, 4 May 2025 19:39:38 +0000 Subject: [PATCH 04/24] feat-72: minor fixes --- backend/src/flashcards/services.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/backend/src/flashcards/services.py b/backend/src/flashcards/services.py index c6774cf..ffb9195 100644 --- a/backend/src/flashcards/services.py +++ b/backend/src/flashcards/services.py @@ -437,8 +437,8 @@ def _increase_usage_quota(session: Session, user_id: uuid.UUID): if not quota: quota = AIUsageQuota(user_id=user_id) session.add(quota) - if datetime.now(timezone.utc) - quota.last_reset_time => timedelta(days=settings.AI_QUOTA_TIME_RANGE_DAYS): + if datetime.now(timezone.utc) - quota.last_reset_time >= timedelta(days=settings.AI_QUOTA_TIME_RANGE_DAYS): quota.usage_count = 0 - quota.last_reset_time = datetime.now() + quota.last_reset_time = datetime.now(timezone.utc) quota.usage_count += 1 session.commit() From 048a10a2674e8003bcebd991080f7579ced7e92c Mon Sep 17 00:00:00 2001 From: Sebastian Date: Sun, 4 May 2025 21:52:21 +0000 Subject: [PATCH 05/24] feat-72: include tests for services --- backend/src/flashcards/services.py | 10 +++-- backend/tests-start.sh | 0 .../tests/flashcards/card/test_services.py | 40 ++++++++++++++++++- 3 files changed, 46 insertions(+), 4 deletions(-) mode change 100644 => 100755 backend/tests-start.sh diff --git a/backend/src/flashcards/services.py b/backend/src/flashcards/services.py index ffb9195..45d4413 100644 --- a/backend/src/flashcards/services.py +++ b/backend/src/flashcards/services.py @@ -428,6 +428,13 @@ def _has_exceeded_usage_quota(session: Session, user_id: uuid.UUID) -> bool: quota = session.exec(statement).first() if not quota: return False + if ( + datetime.now(timezone.utc) - quota.last_reset_time + >= timedelta(days=settings.AI_QUOTA_TIME_RANGE_DAYS) + ): + quota.usage_count = 0 + quota.last_reset_time = datetime.now(timezone.utc) + session.commit() return quota.usage_count >= settings.AI_MAX_USAGE_QUOTA @@ -437,8 +444,5 @@ def _increase_usage_quota(session: Session, user_id: uuid.UUID): if not quota: quota = AIUsageQuota(user_id=user_id) session.add(quota) - if datetime.now(timezone.utc) - quota.last_reset_time >= timedelta(days=settings.AI_QUOTA_TIME_RANGE_DAYS): - quota.usage_count = 0 - quota.last_reset_time = datetime.now(timezone.utc) quota.usage_count += 1 session.commit() diff --git a/backend/tests-start.sh b/backend/tests-start.sh old mode 100644 new mode 100755 diff --git a/backend/tests/flashcards/card/test_services.py b/backend/tests/flashcards/card/test_services.py index 51f7725..d7613b6 100644 --- a/backend/tests/flashcards/card/test_services.py +++ b/backend/tests/flashcards/card/test_services.py @@ -1,8 +1,10 @@ +from datetime import datetime, timedelta, timezone +from unittest.mock import MagicMock, patch import uuid from sqlmodel import Session -from src.flashcards.models import Card, Collection +from src.flashcards.models import AIUsageQuota, Card, Collection from src.flashcards.schemas import CardCreate, CardUpdate from src.flashcards.services import ( create_card, @@ -12,7 +14,9 @@ get_card_with_collection, get_cards, update_card, + is_within_ai_usage_quota ) +from src.users.models import User def test_create_card(db: Session, test_collection: Collection): @@ -152,3 +156,37 @@ def test_delete_card(db: Session, test_collection: Collection, test_card: Card): session=db, card_id=test_card.id, user_id=test_collection.user_id ) assert card is None + + +def test_ai_usage_quota_not_reached_first_time(db: Session, test_user: User): + within_quota = is_within_ai_usage_quota(db, test_user["id"]) + assert within_quota is True + + +def test_ai_usage_quota_not_reached(db: Session, test_user: User): + within_quota = is_within_ai_usage_quota(db, test_user["id"]) + assert within_quota is True + + +def test_ai_usage_quota_reached(test_user: User): + test_session = MagicMock(spec=Session) + mock_quota = MagicMock(spec=AIUsageQuota) + mock_quota.usage_count = 3000 # exagerated for testing + mock_quota.last_reset_time = datetime.now(timezone.utc) + mock_quota.user_id = test_user["id"] + test_session.exec.return_value.first.return_value = mock_quota + + within_quota = is_within_ai_usage_quota(test_session, test_user["id"]) + assert within_quota is False + + +def test_ai_usage_quota_reset(test_user: User): + test_session = MagicMock(spec=Session) + mock_quota = MagicMock(spec=AIUsageQuota) + mock_quota.usage_count = 3000 # exagerated for testing + # exagerated for testing + mock_quota.last_reset_time = datetime.now(timezone.utc) - timedelta(days=700) + test_session.exec.return_value.first.return_value = mock_quota + + within_quota = is_within_ai_usage_quota(test_session, test_user["id"]) + assert within_quota is True From 904542a03991149108c7e4c9c20758b37dcc1a81 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Sun, 4 May 2025 22:32:13 +0000 Subject: [PATCH 06/24] feat-72: include tests for new service call --- backend/src/flashcards/api.py | 11 +++++- backend/src/flashcards/schemas.py | 5 +++ backend/src/flashcards/services.py | 21 +++++++++++ .../tests/flashcards/card/test_services.py | 37 ++++++++++++++++++- 4 files changed, 72 insertions(+), 2 deletions(-) diff --git a/backend/src/flashcards/api.py b/backend/src/flashcards/api.py index ea0ff2c..0dd961c 100644 --- a/backend/src/flashcards/api.py +++ b/backend/src/flashcards/api.py @@ -11,6 +11,7 @@ from . import services from .exceptions import EmptyCollectionError from .schemas import ( + AIUsageQuota, Card, CardCreate, CardList, @@ -24,12 +25,20 @@ PracticeCardResultPatch, PracticeSession, PracticeSessionCreate, - PracticeSessionList, + PracticeSessionList ) router = APIRouter() +@router.get("/aiquota", response_model=AIUsageQuota) +def get_ai_usage_quota( + session: SessionDep, + current_user: CurrentUser, +) -> Any: + return services.get_usage_quota(session, current_user.id) + + @router.get("/collections/", response_model=CollectionList) def read_collections( session: SessionDep, current_user: CurrentUser, skip: int = 0, limit: int = 100 diff --git a/backend/src/flashcards/schemas.py b/backend/src/flashcards/schemas.py index bd4568c..f9b513a 100644 --- a/backend/src/flashcards/schemas.py +++ b/backend/src/flashcards/schemas.py @@ -130,3 +130,8 @@ class PracticeCardListResponse(SQLModel): class PracticeCardResultPatch(SQLModel): is_correct: bool + + +class AIUsageQuota(SQLModel): + percentage_used: int + reset_date: datetime diff --git a/backend/src/flashcards/services.py b/backend/src/flashcards/services.py index 45d4413..67eb38f 100644 --- a/backend/src/flashcards/services.py +++ b/backend/src/flashcards/services.py @@ -21,6 +21,7 @@ PracticeSession ) from .schemas import ( + AIUsageQuota as AIUsageQuotaSchema, AIFlashcardCollection, CardBase, CardCreate, @@ -417,6 +418,26 @@ async def generate_ai_flashcard(prompt: str, provider) -> CardBase: raise AIGenerationError(f"Error processing AI response: {str(e)}") +def get_usage_quota(session: Session, user_id: uuid.UUID) -> AIUsageQuotaSchema: + statement = select(AIUsageQuota).filter_by(user_id=user_id) + quota = session.exec(statement).first() + if not quota: + return AIUsageQuotaSchema( + percentage_used=0, + reset_date=( + datetime.now(timezone.utc) + + timedelta(days=settings.AI_QUOTA_TIME_RANGE_DAYS) + ) + ) + return AIUsageQuotaSchema( + percentage_used=int(quota.usage_count / settings.AI_MAX_USAGE_QUOTA * 100), + reset_date=( + quota.last_reset_time + + timedelta(days=settings.AI_QUOTA_TIME_RANGE_DAYS) + ) + ) + + def is_within_ai_usage_quota(session: Session, user_id: uuid.UUID) -> bool: if _has_exceeded_usage_quota(session, user_id): return False diff --git a/backend/tests/flashcards/card/test_services.py b/backend/tests/flashcards/card/test_services.py index d7613b6..bc48956 100644 --- a/backend/tests/flashcards/card/test_services.py +++ b/backend/tests/flashcards/card/test_services.py @@ -13,10 +13,12 @@ get_card_by_id, get_card_with_collection, get_cards, + get_usage_quota, + is_within_ai_usage_quota, update_card, - is_within_ai_usage_quota ) from src.users.models import User +from src.core.config import settings def test_create_card(db: Session, test_collection: Collection): @@ -190,3 +192,36 @@ def test_ai_usage_quota_reset(test_user: User): within_quota = is_within_ai_usage_quota(test_session, test_user["id"]) assert within_quota is True + + +def test_get_usage_quota_empty(test_user: User): + test_session = MagicMock(spec=Session) + test_session.exec.return_value.first.return_value = None + + ai_usage_quota = get_usage_quota(test_session, test_user["id"]) + assert ai_usage_quota.percentage_used == 0 + assert ( + ai_usage_quota.reset_date + >= datetime.now(timezone.utc) + + timedelta(days=settings.AI_QUOTA_TIME_RANGE_DAYS) + - timedelta(milliseconds=10) # little bit of tolerance + ) + + +def test_get_usage_quota(test_user: User): + test_session = MagicMock(spec=Session) + mock_quota = MagicMock(spec=AIUsageQuota) + mock_quota.usage_count = settings.AI_MAX_USAGE_QUOTA / 2 + mock_quota.last_reset_time = ( + datetime.now(timezone.utc) + - timedelta(days=settings.AI_QUOTA_TIME_RANGE_DAYS / 2) + ) + test_session.exec.return_value.first.return_value = mock_quota + ai_usage_quota = get_usage_quota(test_session, test_user["id"]) + assert ai_usage_quota.percentage_used == 50 + assert ( + ai_usage_quota.reset_date + >= datetime.now(timezone.utc) + + timedelta(days=settings.AI_QUOTA_TIME_RANGE_DAYS / 2) + - timedelta(milliseconds=10) # little bit of tolerance + ) From 76a7eb11d60596a7ea748754d7c5873db4c7049f Mon Sep 17 00:00:00 2001 From: Sebastian Date: Sun, 4 May 2025 22:34:41 +0000 Subject: [PATCH 07/24] feat-72: run ruff formatter and linter --- backend/src/flashcards/api.py | 10 +++++++--- backend/src/flashcards/models.py | 5 ++++- backend/src/flashcards/services.py | 27 +++++++++++---------------- 3 files changed, 22 insertions(+), 20 deletions(-) diff --git a/backend/src/flashcards/api.py b/backend/src/flashcards/api.py index 0dd961c..6c8c6b6 100644 --- a/backend/src/flashcards/api.py +++ b/backend/src/flashcards/api.py @@ -25,7 +25,7 @@ PracticeCardResultPatch, PracticeSession, PracticeSessionCreate, - PracticeSessionList + PracticeSessionList, ) router = APIRouter() @@ -62,7 +62,9 @@ async def create_collection( if collection_in.prompt: try: if not services.is_within_ai_usage_quota(session, current_user.id): - raise HTTPException(status_code=429, detail="Quota for AI usage is reached.") + raise HTTPException( + status_code=429, detail="Quota for AI usage is reached." + ) flashcard_collection = await services.generate_ai_collection( provider, collection_in.prompt ) @@ -157,7 +159,9 @@ async def create_card( raise HTTPException(status_code=404, detail="Collection not found") if card_in.prompt: if not services.is_within_ai_usage_quota(session, current_user.id): - raise HTTPException(status_code=429, detail="Quota for AI usage is reached.") + raise HTTPException( + status_code=429, detail="Quota for AI usage is reached." + ) card_base = await services.generate_ai_flashcard(card_in.prompt, provider) card_in.front = card_base.front card_in.back = card_base.back diff --git a/backend/src/flashcards/models.py b/backend/src/flashcards/models.py index 3d1225f..9a709fb 100644 --- a/backend/src/flashcards/models.py +++ b/backend/src/flashcards/models.py @@ -64,8 +64,11 @@ class PracticeCard(SQLModel, table=True): session: PracticeSession = Relationship(back_populates="practice_cards") card: Card = Relationship(back_populates="practice_cards") + class AIUsageQuota(SQLModel, table=True): id: uuid.UUID | None = Field(default_factory=uuid.uuid4, primary_key=True) user_id: uuid.UUID = Field(foreign_key="user.id", index=True, ondelete="CASCADE") usage_count: int = Field(default=0) - last_reset_time: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) \ No newline at end of file + last_reset_time: datetime = Field( + default_factory=lambda: datetime.now(timezone.utc) + ) diff --git a/backend/src/flashcards/services.py b/backend/src/flashcards/services.py index 67eb38f..00a10e5 100644 --- a/backend/src/flashcards/services.py +++ b/backend/src/flashcards/services.py @@ -13,21 +13,17 @@ from .ai_config import get_card_config, get_flashcard_config from .exceptions import EmptyCollectionError -from .models import ( - AIUsageQuota, - Card, - Collection, - PracticeCard, - PracticeSession -) +from .models import AIUsageQuota, Card, Collection, PracticeCard, PracticeSession from .schemas import ( - AIUsageQuota as AIUsageQuotaSchema, AIFlashcardCollection, CardBase, CardCreate, CardUpdate, CollectionUpdate, ) +from .schemas import ( + AIUsageQuota as AIUsageQuotaSchema, +) def get_collections( @@ -425,16 +421,15 @@ def get_usage_quota(session: Session, user_id: uuid.UUID) -> AIUsageQuotaSchema: return AIUsageQuotaSchema( percentage_used=0, reset_date=( - datetime.now(timezone.utc) + datetime.now(timezone.utc) + timedelta(days=settings.AI_QUOTA_TIME_RANGE_DAYS) - ) + ), ) return AIUsageQuotaSchema( percentage_used=int(quota.usage_count / settings.AI_MAX_USAGE_QUOTA * 100), reset_date=( - quota.last_reset_time - + timedelta(days=settings.AI_QUOTA_TIME_RANGE_DAYS) - ) + quota.last_reset_time + timedelta(days=settings.AI_QUOTA_TIME_RANGE_DAYS) + ), ) @@ -444,14 +439,14 @@ def is_within_ai_usage_quota(session: Session, user_id: uuid.UUID) -> bool: _increase_usage_quota(session, user_id) return True + def _has_exceeded_usage_quota(session: Session, user_id: uuid.UUID) -> bool: statement = select(AIUsageQuota).filter_by(user_id=user_id) quota = session.exec(statement).first() if not quota: return False - if ( - datetime.now(timezone.utc) - quota.last_reset_time - >= timedelta(days=settings.AI_QUOTA_TIME_RANGE_DAYS) + if datetime.now(timezone.utc) - quota.last_reset_time >= timedelta( + days=settings.AI_QUOTA_TIME_RANGE_DAYS ): quota.usage_count = 0 quota.last_reset_time = datetime.now(timezone.utc) From d9518c7a7f591f977b940968a0de5841a7c61151 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Sun, 4 May 2025 23:15:04 +0000 Subject: [PATCH 08/24] feat-72: generate openapi for new route --- frontend/openapi.json | 25 +++++++++++++++++++++++++ frontend/src/client/sdk.gen.ts | 13 +++++++++++++ frontend/src/client/types.gen.ts | 7 +++++++ 3 files changed, 45 insertions(+) diff --git a/frontend/openapi.json b/frontend/openapi.json index b5ed7e5..f1618df 100644 --- a/frontend/openapi.json +++ b/frontend/openapi.json @@ -80,6 +80,22 @@ } } }, + "/api/v1/aiquota": { + "get": { + "tags": ["flashcards"], + "summary": "Get Ai Usage Quota", + "operationId": "flashcards-get_ai_usage_quota", + "responses": { + "200": { + "description": "Successful Response", + "content": { + "application/json": { "schema": { "$ref": "#/components/schemas/AIUsageQuota" } } + } + } + }, + "security": [{ "OAuth2PasswordBearer": [] }] + } + }, "/api/v1/collections/": { "get": { "tags": ["flashcards"], @@ -692,6 +708,15 @@ }, "components": { "schemas": { + "AIUsageQuota": { + "properties": { + "percentage_used": { "type": "integer", "title": "Percentage Used" }, + "reset_date": { "type": "string", "format": "date-time", "title": "Reset Date" } + }, + "type": "object", + "required": ["percentage_used", "reset_date"], + "title": "AIUsageQuota" + }, "Body_login-login_access_token": { "properties": { "grant_type": { diff --git a/frontend/src/client/sdk.gen.ts b/frontend/src/client/sdk.gen.ts index 59e577b..20fc818 100644 --- a/frontend/src/client/sdk.gen.ts +++ b/frontend/src/client/sdk.gen.ts @@ -12,6 +12,7 @@ import type { FlashcardsDeleteCardResponse, FlashcardsDeleteCollectionData, FlashcardsDeleteCollectionResponse, + FlashcardsGetAiUsageQuotaResponse, FlashcardsGetPracticeSessionStatusData, FlashcardsGetPracticeSessionStatusResponse, FlashcardsListPracticeCardsData, @@ -44,6 +45,18 @@ import type { } from './types.gen' export class FlashcardsService { + /** + * Get Ai Usage Quota + * @returns AIUsageQuota Successful Response + * @throws ApiError + */ + public static getAiUsageQuota(): CancelablePromise { + return __request(OpenAPI, { + method: 'GET', + url: '/api/v1/aiquota', + }) + } + /** * Read Collections * @param data The data for the request. diff --git a/frontend/src/client/types.gen.ts b/frontend/src/client/types.gen.ts index 6f98fcb..4131154 100644 --- a/frontend/src/client/types.gen.ts +++ b/frontend/src/client/types.gen.ts @@ -1,5 +1,10 @@ // This file is auto-generated by @hey-api/openapi-ts +export type AIUsageQuota = { + percentage_used: number + reset_date: string +} + export type Body_login_login_access_token = { grant_type?: string | null username: string @@ -155,6 +160,8 @@ export type ValidationError = { type: string } +export type FlashcardsGetAiUsageQuotaResponse = AIUsageQuota + export type FlashcardsReadCollectionsData = { limit?: number skip?: number From 16853996b0becf119adadce56e6430376aac9c30 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Thu, 8 May 2025 18:00:51 +0000 Subject: [PATCH 09/24] feat-72: add state for loading ai usage quota --- frontend/public/locales/en/translation.json | 3 ++ frontend/public/locales/es/translation.json | 3 ++ frontend/public/locales/nl/translation.json | 3 ++ .../components/commonUI/AiPromptDialog.tsx | 13 +++++++ frontend/src/hooks/useAiDialog.ts | 36 +++++++++++++++++++ frontend/src/hooks/useCard.ts | 2 +- 6 files changed, 59 insertions(+), 1 deletion(-) create mode 100644 frontend/src/hooks/useAiDialog.ts diff --git a/frontend/public/locales/en/translation.json b/frontend/public/locales/en/translation.json index 304e480..c1b68a4 100644 --- a/frontend/public/locales/en/translation.json +++ b/frontend/public/locales/en/translation.json @@ -84,6 +84,8 @@ "general": { "actions": { "addCard": "Add Card", + "aiQuotaResetDate": "AI Usage quota reset date", + "aiQuotaUsed": "AI Usage quota used", "generateAiCard": "Add Card with AI", "addCollection": "Add Collection", "backToCollection": "Back to Collection", @@ -120,6 +122,7 @@ "emailIsRequired": "Email is required", "errorCreatingAccount": "Error creating account", "errorLoadingCard": "Error loading card", + "errorLoadingAIUsageQuota": "Error loading AI usage quota", "errorSavingCard": "Error saving card", "invalidCredentials": "Invalid credentials", "invalidName": "Invalid name", diff --git a/frontend/public/locales/es/translation.json b/frontend/public/locales/es/translation.json index 62f764e..ec48978 100644 --- a/frontend/public/locales/es/translation.json +++ b/frontend/public/locales/es/translation.json @@ -84,6 +84,8 @@ "general": { "actions": { "addCard": "Agregar Tarjeta", + "aiQuotaResetDate": "Fecha de reinicio de cuota de IA", + "aiQuotaUsed": "Cuota de IA usada", "generateAiCard": "Agregar tarjeta con IA", "addCollection": "Agregar Colección", "backToCollection": "Volver a la Colección", @@ -120,6 +122,7 @@ "emailIsRequired": "Correo electrónico es requerido", "errorCreatingAccount": "Error al crear la cuenta", "errorLoadingCard": "Error al cargar la tarjeta", + "errorLoadingAIUsageQuota": "Error al cargar cuota de uso de IA", "errorSavingCard": "Error al guardar la tarjeta", "invalidCredentials": "Credenciales inválidas", "invalidName": "Nombre inválido", diff --git a/frontend/public/locales/nl/translation.json b/frontend/public/locales/nl/translation.json index 39a694c..39ffc42 100644 --- a/frontend/public/locales/nl/translation.json +++ b/frontend/public/locales/nl/translation.json @@ -84,6 +84,8 @@ "general": { "actions": { "addCard": "Kaart Toevoegen", + "aiQuotaResetDate": "Datum van herstart van AI-quota", + "aiQuotaUsed": "Gebruikte AI-quota", "generateAiCard": "Kaart toevoegen met AI", "addCollection": "Collectie Toevoegen", "backToCollection": "Terug naar Collectie", @@ -120,6 +122,7 @@ "emailIsRequired": "E-mail is vereist", "errorCreatingAccount": "Fout bij het aanmaken van account", "errorLoadingCard": "Fout bij het laden van kaart", + "errorLoadingAIUsageQuota": "Fout bij het laden van AI-gebruiksquota", "errorSavingCard": "Fout bij het opslaan van kaart", "invalidCredentials": "Ongeldige inloggegevens", "invalidName": "Ongeldige naam", diff --git a/frontend/src/components/commonUI/AiPromptDialog.tsx b/frontend/src/components/commonUI/AiPromptDialog.tsx index c92230a..e90e0ac 100644 --- a/frontend/src/components/commonUI/AiPromptDialog.tsx +++ b/frontend/src/components/commonUI/AiPromptDialog.tsx @@ -11,6 +11,7 @@ import { import { Text } from '@chakra-ui/react' import type { OpenChangeDetails } from 'node_modules/@chakra-ui/react/dist/types/components/dialog/namespace' import { useRef, useState } from 'react' +import { useAiDialog } from '@/hooks/useAiDialog' import { useTranslation } from 'react-i18next' import { BlueButton, RedButton } from '../commonUI/Button' import { DefaultInput } from '../commonUI/Input' @@ -37,6 +38,7 @@ const AiPromptDialog: React.FC = ({ const { t } = useTranslation() const [prompt, setPrompt] = useState('') const closeButtonRef = useRef(null) + const { usageQuota } = useAiDialog() const handleSubmit = () => { if (!prompt.trim() || isLoading) return @@ -85,6 +87,17 @@ const AiPromptDialog: React.FC = ({ {prompt.length}/{MAX_CHARS} + + {`${t('general.actions.aiQuotaResetDate')}: ${new Date(usageQuota.reset_date).toLocaleDateString()}`} + + + {`${t('general.actions.aiQuotaUsed')}: ${usageQuota.percentage_used}%`} + diff --git a/frontend/src/hooks/useAiDialog.ts b/frontend/src/hooks/useAiDialog.ts new file mode 100644 index 0000000..61a4d35 --- /dev/null +++ b/frontend/src/hooks/useAiDialog.ts @@ -0,0 +1,36 @@ +import { toaster } from '@/components/ui/toaster' +import { useEffect, useState } from 'react' +import { useTranslation } from 'react-i18next' +import { FlashcardsService } from '@/client' + +interface AIUsageQuotaData { + reset_date: string + percentage_used: number +} + + +export function useAiDialog() { + const { t } = useTranslation() + const [usageQuota, setUsageQuota] = useState({reset_date: '', percentage_used: 0}) + + useEffect(() => { + const fetchUsageQuota = async () => { + try { + const data = await FlashcardsService.getAiUsageQuota() + if (data) { + setUsageQuota(data) + } + } catch (error) { + toaster.create({ + title: t('general.errors.errorloadingCard'), + type: 'error', + }) + } + } + fetchUsageQuota() + }, []) + + return { + usageQuota + } +} diff --git a/frontend/src/hooks/useCard.ts b/frontend/src/hooks/useCard.ts index c4d4ede..50ae0ff 100644 --- a/frontend/src/hooks/useCard.ts +++ b/frontend/src/hooks/useCard.ts @@ -42,7 +42,7 @@ export function useCard(collectionId: string, cardId?: string) { } } catch (error) { toaster.create({ - title: t('general.errors.errorloadingCard'), + title: t('general.errors.errorLoadingCard'), type: 'error', }) } finally { From 96bbc3981441eba8791e3a59546c20276dc1e85c Mon Sep 17 00:00:00 2001 From: Sebastian Date: Fri, 9 May 2025 22:07:28 +0000 Subject: [PATCH 10/24] feat-72: only allow creation if quota is not 100% --- .../src/components/commonUI/AiPromptDialog.tsx | 2 +- frontend/src/hooks/useAiDialog.ts | 16 ++-------------- 2 files changed, 3 insertions(+), 15 deletions(-) diff --git a/frontend/src/components/commonUI/AiPromptDialog.tsx b/frontend/src/components/commonUI/AiPromptDialog.tsx index e90e0ac..b2b63a9 100644 --- a/frontend/src/components/commonUI/AiPromptDialog.tsx +++ b/frontend/src/components/commonUI/AiPromptDialog.tsx @@ -106,7 +106,7 @@ const AiPromptDialog: React.FC = ({ - + {isLoading ? `${t('general.actions.creating')}...` : t('general.actions.create')} diff --git a/frontend/src/hooks/useAiDialog.ts b/frontend/src/hooks/useAiDialog.ts index 61a4d35..ac9027a 100644 --- a/frontend/src/hooks/useAiDialog.ts +++ b/frontend/src/hooks/useAiDialog.ts @@ -1,6 +1,4 @@ -import { toaster } from '@/components/ui/toaster' import { useEffect, useState } from 'react' -import { useTranslation } from 'react-i18next' import { FlashcardsService } from '@/client' interface AIUsageQuotaData { @@ -10,22 +8,12 @@ interface AIUsageQuotaData { export function useAiDialog() { - const { t } = useTranslation() const [usageQuota, setUsageQuota] = useState({reset_date: '', percentage_used: 0}) useEffect(() => { const fetchUsageQuota = async () => { - try { - const data = await FlashcardsService.getAiUsageQuota() - if (data) { - setUsageQuota(data) - } - } catch (error) { - toaster.create({ - title: t('general.errors.errorloadingCard'), - type: 'error', - }) - } + const data = await FlashcardsService.getAiUsageQuota() + setUsageQuota(data) } fetchUsageQuota() }, []) From ae98d08a085c5d5d903bc80d1f245a935dd3d845 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Fri, 9 May 2025 22:17:44 +0000 Subject: [PATCH 11/24] run formatting --- .../src/components/commonUI/AiPromptDialog.tsx | 15 ++++++++++++--- frontend/src/hooks/useAiDialog.ts | 10 ++++++---- 2 files changed, 18 insertions(+), 7 deletions(-) diff --git a/frontend/src/components/commonUI/AiPromptDialog.tsx b/frontend/src/components/commonUI/AiPromptDialog.tsx index b2b63a9..46adadc 100644 --- a/frontend/src/components/commonUI/AiPromptDialog.tsx +++ b/frontend/src/components/commonUI/AiPromptDialog.tsx @@ -8,10 +8,10 @@ import { DialogRoot, DialogTitle, } from '@/components/ui/dialog' +import { useAiDialog } from '@/hooks/useAiDialog' import { Text } from '@chakra-ui/react' import type { OpenChangeDetails } from 'node_modules/@chakra-ui/react/dist/types/components/dialog/namespace' import { useRef, useState } from 'react' -import { useAiDialog } from '@/hooks/useAiDialog' import { useTranslation } from 'react-i18next' import { BlueButton, RedButton } from '../commonUI/Button' import { DefaultInput } from '../commonUI/Input' @@ -94,7 +94,13 @@ const AiPromptDialog: React.FC = ({ fontSize="xs" textAlign="right" mt={1} - color={usageQuota.percentage_used <= 50 ? 'green.500' : usageQuota.percentage_used <= 80 ? 'yellow.500' : 'red.500'} + color={ + usageQuota.percentage_used <= 50 + ? 'green.500' + : usageQuota.percentage_used <= 80 + ? 'yellow.500' + : 'red.500' + } > {`${t('general.actions.aiQuotaUsed')}: ${usageQuota.percentage_used}%`} @@ -106,7 +112,10 @@ const AiPromptDialog: React.FC = ({ - + {isLoading ? `${t('general.actions.creating')}...` : t('general.actions.create')} diff --git a/frontend/src/hooks/useAiDialog.ts b/frontend/src/hooks/useAiDialog.ts index ac9027a..38a6119 100644 --- a/frontend/src/hooks/useAiDialog.ts +++ b/frontend/src/hooks/useAiDialog.ts @@ -1,14 +1,16 @@ -import { useEffect, useState } from 'react' import { FlashcardsService } from '@/client' +import { useEffect, useState } from 'react' interface AIUsageQuotaData { reset_date: string percentage_used: number } - export function useAiDialog() { - const [usageQuota, setUsageQuota] = useState({reset_date: '', percentage_used: 0}) + const [usageQuota, setUsageQuota] = useState({ + reset_date: '', + percentage_used: 0, + }) useEffect(() => { const fetchUsageQuota = async () => { @@ -19,6 +21,6 @@ export function useAiDialog() { }, []) return { - usageQuota + usageQuota, } } From 94c4aaf11d3bdff3b17e5db47c7e362de0b84a5b Mon Sep 17 00:00:00 2001 From: Sebastian Date: Mon, 12 May 2025 23:04:43 +0000 Subject: [PATCH 12/24] feat: apply suggested changes --- backend/src/core/config.py | 7 ++++ backend/src/flashcards/api.py | 14 ++------ backend/src/flashcards/models.py | 9 ----- backend/src/flashcards/schemas.py | 5 --- backend/src/flashcards/services.py | 58 ++---------------------------- backend/src/users/api.py | 7 +++- backend/src/users/models.py | 19 +++++++++- backend/src/users/schemas.py | 6 ++++ backend/src/users/services.py | 55 ++++++++++++++++++++++++++-- 9 files changed, 95 insertions(+), 85 deletions(-) diff --git a/backend/src/core/config.py b/backend/src/core/config.py index e5a2a4a..0102553 100644 --- a/backend/src/core/config.py +++ b/backend/src/core/config.py @@ -94,6 +94,13 @@ def _enforce_non_default_secrets(self) -> Self: "FIRST_SUPERUSER_PASSWORD", self.FIRST_SUPERUSER_PASSWORD ) + if self.AI_MAX_USAGE_QUOTA is None or self.AI_MAX_USAGE_QUOTA <= 0: + raise ValueError("AI_MAX_USAGE_QUOTA must be set to a positive integer.") + if self.AI_QUOTA_TIME_RANGE_DAYS is None or self.AI_QUOTA_TIME_RANGE_DAYS <= 0: + raise ValueError( + "AI_QUOTA_TIME_RANGE_DAYS must be set to a positive integer." + ) + return self diff --git a/backend/src/flashcards/api.py b/backend/src/flashcards/api.py index 6c8c6b6..bebbc62 100644 --- a/backend/src/flashcards/api.py +++ b/backend/src/flashcards/api.py @@ -7,11 +7,11 @@ from src.ai_models.gemini import GeminiProviderDep from src.ai_models.gemini.exceptions import AIGenerationError from src.auth.services import CurrentUser, SessionDep +from src.users.services import check_and_increment_ai_usage_quota from . import services from .exceptions import EmptyCollectionError from .schemas import ( - AIUsageQuota, Card, CardCreate, CardList, @@ -31,14 +31,6 @@ router = APIRouter() -@router.get("/aiquota", response_model=AIUsageQuota) -def get_ai_usage_quota( - session: SessionDep, - current_user: CurrentUser, -) -> Any: - return services.get_usage_quota(session, current_user.id) - - @router.get("/collections/", response_model=CollectionList) def read_collections( session: SessionDep, current_user: CurrentUser, skip: int = 0, limit: int = 100 @@ -61,7 +53,7 @@ async def create_collection( if collection_in.prompt: try: - if not services.is_within_ai_usage_quota(session, current_user.id): + if not check_and_increment_ai_usage_quota(session, current_user): raise HTTPException( status_code=429, detail="Quota for AI usage is reached." ) @@ -158,7 +150,7 @@ async def create_card( if not access_checked: raise HTTPException(status_code=404, detail="Collection not found") if card_in.prompt: - if not services.is_within_ai_usage_quota(session, current_user.id): + if not check_and_increment_ai_usage_quota(session, current_user): raise HTTPException( status_code=429, detail="Quota for AI usage is reached." ) diff --git a/backend/src/flashcards/models.py b/backend/src/flashcards/models.py index 9a709fb..3f2c95b 100644 --- a/backend/src/flashcards/models.py +++ b/backend/src/flashcards/models.py @@ -63,12 +63,3 @@ class PracticeCard(SQLModel, table=True): updated_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) session: PracticeSession = Relationship(back_populates="practice_cards") card: Card = Relationship(back_populates="practice_cards") - - -class AIUsageQuota(SQLModel, table=True): - id: uuid.UUID | None = Field(default_factory=uuid.uuid4, primary_key=True) - user_id: uuid.UUID = Field(foreign_key="user.id", index=True, ondelete="CASCADE") - usage_count: int = Field(default=0) - last_reset_time: datetime = Field( - default_factory=lambda: datetime.now(timezone.utc) - ) diff --git a/backend/src/flashcards/schemas.py b/backend/src/flashcards/schemas.py index f9b513a..bd4568c 100644 --- a/backend/src/flashcards/schemas.py +++ b/backend/src/flashcards/schemas.py @@ -130,8 +130,3 @@ class PracticeCardListResponse(SQLModel): class PracticeCardResultPatch(SQLModel): is_correct: bool - - -class AIUsageQuota(SQLModel): - percentage_used: int - reset_date: datetime diff --git a/backend/src/flashcards/services.py b/backend/src/flashcards/services.py index 00a10e5..da32518 100644 --- a/backend/src/flashcards/services.py +++ b/backend/src/flashcards/services.py @@ -1,7 +1,7 @@ import json import random import uuid -from datetime import datetime, timedelta, timezone +from datetime import datetime, timezone from typing import Literal from google import genai @@ -9,11 +9,10 @@ from sqlmodel import Session, func, select from src.ai_models.gemini.exceptions import AIGenerationError -from src.core.config import settings from .ai_config import get_card_config, get_flashcard_config from .exceptions import EmptyCollectionError -from .models import AIUsageQuota, Card, Collection, PracticeCard, PracticeSession +from .models import Card, Collection, PracticeCard, PracticeSession from .schemas import ( AIFlashcardCollection, CardBase, @@ -21,9 +20,6 @@ CardUpdate, CollectionUpdate, ) -from .schemas import ( - AIUsageQuota as AIUsageQuotaSchema, -) def get_collections( @@ -412,53 +408,3 @@ async def generate_ai_flashcard(prompt: str, provider) -> CardBase: raise AIGenerationError(f"Invalid AI response format: {str(e)}") except Exception as e: raise AIGenerationError(f"Error processing AI response: {str(e)}") - - -def get_usage_quota(session: Session, user_id: uuid.UUID) -> AIUsageQuotaSchema: - statement = select(AIUsageQuota).filter_by(user_id=user_id) - quota = session.exec(statement).first() - if not quota: - return AIUsageQuotaSchema( - percentage_used=0, - reset_date=( - datetime.now(timezone.utc) - + timedelta(days=settings.AI_QUOTA_TIME_RANGE_DAYS) - ), - ) - return AIUsageQuotaSchema( - percentage_used=int(quota.usage_count / settings.AI_MAX_USAGE_QUOTA * 100), - reset_date=( - quota.last_reset_time + timedelta(days=settings.AI_QUOTA_TIME_RANGE_DAYS) - ), - ) - - -def is_within_ai_usage_quota(session: Session, user_id: uuid.UUID) -> bool: - if _has_exceeded_usage_quota(session, user_id): - return False - _increase_usage_quota(session, user_id) - return True - - -def _has_exceeded_usage_quota(session: Session, user_id: uuid.UUID) -> bool: - statement = select(AIUsageQuota).filter_by(user_id=user_id) - quota = session.exec(statement).first() - if not quota: - return False - if datetime.now(timezone.utc) - quota.last_reset_time >= timedelta( - days=settings.AI_QUOTA_TIME_RANGE_DAYS - ): - quota.usage_count = 0 - quota.last_reset_time = datetime.now(timezone.utc) - session.commit() - return quota.usage_count >= settings.AI_MAX_USAGE_QUOTA - - -def _increase_usage_quota(session: Session, user_id: uuid.UUID): - statement = select(AIUsageQuota).filter_by(user_id=user_id) - quota = session.exec(statement).first() - if not quota: - quota = AIUsageQuota(user_id=user_id) - session.add(quota) - quota.usage_count += 1 - session.commit() diff --git a/backend/src/users/api.py b/backend/src/users/api.py index 4bd7007..f4c6704 100644 --- a/backend/src/users/api.py +++ b/backend/src/users/api.py @@ -4,7 +4,7 @@ from src.auth.services import CurrentUser, SessionDep from src.core.config import settings -from src.users.schemas import UserCreate, UserPublic, UserRegister +from src.users.schemas import AIUsageQuota, UserCreate, UserPublic, UserRegister from . import services @@ -38,3 +38,8 @@ def register_user(session: SessionDep, user_in: UserRegister) -> Any: user_create = UserCreate.model_validate(user_in) user = services.create_user(session=session, user_create=user_create) return user + + +@router.get("/users/me/ai-usage-quota", response_model=AIUsageQuota) +def get_my_ai_usage_quota(current_user: CurrentUser): + return services.get_ai_usage_quota_for_user(current_user) diff --git a/backend/src/users/models.py b/backend/src/users/models.py index 528005a..655afa9 100644 --- a/backend/src/users/models.py +++ b/backend/src/users/models.py @@ -1,7 +1,8 @@ import uuid +from datetime import datetime, timezone from typing import TYPE_CHECKING -from sqlmodel import Field, Relationship +from sqlmodel import Field, Relationship, SQLModel from src.users.schemas import UserBase @@ -22,3 +23,19 @@ class User(UserBase, table=True): cascade_delete=True, sa_relationship_kwargs={"lazy": "selectin"}, ) + ai_usage_quota: "AIUsageQuota" = Relationship( + back_populates="user", + sa_relationship_kwargs={"uselist": False, "lazy": "selectin"}, + ) + + +class AIUsageQuota(SQLModel, table=True): + id: uuid.UUID | None = Field(default_factory=uuid.uuid4, primary_key=True) + user_id: uuid.UUID = Field( + foreign_key="user.id", index=True, unique=True, ondelete="CASCADE" + ) + usage_count: int = Field(default=0) + last_reset_time: datetime = Field( + default_factory=lambda: datetime.now(timezone.utc) + ) + user: "User" = Relationship(back_populates="ai_usage_quota") diff --git a/backend/src/users/schemas.py b/backend/src/users/schemas.py index 5da1ce6..94cfe88 100644 --- a/backend/src/users/schemas.py +++ b/backend/src/users/schemas.py @@ -1,4 +1,5 @@ import uuid +from datetime import datetime from pydantic import EmailStr from sqlmodel import Field, SQLModel @@ -26,3 +27,8 @@ class UserRegister(SQLModel): class UserPublic(UserBase): id: uuid.UUID + + +class AIUsageQuota(SQLModel): + percentage_used: int + reset_date: datetime diff --git a/backend/src/users/services.py b/backend/src/users/services.py index 474948e..9277398 100644 --- a/backend/src/users/services.py +++ b/backend/src/users/services.py @@ -1,11 +1,13 @@ import uuid +from datetime import datetime, timedelta, timezone from typing import Any -from sqlmodel import Session, select +from sqlmodel import Session, select, update from src.auth.services import get_password_hash +from src.core.config import settings from src.users.models import User -from src.users.schemas import UserCreate, UserUpdate +from src.users.schemas import AIUsageQuota, UserCreate, UserUpdate def create_user(*, session: Session, user_create: UserCreate) -> User: @@ -42,3 +44,52 @@ def get_user_by_email(*, session: Session, email: str) -> User | None: statement = select(User).where(User.email == email) session_user = session.exec(statement).first() return session_user + + +def get_ai_usage_quota_for_user(user: User) -> AIUsageQuota: + quota = user.ai_usage_quota + if not quota: + return AIUsageQuota( + percentage_used=0, + reset_date=( + datetime.now(timezone.utc) + + timedelta(days=settings.AI_QUOTA_TIME_RANGE_DAYS) + ), + ) + return AIUsageQuota( + percentage_used=int(quota.usage_count / settings.AI_MAX_USAGE_QUOTA * 100), + reset_date=( + quota.last_reset_time + timedelta(days=settings.AI_QUOTA_TIME_RANGE_DAYS) + ), + ) + + +def check_and_increment_ai_usage_quota(session: Session, user: User) -> bool: + quota = user.ai_usage_quota + now = datetime.now(timezone.utc) + + # 1. If no quota record, create and allow + if not quota: + quota = AIUsageQuota(user_id=user.id, usage_count=1, last_reset_time=now) + session.add(quota) + session.commit() + return True + + # 2. Reset if window expired + if now - quota.last_reset_time >= timedelta(days=settings.AI_QUOTA_TIME_RANGE_DAYS): + quota.usage_count = 0 + quota.last_reset_time = now + session.add(quota) + session.commit() + + # 3. Atomic check and increment + result = session.exec( + update(AIUsageQuota) + .where( + (AIUsageQuota.id == quota.id) + & (AIUsageQuota.usage_count < settings.AI_MAX_USAGE_QUOTA) + ) + .values(usage_count=AIUsageQuota.usage_count + 1) + ) + session.commit() + return result.rowcount > 0 From 11ac736c73a29ff42f398ef975987d1d3df90fc1 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Tue, 13 May 2025 22:57:09 +0000 Subject: [PATCH 13/24] tests: update tests for user services --- backend/src/users/schemas.py | 3 +- backend/src/users/services.py | 25 ++++--- .../tests/flashcards/card/test_services.py | 75 +------------------ backend/tests/users/conftest.py | 20 +++++ backend/tests/users/test_services.py | 44 ++++++++++- 5 files changed, 77 insertions(+), 90 deletions(-) create mode 100644 backend/tests/users/conftest.py diff --git a/backend/src/users/schemas.py b/backend/src/users/schemas.py index 94cfe88..ebc490f 100644 --- a/backend/src/users/schemas.py +++ b/backend/src/users/schemas.py @@ -30,5 +30,6 @@ class UserPublic(UserBase): class AIUsageQuota(SQLModel): - percentage_used: int + usage_count: int + max_usage_allowed: int reset_date: datetime diff --git a/backend/src/users/services.py b/backend/src/users/services.py index 9277398..60a93b2 100644 --- a/backend/src/users/services.py +++ b/backend/src/users/services.py @@ -8,6 +8,7 @@ from src.core.config import settings from src.users.models import User from src.users.schemas import AIUsageQuota, UserCreate, UserUpdate +from src.users.models import AIUsageQuota as AIUsageQuotaModel def create_user(*, session: Session, user_create: UserCreate) -> User: @@ -50,14 +51,16 @@ def get_ai_usage_quota_for_user(user: User) -> AIUsageQuota: quota = user.ai_usage_quota if not quota: return AIUsageQuota( - percentage_used=0, + usage_count=0, + max_usage_allowed=settings.AI_MAX_USAGE_QUOTA, reset_date=( datetime.now(timezone.utc) + timedelta(days=settings.AI_QUOTA_TIME_RANGE_DAYS) ), ) return AIUsageQuota( - percentage_used=int(quota.usage_count / settings.AI_MAX_USAGE_QUOTA * 100), + usage_count=quota.usage_count, + max_usage_allowed=settings.AI_MAX_USAGE_QUOTA, reset_date=( quota.last_reset_time + timedelta(days=settings.AI_QUOTA_TIME_RANGE_DAYS) ), @@ -67,29 +70,27 @@ def get_ai_usage_quota_for_user(user: User) -> AIUsageQuota: def check_and_increment_ai_usage_quota(session: Session, user: User) -> bool: quota = user.ai_usage_quota now = datetime.now(timezone.utc) - - # 1. If no quota record, create and allow if not quota: - quota = AIUsageQuota(user_id=user.id, usage_count=1, last_reset_time=now) + quota = AIUsageQuotaModel(user_id=user.id, usage_count=1, last_reset_time=now) session.add(quota) session.commit() return True - # 2. Reset if window expired + if quota.usage_count > settings.AI_MAX_USAGE_QUOTA: + return False + if now - quota.last_reset_time >= timedelta(days=settings.AI_QUOTA_TIME_RANGE_DAYS): quota.usage_count = 0 quota.last_reset_time = now session.add(quota) session.commit() - - # 3. Atomic check and increment result = session.exec( - update(AIUsageQuota) + update(AIUsageQuotaModel) .where( - (AIUsageQuota.id == quota.id) - & (AIUsageQuota.usage_count < settings.AI_MAX_USAGE_QUOTA) + (AIUsageQuotaModel.id == quota.id) + & (AIUsageQuotaModel.usage_count < settings.AI_MAX_USAGE_QUOTA) ) - .values(usage_count=AIUsageQuota.usage_count + 1) + .values(usage_count=AIUsageQuotaModel.usage_count + 1) ) session.commit() return result.rowcount > 0 diff --git a/backend/tests/flashcards/card/test_services.py b/backend/tests/flashcards/card/test_services.py index bc48956..51f7725 100644 --- a/backend/tests/flashcards/card/test_services.py +++ b/backend/tests/flashcards/card/test_services.py @@ -1,10 +1,8 @@ -from datetime import datetime, timedelta, timezone -from unittest.mock import MagicMock, patch import uuid from sqlmodel import Session -from src.flashcards.models import AIUsageQuota, Card, Collection +from src.flashcards.models import Card, Collection from src.flashcards.schemas import CardCreate, CardUpdate from src.flashcards.services import ( create_card, @@ -13,12 +11,8 @@ get_card_by_id, get_card_with_collection, get_cards, - get_usage_quota, - is_within_ai_usage_quota, update_card, ) -from src.users.models import User -from src.core.config import settings def test_create_card(db: Session, test_collection: Collection): @@ -158,70 +152,3 @@ def test_delete_card(db: Session, test_collection: Collection, test_card: Card): session=db, card_id=test_card.id, user_id=test_collection.user_id ) assert card is None - - -def test_ai_usage_quota_not_reached_first_time(db: Session, test_user: User): - within_quota = is_within_ai_usage_quota(db, test_user["id"]) - assert within_quota is True - - -def test_ai_usage_quota_not_reached(db: Session, test_user: User): - within_quota = is_within_ai_usage_quota(db, test_user["id"]) - assert within_quota is True - - -def test_ai_usage_quota_reached(test_user: User): - test_session = MagicMock(spec=Session) - mock_quota = MagicMock(spec=AIUsageQuota) - mock_quota.usage_count = 3000 # exagerated for testing - mock_quota.last_reset_time = datetime.now(timezone.utc) - mock_quota.user_id = test_user["id"] - test_session.exec.return_value.first.return_value = mock_quota - - within_quota = is_within_ai_usage_quota(test_session, test_user["id"]) - assert within_quota is False - - -def test_ai_usage_quota_reset(test_user: User): - test_session = MagicMock(spec=Session) - mock_quota = MagicMock(spec=AIUsageQuota) - mock_quota.usage_count = 3000 # exagerated for testing - # exagerated for testing - mock_quota.last_reset_time = datetime.now(timezone.utc) - timedelta(days=700) - test_session.exec.return_value.first.return_value = mock_quota - - within_quota = is_within_ai_usage_quota(test_session, test_user["id"]) - assert within_quota is True - - -def test_get_usage_quota_empty(test_user: User): - test_session = MagicMock(spec=Session) - test_session.exec.return_value.first.return_value = None - - ai_usage_quota = get_usage_quota(test_session, test_user["id"]) - assert ai_usage_quota.percentage_used == 0 - assert ( - ai_usage_quota.reset_date - >= datetime.now(timezone.utc) - + timedelta(days=settings.AI_QUOTA_TIME_RANGE_DAYS) - - timedelta(milliseconds=10) # little bit of tolerance - ) - - -def test_get_usage_quota(test_user: User): - test_session = MagicMock(spec=Session) - mock_quota = MagicMock(spec=AIUsageQuota) - mock_quota.usage_count = settings.AI_MAX_USAGE_QUOTA / 2 - mock_quota.last_reset_time = ( - datetime.now(timezone.utc) - - timedelta(days=settings.AI_QUOTA_TIME_RANGE_DAYS / 2) - ) - test_session.exec.return_value.first.return_value = mock_quota - ai_usage_quota = get_usage_quota(test_session, test_user["id"]) - assert ai_usage_quota.percentage_used == 50 - assert ( - ai_usage_quota.reset_date - >= datetime.now(timezone.utc) - + timedelta(days=settings.AI_QUOTA_TIME_RANGE_DAYS / 2) - - timedelta(milliseconds=10) # little bit of tolerance - ) diff --git a/backend/tests/users/conftest.py b/backend/tests/users/conftest.py new file mode 100644 index 0000000..e3a4b23 --- /dev/null +++ b/backend/tests/users/conftest.py @@ -0,0 +1,20 @@ +from typing import Any + +import pytest +from sqlmodel import Session + +from src.users.schemas import UserCreate +from src.users.services import create_user +from tests.utils.utils import random_email, random_lower_string + + +@pytest.fixture +def test_user(db: Session) -> dict[str, Any]: + email = random_email() + password = random_lower_string() + full_name = random_lower_string() + + user_in = UserCreate(email=email, password=password, full_name=full_name) + user = create_user(session=db, user_create=user_in) + + return user \ No newline at end of file diff --git a/backend/tests/users/test_services.py b/backend/tests/users/test_services.py index 970b4b4..25f5839 100644 --- a/backend/tests/users/test_services.py +++ b/backend/tests/users/test_services.py @@ -1,12 +1,14 @@ +from datetime import datetime, timedelta, timezone +from unittest.mock import MagicMock, patch from fastapi.encoders import jsonable_encoder from sqlmodel import Session from src.auth.services import verify_password -from src.users.models import User +from src.users.models import AIUsageQuota, User from src.users.schemas import UserCreate, UserUpdate -from src.users.services import create_user, get_user_by_email, update_user +from src.users.services import check_and_increment_ai_usage_quota, create_user, get_ai_usage_quota_for_user, get_user_by_email, update_user from tests.utils.utils import random_email, random_lower_string - +from src.core.config import settings def test_create_user(db: Session) -> None: email = random_email() @@ -83,3 +85,39 @@ def test_update_user(db: Session) -> None: assert user_2 assert user.email == user_2.email assert verify_password(new_password, user_2.hashed_password) + + +def test_get_ai_usage_quota_for_user_no_quota(test_user): + test_user.ai_usage_quota = None + with patch('src.users.services.datetime') as mock_datetime: + mock_datetime.now.return_value = datetime(2025, 1, 1, tzinfo=timezone.utc) + quota = get_ai_usage_quota_for_user(test_user) + assert quota.usage_count == 0 + assert quota.reset_date == datetime(2025, 1, 31, tzinfo=timezone.utc) + +def test_get_ai_usage_quota_for_user_with_quota(test_user): + test_user.ai_usage_quota = AIUsageQuota(usage_count=50, last_reset_time=datetime(2025, 1, 1, tzinfo=timezone.utc)) + with patch('src.users.services.settings.AI_MAX_USAGE_QUOTA', 100): + with patch('src.users.services.datetime') as mock_datetime: + mock_datetime.now.return_value = datetime(2025, 1, 1, tzinfo=timezone.utc) + quota = get_ai_usage_quota_for_user(test_user) + assert quota.usage_count == 50 + assert quota.reset_date == datetime(2025, 1, 31, tzinfo=timezone.utc) + +def test_check_and_increment_ai_usage_quota_first_time(db, test_user): + within_quota = check_and_increment_ai_usage_quota(db, test_user) + assert within_quota is True + +def test_check_and_increment_ai_usage_quota_max_count_reached(db, test_user): + test_user.ai_usage_quota = AIUsageQuota(usage_count=30000, last_reset_time=datetime.now(timezone.utc)) + within_quota = check_and_increment_ai_usage_quota(db, test_user) + assert within_quota is False + +def test_check_and_increment_ai_usage_quota_reset_count(db, test_user): + with patch('src.users.services.datetime') as mock_datetime: + mock_datetime.now.return_value = datetime(2025, 2, 1, tzinfo=timezone.utc) + test_user.ai_usage_quota = AIUsageQuota(usage_count=1, last_reset_time=datetime(2025, 1, 1, tzinfo=timezone.utc)) + within_quota = check_and_increment_ai_usage_quota(db, test_user) + assert within_quota is True + assert test_user.ai_usage_quota.usage_count == 1 + assert test_user.ai_usage_quota.last_reset_time == datetime(2025, 2, 1, tzinfo=timezone.utc) From 2097f5b9cfa57bbce1d1578fc7b5445484481ba3 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 May 2025 23:39:54 +0000 Subject: [PATCH 14/24] fix more tests --- backend/tests/flashcards/card/test_api.py | 28 ++++---- .../tests/flashcards/collection/test_api.py | 72 ++++++++++--------- 2 files changed, 56 insertions(+), 44 deletions(-) diff --git a/backend/tests/flashcards/card/test_api.py b/backend/tests/flashcards/card/test_api.py index d4222bc..daec27e 100644 --- a/backend/tests/flashcards/card/test_api.py +++ b/backend/tests/flashcards/card/test_api.py @@ -394,18 +394,22 @@ def test_create_card_with_prompt_ai( with patch( "src.flashcards.services.generate_ai_flashcard", new_callable=AsyncMock ) as mock_ai: - mock_ai.return_value = type("Card", (), ai_card)() - card_data = {"prompt": prompt, "front": "", "back": ""} - rsp = client.post( - f"{settings.API_V1_STR}/collections/{collection_id}/cards/", - json=card_data, - headers=normal_user_token_headers, - ) - assert rsp.status_code == 200 - content = rsp.json() - assert content["front"] == ai_card["front"] - assert content["back"] == ai_card["back"] - mock_ai.assert_called_once_with(prompt, ANY) + with patch( + "src.flashcards.api.check_and_increment_ai_usage_quota" + ) as mock_quota_check: + mock_quota_check.return_value = True + mock_ai.return_value = type("Card", (), ai_card)() + card_data = {"prompt": prompt, "front": "", "back": ""} + rsp = client.post( + f"{settings.API_V1_STR}/collections/{collection_id}/cards/", + json=card_data, + headers=normal_user_token_headers, + ) + assert rsp.status_code == 200 + content = rsp.json() + assert content["front"] == ai_card["front"] + assert content["back"] == ai_card["back"] + mock_ai.assert_called_once_with(prompt, ANY) def test_create_card_with_prompt_too_long( diff --git a/backend/tests/flashcards/collection/test_api.py b/backend/tests/flashcards/collection/test_api.py index d2bdff0..63219b7 100644 --- a/backend/tests/flashcards/collection/test_api.py +++ b/backend/tests/flashcards/collection/test_api.py @@ -83,25 +83,29 @@ def test_create_collection_with_prompt( with patch( "src.flashcards.services.generate_ai_collection", new_callable=AsyncMock ) as mock_ai_generate: - mock_ai_generate.return_value = mock_collection - - rsp = client.post( - f"{settings.API_V1_STR}/collections/", - json=collection_data.model_dump(), - headers=normal_user_token_headers, - ) - - assert rsp.status_code == 200 - content = rsp.json() - assert content["name"] == collection_data.name - assert "id" in content - assert isinstance(content["id"], str) - assert len(content["cards"]) == len(mock_collection.cards) - for i, card in enumerate(mock_collection.cards): - assert content["cards"][i]["front"] == card.front - assert content["cards"][i]["back"] == card.back - - mock_ai_generate.assert_called_once() + with patch( + "src.flashcards.api.check_and_increment_ai_usage_quota" + ) as mock_quota_check: + mock_ai_generate.return_value = mock_collection + mock_quota_check.return_value = True + + rsp = client.post( + f"{settings.API_V1_STR}/collections/", + json=collection_data.model_dump(), + headers=normal_user_token_headers, + ) + + assert rsp.status_code == 200 + content = rsp.json() + assert content["name"] == collection_data.name + assert "id" in content + assert isinstance(content["id"], str) + assert len(content["cards"]) == len(mock_collection.cards) + for i, card in enumerate(mock_collection.cards): + assert content["cards"][i]["front"] == card.front + assert content["cards"][i]["back"] == card.back + + mock_ai_generate.assert_called_once() def test_create_collection_with_ai_generation_error( @@ -114,19 +118,23 @@ def test_create_collection_with_ai_generation_error( with patch( "src.flashcards.services.generate_ai_collection", new_callable=AsyncMock ) as mock_ai_generate: - err_msg = "AI service is unavailable" - mock_ai_generate.side_effect = AIGenerationError(err_msg) - - rsp = client.post( - f"{settings.API_V1_STR}/collections/", - json=collection_data.model_dump(), - headers=normal_user_token_headers, - ) - - assert rsp.status_code == 500 - content = rsp.json() - assert "detail" in content - assert err_msg in content["detail"] + with patch( + "src.flashcards.api.check_and_increment_ai_usage_quota" + ) as mock_quota_check: + err_msg = "AI service is unavailable" + mock_ai_generate.side_effect = AIGenerationError(err_msg) + mock_quota_check.return_value = True + + rsp = client.post( + f"{settings.API_V1_STR}/collections/", + json=collection_data.model_dump(), + headers=normal_user_token_headers, + ) + + assert rsp.status_code == 500 + content = rsp.json() + assert "detail" in content + assert err_msg in content["detail"] def test_read_collection( From be5539c431de28b31450cfd39e0dbc5959bbf07a Mon Sep 17 00:00:00 2001 From: Sebastian Date: Thu, 15 May 2025 22:40:44 +0000 Subject: [PATCH 15/24] fix linting --- backend/src/users/services.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/backend/src/users/services.py b/backend/src/users/services.py index 60a93b2..d791800 100644 --- a/backend/src/users/services.py +++ b/backend/src/users/services.py @@ -6,9 +6,9 @@ from src.auth.services import get_password_hash from src.core.config import settings +from src.users.models import AIUsageQuota as AIUsageQuotaModel from src.users.models import User from src.users.schemas import AIUsageQuota, UserCreate, UserUpdate -from src.users.models import AIUsageQuota as AIUsageQuotaModel def create_user(*, session: Session, user_create: UserCreate) -> User: From fd6962998b4aff07d80f1c144acd45e49ecf7310 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Thu, 15 May 2025 22:42:42 +0000 Subject: [PATCH 16/24] run client and regenerate openapi spec --- frontend/openapi.json | 13 +++++++------ frontend/src/client/sdk.gen.ts | 26 +++++++++++++------------- frontend/src/client/types.gen.ts | 7 ++++--- 3 files changed, 24 insertions(+), 22 deletions(-) diff --git a/frontend/openapi.json b/frontend/openapi.json index f1618df..536d565 100644 --- a/frontend/openapi.json +++ b/frontend/openapi.json @@ -80,11 +80,11 @@ } } }, - "/api/v1/aiquota": { + "/api/v1/users/users/me/ai-usage-quota": { "get": { - "tags": ["flashcards"], - "summary": "Get Ai Usage Quota", - "operationId": "flashcards-get_ai_usage_quota", + "tags": ["users"], + "summary": "Get My Ai Usage Quota", + "operationId": "users-get_my_ai_usage_quota", "responses": { "200": { "description": "Successful Response", @@ -710,11 +710,12 @@ "schemas": { "AIUsageQuota": { "properties": { - "percentage_used": { "type": "integer", "title": "Percentage Used" }, + "usage_count": { "type": "integer", "title": "Usage Count" }, + "max_usage_allowed": { "type": "integer", "title": "Max Usage Allowed" }, "reset_date": { "type": "string", "format": "date-time", "title": "Reset Date" } }, "type": "object", - "required": ["percentage_used", "reset_date"], + "required": ["usage_count", "max_usage_allowed", "reset_date"], "title": "AIUsageQuota" }, "Body_login-login_access_token": { diff --git a/frontend/src/client/sdk.gen.ts b/frontend/src/client/sdk.gen.ts index 20fc818..ea6a713 100644 --- a/frontend/src/client/sdk.gen.ts +++ b/frontend/src/client/sdk.gen.ts @@ -12,7 +12,6 @@ import type { FlashcardsDeleteCardResponse, FlashcardsDeleteCollectionData, FlashcardsDeleteCollectionResponse, - FlashcardsGetAiUsageQuotaResponse, FlashcardsGetPracticeSessionStatusData, FlashcardsGetPracticeSessionStatusResponse, FlashcardsListPracticeCardsData, @@ -39,24 +38,13 @@ import type { LoginLoginAccessTokenResponse, StatsGetCollectionStatisticsEndpointData, StatsGetCollectionStatisticsEndpointResponse, + UsersGetMyAiUsageQuotaResponse, UsersReadUserMeResponse, UsersRegisterUserData, UsersRegisterUserResponse, } from './types.gen' export class FlashcardsService { - /** - * Get Ai Usage Quota - * @returns AIUsageQuota Successful Response - * @throws ApiError - */ - public static getAiUsageQuota(): CancelablePromise { - return __request(OpenAPI, { - method: 'GET', - url: '/api/v1/aiquota', - }) - } - /** * Read Collections * @param data The data for the request. @@ -515,4 +503,16 @@ export class UsersService { }, }) } + + /** + * Get My Ai Usage Quota + * @returns AIUsageQuota Successful Response + * @throws ApiError + */ + public static getMyAiUsageQuota(): CancelablePromise { + return __request(OpenAPI, { + method: 'GET', + url: '/api/v1/users/users/me/ai-usage-quota', + }) + } } diff --git a/frontend/src/client/types.gen.ts b/frontend/src/client/types.gen.ts index 4131154..d260da6 100644 --- a/frontend/src/client/types.gen.ts +++ b/frontend/src/client/types.gen.ts @@ -1,7 +1,8 @@ // This file is auto-generated by @hey-api/openapi-ts export type AIUsageQuota = { - percentage_used: number + usage_count: number + max_usage_allowed: number reset_date: string } @@ -160,8 +161,6 @@ export type ValidationError = { type: string } -export type FlashcardsGetAiUsageQuotaResponse = AIUsageQuota - export type FlashcardsReadCollectionsData = { limit?: number skip?: number @@ -290,3 +289,5 @@ export type UsersRegisterUserData = { } export type UsersRegisterUserResponse = UserPublic + +export type UsersGetMyAiUsageQuotaResponse = AIUsageQuota From 9ab44433f76472ff0111d3fc16ebf9bdcc5fb915 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Thu, 15 May 2025 23:14:40 +0000 Subject: [PATCH 17/24] update cheking function previous implementation does not increase usage. Taking advantage of optimistic locking for usage increment --- backend/src/users/services.py | 22 +++++++++---------- frontend/public/locales/en/translation.json | 2 +- frontend/public/locales/es/translation.json | 2 +- frontend/public/locales/nl/translation.json | 2 +- .../components/commonUI/AiPromptDialog.tsx | 12 +++------- frontend/src/hooks/useAiDialog.ts | 14 +++++------- 6 files changed, 21 insertions(+), 33 deletions(-) diff --git a/backend/src/users/services.py b/backend/src/users/services.py index d791800..2864503 100644 --- a/backend/src/users/services.py +++ b/backend/src/users/services.py @@ -76,21 +76,19 @@ def check_and_increment_ai_usage_quota(session: Session, user: User) -> bool: session.commit() return True - if quota.usage_count > settings.AI_MAX_USAGE_QUOTA: - return False - if now - quota.last_reset_time >= timedelta(days=settings.AI_QUOTA_TIME_RANGE_DAYS): quota.usage_count = 0 quota.last_reset_time = now session.add(quota) session.commit() - result = session.exec( - update(AIUsageQuotaModel) - .where( - (AIUsageQuotaModel.id == quota.id) - & (AIUsageQuotaModel.usage_count < settings.AI_MAX_USAGE_QUOTA) - ) - .values(usage_count=AIUsageQuotaModel.usage_count + 1) + result = session.exec( + update(AIUsageQuotaModel) + .where( + (AIUsageQuotaModel.id == quota.id) + & (AIUsageQuotaModel.usage_count <= settings.AI_MAX_USAGE_QUOTA) ) - session.commit() - return result.rowcount > 0 + .values(usage_count=AIUsageQuotaModel.usage_count + 1) + ) + session.commit() + return result.rowcount > 0 + diff --git a/frontend/public/locales/en/translation.json b/frontend/public/locales/en/translation.json index c1b68a4..ab4ecec 100644 --- a/frontend/public/locales/en/translation.json +++ b/frontend/public/locales/en/translation.json @@ -85,7 +85,7 @@ "actions": { "addCard": "Add Card", "aiQuotaResetDate": "AI Usage quota reset date", - "aiQuotaUsed": "AI Usage quota used", + "aiUsageCount": "AI creation count", "generateAiCard": "Add Card with AI", "addCollection": "Add Collection", "backToCollection": "Back to Collection", diff --git a/frontend/public/locales/es/translation.json b/frontend/public/locales/es/translation.json index ec48978..1ea0475 100644 --- a/frontend/public/locales/es/translation.json +++ b/frontend/public/locales/es/translation.json @@ -85,7 +85,7 @@ "actions": { "addCard": "Agregar Tarjeta", "aiQuotaResetDate": "Fecha de reinicio de cuota de IA", - "aiQuotaUsed": "Cuota de IA usada", + "aiUsageCount": "Uso de IA", "generateAiCard": "Agregar tarjeta con IA", "addCollection": "Agregar Colección", "backToCollection": "Volver a la Colección", diff --git a/frontend/public/locales/nl/translation.json b/frontend/public/locales/nl/translation.json index 39ffc42..a36b00e 100644 --- a/frontend/public/locales/nl/translation.json +++ b/frontend/public/locales/nl/translation.json @@ -85,7 +85,7 @@ "actions": { "addCard": "Kaart Toevoegen", "aiQuotaResetDate": "Datum van herstart van AI-quota", - "aiQuotaUsed": "Gebruikte AI-quota", + "aiUsageCount": "Gebruik van AI", "generateAiCard": "Kaart toevoegen met AI", "addCollection": "Collectie Toevoegen", "backToCollection": "Terug naar Collectie", diff --git a/frontend/src/components/commonUI/AiPromptDialog.tsx b/frontend/src/components/commonUI/AiPromptDialog.tsx index 46adadc..807850f 100644 --- a/frontend/src/components/commonUI/AiPromptDialog.tsx +++ b/frontend/src/components/commonUI/AiPromptDialog.tsx @@ -94,15 +94,9 @@ const AiPromptDialog: React.FC = ({ fontSize="xs" textAlign="right" mt={1} - color={ - usageQuota.percentage_used <= 50 - ? 'green.500' - : usageQuota.percentage_used <= 80 - ? 'yellow.500' - : 'red.500' - } + color="white.500" > - {`${t('general.actions.aiQuotaUsed')}: ${usageQuota.percentage_used}%`} + {`AI Usage: ${usageQuota.usage_count}/${usageQuota.max_usage_allowed}`} @@ -114,7 +108,7 @@ const AiPromptDialog: React.FC = ({ {isLoading ? `${t('general.actions.creating')}...` : t('general.actions.create')} diff --git a/frontend/src/hooks/useAiDialog.ts b/frontend/src/hooks/useAiDialog.ts index 38a6119..f6f5a89 100644 --- a/frontend/src/hooks/useAiDialog.ts +++ b/frontend/src/hooks/useAiDialog.ts @@ -1,20 +1,16 @@ -import { FlashcardsService } from '@/client' +import { AIUsageQuota, UsersService } from '@/client' import { useEffect, useState } from 'react' -interface AIUsageQuotaData { - reset_date: string - percentage_used: number -} - export function useAiDialog() { - const [usageQuota, setUsageQuota] = useState({ + const [usageQuota, setUsageQuota] = useState({ reset_date: '', - percentage_used: 0, + usage_count: 0, + max_usage_allowed: 0 }) useEffect(() => { const fetchUsageQuota = async () => { - const data = await FlashcardsService.getAiUsageQuota() + const data = await UsersService.getMyAiUsageQuota() setUsageQuota(data) } fetchUsageQuota() From 21d9fc7ddb05804e36795150570c77d32cc8d026 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Thu, 15 May 2025 23:20:12 +0000 Subject: [PATCH 18/24] fix linting --- backend/src/users/services.py | 2 +- .../src/components/commonUI/AiPromptDialog.tsx | 15 +++++++-------- frontend/src/hooks/useAiDialog.ts | 4 ++-- 3 files changed, 10 insertions(+), 11 deletions(-) diff --git a/backend/src/users/services.py b/backend/src/users/services.py index 2864503..d79bdfc 100644 --- a/backend/src/users/services.py +++ b/backend/src/users/services.py @@ -91,4 +91,4 @@ def check_and_increment_ai_usage_quota(session: Session, user: User) -> bool: ) session.commit() return result.rowcount > 0 - + diff --git a/frontend/src/components/commonUI/AiPromptDialog.tsx b/frontend/src/components/commonUI/AiPromptDialog.tsx index 807850f..0d44e37 100644 --- a/frontend/src/components/commonUI/AiPromptDialog.tsx +++ b/frontend/src/components/commonUI/AiPromptDialog.tsx @@ -90,13 +90,8 @@ const AiPromptDialog: React.FC = ({ {`${t('general.actions.aiQuotaResetDate')}: ${new Date(usageQuota.reset_date).toLocaleDateString()}`} - - {`AI Usage: ${usageQuota.usage_count}/${usageQuota.max_usage_allowed}`} + + {`${t('general.actions.aiUsageCount')}: ${usageQuota.usage_count}/${usageQuota.max_usage_allowed}`} @@ -108,7 +103,11 @@ const AiPromptDialog: React.FC = ({ {isLoading ? `${t('general.actions.creating')}...` : t('general.actions.create')} diff --git a/frontend/src/hooks/useAiDialog.ts b/frontend/src/hooks/useAiDialog.ts index f6f5a89..4544fdc 100644 --- a/frontend/src/hooks/useAiDialog.ts +++ b/frontend/src/hooks/useAiDialog.ts @@ -1,11 +1,11 @@ -import { AIUsageQuota, UsersService } from '@/client' +import { type AIUsageQuota, UsersService } from '@/client' import { useEffect, useState } from 'react' export function useAiDialog() { const [usageQuota, setUsageQuota] = useState({ reset_date: '', usage_count: 0, - max_usage_allowed: 0 + max_usage_allowed: 0, }) useEffect(() => { From 9e91cf5b6eb9e635f2d5af9b81852df89601e40d Mon Sep 17 00:00:00 2001 From: Sebastian Date: Thu, 15 May 2025 23:26:45 +0000 Subject: [PATCH 19/24] run formatter --- backend/src/users/services.py | 1 - 1 file changed, 1 deletion(-) diff --git a/backend/src/users/services.py b/backend/src/users/services.py index d79bdfc..548af22 100644 --- a/backend/src/users/services.py +++ b/backend/src/users/services.py @@ -91,4 +91,3 @@ def check_and_increment_ai_usage_quota(session: Session, user: User) -> bool: ) session.commit() return result.rowcount > 0 - From 72f2408f6a9216bc2a4f499bacae9e14bdb3339f Mon Sep 17 00:00:00 2001 From: Sebastian Date: Fri, 16 May 2025 12:54:55 +0000 Subject: [PATCH 20/24] remove unneeded hook for ai prompt dialog, only one query is needed --- frontend/public/locales/en/translation.json | 2 +- .../components/commonUI/AiPromptDialog.tsx | 15 +++++++++++-- frontend/src/hooks/useAiDialog.ts | 22 ------------------- 3 files changed, 14 insertions(+), 25 deletions(-) delete mode 100644 frontend/src/hooks/useAiDialog.ts diff --git a/frontend/public/locales/en/translation.json b/frontend/public/locales/en/translation.json index ab4ecec..677d128 100644 --- a/frontend/public/locales/en/translation.json +++ b/frontend/public/locales/en/translation.json @@ -85,7 +85,7 @@ "actions": { "addCard": "Add Card", "aiQuotaResetDate": "AI Usage quota reset date", - "aiUsageCount": "AI creation count", + "aiUsageCount": "AI usage count", "generateAiCard": "Add Card with AI", "addCollection": "Add Collection", "backToCollection": "Back to Collection", diff --git a/frontend/src/components/commonUI/AiPromptDialog.tsx b/frontend/src/components/commonUI/AiPromptDialog.tsx index 0d44e37..6b627cb 100644 --- a/frontend/src/components/commonUI/AiPromptDialog.tsx +++ b/frontend/src/components/commonUI/AiPromptDialog.tsx @@ -1,3 +1,5 @@ +import { UsersService } from '@/client' +import type { AIUsageQuota } from '@/client/types.gen' import { DialogActionTrigger, DialogBody, @@ -8,8 +10,8 @@ import { DialogRoot, DialogTitle, } from '@/components/ui/dialog' -import { useAiDialog } from '@/hooks/useAiDialog' import { Text } from '@chakra-ui/react' +import { useQuery } from '@tanstack/react-query' import type { OpenChangeDetails } from 'node_modules/@chakra-ui/react/dist/types/components/dialog/namespace' import { useRef, useState } from 'react' import { useTranslation } from 'react-i18next' @@ -38,7 +40,16 @@ const AiPromptDialog: React.FC = ({ const { t } = useTranslation() const [prompt, setPrompt] = useState('') const closeButtonRef = useRef(null) - const { usageQuota } = useAiDialog() + const { data } = useQuery({ + queryKey: ['usageQuota'], + queryFn: UsersService.getMyAiUsageQuota, + enabled: isOpen, + }) + const usageQuota: AIUsageQuota = data || { + max_usage_allowed: 0, + usage_count: 0, + reset_date: new Date().toDateString(), + } const handleSubmit = () => { if (!prompt.trim() || isLoading) return diff --git a/frontend/src/hooks/useAiDialog.ts b/frontend/src/hooks/useAiDialog.ts deleted file mode 100644 index 4544fdc..0000000 --- a/frontend/src/hooks/useAiDialog.ts +++ /dev/null @@ -1,22 +0,0 @@ -import { type AIUsageQuota, UsersService } from '@/client' -import { useEffect, useState } from 'react' - -export function useAiDialog() { - const [usageQuota, setUsageQuota] = useState({ - reset_date: '', - usage_count: 0, - max_usage_allowed: 0, - }) - - useEffect(() => { - const fetchUsageQuota = async () => { - const data = await UsersService.getMyAiUsageQuota() - setUsageQuota(data) - } - fetchUsageQuota() - }, []) - - return { - usageQuota, - } -} From 8ef3ed49d9170fa51de8e2c2ee965d2a5858e1d5 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Fri, 16 May 2025 13:04:59 +0000 Subject: [PATCH 21/24] change UI so that usage left is shown --- frontend/public/locales/en/translation.json | 2 +- frontend/public/locales/es/translation.json | 2 +- frontend/public/locales/nl/translation.json | 2 +- frontend/src/components/commonUI/AiPromptDialog.tsx | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/frontend/public/locales/en/translation.json b/frontend/public/locales/en/translation.json index 677d128..77bdd87 100644 --- a/frontend/public/locales/en/translation.json +++ b/frontend/public/locales/en/translation.json @@ -85,7 +85,7 @@ "actions": { "addCard": "Add Card", "aiQuotaResetDate": "AI Usage quota reset date", - "aiUsageCount": "AI usage count", + "aiUsageLeft": "AI usage left", "generateAiCard": "Add Card with AI", "addCollection": "Add Collection", "backToCollection": "Back to Collection", diff --git a/frontend/public/locales/es/translation.json b/frontend/public/locales/es/translation.json index 1ea0475..7108d1b 100644 --- a/frontend/public/locales/es/translation.json +++ b/frontend/public/locales/es/translation.json @@ -85,7 +85,7 @@ "actions": { "addCard": "Agregar Tarjeta", "aiQuotaResetDate": "Fecha de reinicio de cuota de IA", - "aiUsageCount": "Uso de IA", + "aiUsageLeft": "Uso de IA restante", "generateAiCard": "Agregar tarjeta con IA", "addCollection": "Agregar Colección", "backToCollection": "Volver a la Colección", diff --git a/frontend/public/locales/nl/translation.json b/frontend/public/locales/nl/translation.json index a36b00e..5f54f1d 100644 --- a/frontend/public/locales/nl/translation.json +++ b/frontend/public/locales/nl/translation.json @@ -85,7 +85,7 @@ "actions": { "addCard": "Kaart Toevoegen", "aiQuotaResetDate": "Datum van herstart van AI-quota", - "aiUsageCount": "Gebruik van AI", + "aiUsageLeft": "Overgebleven AI-gebruik", "generateAiCard": "Kaart toevoegen met AI", "addCollection": "Collectie Toevoegen", "backToCollection": "Terug naar Collectie", diff --git a/frontend/src/components/commonUI/AiPromptDialog.tsx b/frontend/src/components/commonUI/AiPromptDialog.tsx index 6b627cb..b52581b 100644 --- a/frontend/src/components/commonUI/AiPromptDialog.tsx +++ b/frontend/src/components/commonUI/AiPromptDialog.tsx @@ -102,7 +102,7 @@ const AiPromptDialog: React.FC = ({ {`${t('general.actions.aiQuotaResetDate')}: ${new Date(usageQuota.reset_date).toLocaleDateString()}`} - {`${t('general.actions.aiUsageCount')}: ${usageQuota.usage_count}/${usageQuota.max_usage_allowed}`} + {`${t('general.actions.aiUsageLeft')}: ${usageQuota.max_usage_allowed - usageQuota.usage_count}`} From ba9109c0038fe4ef9af670fd78538096849d90a2 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Thu, 22 May 2025 12:59:35 +0000 Subject: [PATCH 22/24] feat: apply suggested changes --- backend/src/core/config.py | 4 ++-- backend/src/flashcards/api.py | 8 ++++++-- frontend/src/components/commonUI/AiPromptDialog.tsx | 2 +- 3 files changed, 9 insertions(+), 5 deletions(-) diff --git a/backend/src/core/config.py b/backend/src/core/config.py index 0102553..6062b3d 100644 --- a/backend/src/core/config.py +++ b/backend/src/core/config.py @@ -64,8 +64,8 @@ def SQLALCHEMY_DATABASE_URI(self) -> PostgresDsn: AI_API_KEY: str | None = None AI_MODEL: str | None = None - AI_MAX_USAGE_QUOTA: int | None = None - AI_QUOTA_TIME_RANGE_DAYS: int | None = None + AI_MAX_USAGE_QUOTA: int = 30 + AI_QUOTA_TIME_RANGE_DAYS: int = 1 COLLECTION_GENERATION_PROMPT: str | None = None CARD_GENERATION_PROMPT: str | None = None diff --git a/backend/src/flashcards/api.py b/backend/src/flashcards/api.py index bebbc62..96b01fa 100644 --- a/backend/src/flashcards/api.py +++ b/backend/src/flashcards/api.py @@ -53,7 +53,9 @@ async def create_collection( if collection_in.prompt: try: - if not check_and_increment_ai_usage_quota(session, current_user): + if not await asyncio.to_thread( + lambda: check_and_increment_ai_usage_quota(session, current_user) + ): raise HTTPException( status_code=429, detail="Quota for AI usage is reached." ) @@ -150,7 +152,9 @@ async def create_card( if not access_checked: raise HTTPException(status_code=404, detail="Collection not found") if card_in.prompt: - if not check_and_increment_ai_usage_quota(session, current_user): + if not await asyncio.to_thread( + lambda: check_and_increment_ai_usage_quota(session, current_user) + ): raise HTTPException( status_code=429, detail="Quota for AI usage is reached." ) diff --git a/frontend/src/components/commonUI/AiPromptDialog.tsx b/frontend/src/components/commonUI/AiPromptDialog.tsx index b52581b..f5575dc 100644 --- a/frontend/src/components/commonUI/AiPromptDialog.tsx +++ b/frontend/src/components/commonUI/AiPromptDialog.tsx @@ -48,7 +48,7 @@ const AiPromptDialog: React.FC = ({ const usageQuota: AIUsageQuota = data || { max_usage_allowed: 0, usage_count: 0, - reset_date: new Date().toDateString(), + reset_date: new Date().toLocaleDateString(), } const handleSubmit = () => { From 373e0b007a4996cf93e543e6b4b73faa10d61c87 Mon Sep 17 00:00:00 2001 From: 0010aor <4ndres.or@gmail.com> Date: Tue, 27 May 2025 10:14:55 +0200 Subject: [PATCH 23/24] fix: refactor check_and_increment_ai_usage_quota --- backend/src/users/services.py | 45 +++++++---- backend/tests/scripts/test_test_pre_start.py | 6 +- backend/tests/users/conftest.py | 2 +- backend/tests/users/test_services.py | 82 +++++++++++++++++--- 4 files changed, 108 insertions(+), 27 deletions(-) diff --git a/backend/src/users/services.py b/backend/src/users/services.py index 548af22..0180a68 100644 --- a/backend/src/users/services.py +++ b/backend/src/users/services.py @@ -3,6 +3,7 @@ from typing import Any from sqlmodel import Session, select, update +from sqlalchemy.exc import IntegrityError from src.auth.services import get_password_hash from src.core.config import settings @@ -68,26 +69,44 @@ def get_ai_usage_quota_for_user(user: User) -> AIUsageQuota: def check_and_increment_ai_usage_quota(session: Session, user: User) -> bool: - quota = user.ai_usage_quota now = datetime.now(timezone.utc) - if not quota: - quota = AIUsageQuotaModel(user_id=user.id, usage_count=1, last_reset_time=now) - session.add(quota) + reset_threshold = now - timedelta(days=settings.AI_QUOTA_TIME_RANGE_DAYS) + + if not user.ai_usage_quota: + try: + quota = AIUsageQuotaModel( + user_id=user.id, usage_count=1, last_reset_time=now + ) + session.add(quota) + session.commit() + return True + except IntegrityError: + session.rollback() + + session.refresh(user) + + result_reset = session.exec( + update(AIUsageQuotaModel) + .where( + (AIUsageQuotaModel.user_id == user.id) + & (AIUsageQuotaModel.last_reset_time <= reset_threshold) + ) + .values(usage_count=1, last_reset_time=now) + ) + + if result_reset.rowcount > 0: session.commit() return True - if now - quota.last_reset_time >= timedelta(days=settings.AI_QUOTA_TIME_RANGE_DAYS): - quota.usage_count = 0 - quota.last_reset_time = now - session.add(quota) - session.commit() - result = session.exec( + result_increment = session.exec( update(AIUsageQuotaModel) .where( - (AIUsageQuotaModel.id == quota.id) - & (AIUsageQuotaModel.usage_count <= settings.AI_MAX_USAGE_QUOTA) + (AIUsageQuotaModel.user_id == user.id) + & (AIUsageQuotaModel.last_reset_time > reset_threshold) + & (AIUsageQuotaModel.usage_count < settings.AI_MAX_USAGE_QUOTA) ) .values(usage_count=AIUsageQuotaModel.usage_count + 1) ) + session.commit() - return result.rowcount > 0 + return result_increment.rowcount > 0 diff --git a/backend/tests/scripts/test_test_pre_start.py b/backend/tests/scripts/test_test_pre_start.py index 1dd33da..5baeb25 100644 --- a/backend/tests/scripts/test_test_pre_start.py +++ b/backend/tests/scripts/test_test_pre_start.py @@ -22,6 +22,6 @@ def test_init_successful_connection() -> None: except Exception: connection_successful = False - assert ( - connection_successful - ), "The database connection should be successful and not raise an exception." + assert connection_successful, ( + "The database connection should be successful and not raise an exception." + ) diff --git a/backend/tests/users/conftest.py b/backend/tests/users/conftest.py index e3a4b23..511080b 100644 --- a/backend/tests/users/conftest.py +++ b/backend/tests/users/conftest.py @@ -17,4 +17,4 @@ def test_user(db: Session) -> dict[str, Any]: user_in = UserCreate(email=email, password=password, full_name=full_name) user = create_user(session=db, user_create=user_in) - return user \ No newline at end of file + return user diff --git a/backend/tests/users/test_services.py b/backend/tests/users/test_services.py index 25f5839..4c72942 100644 --- a/backend/tests/users/test_services.py +++ b/backend/tests/users/test_services.py @@ -6,10 +6,17 @@ from src.auth.services import verify_password from src.users.models import AIUsageQuota, User from src.users.schemas import UserCreate, UserUpdate -from src.users.services import check_and_increment_ai_usage_quota, create_user, get_ai_usage_quota_for_user, get_user_by_email, update_user +from src.users.services import ( + check_and_increment_ai_usage_quota, + create_user, + get_ai_usage_quota_for_user, + get_user_by_email, + update_user, +) from tests.utils.utils import random_email, random_lower_string from src.core.config import settings + def test_create_user(db: Session) -> None: email = random_email() password = random_lower_string() @@ -89,35 +96,90 @@ def test_update_user(db: Session) -> None: def test_get_ai_usage_quota_for_user_no_quota(test_user): test_user.ai_usage_quota = None - with patch('src.users.services.datetime') as mock_datetime: + with patch("src.users.services.datetime") as mock_datetime: mock_datetime.now.return_value = datetime(2025, 1, 1, tzinfo=timezone.utc) quota = get_ai_usage_quota_for_user(test_user) assert quota.usage_count == 0 assert quota.reset_date == datetime(2025, 1, 31, tzinfo=timezone.utc) + def test_get_ai_usage_quota_for_user_with_quota(test_user): - test_user.ai_usage_quota = AIUsageQuota(usage_count=50, last_reset_time=datetime(2025, 1, 1, tzinfo=timezone.utc)) - with patch('src.users.services.settings.AI_MAX_USAGE_QUOTA', 100): - with patch('src.users.services.datetime') as mock_datetime: + test_user.ai_usage_quota = AIUsageQuota( + usage_count=50, last_reset_time=datetime(2025, 1, 1, tzinfo=timezone.utc) + ) + with patch("src.users.services.settings.AI_MAX_USAGE_QUOTA", 100): + with patch("src.users.services.datetime") as mock_datetime: mock_datetime.now.return_value = datetime(2025, 1, 1, tzinfo=timezone.utc) quota = get_ai_usage_quota_for_user(test_user) assert quota.usage_count == 50 assert quota.reset_date == datetime(2025, 1, 31, tzinfo=timezone.utc) + def test_check_and_increment_ai_usage_quota_first_time(db, test_user): within_quota = check_and_increment_ai_usage_quota(db, test_user) assert within_quota is True + def test_check_and_increment_ai_usage_quota_max_count_reached(db, test_user): - test_user.ai_usage_quota = AIUsageQuota(usage_count=30000, last_reset_time=datetime.now(timezone.utc)) + from src.users.models import AIUsageQuota as AIUsageQuotaModel + + quota = AIUsageQuotaModel( + user_id=test_user.id, + usage_count=settings.AI_MAX_USAGE_QUOTA, + last_reset_time=datetime.now(timezone.utc), + ) + db.add(quota) + db.commit() + db.refresh(test_user) + within_quota = check_and_increment_ai_usage_quota(db, test_user) assert within_quota is False + def test_check_and_increment_ai_usage_quota_reset_count(db, test_user): - with patch('src.users.services.datetime') as mock_datetime: + from src.users.models import AIUsageQuota as AIUsageQuotaModel + + old_reset_time = datetime(2025, 1, 1, tzinfo=timezone.utc) + quota = AIUsageQuotaModel( + user_id=test_user.id, + usage_count=10, + last_reset_time=old_reset_time, + ) + db.add(quota) + db.commit() + db.refresh(test_user) + + with patch("src.users.services.datetime") as mock_datetime: mock_datetime.now.return_value = datetime(2025, 2, 1, tzinfo=timezone.utc) - test_user.ai_usage_quota = AIUsageQuota(usage_count=1, last_reset_time=datetime(2025, 1, 1, tzinfo=timezone.utc)) within_quota = check_and_increment_ai_usage_quota(db, test_user) assert within_quota is True - assert test_user.ai_usage_quota.usage_count == 1 - assert test_user.ai_usage_quota.last_reset_time == datetime(2025, 2, 1, tzinfo=timezone.utc) + + db.refresh(quota) + assert quota.usage_count == 1 + assert quota.last_reset_time == datetime(2025, 2, 1, tzinfo=timezone.utc) + + +def test_check_and_increment_ai_usage_quota_near_limit(db, test_user): + """Test that quota checking works correctly near the limit.""" + from src.users.models import AIUsageQuota as AIUsageQuotaModel + + quota = AIUsageQuotaModel( + user_id=test_user.id, + usage_count=settings.AI_MAX_USAGE_QUOTA - 1, + last_reset_time=datetime.now(timezone.utc), + ) + db.add(quota) + db.commit() + db.refresh(test_user) + + within_quota = check_and_increment_ai_usage_quota(db, test_user) + assert within_quota is True + + db.refresh(quota) + assert quota.usage_count == settings.AI_MAX_USAGE_QUOTA + + within_quota = check_and_increment_ai_usage_quota(db, test_user) + assert within_quota is False + + db.refresh(quota) + assert quota.usage_count == settings.AI_MAX_USAGE_QUOTA From 9db02b1c0946be5fb26b2f0d5316cc138298d1a7 Mon Sep 17 00:00:00 2001 From: 0010aor <4ndres.or@gmail.com> Date: Tue, 27 May 2025 10:19:55 +0200 Subject: [PATCH 24/24] refactor: format backend files --- backend/src/users/services.py | 2 +- backend/tests/scripts/test_test_pre_start.py | 6 +++--- backend/tests/users/test_services.py | 7 ++++--- 3 files changed, 8 insertions(+), 7 deletions(-) diff --git a/backend/src/users/services.py b/backend/src/users/services.py index 0180a68..f6acb6a 100644 --- a/backend/src/users/services.py +++ b/backend/src/users/services.py @@ -2,8 +2,8 @@ from datetime import datetime, timedelta, timezone from typing import Any -from sqlmodel import Session, select, update from sqlalchemy.exc import IntegrityError +from sqlmodel import Session, select, update from src.auth.services import get_password_hash from src.core.config import settings diff --git a/backend/tests/scripts/test_test_pre_start.py b/backend/tests/scripts/test_test_pre_start.py index 5baeb25..1dd33da 100644 --- a/backend/tests/scripts/test_test_pre_start.py +++ b/backend/tests/scripts/test_test_pre_start.py @@ -22,6 +22,6 @@ def test_init_successful_connection() -> None: except Exception: connection_successful = False - assert connection_successful, ( - "The database connection should be successful and not raise an exception." - ) + assert ( + connection_successful + ), "The database connection should be successful and not raise an exception." diff --git a/backend/tests/users/test_services.py b/backend/tests/users/test_services.py index 4c72942..f3a17a1 100644 --- a/backend/tests/users/test_services.py +++ b/backend/tests/users/test_services.py @@ -1,9 +1,11 @@ -from datetime import datetime, timedelta, timezone -from unittest.mock import MagicMock, patch +from datetime import datetime, timezone +from unittest.mock import patch + from fastapi.encoders import jsonable_encoder from sqlmodel import Session from src.auth.services import verify_password +from src.core.config import settings from src.users.models import AIUsageQuota, User from src.users.schemas import UserCreate, UserUpdate from src.users.services import ( @@ -14,7 +16,6 @@ update_user, ) from tests.utils.utils import random_email, random_lower_string -from src.core.config import settings def test_create_user(db: Session) -> None: