Update main.py

main.py

@@ -1,11 +1,16 @@
 import os
 import json
+import csv
+from io import StringIO
+from functools import wraps
+from datetime import datetime
 from flask import Flask, request, jsonify
 from flask_cors import CORS
 from dotenv import load_dotenv
 
-# Firebase
-from firebase_admin import credentials, firestore, storage, initialize_app
+# Firebase Admin SDK
+import firebase_admin
+from firebase_admin import credentials, auth, firestore, storage, db as firebase_db, initialize_app
 
 # Exa.ai
 from exa_py import Exa
@@ -13,129 +18,249 @@ from exa_py import Exa
 # Google GenAI (Gemini)
 from google import genai
 
-# (Optional) Faiss for behavioral embeddings 
-import faiss
-
 # ─── Load environment ───────────────────────────────────────────────────────────
 load_dotenv()
-EXA_API_KEY    = os.getenv("EXA_API_KEY")
-GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
-FIREBASE_JSON  = os.getenv("FIREBASE")
-STORAGE_BUCKET = os.getenv("Firebase_Storage")
+EXA_API_KEY      = os.getenv("EXA_API_KEY")
+GEMINI_API_KEY   = os.getenv("GEMINI_API_KEY")
+FIREBASE_JSON    = os.getenv("FIREBASE")
+FIREBASE_DB_URL  = os.getenv("FIREBASE_DB_URL")
+STORAGE_BUCKET   = os.getenv("Firebase_Storage")
 
-if not (EXA_API_KEY and GEMINI_API_KEY and FIREBASE_JSON and STORAGE_BUCKET):
-    raise RuntimeError("Missing one or more required env vars: EXA_API_KEY, GEMINI_API_KEY, FIREBASE, Firebase_Storage")
+if not (EXA_API_KEY and GEMINI_API_KEY and FIREBASE_JSON and FIREBASE_DB_URL and STORAGE_BUCKET):
+    raise RuntimeError("Missing one or more required env vars: EXA_API_KEY, GEMINI_API_KEY, FIREBASE, FIREBASE_DB_URL, Firebase_Storage")
 
-# ─── Initialize Firebase ───────────────────────────────────────────────────────
+# ─── Initialize Firebase (Auth, Firestore, Realtime DB, Storage) ─────────────
 cred = credentials.Certificate(json.loads(FIREBASE_JSON))
-initialize_app(cred, {"storageBucket": STORAGE_BUCKET})
+initialize_app(cred, {
+    "storageBucket": STORAGE_BUCKET,
+    "databaseURL": FIREBASE_DB_URL
+})
 fs     = firestore.client()
 bucket = storage.bucket()
 
+# ─── Ensure dummy admin exists ─────────────────────────────────────────────────
+try:
+    admin_user = auth.get_user_by_email("rairorr@gmail.com")
+except firebase_admin._auth_utils.UserNotFoundError:
+    admin_user = auth.create_user(
+        email="rairorr@gmail.com",
+        email_verified=True,
+        password="123456",
+        display_name="Admin"
+    )
+auth.set_custom_user_claims(admin_user.uid, {"admin": True})
+
 # ─── Initialize Exa.ai ─────────────────────────────────────────────────────────
-exa = Exa(EXA_API_KEY)  #  [oai_citation:0‡docs.exa.ai](https://docs.exa.ai/reference/quickstart)
+exa = Exa(EXA_API_KEY)
 
 # ─── Initialize Gemini Client ─────────────────────────────────────────────────
-client = genai.Client(api_key=GEMINI_API_KEY)  #  [oai_citation:1‡Google Cloud](https://cloud.google.com/vertex-ai/generative-ai/docs/sdks/overview?utm_source=chatgpt.com)
+client = genai.Client(api_key=GEMINI_API_KEY)
 MODEL = "gemini-2.0-flash-001"
 
-# ─── (Optional) Setup a Faiss index for future behavioral targeting ─────────────
-# dim = 1536  # e.g., if you use a 1536‑dim embedding model
-# faiss_index = faiss.IndexFlatL2(dim)
-
 # ─── Flask App ────────────────────────────────────────────────────────────────
 app = Flask(__name__)
 CORS(app)
 
