Update utils/memory_utils.py

utils/memory_utils.py

@@ -4,7 +4,7 @@ import json
 import time
 import shutil
 import gradio as gr
-from huggingface_hub import HfApi
+from huggingface_hub import HfApi, snapshot_download
 
 # LlamaIndex Imports
 try:
@@ -13,7 +13,6 @@ except ImportError:
     pass
 
 # Local Imports setup
-# فرض بر این است که این فایل در utils/ است و memory_bank یک پوشه عقب‌تر است
 bank_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), '../memory_bank')
 sys.path.append(bank_path)
 
@@ -28,20 +27,25 @@ except ImportError:
 # 1. تنظیمات مسیر و هاب (Cloud Config)
 # ==========================================
 
-# !!! مهم: این قسمت را با نام کاربری و نام دیتاست خودتان پر کنید !!!
-# مثال: "ali-b/my-bot-memory"
-REPO_ID = "Keyvan1986/Emma-memory-storage" 
+# !!! مهم: نام کاربری و دیتاست خود را اینجا وارد کنید !!!
+REPO_ID = "YOUR_USERNAME/YOUR_DATASET_NAME"
 REPO_TYPE = "dataset"
 
-# دریافت توکن از متغیرهای محیطی (Secrets)
+# دریافت توکن
 HF_TOKEN = os.environ.get("HF_TOKEN")
 
 CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
 BASE_DIR = os.path.dirname(CURRENT_DIR) 
 MEMORIES_DIR = os.path.join(BASE_DIR, "memories")
+
+# فایل اصلی JSON
 MEMORY_FILE_NAME = "update_memory_0512_eng.json"
 MEMORY_FILE_PATH = os.path.join(MEMORIES_DIR, MEMORY_FILE_NAME)
 
+# پوشه ایندکس‌های LlamaIndex
+MEMORY_INDEX_DIR_NAME = "memory_index"
+MEMORY_INDEX_PATH = os.path.join(MEMORIES_DIR, MEMORY_INDEX_DIR_NAME)
+
 # اطمینان از وجود پوشه در زمان ایمپورت
 os.makedirs(MEMORIES_DIR, exist_ok=True)
 
@@ -52,44 +56,61 @@ os.makedirs(MEMORIES_DIR, exist_ok=True)
 
 def push_to_hub():
     """
-    فایل حافظه را به دیتاست هاگینگ‌فیس آپلود می‌کند.
+    1. فایل JSON حافظه را آپلود می‌کند.
+    2. پوشه memory_index (شامل وکتورها) را آپلود می‌کند.
     """
     if not HF_TOKEN:
-        print("⚠️ Warning: HF_TOKEN not found. Memory is NOT saved to the cloud.")
+        print("⚠️ Warning: HF_TOKEN not found. Cloud sync skipped.")
         return
     
     if "YOUR_USERNAME" in REPO_ID:
-        print("⚠️ Warning: REPO_ID is not set in memory_utils.py. Cloud sync skipped.")
+        print("⚠️ Warning: REPO_ID is not set. Cloud sync skipped.")
         return
 
     try:
         api = HfApi(token=HF_TOKEN)
-        # آپلود فایل
+        
+        # A) آپلود فایل JSON اصلی
+        print("☁️ Uploading JSON memory file...")
         api.upload_file(
             path_or_fileobj=MEMORY_FILE_PATH,
             path_in_repo=MEMORY_FILE_NAME,
             repo_id=REPO_ID,
             repo_type=REPO_TYPE,
-            commit_message=f"Auto-save: {time.strftime('%Y-%m-%d %H:%M:%S')}"
+            commit_message=f"Auto-save JSON: {time.strftime('%Y-%m-%d %H:%M:%S')}"
         )
