Custom Cache cache

app/cache.py

@@ -0,0 +1,193 @@
+import time
+import threading
+import logging
+import boto3
+from botocore.exceptions import ClientError, NoCredentialsError, PartialCredentialsError
+import os
+
+from dotenv import load_dotenv
+
+logger = logging.getLogger(__name__)
+
+load_dotenv()
+
+AWS_ACCESS_KEY = os.getenv('AWS_ACCESS_KEY')
+AWS_SECRET_KEY = os.getenv('AWS_SECRET_KEY')
+REGION = os.getenv('AWS_REGION')
+
+def upload_file_to_s3(filename):
+    # user_id = users\to_upload\da43883b-7e12-425f-b498-ac8665dd23fe.pkl extract everything before the first dot and the last \
+    user_id = filename.split('\\')[-1].split('.')[0]
+    # user_id = filename.split('.')[0]
+    print(user_id)
+    function_name = upload_file_to_s3.__name__
+    logger.info(f"Uploading file {filename} to S3", extra={'user_id': user_id, 'endpoint': function_name})
+    bucket = 'core-ai-assets'
+    try:
+        if (AWS_ACCESS_KEY and AWS_SECRET_KEY):
+            session = boto3.session.Session(aws_access_key_id=AWS_ACCESS_KEY, aws_secret_access_key=AWS_SECRET_KEY, region_name=REGION)
+        else:
+            session = boto3.session.Session()
+        s3_client = session.client('s3')
+        with open(filename, "rb") as f:
+            ## Upload to Production Folder
+            s3_client.upload_fileobj(f, bucket, f'staging/users/{user_id}.pkl')
+        logger.info(f"File {filename} uploaded successfully to S3", extra={'user_id': user_id, 'endpoint': function_name})
+        
+        os.remove(filename)
+
+        # force_file_move(os.path.join('users', 'to_upload', filename), os.path.join('users', 'data', filename))
+        return True
+    except (FileNotFoundError, NoCredentialsError, PartialCredentialsError) as e:
+        logger.error(f"S3 upload failed for {filename}: {e}", extra={'user_id': user_id, 'endpoint': function_name})
+        return False
+
+class CustomTTLCache:
+    def __init__(self, ttl=60, cleanup_interval=10):
+        """
+        ttl: Time to live (in seconds) for each item
+        cleanup_interval: Interval at which background thread checks for expired items
+        """
+        self.ttl = ttl
+        self.data = {}  # key -> (value, expiration_time)
+        self.lock = threading.Lock()
+        self.cleanup_interval = cleanup_interval
+        self.running = True
+        self.cleanup_thread = threading.Thread(target=self._cleanup_task, daemon=True) # for periodic cleanup
+        self.cleanup_thread.start()
+
+    def __setitem__(self, key, value):
+        """
+        For item assignment i.e. cache[key] = value
+        """
+        self.set(key, value)
+
+    def set(self, key, value):
+        # set expiration time as current time + ttl
+        expire_at = time.time() + self.ttl
+        with self.lock:
+            self.data[key] = (value, expire_at)
+
+    def get(self, key, default=None):
+        """
+        Get a value from the cache. Returns default if key not found or expired.
+        Resets the TTL if the item is successfully accessed.
+        """
+        with self.lock:
+            entry = self.data.get(key)
+            if not entry:
+                return default
+            value, expire_at = entry
+            now = time.time()
+            if now > expire_at:
+                # item expired, evict and upload
+                self._evict_item(key, value)
+                return default
+            # refresh this data's TTL since it's accessed
+            self.data[key] = (value, now + self.ttl)
+            return value
+
+    def __contains__(self, key):
+        """
+        Chheck if a key is in the cache and not expired.
+        If not expired, reset the TTL.
+        """
+        with self.lock:
+            entry = self.data.get(key)
+            if not entry:
+                return False
+            value, expire_at = entry
+            now = time.time()
+            if now > expire_at:
+                self._evict_item(key, value)
+                return False
+            # refresh TTL
+            self.data[key] = (value, now + self.ttl)
+            return True
+
+    def __getitem__(self, key):
+        """
+        Retrieve an item i.e. user_cache[key].
+        Raises KeyError if not found or expired.
+        Resets the TTL if the item is successfully accessed.
+        """
+        with self.lock:
+            if key not in self.data:
+                raise KeyError(key)
+            value, expire_at = self.data[key]
+            now = time.time()
+            if now > expire_at:
+                self._evict_item(key, value)
+                raise KeyError(key)
+            self.data[key] = (value, now + self.ttl)
+            return value
+
+    def pop(self, key, default=None):
+        """
+        Remove an item from the cache and return its value.
+        Upload to S3 if it hasn't expired yet.
+        """
+        with self.lock:
+            entry = self.data.pop(key, None)
+            if entry:
+                value, expire_at = entry
+                # upload before removal
+                self._upload_value(key, value)
+                return value
+            return default
+
+    def close(self):
+        """
+        Stop the background cleanup thread. 
+        Technically we should call this when done using the cache but I guess we are never done using it lol.
+        """
+        self.running = False
+        self.cleanup_thread.join()
+
+    def _cleanup_task(self):
+        """
+        Background task that periodically checks for expired items and removes them.
+        """
+        while self.running:
+            now = time.time()
+            keys_to_remove = []
+            with self.lock:
+                for k, (v, exp) in list(self.data.items()):
+                    if now > exp:
+                        keys_to_remove.append(k)
+            
+            # Evict expired items outside the lock to handle uploading
+            for k in keys_to_remove:
+                with self.lock:
+                    if k in self.data:
+                        value, _ = self.data.pop(k)
+                        self._evict_item(k, value)
+
+            time.sleep(self.cleanup_interval)
+
+    def _evict_item(self, key, value):
+        """
+        Called internally when an item expires. upload to S3 before removal.
+        """
+        self._upload_value(key, value)
+
+    def _upload_value(self, key, value):
+        """
+        Saves the user's data and uploads the resulting file to S3.
+        """
+        try:
+            filename = value.save_user()
+            logger.info(
+                f"FILENAME = {filename}",
+                extra={'user_id': key, 'endpoint': 'cache_eviction'}
+            )
+            upload_file_to_s3(filename)
+            logger.info(
+                f"User {key} saved to S3 during cache eviction",
+                extra={'user_id': key, 'endpoint': 'cache_eviction'}
+            )
+        except Exception as e:
+            logger.error(
+                f"Failed to save user {key} to S3 during cache eviction: {e}",
+                extra={'user_id': key, 'endpoint': 'cache_eviction'}
+            )

