batch_import.py

#!/usr/bin/env python3
"""
批量導入現有記憶到 Qdrant 向量庫
將所有核心記憶文件分塊 → embedding → 寫入 openclaw_mem

v2.0 — 使用 memory_utils 共享模塊
修復：UUID 碰撞 Bug / 斷點續傳 / Markdown-Aware 分塊
"""

import os
import sys
import json
from datetime import datetime
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent))

from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct
from sentence_transformers import SentenceTransformer

from memory_utils import (
    chunk_text, get_tags_from_path, collect_files,
    generate_content_id,
    WORKSPACE, MEMORY_DIR, MEMORY_MD, PROJECT_ROOT,
)

# ── Config ──────────────────────────────────────────
QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333")
MODEL_NAME = os.getenv("EMBEDDING_MODEL", "BAAI/bge-m3")
COLLECTION = os.getenv("COLLECTION", "openclaw_mem")

# ── Checkpoint file for resume support ──
CHECKPOINT_FILE = PROJECT_ROOT / "batch_import_checkpoint.json"


def load_checkpoint() -> dict:
    """Load import checkpoint for resume support"""
    if CHECKPOINT_FILE.exists():
        try:
            return json.loads(CHECKPOINT_FILE.read_text())
        except json.JSONDecodeError:
            return {}
    return {"imported_ids": [], "last_file": None, "last_chunk_index": 0}


def save_checkpoint(checkpoint: dict):
    """Save import checkpoint (atomic)"""
    from memory_utils import atomic_write_json
    atomic_write_json(CHECKPOINT_FILE, checkpoint)


def main():
    print("=" * 60)
    print("📦 向量記憶批量導入工具 v2.0")
    print("=" * 60)

    # ── Load model ──
    print(f"\n🔄 載入 Embedding 模型: {MODEL_NAME}...")
    model = SentenceTransformer(MODEL_NAME, device="mps")
    print(f"✅ 模型已載入")

    # ── Connect Qdrant ──
    print(f"\n🔄 連接 Qdrant: {QDRANT_URL}")
    client = QdrantClient(url=QDRANT_URL)
    print(f"✅ 已連接")

    # ── Ensure collection ──
    collections = [c.name for c in client.get_collections().collections]
    if COLLECTION not in collections:
        print(f"\n📦 創建 collection: {COLLECTION}")
        client.create_collection(
            collection_name=COLLECTION,
            vectors_config=VectorParams(
                size=model.get_sentence_embedding_dimension(),
                distance=Distance.COSINE,
                on_disk=True
            )
        )

    # Get current count
    stats = client.get_collection(COLLECTION)
    existing_count = stats.points_count
    print(f"\n📊 {COLLECTION} 現有: {existing_count} 條記憶")

    # ── Load checkpoint ──
    checkpoint = load_checkpoint()
    imported_ids = set(checkpoint.get("imported_ids", []))
    if imported_ids:
        print(f"🔄 發現斷點：已導入 {len(imported_ids)} 條，將從上次位置繼續")

    # ── Collect files ──
    print(f"\n🔍 掃描記憶文件...")
    files = collect_files()
    print(f"   找到 {len(files)} 個文件")

    # ── Chunk all files ──
    print(f"\n📝 分塊處理（Markdown-Aware）...")
    all_chunks = []
    skipped_files = 0
    for f in files:
        try:
            content = f.read_text(encoding="utf-8")
        except Exception as e:
            print(f"  ⚠️  讀取失敗: {f.name}: {e}")
            skipped_files += 1
            continue

        rel_path = str(f.relative_to(WORKSPACE))
        chunks = chunk_text(content, rel_path, markdown_aware=True)

        if not chunks:
            continue

        tags = get_tags_from_path(f)
        for c in chunks:
            c["tags"] = tags
            c["filename"] = f.name

        all_chunks.extend(chunks)

    print(f"   共 {len(all_chunks)} 個 chunks（跳過 {skipped_files} 個文件）")

    if not all_chunks:
        print("\n❌ 沒有可導入的內容")
        return

    # ── Embed & Import (with checkpoint resume) ──
    print(f"\n🧠 開始 Embedding + 導入...")
    BATCH_SIZE = 50
    total_imported = 0
    skipped_existing = 0

    for i in range(0, len(all_chunks), BATCH_SIZE):
        batch = all_chunks[i:i + BATCH_SIZE]

        # Generate embeddings
        texts = [c["content"] for c in batch]
        embeddings = model.encode(texts, normalize_embeddings=True, show_progress_bar=False)

        # Create points — FIX: 使用 chunk["source"] 而非外層變量
        points = []
        for j, chunk in enumerate(batch):
            # FIX: UUID 使用 chunk 自己的 source（不是外層 rel_path）
            content_id = generate_content_id(chunk["source"], chunk["content"], prefix="batch")

            # Skip if already imported (checkpoint resume)
            if content_id in imported_ids:
                skipped_existing += 1
                continue

            point = PointStruct(
                id=content_id,
                vector=embeddings[j].tolist(),
                payload={
                    "content": chunk["content"],
                    "source": chunk["source"],
                    "filename": chunk.get("filename", ""),
                    "tags": chunk.get("tags", []),
                    "section_path": chunk.get("section_path", ""),
                    "char_length": chunk.get("char_length", len(chunk["content"])),
                    "imported_at": datetime.now().isoformat(),
                    "access_count": 0,
                    "last_accessed": datetime.now().isoformat(),
                }
            )
            points.append(point)

        # Upsert to Qdrant
        if points:
            try:
                client.upsert(collection_name=COLLECTION, points=points)
                total_imported += len(points)

                # Update checkpoint
                for p in points:
                    imported_ids.add(p.id)
                save_checkpoint({
                    "imported_ids": list(imported_ids),
                    "last_file": batch[-1].get("source", ""),
                    "last_chunk_index": i + len(batch),
                })

                pct = min(100, int((i + BATCH_SIZE) / len(all_chunks) * 100))
                print(f"   [{pct:3d}%] 已導入 {total_imported}/{len(all_chunks)} chunks...", end="\r")
            except Exception as e:
                print(f"\n  ❌ Batch {i}-{i + BATCH_SIZE} 導入失敗: {e}")

    # ── Cleanup checkpoint after success ──
    if total_imported + skipped_existing >= len(all_chunks):
        CHECKPOINT_FILE.unlink(missing_ok=True)
        print("\n🧹 導入完成，清除斷點文件")

    # ── Verify ──
    print(f"\n\n✅ 導入完成！")
    stats = client.get_collection(COLLECTION)
    final_count = stats.points_count
    print(f"\n📊 {COLLECTION} 總計: {final_count} 條記憶（新增 {final_count - existing_count} 條）")
    if skipped_existing:
        print(f"   跳過已存在: {skipped_existing} 條（斷點續傳）")

    # ── Breakdown ──
    print(f"\n📂 來源分布:")
    source_counts = {}
    for c in all_chunks:
        src = c["source"].split("/")[0] if "/" in c["source"] else c["source"]
        source_counts[src] = source_counts.get(src, 0) + 1
    for src, count in sorted(source_counts.items(), key=lambda x: -x[1]):
        print(f"   {src}: {count} chunks")


if __name__ == "__main__":
    main()