-        print(f"☁️ Synced to Hub: {REPO_ID}")
+
+        # B) آپلود پوشه ایندکس‌ها (Vector Stores)
+        if os.path.exists(MEMORY_INDEX_PATH):
+            print("☁️ Uploading Index folder (Vectors)...")
+            api.upload_folder(
+                folder_path=MEMORY_INDEX_PATH,
+                path_in_repo=MEMORY_INDEX_DIR_NAME, # در دیتاست هم داخل پوشه memory_index می‌رود
+                repo_id=REPO_ID,
+                repo_type=REPO_TYPE,
+                commit_message=f"Auto-save Indices: {time.strftime('%Y-%m-%d %H:%M:%S')}"
+            )
+            
+        print(f"✅ Synced all data to Hub: {REPO_ID}")
+
     except Exception as e:
         print(f"❌ Cloud Sync Error: {e}")
 
 def pull_from_hub():
     """
-    آخرین نسخه حافظه را دانلود می‌کند.
+    1. فایل JSON را دانلود می‌کند.
+    2. پوشه memory_index را دانلود می‌کند.
     """
     if not HF_TOKEN or "YOUR_USERNAME" in REPO_ID:
         return
 
     try:
         api = HfApi(token=HF_TOKEN)
-        # بررسی وجود فایل
         files = api.list_repo_files(repo_id=REPO_ID, repo_type=REPO_TYPE)
         
+        # A) دانلود فایل JSON
         if MEMORY_FILE_NAME in files:
-            print(f"📥 Downloading memory from Hub...")
+            print(f"📥 Downloading JSON memory...")
             api.hf_hub_download(
                 repo_id=REPO_ID,
                 repo_type=REPO_TYPE,
@@ -97,9 +118,26 @@ def pull_from_hub():
                 local_dir=MEMORIES_DIR,
                 force_download=True
             )
-            print("✅ Download complete.")
         else:
-            print("🆕 Memory file not found on Hub. Starting fresh.")
+            print("🆕 Memory JSON not found on Hub. Creating fresh.")
+
+        # B) دانلود پوشه ایندکس‌ها (اگر وجود داشته باشد)
+        # بررسی می‌کنیم آیا پوشه memory_index در فایل‌های ریپو هست یا نه
+        has_index = any(f.startswith(MEMORY_INDEX_DIR_NAME) for f in files)
+        
+        if has_index:
+            print(f"📥 Downloading Index folder (Vectors)...")
+            snapshot_download(
+                repo_id=REPO_ID,
+                repo_type=REPO_TYPE,
+                local_dir=MEMORIES_DIR,
+                allow_patterns=f"{MEMORY_INDEX_DIR_NAME}/**", # فقط پوشه ایندکس را بگیر
+                force_download=True # برای اطمینان از بروزرسانی
+            )
+            print("✅ Indices downloaded.")
+        else:
+            print("🆕 No existing indices on Hub. Will build locally.")
+
     except Exception as e:
         print(f"⚠️ Cloud Download Error: {e}")
 
@@ -109,19 +147,18 @@ def pull_from_hub():
 
 def save_memory(data):
     """
-    1. ذخیره در دیسک محلی (Atomic Write)
-    2. ارسال به فضای ابری (Push to Hub)
+    ذخیره اتمیک JSON + همگام‌سازی کامل با ابر
     """
     try:
         os.makedirs(MEMORIES_DIR, exist_ok=True)
         
-        # 1. Local Save
+        # 1. Local Save (JSON)
         temp_file = MEMORY_FILE_PATH + ".tmp"
         with open(temp_file, "w", encoding="utf-8") as f:
             json.dump(data, f, ensure_ascii=False, indent=4)
         os.replace(temp_file, MEMORY_FILE_PATH)
         
-        # 2. Cloud Save
+        # 2. Cloud Save (JSON + Folder)
         push_to_hub()
         
     except Exception as e:
@@ -130,17 +167,17 @@ def save_memory(data):
 
 def load_memory():
     """
-    1. تلاش برای دانلود از فضای ابری
-    2. خواندن از دیسک محلی
-    3. اگر فایل نبود، ساخت فایل خالی
+    لود کردن (همراه با دانلود اولیه از ابر)
     """
-    # 1. Cloud Pull
+    # 1. Cloud Pull (JSON + Indices)
     pull_from_hub()
 
     # 2. Local Read / Create
     if not os.path.exists(MEMORY_FILE_PATH):
