Feat: implement dataset-based data persistence

Dockerfile

@@ -57,9 +57,16 @@ COPY --from=builder --chown=nextjs:nodejs /app/.next/static ./.next/static
 # Create directory if it doesn't exist
 # We copy existing stores so the demo works out of the box
 COPY --from=builder --chown=nextjs:nodejs /app/vector_store ./vector_store
+# Copy database if it exists
+COPY --from=builder --chown=nextjs:nodejs /app/rag-kb.db ./rag-kb.db
 # Copy source documents
 COPY --from=builder --chown=nextjs:nodejs /app/data ./data
 
+# Copy scripts for data sync
+COPY --chown=nextjs:nodejs scripts/download_data.py ./scripts/
+COPY --chown=nextjs:nodejs scripts/start.sh ./scripts/
+RUN chmod +x ./scripts/start.sh
+
 USER nextjs
 
 EXPOSE 7860
@@ -67,4 +74,4 @@ EXPOSE 7860
 ENV PORT 7860
 ENV HOSTNAME "0.0.0.0"
 
-CMD ["node", "server.js"]
+CMD ["./scripts/start.sh"]

scripts/download_data.py

@@ -0,0 +1,51 @@
+import os
+import shutil
+import sys
+from huggingface_hub import snapshot_download
+
+def main():
+    repo_id = os.environ.get("HF_DATASET_REPO")
+    if not repo_id:
+        print("HF_DATASET_REPO environment variable not set. Skipping data download.")
+        return
+
+    print(f"Downloading data from dataset: {repo_id}...")
+    try:
+        # Download to a temporary directory
+        local_dir = snapshot_download(repo_id=repo_id, repo_type="dataset", allow_patterns=["rag-kb.db", "vector_store/*"])
+        
+        print(f"Data downloaded to: {local_dir}")
+        
+        # Move rag-kb.db
+        db_source = os.path.join(local_dir, "rag-kb.db")
+        if os.path.exists(db_source):
+            if os.path.exists("rag-kb.db"):
+                print("Backing up existing rag-kb.db...")
+                shutil.move("rag-kb.db", "rag-kb.db.bak")
+            
+            print("Restoring rag-kb.db...")
+            shutil.copy(db_source, ".")
+        else:
+            print("Warning: rag-kb.db not found in dataset.")
+
+        # Move vector_store
+        vs_source = os.path.join(local_dir, "vector_store")
+        if os.path.exists(vs_source):
+            if os.path.exists("vector_store"):
+                print("Removing existing vector_store...")
+                shutil.rmtree("vector_store")
+            
+            print("Restoring vector_store...")
+            shutil.copytree(vs_source, "vector_store")
+        else:
+            print("Warning: vector_store not found in dataset.")
+
+        print("✅ Data restoration completed.")
+
+    except Exception as e:
+        print(f"Error downloading data: {e}")
+        # Don't exit with error, allow app to start empty if download fails
+        pass
+
+if __name__ == "__main__":
+    main()

scripts/start.sh

@@ -0,0 +1,13 @@
+#!/bin/bash
+
+# Try to download data if configured
+if [ -n "$HF_DATASET_REPO" ]; then
+    echo "Attempting to restore data from Hugging Face Dataset..."
+    python3 scripts/download_data.py
+else
+    echo "HF_DATASET_REPO not set. Starting with local/empty data."
+fi
+
+# Start the application
+echo "Starting Next.js server..."
+exec node server.js

scripts/upload_data.py

@@ -0,0 +1,75 @@
+import os
+import sys
+from huggingface_hub import HfApi, create_repo
+
+def main():
+    print("=== Rag-KB Data Uploader ===")
+    print("This script helps you upload your local database and vector store to a Hugging Face Dataset.")
+    print("This ensures your data persists across Space deployments.\n")
+
+    # Get Token
+    token = os.environ.get("HF_TOKEN")
+    if not token:
+        token = input("Please enter your Hugging Face Write Token: ").strip()
+    
+    if not token:
+        print("Error: Token is required.")
+        sys.exit(1)
+
+    # Get Repo ID
+    default_repo = "duqing2026/rag-kb-data"
+    repo_id = input(f"Enter Dataset Repo ID (default: {default_repo}): ").strip() or default_repo
+
+    api = HfApi(token=token)
+
+    # Create Repo if not exists
+    print(f"\nChecking/Creating dataset repository: {repo_id}...")
+    try:
+        create_repo(repo_id=repo_id, repo_type="dataset", exist_ok=True, token=token, private=True)
+        print("Repository ready.")
+    except Exception as e:
+        print(f"Error creating repository: {e}")
+        sys.exit(1)
+
+    # Upload files
+    files_to_upload = [
+        ("rag-kb.db", "rag-kb.db"),
+    ]
+    
+    # Add vector store files
+    vector_store_path = "vector_store"
+    if os.path.exists(vector_store_path):
+        for root, dirs, files in os.walk(vector_store_path):
+            for file in files:
+                local_path = os.path.join(root, file)
+                # Keep the directory structure relative to vector_store
+                # We want vector_store/xxx in the repo
+                repo_path = local_path
+                files_to_upload.append((local_path, repo_path))
+    else:
+        print("Warning: vector_store directory not found.")
+
+    print(f"\nStarting upload of {len(files_to_upload)} files...")
+    
+    for local_path, repo_path in files_to_upload:
+        if os.path.exists(local_path):
+            print(f"Uploading {local_path} -> {repo_path}...")
+            try:
+                api.upload_file(
+                    path_or_fileobj=local_path,
+                    path_in_repo=repo_path,
+                    repo_id=repo_id,
+                    repo_type="dataset",
+                    token=token
+                )
+            except Exception as e:
+                print(f"Failed to upload {local_path}: {e}")
+        else:
+            print(f"Skipping {local_path} (not found)")
+
+    print("\n✅ Upload completed successfully!")
+    print(f"Please set the following environment variable in your Hugging Face Space settings:")
+    print(f"HF_DATASET_REPO={repo_id}")
+
+if __name__ == "__main__":
+    main()