create_db.py

"""
CodeLens – vector database builder.

Scans a codebase, chunks files by language, embeds them with
HuggingFace nomic-embed-text-v1.5, and stores them in ChromaDB.

Usage:
    python create_db.py [--repo-path ./my-project] [--db-path ./chroma_db]

Fixes vs original:
- REPO_PATH read from env / CLI arg (no hardcoded absolute path)
- Bare except replaced with explicit error logging
- torch.mps.empty_cache() gated by MPS availability check
- torch removed as a hard dependency (only imported if MPS available)
- Incremental mode: skip rebuild if DB already exists unless --force
"""
import os
import shutil
import glob
import gc
import argparse
from pathlib import Path

from langchain_text_splitters import Language, RecursiveCharacterTextSplitter
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_chroma import Chroma
from langchain_core.documents import Document

from config import DB_PATH, REPO_PATH, EMBED_MODEL, BATCH_SIZE, CHUNK_SIZE, CHUNK_OVERLAP, FILE_TYPES

# Language enum mapping
LANG_MAP = {
    "python":   Language.PYTHON,
    "js":       Language.JS,
    "ts":       Language.TS,
    "html":     Language.HTML,
    "markdown": Language.MARKDOWN,
}


def _clear_mps_cache():
    """Release Apple MPS GPU memory if available — no-op everywhere else."""
    try:
        import torch
        if torch.backends.mps.is_available():
            torch.mps.empty_cache()
    except ImportError:
        pass


def load_files(extension: str, repo_path: str) -> list[Document]:
    """Load all files with the given extension from repo_path."""
    docs = []
    pattern = f"{repo_path}/**/*{extension}"
    for path in glob.glob(pattern, recursive=True):
        # Skip hidden dirs and virtual environments
        parts = Path(path).parts
        if any(p.startswith(".") or p in ("__pycache__", "node_modules", "venv", ".venv", "chroma_db") for p in parts):
            continue
        try:
            with open(path, "r", encoding="utf-8", errors="ignore") as f:
                text = f.read()
            if text.strip():
                docs.append(Document(page_content=text, metadata={"source": path}))
        except (PermissionError, OSError) as exc:
            print(f"  Skipping {path}: {exc}")
    return docs


def main(repo_path: str = None, db_path: str = None, force: bool = False):
    """
    Build (or rebuild) the ChromaDB vector database from a codebase.

    Args:
        repo_path: Root directory of the codebase to index.
        db_path:   Where to store the ChromaDB database.
        force:     If True, wipe and rebuild even if DB already exists.
    """
    repo_path = repo_path or REPO_PATH
    db_path   = db_path   or DB_PATH

    if not os.path.exists(repo_path):
        raise FileNotFoundError(f"Repo path not found: {repo_path}")

    if os.path.exists(db_path) and not force:
        print(f"Database already exists at {db_path}. Use --force to rebuild.")
        return

    if os.path.exists(db_path):
        shutil.rmtree(db_path)

    print(f"Scanning {repo_path}...")
    all_chunks = []

    for ext, lang_key in FILE_TYPES.items():
        docs = load_files(ext, repo_path)
        if not docs:
            continue

        lang_type = LANG_MAP.get(lang_key) if lang_key else None
        if lang_type:
            splitter = RecursiveCharacterTextSplitter.from_language(
                language=lang_type,
                chunk_size=CHUNK_SIZE,
                chunk_overlap=CHUNK_OVERLAP,
            )
        else:
            splitter = RecursiveCharacterTextSplitter(
                chunk_size=CHUNK_SIZE,
                chunk_overlap=CHUNK_OVERLAP,
            )

        chunks = splitter.split_documents(docs)
        print(f"  {ext}: {len(docs)} files → {len(chunks)} chunks")
        all_chunks.extend(chunks)

    if not all_chunks:
        print("No files found. Check REPO_PATH.")
        return

    print(f"\nEmbedding {len(all_chunks)} chunks (batch size {BATCH_SIZE})...")

    embeddings = HuggingFaceEmbeddings(
        model_name=EMBED_MODEL,
        model_kwargs={"device": "cpu", "trust_remote_code": True},
    )

    db = Chroma(persist_directory=db_path, embedding_function=embeddings)
    total_batches = (len(all_chunks) + BATCH_SIZE - 1) // BATCH_SIZE

    for i in range(0, len(all_chunks), BATCH_SIZE):
        batch = all_chunks[i: i + BATCH_SIZE]
        batch_num = i // BATCH_SIZE + 1
        try:
            db.add_documents(batch)
            print(f"  Saved batch {batch_num}/{total_batches}")
            gc.collect()
            _clear_mps_cache()
        except Exception as exc:
            print(f"  Error saving batch {batch_num}: {exc}")

    print(f"\nDatabase built at {db_path} ({len(all_chunks)} chunks indexed).")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Build CodeLens vector database")
    parser.add_argument("--repo-path", default=None, help="Path to the codebase to index")
    parser.add_argument("--db-path",   default=None, help="Where to store ChromaDB")
    parser.add_argument("--force",     action="store_true", help="Wipe and rebuild existing DB")
    args = parser.parse_args()
    main(repo_path=args.repo_path, db_path=args.db_path, force=args.force)