broadfield-dev/memvid
0
1import os2from flask import Flask, render_template, request, jsonify, stream_with_context, Response3from memvid_sdk import create, open as open_memvid4from huggingface_hub import hf_hub_download, upload_file, HfApi5 6app = Flask(__name__)7 8# CONFIGURATION9FILENAME = "knowledge.mv2"10HF_TOKEN = os.environ.get("HF_TOKEN")11DATASET_NAME = "memvid-storage" 12 13# Global variables14db = None15DB_PATH = os.path.abspath(FILENAME)16DATASET_REPO_ID = None17 18def get_repo_id():19 """Helper to dynamically resolve 'username/dataset_name'"""20 global DATASET_REPO_ID21 if DATASET_REPO_ID:22 return DATASET_REPO_ID23 24 if HF_TOKEN:25 try:26 api = HfApi(token=HF_TOKEN)27 username = api.whoami()['name']28 DATASET_REPO_ID = f"{username}/{DATASET_NAME}"29 return DATASET_REPO_ID30 except Exception as e:31 print(f"⚠️ Error getting username: {e}")32 return None33 return None34 35def init_db():36 """37 1. Ensure Dataset Exists.38 2. Try to download existing DB.39 3. Initialize Memvid.40 """41 global db, DATASET_REPO_ID42 43 repo_id = get_repo_id()44 45 # 1. Sync / Setup Cloud Storage46 if HF_TOKEN and repo_id:47 print(f"🔄 Checking cloud storage at {repo_id}...")48 api = HfApi(token=HF_TOKEN)49 50 try:51 # Create the repo if it doesn't exist52 api.create_repo(repo_id=repo_id, repo_type="dataset", exist_ok=True)53 54 # Check for file existence55 files = api.list_repo_files(repo_id=repo_id, repo_type="dataset")56 57 if FILENAME in files:58 downloaded_path = hf_hub_download(59 repo_id=repo_id,60 filename=FILENAME,61 repo_type="dataset",62 token=HF_TOKEN,63 local_dir=".", 64 local_dir_use_symlinks=False65 )66 print(f"✅ Downloaded database to {downloaded_path}")67 else:68 print("⚠️ Database file not found in repo. A new one will be created and synced.")69 70 except Exception as e:71 print(f"⚠️ Cloud sync warning: {e}")72 73 try:74 if os.path.exists(DB_PATH):75 db = open_memvid(DB_PATH, read_only=False) 76 else:77 db = create(DB_PATH) 78 79 except ImportError:80 from memvid_sdk import Memvid81 if os.path.exists(DB_PATH):82 db = Memvid()83 db.open(DB_PATH) 84 else:85 db = Memvid()86 db.create(DB_PATH)87 88def sync_to_hub():89 """Uploads the local .mv2 file back to Hugging Face"""90 repo_id = get_repo_id()91 92 if not HF_TOKEN or not repo_id:93 print("⚠️ No HF_TOKEN or Repo ID found. Skipping sync.")94 return95 96 try:97 print("☁️ Syncing to Hub...")98 upload_file(99 path_or_fileobj=DB_PATH,100 path_in_repo=FILENAME,101 repo_id=repo_id,102 repo_type="dataset",103 token=HF_TOKEN,104 commit_message="Memvid: Auto-save memory update"105 )106 print("✅ Sync complete.")107 except Exception as e:108 print(f"❌ Sync failed: {e}")109 110# Initialize on startup111init_db()112 113@app.route('/')114def index():115 return render_template('index.html')116 117@app.route('/add', methods=['POST'])118def add_memory():119 # 1. Setup Validation120 global db121 content = request.form.get('content')122 123 if not content:124 return jsonify({"error": "No content provided"}), 400125 126 # 2. Define the Stream Generator127 def generate():128 try:129 # Step A: Re-init if needed inside the stream130 global db131 if not db:132 init_db()133 if not db:134 yield '{"status": "error", "message": "Database init failed"}\n'135 return136 137 # Step B: Database Put138 yield '{"status": "processing", "message": "Ingesting content..."}\n'139 140 payload = {141 "text": content,142 "labels": ["web-entry"], 143 "title": "User Memory"144 }145 db.put(payload)146 147 # Step C: Flush to Disk148 yield '{"status": "processing", "message": "Flushing to disk..."}\n'149 del db150 db = None151 152 # Step D: Sync153 yield '{"status": "processing", "message": "Syncing to cloud (this may take a moment)..."}\n'154 sync_to_hub()155 156 # Step E: Reload157 yield '{"status": "processing", "message": "Reloading index..."}\n'158 init_db()159 160 # Final Success Message161 yield '{"status": "success", "message": "Memory added and synced."}\n'162 163 except Exception as e:164 # Capture any errors during the process165 yield f'{{"status": "error", "message": "{str(e)}"}}\n'166 167 # 3. Return the Stream168 return Response(stream_with_context(generate()), mimetype='application/x-ndjson')169 170@app.route('/search', methods=['POST'])171def search_memory():172 if not db:173 return jsonify({"error": "Database not initialized"}), 500174 175 query = request.form.get('query')176 if not query:177 return jsonify({"error": "No query provided"}), 400178 179 try:180 # 1. Search181 response = db.find(query)182 183 # 2. Parse & Clean184 clean_results = []185 hits = response.get('hits', [])186 187 for hit in hits:188 score = hit.get('score', 0.0)189 if score < 0.65: continue190 191 # --- CLEANING LOGIC ---192 raw_snippet = hit.get('snippet', '')193 194 lines = raw_snippet.split('\n')195 content_lines = [196 line for line in lines 197 if not line.strip().startswith(('title:', 'tags:', 'labels:', 'extractous_metadata:'))198 ]199 clean_text = "\n".join(content_lines).strip()200 201 tags = hit.get('tags', [])202 labels = hit.get('labels', [])203 204 clean_results.append({205 "title": hit.get('title') or "Untitled Memory",206 "text": clean_text, 207 "tags": tags, 208 "labels": labels, 209 "date": hit.get('created_at', ''),210 "score": f"{score:.2f}"211 })212 213 return jsonify({"success": True, "results": clean_results})214 except Exception as e:215 return jsonify({"error": str(e)}), 500216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231if __name__ == '__main__':232 app.run(host='0.0.0.0', port=7860)