Add checkpoint resume + batch upsert to handle quota limits

src/vector_store.py

@@ -42,28 +42,65 @@ def get_or_create_collection(client):
     return client.get_collection(COLLECTION_NAME)
 
 
+CHECKPOINT_PATH = "embed_checkpoint.pkl"
+
+
 def add_chunks(chunks: list[dict]):
-    """Embed và lưu chunks vào Qdrant + build BM25 index."""
+    """Embed và lưu chunks vào Qdrant + build BM25 index. Có checkpoint để resume khi lỗi quota."""
     client = get_client()
     get_or_create_collection(client)
 
     texts = [c["text"] for c in chunks]
     metadatas = [c["metadata"] for c in chunks]
 
-    print(f"Embedding {len(chunks)} chunks...")
-    all_embeddings = []
-    for i in tqdm(range(0, len(chunks), BATCH_SIZE)):
+    # Load checkpoint nếu có
+    done_batches: dict[int, list] = {}
+    if Path(CHECKPOINT_PATH).exists():
+        with open(CHECKPOINT_PATH, "rb") as f:
+            done_batches = pickle.load(f)
+        print(f"Resume từ checkpoint: đã xong {len(done_batches)} batches (~{len(done_batches) * BATCH_SIZE} chunks)")
+
+    pending = [i for i in range(0, len(chunks), BATCH_SIZE) if i not in done_batches]
+    print(f"Embedding {len(chunks)} chunks... ({len(pending)} batches còn lại)")
+
+    for i in tqdm(pending):
         batch_texts = texts[i:i + BATCH_SIZE]
-        embeddings = embed_texts(batch_texts)
-        all_embeddings.extend(embeddings)
+
+        for attempt in range(3):
+            try:
+                embeddings = embed_texts(batch_texts)
+                break
+            except Exception as e:
+                if "429" in str(e) or "RESOURCE_EXHAUSTED" in str(e):
+                    wait = 60 * (attempt + 1)
+                    print(f"\n[Quota] Lỗi rate limit, đợi {wait}s rồi thử lại (lần {attempt + 1}/3)...")
+                    time.sleep(wait)
+                    if attempt == 2:
+                        print("[Quota] Hết quota ngày hôm nay. Chạy lại vào ngày mai, sẽ resume từ checkpoint.")
+                        raise
+                else:
+                    raise
+
+        done_batches[i] = embeddings
+        with open(CHECKPOINT_PATH, "wb") as f:
+            pickle.dump(done_batches, f)
+
         time.sleep(BATCH_SIZE * 1.5)
 
-    # Lưu vào Qdrant
-    points = [
-        PointStruct(id=i, vector=emb, payload={"text": texts[i], **metadatas[i]})
-        for i, emb in enumerate(all_embeddings)
-    ]
-    client.upsert(collection_name=COLLECTION_NAME, points=points)
+    # Ghép tất cả embeddings theo thứ tự
+    all_embeddings = []
+    for i in range(0, len(chunks), BATCH_SIZE):
+        all_embeddings.extend(done_batches[i])
+
+    # Lưu vào Qdrant theo batch để tránh timeout
+    UPSERT_BATCH = 200
+    print(f"Upserting {len(all_embeddings)} points vào Qdrant...")
+    for i in tqdm(range(0, len(all_embeddings), UPSERT_BATCH)):
+        batch_points = [
+            PointStruct(id=i + j, vector=all_embeddings[i + j], payload={"text": texts[i + j], **metadatas[i + j]})
+            for j in range(min(UPSERT_BATCH, len(all_embeddings) - i))
+        ]
+        client.upsert(collection_name=COLLECTION_NAME, points=batch_points)
 
     # Build BM25 index
     print("Building BM25 index...")
@@ -72,6 +109,8 @@ def add_chunks(chunks: list[dict]):
     with open(BM25_PATH, "wb") as f:
         pickle.dump({"bm25": bm25, "texts": texts, "metadatas": metadatas}, f)
 
+    # Xóa checkpoint sau khi hoàn thành
+    Path(CHECKPOINT_PATH).unlink(missing_ok=True)
     print(f"Đã lưu {len(chunks)} chunks vào Qdrant + BM25")