Update app.py

app.py

@@ -1,30 +1,95 @@
 import os
 from flask import Flask, render_template, request, jsonify
-# Importing Memvid based on the SDK package name
-from memvid_sdk import Memvid 
+from memvid_sdk import Memvid
+from huggingface_hub import hf_hub_download, upload_file, HfApi
 
 app = Flask(__name__)
 
 # CONFIGURATION
-# We store the .mv2 file in a writable path. 
-# On HF Spaces, /tmp is writable, or you can use the current directory 
-# if you don't need persistence across reboots.
-DB_PATH = "knowledge.mv2"
-
-# Initialize Memvid
-# We use a global variable to hold the reference
-try:
-    # Attempt to open or create. 
-    # Python SDK usually follows: Memvid(path) or Memvid.create(path)
-    # Adjust this line based on exact SDK docs.
-    if os.path.exists(DB_PATH):
-        db = Memvid.open(DB_PATH)
-    else:
-        db = Memvid.create(DB_PATH)
-    print(f"Memvid loaded at {DB_PATH}")
-except Exception as e:
-    print(f"Error initializing Memvid: {e}")
-    db = None
+# Format: "username/dataset-name"
+DATASET_REPO_ID = "YOUR_USERNAME/memvid-storage" 
+FILENAME = "knowledge.mv2"
+HF_TOKEN = os.environ.get("HF_TOKEN")
+
+# Global DB reference
+db = None
+DB_PATH = os.path.abspath(FILENAME)
+
+def init_db():
+    """
+    1. Try to download existing DB from HF Hub.
+    2. If not found, create a new one locally.
+    3. Load Memvid.
+    """
+    global db
+    
+    # 1. Try to sync from Hub
+    if HF_TOKEN:
+        print(f"🔄 Attempting to download {FILENAME} from {DATASET_REPO_ID}...")
+        try:
+            # This downloads the file to the local cache and returns the path
+            # We copy it or link it to our working dir if needed, but usually 
+            # we just want it in the current directory for Memvid to write to.
+            
+            # Check if file exists in repo first
+            api = HfApi(token=HF_TOKEN
+            username = api.whoami()
+            print(username)
+            print(username.name)
+            files = api.list_repo_files(repo_id=DATASET_REPO_ID, repo_type="dataset")
+            
+            if FILENAME in files:
+                downloaded_path = hf_hub_download(
+                    repo_id=DATASET_REPO_ID,
+                    filename=FILENAME,
+                    repo_type="dataset",
+                    token=HF_TOKEN,
+                    local_dir=".", # Download to current directory
+                    local_dir_use_symlinks=False # We need the actual file to write to it
+                )
+                print(f"✅ Downloaded database to {downloaded_path}")
+            else:
+                print("⚠️ Database file not found in repo. Creating new one.")
+        except Exception as e:
+            print(f"⚠️ Could not download from Hub (might be first run): {e}")
+
+    # 2. Open or Create Memvid
+    try:
+        # Memvid 2.0 pattern:
+        # If file exists, open it. If not, create it.
+        if os.path.exists(DB_PATH):
+            db = Memvid.open(DB_PATH)
+            print(f"📂 Memvid opened at {DB_PATH}")
+        else:
+            db = Memvid.create(DB_PATH)
+            print(f"✨ New Memvid created at {DB_PATH}")
+            
+    except Exception as e:
+        print(f"❌ CRITICAL ERROR initializing Memvid: {e}")
+        db = None
+
+def sync_to_hub():
+    """Uploads the local .mv2 file back to Hugging Face"""
+    if not HF_TOKEN:
+        print("⚠️ No HF_TOKEN found. Skipping sync.")
+        return
+
+    try:
+        print("☁️ Syncing to Hub...")
+        upload_file(
+            path_or_fileobj=DB_PATH,
+            path_in_repo=FILENAME,
+            repo_id=DATASET_REPO_ID,
+            repo_type="dataset",
+            token=HF_TOKEN,
+            commit_message="Memvid: Auto-save memory update"
+        )
+        print("✅ Sync complete.")
+    except Exception as e:
+        print(f"❌ Sync failed: {e}")
+
+# Initialize on startup
+init_db()
 
 @app.route('/')
 def index():
@@ -32,25 +97,34 @@ def index():
 
 @app.route('/add', methods=['POST'])
 def add_memory():
+    global db
     if not db:
-        return jsonify({"error": "Database not initialized"}), 500
+        # Try to re-init if it failed before
+        init_db()
+        if not db:
+            return jsonify({"error": "Database could not be initialized. Check logs."}), 500
 
     content = request.form.get('content')
-    tags = request.form.get('tags', '')
-    
+    tags = request.form.get('tags', '') # Not used in basic put, but good for expansion
+
     if not content:
         return jsonify({"error": "No content provided"}), 400
 
     try:
-        # Assuming the Python SDK has a .put() or .add() method
-        # and accepts string content. 
-        # We simulate "PutOptions" via kwargs or a dict.
-        db.put(content, tags=tags)
+        # Add the memory
+        # Note: Check if your SDK version requires explicit transaction/commit
+        # Some versions use db.put(content), others db.add(content)
+        db.put(content) 
         
-        # Depending on the SDK, you might need an explicit commit
-        # db.commit() 
+        # IMPORTANT: Force a commit/flush to disk before uploading
+        # If the SDK has a .commit() or .flush(), call it here.
+        if hasattr(db, 'commit'):
+            db.commit()
+
+        # Sync back to cloud
+        sync_to_hub()
         
-        return jsonify({"success": True, "message": "Memory added successfully."})
+        return jsonify({"success": True, "message": "Memory added and synced to cloud."})
     except Exception as e:
         return jsonify({"error": str(e)}), 500
 
@@ -60,21 +134,19 @@ def search_memory():
         return jsonify({"error": "Database not initialized"}), 500
 
     query = request.form.get('query')
-    
     if not query:
         return jsonify({"error": "No query provided"}), 400
 
     try:
-        # Perform search. Assuming .search returns a list of result objects
+        # Search
         results = db.search(query, top_k=5)
         
-        # Format results for JSON
-        # Assuming results have .text and .score attributes
+        # Transform results based on SDK object structure
         formatted_results = []
         for hit in results:
             formatted_results.append({
-                "text": hit.text,
-                # "score": hit.score # specific attributes depend on SDK
+                "text": hit.text, 
+                # "score": hit.score # Uncomment if available
             })
             
         return jsonify({"success": True, "results": formatted_results})
@@ -82,5 +154,4 @@ def search_memory():
         return jsonify({"error": str(e)}), 500
 
 if __name__ == '__main__':
-    # Hugging Face Spaces runs on port 7860 by default
     app.run(host='0.0.0.0', port=7860)