-        print(f"⚠️ Memory file missing locally. Creating new at: {MEMORY_FILE_PATH}")
-        save_memory({}) # فایل خالی می‌سازد و آپلود می‌کند
+        print(f"⚠️ Memory file missing locally. Creating new.")
+        # فایل خالی می‌سازیم اما هنوز آپلود نمی‌کنیم تا وقتی دیتایی بیاید
+        with open(MEMORY_FILE_PATH, "w", encoding="utf-8") as f:
+            json.dump({}, f)
         return {}
 
     try:
@@ -159,10 +196,9 @@ def load_memory():
 # ==========================================
 
 def enter_name_llamaindex(name, memory_data, data_args, update_memory_index=True):
-    """
-    کاربر را ایجاد کرده و فایل را فوراً ذخیره می‌کند.
-    """
     is_new_user = False
+    
+    # اطمینان از وجود ساختار کاربر در دیکشنری
     if name not in memory_data:
         print(f"🆕 Creating new user profile for: {name}")
         memory_data[name] = {
@@ -173,38 +209,45 @@ def enter_name_llamaindex(name, memory_data, data_args, update_memory_index=True
             "profile": {}
         }
         is_new_user = True
+        # اگر کاربر جدید است، همان اول فایل جیسون را ذخیره کن
+        save_memory(memory_data)
     
     user_memory = memory_data[name]
 
-    # اگر کاربر جدید است، سریعاً ذخیره و سینک کن
-    if is_new_user:
-        save_memory(memory_data)
-
-    # 2. مدیریت ایندکس‌های وکتوری (LlamaIndex)
+    # --- مدیریت ایندکس‌های LlamaIndex ---
     sessions_memory = None
     episodic_memory = None
     semantic_memory = None
 
-    base_index_dir = os.path.join(MEMORIES_DIR, "memory_index", "llamaindex", name)
+    base_index_dir = os.path.join(MEMORY_INDEX_PATH, "llamaindex", name)
     
     sessions_path = os.path.join(base_index_dir, "sessions_index.json")
     episodic_path = os.path.join(base_index_dir, "episodic_memory_index.json")
     semantic_path = os.path.join(base_index_dir, "semantic_memory_index.json")
 
+    # بررسی اینکه آیا فایل‌های ایندکس به صورت لوکال وجود دارند (یا از ابر دانلود شده‌اند)
     indices_exist = (
         os.path.exists(sessions_path) and 
         os.path.exists(episodic_path) and 
         os.path.exists(semantic_path)
     )
 
+    # اگر ایندکس وجود ندارد یا درخواست آپدیت داریم، بسازیم
     if update_memory_index or not indices_exist:
         try:
+            print(f"⚙️ Building/Updating indices locally for {name}...")
             os.makedirs(base_index_dir, exist_ok=True)
             build_memory_index(memory_data, data_args, name=name)
+            
+            # **نکته مهم:** الان که فایل‌های ایندکس ساخته شدند، باید آنها را به ابر بفرستیم
+            # تا اگر سرور ریستارت شد، این فایل‌ها نپرند.
+            print("☁️ Syncing new indices to cloud...")
+            push_to_hub()
+            
         except Exception as e:
             print(f"Warning: Error building memory index (skipping): {e}")
 
-    # بارگذاری ایندکس‌ها
+    # بارگذاری ایندکس‌ها در حافظه رم (RAM) برای استفاده مدل
     if os.path.exists(sessions_path):
         try:
             storage_context = StorageContext.from_defaults(persist_dir=sessions_path)
@@ -227,11 +270,16 @@ def enter_name_llamaindex(name, memory_data, data_args, update_memory_index=True
 
 
 def summarize_memory_event_personality(data_args, memory_data, user_name):
-    # ذخیره قبل از عملیات سنگین
+    # قبل از عملیات سنگین، وضعیت فعلی را ذخیره کن
     save_memory(memory_data)
     
     try:
+        # این تابع احتمالا فایل‌های ایندکس را تغییر می‌دهد
         updated_memory = summarize_memory(MEMORY_FILE_PATH, user_name, language='en')
+        
+        # چون خلاصه سازی ممکن است ایندکس‌ها را تغییر دهد، یک سینک نهایی انجام می‌دهیم
+        push_to_hub()
+        
         return updated_memory.get(user_name, {})
     except Exception as e:
         print(f"Error in summarize_memory: {e}")