app/main.py

@@ -587,20 +587,7 @@ def create_user(request: CreateUserItem, api_key: str = Security(get_api_key)):
             )
 
         user = User(request.user_id, user_info, client, GENERAL_ASSISTANT)
-        save = user.save_user()
-        
 
-        if save:
-            print_log("INFO",f"Created pickle file for user", extra={"user_id": request.user_id, "endpoint": "/create_user"})
-            logger.info(f"Created pickle file for user", extra={"user_id": request.user_id, "endpoint": "/create_user"})
-        else:
-            print_log("ERROR",f"Failed to create (user.save_user()) pickle file", extra={"user_id": request.user_id, "endpoint": "/create_user"})
-            logger.error(f"Failed to create (user.save_user()) pickle file", extra={"user_id": request.user_id, "endpoint": "/create_user"})
-            raise HTTPException(
-                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
-                detail="Failed to create user pickle file"
-            )
-        
         # create memento folder for user
         folder_path = os.path.join("mementos", "to_upload", request.user_id)
 
@@ -612,7 +599,6 @@ def create_user(request: CreateUserItem, api_key: str = Security(get_api_key)):
 
         # upload user pickle file to s3 bucket 
         try:
-            # value.save_user()
             add_to_cache(user)
             pop_cache(request.user_id)
             upload = True
