added scheduler, logging and optimization updates

app.py

@@ -12,7 +12,17 @@ from torch_geometric.nn import SAGEConv
 from supabase import create_client
 from fastapi import FastAPI, HTTPException, Query
 from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import StreamingResponse
 from dotenv import load_dotenv
+import json
+from apscheduler.schedulers.background import BackgroundScheduler
+from apscheduler.triggers.cron import CronTrigger
+import logging
+import uvicorn
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 load_dotenv()
 
@@ -46,24 +56,32 @@ def get_supabase_client():
 def load_and_preprocess_data():
     supabase = get_supabase_client()
 
-    profiles = pd.DataFrame(supabase.table('profiles').select('id').execute().data)
-    posts = pd.DataFrame(supabase.table('posts').select('id, author').execute().data)
-    likes = pd.DataFrame(supabase.table('likes').select('user_id, post_id').execute().data)
+    def fetch_table(table, columns, chunk_size=1000):
+        offset = 0
+        all_data = []
+        while True:
+            response = supabase.table(table).select(columns).range(offset, offset + chunk_size - 1).execute()
+            data = response.data
+            if not data:
+                break
+            all_data.extend(data)
+            offset += chunk_size
+        return all_data
+
+    profiles = fetch_table('profiles', 'id')
+    posts = fetch_table('posts', 'id, author')
+    likes = fetch_table('likes', 'user_id, post_id')
 
     bipartite = nx.DiGraph()
 
-    user_set = set(posts['author']) | set(likes['user_id'])
-    post_set = set(posts['id'])
+    user_set = {p['author'] for p in posts} | {l['user_id'] for l in likes}
+    post_set = {p['id'] for p in posts}
 
-    for user in user_set:
-        bipartite.add_node(user, type='user')
-    for post in post_set:
-        bipartite.add_node(post, type='post')
+    bipartite.add_nodes_from(user_set, type='user')
+    bipartite.add_nodes_from(post_set, type='post')
 
-    for _, row in posts.iterrows():
-        bipartite.add_edge(row['author'], row['id'])
-    for _, row in likes.iterrows():
-        bipartite.add_edge(row['user_id'], row['post_id'])
+    bipartite.add_edges_from((p['author'], p['id']) for p in posts)
+    bipartite.add_edges_from((l['user_id'], l['post_id']) for l in likes)
 
     return bipartite
 
@@ -116,54 +134,73 @@ def train_model(model, data, pos_edges, neg_edges, epochs=200):
 
 def rebuild_model():
     global G, features, user_nodes, post_nodes, node2idx, pyg_data, trained_model
-
-    G = load_and_preprocess_data()
-    user_nodes = sorted(n for n, attr in G.nodes(data=True) if attr['type'] == 'user')
-    post_nodes = sorted(n for n, attr in G.nodes(data=True) if attr['type'] == 'post')
-
-    all_nodes = user_nodes + post_nodes
-    node2idx = {node: i for i, node in enumerate(all_nodes)}
-
-    features = torch.eye(len(all_nodes))
-    pyg_data = from_networkx(G)
-    pyg_data.x = features
-
-    pos_edges, neg_edges = prepare_training_data(G, node2idx, user_nodes, post_nodes)
-
-    input_dim = features.shape[1]
-    model = GraphRecommender(input_dim)
-    trained_model = train_model(model, pyg_data, pos_edges, neg_edges)
-
-def get_recommendations(user_id, model, data, G, user_nodes, post_nodes, node2idx, top_k=10):
+    logger.info("Starting model rebuild at 3:30 AM")
+    try:
+        G = load_and_preprocess_data()
+        user_nodes = sorted(n for n, attr in G.nodes(data=True) if attr['type'] == 'user')
+        post_nodes = sorted(n for n, attr in G.nodes(data=True) if attr['type'] == 'post')
+
+        all_nodes = user_nodes + post_nodes
+        node2idx = {node: i for i, node in enumerate(all_nodes)}
+
+        features = torch.sparse_coo_tensor(
+            torch.arange(len(all_nodes)).repeat(2, 1),
+            torch.ones(len(all_nodes)),
+            (len(all_nodes), len(all_nodes))
+        )
+        pyg_data = from_networkx(G)
+        pyg_data.x = features
+
+        pos_edges, neg_edges = prepare_training_data(G, node2idx, user_nodes, post_nodes)
+
+        input_dim = features.shape[1]
+        model = GraphRecommender(input_dim)
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        model = model.to(device)
+        pyg_data = pyg_data.to(device)
+        pos_edges = pos_edges.to(device)
+        neg_edges = neg_edges.to(device)
+
+        trained_model = train_model(model, pyg_data, pos_edges, neg_edges)
+        trained_model = trained_model.to('cpu')
+        logger.info("Model rebuild completed successfully")
+    except Exception as e:
+        logger.error(f"Error during model rebuild: {str(e)}")
+        raise
+
+def get_recommendations(user_id, model, data, G, user_nodes, post_nodes, node2idx):
     if user_id not in user_nodes:
         return []
 
     user_idx = node2idx[user_id]
