lambdagent/agentexample/qaagent67lambda/scripts/build_vector_index.py

#!/usr/bin/env python3
"""
Build vector index for RAG v2 hybrid search.
Embeds all chunks via local Ollama bge-m3, saves as numpy .npy + metadata pickle.
Supports checkpoint/resume for large corpora.

Usage:
  python build_vector_index.py [--base-dir the knowledge base directory]
"""
import json
import os
import pickle
import sys
import time
import urllib.request
from pathlib import Path

OLLAMA_EMBED_URL = os.environ.get('OLLAMA_EMBED_URL', 'http://127.0.0.1:11435/api/embed')
EMBED_MODEL = os.environ.get('EMBED_MODEL', 'bge-m3')
BATCH_SIZE = 32
CHECKPOINT_INTERVAL = 500  # Save every N chunks


def embed_batch(texts, model=None):
    """Embed a batch of texts via Ollama"""
    if model is None:
        model = EMBED_MODEL
    body = json.dumps({
        "model": model,
        "input": texts,
    }).encode('utf-8')
    req = urllib.request.Request(
        OLLAMA_EMBED_URL,
        data=body,
        headers={"Content-Type": "application/json"},
        method="POST"
    )
    with urllib.request.urlopen(req, timeout=300) as resp:
        data = json.loads(resp.read())
        return data.get('embeddings', [])


def main():
    import argparse

    # Default from config.py
    try:
        from config import cfg as _app_cfg
        default_base = str(_app_cfg.base_dir)
    except ImportError:
        default_base = os.environ.get('KNOWLEDGE_DIR', str(Path(__file__).parent.parent / 'knowledge'))

    parser = argparse.ArgumentParser()
    parser.add_argument('--base-dir', default=default_base)
    args = parser.parse_args()

    base = Path(args.base_dir)
    index_path = base / 'rag_index.json'
    vec_path = base / 'rag_vectors_v2.npy'
    meta_path = base / 'rag_vectors_meta_v2.pkl'
    checkpoint_path = base / 'rag_vectors_checkpoint.pkl'

    print("=" * 60)
    print("Building Vector Index (bge-m3 via Ollama)")
    print("=" * 60)

    # Load chunks from existing index
    if not index_path.exists():
        print(f"ERROR: {index_path} not found. Run rebuild_index.py first.")
        sys.exit(1)

    print(f"Loading chunks from {index_path}...")
    t0 = time.time()
    with open(index_path) as f:
        data = json.load(f)
    chunks = data['chunks']
    print(f"  Loaded {len(chunks)} chunks in {time.time()-t0:.1f}s")

    # Check for checkpoint (resume support)
    all_embeddings = []
    all_meta = []
    start_idx = 0

    if checkpoint_path.exists():
        print("Found checkpoint, resuming...")
        with open(checkpoint_path, 'rb') as f:
            cp = pickle.load(f)
        all_embeddings = cp['embeddings']
        all_meta = cp['meta']
        start_idx = cp['next_idx']
        print(f"  Resumed from index {start_idx} ({len(all_embeddings)} embeddings done)")

    # Test embedding API
    print("\nTesting Ollama embedding API...")
    test_vecs = embed_batch(["test"])
    if not test_vecs:
        print("ERROR: Ollama embedding API not available")
        sys.exit(1)
    dim = len(test_vecs[0])
    print(f"  OK, embedding dim: {dim}")

    # Embed all chunks
    total = len(chunks)
    remaining = total - start_idx
    print(f"\nEmbedding {remaining} chunks (batch_size={BATCH_SIZE})...")
    print(f"Estimated time: ~{remaining / BATCH_SIZE * 0.5 / 60:.1f} minutes")

    batch_texts = []
    batch_indices = []
    errors = 0
    embed_start = time.time()

    for i in range(start_idx, total):
        chunk = chunks[i]
        text = chunk['text'][:2000]  # bge-m3 max ~8192 tokens, truncate for safety

        batch_texts.append(text)
        batch_indices.append(i)

        if len(batch_texts) >= BATCH_SIZE or i == total - 1:
            try:
                vecs = embed_batch(batch_texts)
                if len(vecs) == len(batch_texts):
                    all_embeddings.extend(vecs)
                    for j, idx in enumerate(batch_indices):
                        c = chunks[idx]
                        all_meta.append({
                            'source': c['source'],
                            'chunk_id': c['chunk_id'],
                            'text': c['text'],
                        })
                else:
                    # Partial result, embed one by one
                    for t in batch_texts:
                        v = embed_batch([t])
                        if v:
                            all_embeddings.append(v[0])
                        else:
                            all_embeddings.append([0.0] * dim)
                            errors += 1
                    for idx in batch_indices:
                        c = chunks[idx]
                        all_meta.append({
                            'source': c['source'],
                            'chunk_id': c['chunk_id'],
                            'text': c['text'],
                        })
            except Exception as e:
                errors += 1
                # Fill with zeros for failed batch
                for _ in batch_texts:
                    all_embeddings.append([0.0] * dim)
                for idx in batch_indices:
                    c = chunks[idx]
                    all_meta.append({
                        'source': c['source'],
                        'chunk_id': c['chunk_id'],
                        'text': c['text'],
                    })
                if errors <= 3:
                    print(f"  ERROR at batch {i}: {e}")

            batch_texts = []
            batch_indices = []

            # Progress
            done = len(all_embeddings)
            if done % (CHECKPOINT_INTERVAL) < BATCH_SIZE or i == total - 1:
                elapsed = time.time() - embed_start
                rate = (done - start_idx) / max(elapsed, 1)
                eta = (total - done) / max(rate, 0.01) / 60
                print(f"  [{done}/{total}] {rate:.1f} chunks/s, ETA: {eta:.1f}min, errors: {errors}")

                # Save checkpoint
                with open(checkpoint_path, 'wb') as f:
                    pickle.dump({
                        'embeddings': all_embeddings,
                        'meta': all_meta,
                        'next_idx': i + 1,
                    }, f, protocol=pickle.HIGHEST_PROTOCOL)

    # Save final index
    print(f"\nSaving vector index...")
    try:
        import numpy as np
        vectors = np.array(all_embeddings, dtype=np.float32)
        np.save(str(vec_path), vectors)
        print(f"  Vectors: {vec_path} ({vectors.shape}, {vec_path.stat().st_size/1024/1024:.1f} MB)")
    except ImportError:
        # Fallback: save as pickle
        vec_path = vec_path.with_suffix('.pkl')
        with open(vec_path, 'wb') as f:
            pickle.dump(all_embeddings, f, protocol=pickle.HIGHEST_PROTOCOL)
        print(f"  Vectors (pickle): {vec_path}")

    with open(meta_path, 'wb') as f:
        pickle.dump(all_meta, f, protocol=pickle.HIGHEST_PROTOCOL)
    print(f"  Meta: {meta_path} ({meta_path.stat().st_size/1024/1024:.1f} MB)")

    # Cleanup checkpoint
    if checkpoint_path.exists():
        checkpoint_path.unlink()
        print("  Checkpoint removed")

    total_time = time.time() - embed_start
    print(f"\nDone! {len(all_embeddings)} vectors in {total_time/60:.1f} minutes, {errors} errors")


if __name__ == '__main__':
    main()