update: v2 daily limit and previous recommendations

src/__init__.py

@@ -13,8 +13,8 @@
     get_user_collection,
     append_recommendations_to_db,
     get_user_previous_recommendations,
-    get_user_recommendation_by_id
-    # check_daily_limit
+    get_user_recommendation_by_id,
+    check_and_update_daily_limit,
 )
 from .db import (
     recommend,

src/api.py

@@ -17,7 +17,11 @@
 from .user_data import get_repos
 from src.db import (recommend, 
                     get_topic_based_recommendations)
-from src.models import User, GithubUser, get_user_collection, append_recommendations_to_db, get_user_previous_recommendations, get_user_recommendation_by_id
+from src.models import (User, 
+                        GithubUser, 
+                        get_user_collection, 
+                        append_recommendations_to_db, get_user_previous_recommendations, 
+                        get_user_recommendation_by_id, check_and_update_daily_limit)
 
 # load_dotenv()
 logger = logging.getLogger(__name__)
@@ -189,10 +193,15 @@ async def get_recommendations(request: Request, current_user: dict = Depends(get
         languages = body.get("languages", [])
 
         # TODO: v2
-        # try: 
-        #     check_daily_limit(username)
-        # except ValueError as e:
-        #     raise HTTPException(status_code=403, detail=str(e))
+        limit_updated = await check_and_update_daily_limit(username)
+
+        if not limit_updated:
+            logger.warning(f"Daily limit exceeded for user: {username}")
+            return {
+                "success": False,
+                "message": "Reached your daily limit",
+                "recommendations": []
+            }
 
         urls = []
         user = User(username=current_user["username"], access_token=current_user["access_token"],extra_topics=extra_topics, languages=languages)

src/models.py

@@ -1,6 +1,7 @@
 from pydantic import BaseModel
 from typing import Optional, List
 from motor.motor_asyncio import AsyncIOMotorClient
+from datetime import datetime, timedelta
 import os
 import uuid
 import logging
@@ -153,29 +154,56 @@ async def get_user_previous_recommendations(username: str) -> list:
 
 
 # TODO: v2
-# def update_daily_limit_to_all_users():
-#     try:
-#         user_collection = db['users']
-#         user_collection.update_many({}, {"$set": {"daily_limit": 2}})
-#     except Exception as e:
-#         logger.error(f"Failed to update daily limit: {str(e)}")
-#         raise ValueError("Failed to update daily limit")
-
-
-# def check_and_update_daily_limit(username: str):
-#     user_collection = db['users']
-#     result = user_collection.find_one_and_update(
-#         {"username": username, "daily_limit": {"$gt": 0}},
-#         {"$inc": {"daily_limit": -1}},
-#         return_document=True
-#     )
-
-#     if not result:
-#         logger.warning(f"Daily limit exceeded for user: {username}")
-#         raise ValueError("Daily limit exceeded")
-
-#     return result
+async def update_daily_limit(username: str) -> bool:
+    try:
+        user_recommendations_coll = db['user_recommendations']
+        now = datetime.now()
+        start_of_day = now.replace(hour=0, minute=0, second=0, microsecond=0)
 
+        # Fetch the user's recommendations
+        user_recommendations = await user_recommendations_coll.find_one({"username": username})
+
+        # Initialize if no record exists
+        if not user_recommendations:
+            # Initialize user recommendations with a daily limit
+            await user_recommendations_coll.insert_one({
+                "username": username,
+                "recommendation_refs": [],
+                "last_updated": start_of_day,
+                "daily_limit": 2
+            })
+            return True  # Initial limit set
+
+        last_updated = user_recommendations.get("last_updated", start_of_day)
+        daily_limit = user_recommendations.get("daily_limit", 2)
+
+        # Check if the limit needs to be reset
+        if now - last_updated >= timedelta(days=1):
+            # Reset daily limit and update the timestamp
+            await user_recommendations_coll.update_one(
+                {"username": username},
+                {"$set": {"daily_limit": 2, "last_updated": start_of_day}}
+            )
+            return True  # Limit reset
+
+        # If the limit is zero or less, return False
+        if daily_limit <= 0:
+            return False
+
+        # Decrease the limit if it's still valid
+        await user_recommendations_coll.update_one(
+            {"username": username},
+            {"$inc": {"daily_limit": -1}}
+        )
+
+        return True
+    except Exception as e:
+        logger.error(f"Failed to update daily limit: {str(e)}")
+        return False
+
+async def check_and_update_daily_limit(username: str) -> bool:
+    return await update_daily_limit(username)
+
 async def main():
     recommendations = await get_user_previous_recommendations("Hk669")
     print(recommendations)