+# ─── Helpers ──────────────────────────────────────────────────────────────────
+def verify_id_token(f):
+    @wraps(f)
+    def wrapper(*args, **kwargs):
+        auth_header = request.headers.get("Authorization", "")
+        if not auth_header.startswith("Bearer "):
+            return jsonify({"error": "Missing or invalid Authorization header"}), 401
+        id_token = auth_header.split(" ")[1]
+        try:
+            decoded = auth.verify_id_token(id_token)
+            request.user = decoded
+        except Exception:
+            return jsonify({"error": "Invalid or expired token"}), 401
+        return f(*args, **kwargs)
+    return wrapper
+
+def require_admin(f):
+    @wraps(f)
+    def wrapper(*args, **kwargs):
+        if not getattr(request, "user", None) or not request.user.get("admin", False):
+            return jsonify({"error": "Admin privileges required"}), 403
+        return f(*args, **kwargs)
+    return wrapper
+
+def increment_query_count(ip):
+    doc = fs.collection("ip_queries").document(ip)
+    snap = doc.get()
+    if snap.exists:
+        count = snap.get("count") + 1
+        doc.update({"count": count})
+    else:
+        count = 1
+        doc.set({"count": 1})
+    return count
+
+def store_chat(ip, message, response):
+    """ Store each chat message/response pair in Realtime Database """
+    ref = firebase_db.reference("chats")
+    ref.push({
+        "ip": ip,
+        "message": message,
+        "response": response,
+        "timestamp": datetime.utcnow().isoformat()
+    })
+
+# ─── Routes ───────────────────────────────────────────────────────────────────
+
 @app.route("/chat", methods=["POST"])
 def chat():
-    payload      = request.get_json()
-    user_message = payload.get("message", "").strip()
-    user_ip      = request.remote_addr or "0.0.0.0"
+    user_ip    = request.remote_addr or "0.0.0.0"
+    count      = increment_query_count(user_ip)
+    need_login = count > 5
+
+    # optional user auth
+    token     = request.headers.get("Authorization", "")
+    user_info = None
+    if token.startswith("Bearer "):
+        try:
+            user_info = auth.verify_id_token(token.split(" ")[1])
+        except:
+            user_info = None
+
+    if need_login and not user_info:
+        return jsonify({
+            "error": "Please log in to continue after 5 free queries",
+            "login_required": True
+        }), 403
+
+    data         = request.get_json()
+    user_message = data.get("message", "").strip()
 
-    # 1️⃣ Classify intent
+    # 1) Classify intent
     classify_prompt = (
-        "Categorize the user's intent into one or both:\n"
-        " - self-help\n"
-        " - product_search\n"
-        "Return as a JSON list, e.g.: [\"self-help\"] or [\"product_search\"] or [\"self-help\",\"product_search\"]\n\n"
-        f"User message: \"{user_message}\""
-    )
-    classify_resp = client.models.generate_content(
-        model=MODEL,
-        contents=classify_prompt
+        "Categorize the user's intent as JSON list of one or both:\n"
+        "['self-help','product_search']\n\n"
+        f"User: \"{user_message}\""
     )
+    classify_resp = client.models.generate_content(model=MODEL, contents=classify_prompt)
     try:
         intents = json.loads(classify_resp.text)
-    except json.JSONDecodeError:
-        # fallback: treat as a catch‑all
+    except:
         intents = ["self-help"]
 
-    response_parts = []
+    parts = []
 
-    # 2️⃣ If self‑help requested, generate guidance
+    # 2) Self-help
     if "self-help" in intents:
-        help_prompt = (
-            "You are a friendly assistant. Provide clear, step‑by‑step guidance for:\n\n"
-            f"\"{user_message}\""
-        )
-        help_resp = client.models.generate_content(
-            model=MODEL,
-            contents=help_prompt
-        )
-        response_parts.append(help_resp.text.strip())
-
-        # Also dynamically lookup “what to buy” where applicable
-        # e.g. for “I want to bake a cake” → search “bake a cake ingredients”
-        search_query = user_message + " shop links for supplies"
-        exa_results = exa.search_and_contents(
-            search_query,
-            type="auto",
-            text=True
-        )
-        # take top 3 links
-        links = [r.url for r in exa_results.results[:3]]
-        suggestion_text = (
-            "If you need to pick up supplies, here are some helpful links:\n" +
-            "\n".join(f"- {url}" for url in links)
-        )
-        response_parts.append(suggestion_text)
+        help_prompt = f"You are a helpful assistant. Give step-by-step guidance for: \"{user_message}\""
+        help_resp   = client.models.generate_content(model=MODEL, contents=help_prompt)
+        parts.append(help_resp.text.strip())
 
-    # 3️⃣ If pure product_search requested, search Exa and wrap conversationally
+        # dynamic product links
+        search_q = user_message + " ingredients"
+        exa_res  = exa.search_and_contents(search_q, type="auto", text=True)
+        links    = [r.url for r in exa_res.results[:3]]
+        if links:
+            parts.append("Here are some useful links to get supplies:\n" + "\n".join(f"- {u}" for u in links))
+
+    # 3) Product search only
     if "product_search" in intents and "self-help" not in intents:
-        exa_results = exa.search_and_contents(
-            user_message,
-            type="auto",
-            text=True
-        )
-        links = [r.url for r in exa_results.results[:5]]
+        exa_res    = exa.search_and_contents(user_message, type="auto", text=True)
+        links      = [r.url for r in exa_res.results[:5]]
         rec_prompt = (
-            "You are a helpful shopping assistant. A user is looking for:\n\n"
-            f"\"{user_message}\"\n\n"
-            "From these links:\n" +
-            "\n".join(links) +
-            "\n\n"
-            "Suggest the most relevant products in a friendly, conversational way. if the user does not provide specifics ie budget, brand etc you provide the suggestions. dont prompt for more information. also always return the links."
+            "You are a shopping assistant. Suggest products for:\n"
+            f"\"{user_message}\"\n"
+            "Always include links:\n" + "\n".join(links)
         )
-        rec_resp = client.models.generate_content(
-            model=MODEL,
-            contents=rec_prompt
-        )
-        response_parts.append(rec_resp.text.strip())
+        rec_resp   = client.models.generate_content(model=MODEL, contents=rec_prompt)
+        parts.append(rec_resp.text.strip())
 
-    # 4️⃣ Fallback if nothing matched
-    if not response_parts:
-        default_resp = client.models.generate_content(
-            model=MODEL,
-            contents=f"Assist the user with: \"{user_message}\""
-        )
-        response_parts.append(default_resp.text.strip())
+    # 4) Fallback
+    if not parts:
+        default    = client.models.generate_content(model=MODEL, contents=f"Help the user with: \"{user_message}\"")
+        parts.append(default.text.strip())
+
+    final_response = "\n\n".join(parts)
+
+    # Store in Realtime Database
+    store_chat(user_ip, user_message, final_response)
+
+    return jsonify({"response": final_response})
+
+@app.route("/auth/signup", methods=["POST"])
+def signup():
+    data  = request.get_json()
+    email = data.get("email")
+    pwd   = data.get("password")
+    try:
+        user = auth.create_user(email=email, password=pwd)
+        fs.collection("users").document(user.uid).set({
+            "preferences": {},
+            "email": email
+        })
+        return jsonify({"message": "User created"}), 201
+    except Exception as e:
+        return jsonify({"error": str(e)}), 400
+
+@app.route("/dashboard", methods=["GET"])
+@verify_id_token
+def user_dashboard():
+    uid = request.user["uid"]
+    doc = fs.collection("users").document(uid).get()
+    if not doc.exists:
+        return jsonify({"error": "User not found"}), 404
+    data = doc.to_dict()
+    return jsonify({
+        "email": request.user.get("email"),
+        "preferences": data.get("preferences", {})
+    })
+
+@app.route("/dashboard/preferences", methods=["POST"])
+@verify_id_token
+def update_preferences():
+    uid   = request.user["uid"]
+    prefs = request.get_json().get("preferences", {})
+    fs.collection("users").document(uid).update({"preferences": prefs})
+    return jsonify({"message": "Preferences updated"})
+
+# ── Admin Routes ──────────────────────────────────────────────────────────────
+
+@app.route("/admin/upload_links", methods=["POST"])
+@verify_id_token
+@require_admin
+def upload_links():
+    if "file" not in request.files:
+        return jsonify({"error": "CSV file required"}), 400
+    f      = request.files["file"]
+    stream = StringIO(f.stream.read().decode("utf-8"))
+    reader = csv.DictReader(stream)
+    batch  = fs.batch()
+    for row in reader:
+        doc = fs.collection("sponsoredLinks").document()
+        batch.set(doc, {"keyword": row.get("keyword"), "url": row.get("url")})
+    batch.commit()
+    return jsonify({"message": "Links uploaded from CSV"})
 
-    # 5️⃣ Combine and return
-    final_text = "\n\n".join(response_parts)
-    return jsonify({"response": final_text})
+@app.route("/admin/add_link", methods=["POST"])
+@verify_id_token
+@require_admin
+def add_link():
+    data    = request.get_json()
+    keyword = data.get("keyword")
+    url     = data.get("url")
+    if not (keyword and url):
+        return jsonify({"error": "Both 'keyword' and 'url' are required"}), 400
+    fs.collection("sponsoredLinks").add({"keyword": keyword, "url": url})
+    return jsonify({"message": "Link added successfully"})
 
+@app.route("/admin/stats", methods=["GET"])
+@verify_id_token
+@require_admin
+def stats():
+    users        = auth.list_users().users
+    total_users  = len(users)
+    ip_docs      = fs.collection("ip_queries").stream()
+    total_queries = sum(doc.get("count", 0) for doc in ip_docs)
+    return jsonify({
+        "total_users": total_users,
+        "total_queries": total_queries
+    })
 
 if __name__ == "__main__":
-    # Host on all interfaces, port 7860
     app.run(host="0.0.0.0", port=7860, debug=True)