@@ -622,7 +608,7 @@ def create_user(request: CreateUserItem, api_key: str = Security(get_api_key)):
         if upload == True:
             print_log("INFO",f"Successfully created user", extra={"user_id": request.user_id, "endpoint": "/create_user"})
             logger.info(f"Successfully created user", extra={"user_id": request.user_id, "endpoint": "/create_user"})
-            return {"message": f"[OK] User created: {user.user_id}"} 
+            return {"message": {"info": f"[OK] User created: {user}", "messages": user.get_messages()}} 
         else:
             print_log("ERROR",f"Failed to upload user pickle to S3", extra={"user_id": request.user_id, "endpoint": "/create_user"})
             logger.error(f"Failed to upload user pickle to S3", extra={"user_id": request.user_id, "endpoint": "/create_user"})

app/utils.py

@@ -23,6 +23,8 @@ import threading
 import time
 import uuid
 
+from app.cache import CustomTTLCache, upload_file_to_s3
+
 load_dotenv()
 
 # Environment Variables for API Keys
@@ -37,29 +39,8 @@ REGION = os.getenv('AWS_REGION')
 
 logger = logging.getLogger(__name__)
 
-class AutoSaveTTLCache(TTLCache):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.lock = threading.Lock()
-
-    def pop(self, key, *args):
-        with self.lock:
-            if key in self:
-                value = self[key]
-                # Save to S3 before removing
-                filename = f'{key}.pkl'
-                try:
-                    value.save_user()
-                    upload_file_to_s3(filename)
-                    logger.info(f"User {key} saved to S3 during cache eviction", 
-                               extra={'user_id': key, 'endpoint': 'cache_eviction'})
-                except Exception as e:
-                    logger.error(f"Failed to save user {key} to S3 during cache eviction: {e}", 
-                               extra={'user_id': key, 'endpoint': 'cache_eviction'})
-            return super().pop(key, *args)
-
 # Replace the simple TTLCache with our custom implementation
-user_cache = AutoSaveTTLCache(maxsize=100, ttl=1200)  # 20 minutes TTL
+user_cache = CustomTTLCache(ttl=120, cleanup_interval=30) # 2 minutes TTL
 
 def force_file_move(source, destination):
     function_name = force_file_move.__name__
@@ -79,45 +60,41 @@ def force_file_move(source, destination):
 def get_user(user_id):
     function_name = get_user.__name__
     logger.info(f"Fetching user {user_id}", extra={'user_id': user_id, 'endpoint': function_name})
-    
-    try:
-        if user_id in user_cache:
-            logger.info(f"User {user_id} found in cache", extra={'user_id': user_id, 'endpoint': function_name})
-            return user_cache[user_id]
-    except KeyError:
-        # TTLCache automatically removes expired entries
-        pass
-
-    client = OpenAI(api_key=os.getenv('OPENAI_API_KEY'))
-    user_file = os.path.join('users', 'data', f'{user_id}.pkl')
-    # if os.path.exists(user_file):
-    #     with open(user_file, 'rb') as f:
-    #         user = pickle.load(f)
-    #         user.client = client
-    #         user.conversations.client = client
-    #     with cache_lock:
-    #         user_cache[user_id] = user
-    #     return user
-    logger.warning(f"User {user_id} not found locally. Attempting to download from S3", extra={'user_id': user_id, 'endpoint': function_name})
-    download = download_file_from_s3(f'{user_id}.pkl', 'core-ai-assets')
-    logger.info(f"Download success: {download}", extra={'user_id': user_id, 'endpoint': function_name})
-    if (download):
-        with open(user_file, 'rb') as f:
-            user = pickle.load(f)
-            user.client = client
-            user.conversations.client = client
-        user_cache[user_id] = user  # No need for lock here
-        os.remove(user_file)
-        logger.info(f"User {user_id} loaded successfully from S3", extra={'user_id': user_id, 'endpoint': function_name})
-        return user
+    logger.info(f"[CACHE]: {user_cache}", extra={'user_id': user_id, 'endpoint': function_name})
+    if user_id in user_cache:
+        logger.info(f"User {user_id} found in cache", extra={'user_id': user_id, 'endpoint': function_name})
+        return user_cache[user_id]
     else:
