Nomearod/agentbench

1"""Ingest documents into the hybrid vector store.2 3Usage:4    python scripts/ingest.py --config configs/tasks/tech_docs.yaml5    python scripts/ingest.py --doc-dir data/tech_docs/ --store-path .cache/store6"""7 8from __future__ import annotations9 10import argparse11import sys12from pathlib import Path13 14# Ensure the package is importable when running as a script15sys.path.insert(0, str(Path(__file__).resolve().parent.parent))16 17from agent_bench.rag.chunker import chunk_text18from agent_bench.rag.embedder import Embedder19from agent_bench.rag.store import HybridStore20 21 22def ingest(23    doc_dir: str,24    store_path: str,25    chunk_strategy: str = "recursive",26    chunk_size: int = 512,27    chunk_overlap: int = 64,28    model_name: str = "all-MiniLM-L6-v2",29    cache_dir: str = ".cache/embeddings",30) -> None:31    """Ingest all markdown files from doc_dir into a HybridStore."""32    doc_path = Path(doc_dir)33    if not doc_path.exists():34        print(f"Error: document directory {doc_dir} does not exist")35        sys.exit(1)36 37    # Exclude curation metadata files that live alongside corpus content.38    # SOURCES.md and QUESTION_PLAN.md are version-controlled curation39    # artifacts, not corpus content.40    _EXCLUDED = {"SOURCES.md", "QUESTION_PLAN.md", "README.md"}41    md_files = sorted(f for f in doc_path.glob("*.md") if f.name not in _EXCLUDED)42    if not md_files:43        print(f"Error: no markdown files found in {doc_dir}")44        sys.exit(1)45 46    print(f"Found {len(md_files)} markdown files in {doc_dir}")47 48    # Chunk all documents49    all_chunks = []50    for md_file in md_files:51        text = md_file.read_text(encoding="utf-8")52        source = md_file.name  # bare filename53        chunks = chunk_text(54            text, source, strategy=chunk_strategy, chunk_size=chunk_size, chunk_overlap=chunk_overlap55        )56        print(f"  {source}: {len(chunks)} chunks")57        all_chunks.extend(chunks)58 59    print(f"Total chunks: {len(all_chunks)}")60 61    # Embed62    print(f"Embedding with {model_name}...")63    embedder = Embedder(model_name=model_name, cache_dir=cache_dir)64    texts = [c.content for c in all_chunks]65    embeddings = embedder.embed_batch(texts)66    print(f"Embeddings shape: {embeddings.shape}")67 68    # Store69    store = HybridStore(dimension=embeddings.shape[1])70    store.add(all_chunks, embeddings)71    store.save(store_path)72 73    stats = store.stats()74    print(f"Store saved to {store_path}")75    print(f"  Chunks: {stats.total_chunks}")76    print(f"  FAISS index size: {stats.faiss_index_size}")77    print(f"  Unique sources: {stats.unique_sources}")78 79 80def main() -> None:81    parser = argparse.ArgumentParser(description="Ingest documents into vector store")82    parser.add_argument("--doc-dir", default="data/tech_docs/", help="Document directory")83    parser.add_argument("--store-path", default=".cache/store", help="Store output path")84    parser.add_argument("--chunk-strategy", default="recursive", choices=["recursive", "fixed"])85    parser.add_argument("--chunk-size", type=int, default=512)86    parser.add_argument("--chunk-overlap", type=int, default=64)87    parser.add_argument("--model", default="all-MiniLM-L6-v2", help="Embedding model name")88    parser.add_argument("--cache-dir", default=".cache/embeddings", help="Embedding cache dir")89    parser.add_argument(90        "--config", default=None, help="Task config YAML (overrides other args for doc-dir)"91    )92    args = parser.parse_args()93 94    doc_dir = args.doc_dir95    if args.config:96        from agent_bench.core.config import load_task_config97 98        task = load_task_config(Path(args.config).stem, path=Path(args.config))99        doc_dir = task.document_dir100 101    ingest(102        doc_dir=doc_dir,103        store_path=args.store_path,104        chunk_strategy=args.chunk_strategy,105        chunk_size=args.chunk_size,106        chunk_overlap=args.chunk_overlap,107        model_name=args.model,108        cache_dir=args.cache_dir,109    )110 111 112if __name__ == "__main__":113    main()114