-
     user_interacted = {v for _, v in G.out_edges(user_id) if G.nodes[v]['type'] == 'post'}
+    post_indices = [node2idx[p] for p in post_nodes if p not in user_interacted]
 
     with torch.no_grad():
         embeddings = model(data.x, data.edge_index)
-    user_embed = embeddings[user_idx]
+        user_embed = embeddings[user_idx].unsqueeze(0)
+        post_embeds = embeddings[post_indices]
 
-    scores = [(post, torch.dot(user_embed, embeddings[node2idx[post]]).item()) for post in post_nodes if post not in user_interacted]
-    scores = sorted(scores, key=lambda x: x[1], reverse=True)
+        scores = torch.matmul(user_embed, post_embeds.T).squeeze(0)
 
-    # Return the top_k post IDs
-    recommended_post_ids = [post for post, _ in scores[:top_k]]
-    return recommended_post_ids
+    post_scores = [(post_nodes[i], score.item()) for i, score in zip(post_indices, scores)]
+    post_scores = sorted(post_scores, key=lambda x: x[1], reverse=True)
 
-def fetch_full_post_records(post_ids):
-    """Fetch full post records from Supabase for the given post IDs."""
+    return [{"post_id": post, "score": score} for post, score in post_scores]
+
+def fetch_full_post_records(post_ids, batch_size=1000):
     supabase = get_supabase_client()
     if not post_ids:
         return []
 
-    response = supabase.table('posts').select('*').in_('id', post_ids).execute()
-    records = response.data
-
-    for record in records:
-        record['type'] = 'post'
+    records = []
+    for i in range(0, len(post_ids), batch_size):
+        batch_ids = post_ids[i:i + batch_size]
+        response = supabase.table('posts').select('*').in_('id', batch_ids).execute()
+        batch_records = response.data
+        for record in batch_records:
+            record['type'] = 'post'
+        records.extend(batch_records)
 
     return records
 
@@ -177,16 +214,55 @@ async def get_recommendations_handler(user_id: str = Query(...)):
     if trained_model is None:
         raise HTTPException(status_code=500, detail="Model not initialized, please rebuild first.")
 
-    # Get recommended post IDs
-    recommended_post_ids = get_recommendations(user_id, trained_model, pyg_data, G, user_nodes, post_nodes, node2idx)
-    
-    # Fetch full post records for the recommended post IDs
-    full_post_records = fetch_full_post_records(recommended_post_ids)
-    
-    return {"status": "success", "recommendations": full_post_records}
+    recommended_posts = get_recommendations(user_id, trained_model, pyg_data, G, user_nodes, post_nodes, node2idx)
+    if not recommended_posts:
+        return {"status": "success", "recommendations": []}
+
+    post_ids = [post["post_id"] for post in recommended_posts]
+    full_post_records = fetch_full_post_records(post_ids)
+
+    post_dict = {post["id"]: post for post in full_post_records}
+    ordered_recommendations = []
+    for post in recommended_posts:
+        post_id = post["post_id"]
+        if post_id in post_dict:
+            post_record = post_dict[post_id]
+            post_record["score"] = post["score"]
+            ordered_recommendations.append(post_record)
+
+    def generate():
+        yield '{"status": "success", "recommendations": ['
+        for i, rec in enumerate(ordered_recommendations):
+            yield json.dumps(rec)
+            if i < len(ordered_recommendations) - 1:
+                yield ','
+        yield ']}'
+
+    return StreamingResponse(generate(), media_type="application/json")
 
 @app.get("/")
 async def health_check():
     return {"status": "success", "message": "Service operational"}
 
-rebuild_model()
+scheduler = BackgroundScheduler(timezone="PST")
+
+scheduler.add_job(
+    rebuild_model,
+    trigger=CronTrigger(hour=3, minute=30),
+    id='daily_model_rebuild',
+    replace_existing=True
+)
+
+@app.on_event("startup")
+async def startup_event():
+    rebuild_model()
+    scheduler.start()
+    logger.info("Scheduler started, model will rebuild daily at 3:30 AM")
+
+@app.on_event("shutdown")
+async def shutdown_event():
+    scheduler.shutdown()
+    logger.info("Scheduler shut down")
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=8000)

requirements.txt

@@ -6,4 +6,5 @@ torch_geometric
 supabase
 fastapi
 python-dotenv
-uvicorn
+uvicorn
+apscheduler