-        logger.error(f"User {user_id} pickle does not exist in S3", extra={'user_id': user_id, 'endpoint': function_name})
-        # check if user_info exists
-        user_info = get_user_info(user_id)
-        if (user_info):
-            # user has done onboarding but pickle file not created
-            raise ReferenceError(f"User {user_id} pickle still being created")
-        raise LookupError(f"User [{user_id}] has not onboarded yet")
+        client = OpenAI(api_key=os.getenv('OPENAI_API_KEY'))
+        user_file = os.path.join('users', 'data', f'{user_id}.pkl')
+        # if os.path.exists(user_file):
+        #     with open(user_file, 'rb') as f:
+        #         user = pickle.load(f)
+        #         user.client = client
+        #         user.conversations.client = client
+        #     with cache_lock:
+        #         user_cache[user_id] = user
+        #     return user
+        logger.warning(f"User {user_id} not found locally. Attempting to download from S3", extra={'user_id': user_id, 'endpoint': function_name})
+        download = download_file_from_s3(f'{user_id}.pkl', 'core-ai-assets')
+        logger.info(f"Download success: {download}", extra={'user_id': user_id, 'endpoint': function_name})
+        if (download):
+            with open(user_file, 'rb') as f:
+                user = pickle.load(f)
+                user.client = client
+                user.conversations.client = client
+            user_cache[user_id] = user  # No need for lock here
+            os.remove(user_file)
+            logger.info(f"User {user_id} loaded successfully from S3", extra={'user_id': user_id, 'endpoint': function_name})
+            return user
+        else:
+            logger.error(f"User {user_id} pickle does not exist in S3", extra={'user_id': user_id, 'endpoint': function_name})
+            # check if user_info exists
+            user_info = get_user_info(user_id)
+            if (user_info):
+                # user has done onboarding but pickle file not created
+                raise ReferenceError(f"User {user_id} pickle still being created")
+            raise LookupError(f"User [{user_id}] has not onboarded yet")
 
 def get_life_score(user_id):
     function_name = get_life_score.__name__
@@ -1194,30 +1171,6 @@ def get_growth_guide_session(user_id, session_id):
         return None
     
 
-def upload_file_to_s3(filename):
-    user_id = filename.split('.')[0]
-    function_name = upload_file_to_s3.__name__
-    logger.info(f"Uploading file {filename} to S3", extra={'user_id': user_id, 'endpoint': function_name})
-    bucket = 'core-ai-assets'
-    try:
-        if (AWS_ACCESS_KEY and AWS_SECRET_KEY):
-            session = boto3.session.Session(aws_access_key_id=AWS_ACCESS_KEY, aws_secret_access_key=AWS_SECRET_KEY, region_name=REGION)
-        else:
-            session = boto3.session.Session()
-        s3_client = session.client('s3')
-        with open(os.path.join('users', 'to_upload', filename), "rb") as f:
-            ## Upload to Production Folder
-            s3_client.upload_fileobj(f, bucket, f'staging/users/{filename}')
-        logger.info(f"File {filename} uploaded successfully to S3", extra={'user_id': user_id, 'endpoint': function_name})
-        
-        os.remove(os.path.join('users', 'to_upload', filename))
-
-        # force_file_move(os.path.join('users', 'to_upload', filename), os.path.join('users', 'data', filename))
-        return True
-    except (FileNotFoundError, NoCredentialsError, PartialCredentialsError) as e:
-        logger.error(f"S3 upload failed for {filename}: {e}", extra={'user_id': user_id, 'endpoint': function_name})
-        return False
-
 def download_file_from_s3(filename, bucket):
     user_id = filename.split('.')[0]
     function_name = download_file_from_s3.__name__