Quick and dirty RAG skill with Claude

"""
RAG Indexer
-----------
Indexes a set of documents into a local ChromaDB database for semantic search.
Designed for use with Claude Code or any LLM workflow where you want to search
a large corpus without stuffing it into the context window.

Two collection types out of the box:
  - docs:    Markdown files, chunked by heading
  - code:    Python files, chunked by class and function/method

Usage:
    First run (build the DB):
        python rag.py --index

    Query:
        python rag.py --query "topic or concept" --collection docs
        python rag.py --query "topic or concept" --collection code
        python rag.py --query "topic or concept" --collection all

    Reindex a single file after edits:
        python rag.py --reindex path/to/file.py

Dependencies:
    pip install chromadb sentence-transformers tqdm
"""

import ast
import argparse
import re
from pathlib import Path

import chromadb
from chromadb.utils.embedding_functions import SentenceTransformerEmbeddingFunction
from tqdm import tqdm

import sys
sys.stdout.reconfigure(encoding="utf-8")
sys.stderr.reconfigure(encoding="utf-8")

# ── Config ────────────────────────────────────────────────────────────────────

ROOT      = Path(__file__).parent.parent        # one level up from this script
DB_PATH   = Path(__file__).parent / "rag_db"    # where ChromaDB persists to disk
DOCS_DIR  = Path(__file__).parent               # folder of .md files (or a single .md)
CODE_DIR  = ROOT / "rag"                        # folder of .py files

# Embedding model. Delete db and reindex when switching.
#   "all-MiniLM-L6-v2"   — fast, small (~90MB), good baseline
#   "all-MiniLM-L12-v2"  — better quality, still small (~130MB)
#   "all-mpnet-base-v2"  — best quality, larger (~420MB)
EMBEDDING_MODEL = "all-MiniLM-L12-v2"

FALLBACK_WORDS   = 400
FALLBACK_OVERLAP = 50

ALL_COLLECTIONS  = ["docs", "code"]


def get_embedding_function():
    return SentenceTransformerEmbeddingFunction(model_name=EMBEDDING_MODEL)


# ── Chunking ──────────────────────────────────────────────────────────────────

def word_split(text, chunk_size=FALLBACK_WORDS, overlap=FALLBACK_OVERLAP):
    """Split text into overlapping word-count chunks."""
    words = text.split()
    chunks = []
    i = 0
    while i < len(words):
        chunks.append(" ".join(words[i : i + chunk_size]))
        i += chunk_size - overlap
    return chunks


def chunk_markdown(text, source_name="doc"):
    """
    Split markdown by headings. Sections longer than FALLBACK_WORDS are
    further split, with the heading prepended to each sub-chunk.
    """
    sections = re.split(r"\n(?=#{1,6} )", text)
    chunks, metadatas = [], []

    for section in sections:
        section = section.strip()
        if not section or len(section) < 50:
            continue

        lines = section.split("\n")
        heading = lines[0].lstrip("#").strip() if lines[0].startswith("#") else "Preamble"

        if len(section.split()) <= FALLBACK_WORDS:
            chunks.append(section)
            metadatas.append({"source": source_name, "section": heading})
        else:
            for idx, sub in enumerate(word_split(section)):
                chunks.append(f"[{heading} — part {idx + 1}]\n{sub}")
                metadatas.append({"source": source_name, "section": heading, "part": idx + 1})

    return chunks, metadatas


def chunk_python(text, filename):
    """
    Parse a Python file with ast and extract top-level classes and functions
    as individual chunks. Methods are included within their parent class chunk.
    Falls back to raw word splitting if the file can't be parsed.
    """
    chunks, metadatas = [], []

    try:
        tree = ast.parse(text)
    except SyntaxError:
        for idx, sub in enumerate(word_split(text)):
            chunks.append(sub)
            metadatas.append({"source": filename, "section": f"chunk_{idx}"})
        return chunks, metadatas

    lines = text.splitlines()

    for node in ast.iter_child_nodes(tree):
        if not isinstance(node, (ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef)):
            continue

        start = node.lineno - 1
        end   = node.end_lineno
        block = "\n".join(lines[start:end]).strip()

        if isinstance(node, ast.ClassDef):
            label     = f"class {node.name}"
            docstring = ast.get_docstring(node) or ""
            methods   = [
                n.name for n in ast.iter_child_nodes(node)
                if isinstance(n, (ast.FunctionDef, ast.AsyncFunctionDef))
            ]
            summary   = f"class {node.name}:\n"
            if docstring:
                summary += f'    """{docstring}"""\n'
            if methods:
                summary += f"    # methods: {', '.join(methods)}\n"
        else:
            label   = f"def {node.name}"
            summary = f"# {label}\n"

        if len(block.split()) <= FALLBACK_WORDS:
            chunks.append(block)
            metadatas.append({"source": filename, "section": label})
        else:
            for idx, sub in enumerate(word_split(block)):
                chunks.append(f"{summary}\n{sub}")
                metadatas.append({"source": filename, "section": label, "part": idx + 1})

    return chunks, metadatas


# ── Indexing ──────────────────────────────────────────────────────────────────

def get_client():
    return chromadb.PersistentClient(path=str(DB_PATH))


def get_collection(client, name):
    return client.get_or_create_collection(name, embedding_function=get_embedding_function())


def clear_collection(collection):
    existing = collection.get()
    if existing["ids"]:
        collection.delete(ids=existing["ids"])


def index_docs(client):
    docs_path = Path(DOCS_DIR)
    if not docs_path.exists():
        print(f"Docs path not found: {docs_path}")
        return

    collection = get_collection(client, "docs")
    clear_collection(collection)

    md_files = [docs_path] if docs_path.is_file() else list(docs_path.rglob("*.md"))
    if not md_files:
        print("No .md files found.")
        return

    total = 0
    for filepath in tqdm(md_files, desc="Indexing docs"):
        text = filepath.read_text(encoding="utf-8")
        chunks, metadatas = chunk_markdown(text, source_name=str(filepath.resolve()))
        if chunks:
            ids = [f"doc_{filepath.stem}_{i}" for i in range(len(chunks))]
            collection.add(documents=chunks, ids=ids, metadatas=metadatas)
            total += len(chunks)

    print(f"  → {total} doc chunks added")


def index_code(client):
    code_path = Path(CODE_DIR)
    if not code_path.exists():
        print(f"Code directory not found: {code_path}")
        return

    collection = get_collection(client, "code")
    clear_collection(collection)

    py_files = list(code_path.rglob("*.py"))
    print(f"Found {len(py_files)} .py files in {code_path}")
    if not py_files:
        return

    total = 0
    for filepath in tqdm(py_files, desc="Indexing code"):
        text = filepath.read_text(encoding="utf-8")
        chunks, metadatas = chunk_python(text, str(filepath.resolve()))
        if chunks:
            ids = [f"{filepath.stem}_{i}" for i in range(len(chunks))]
            collection.add(documents=chunks, ids=ids, metadatas=metadatas)
            total += len(chunks)

    print(f"  → {total} code chunks added")


def reindex_file(client, filepath_str):
    """Reindex a single file without rebuilding the whole collection."""
    filepath  = Path(filepath_str).resolve()
    full_path = str(filepath)
    text      = filepath.read_text(encoding="utf-8")
    suffix    = filepath.suffix.lower()

    if suffix == ".md":
        chunks, metadatas = chunk_markdown(text, source_name=full_path)
        col_name = "docs"
    elif suffix == ".py":
        chunks, metadatas = chunk_python(text, full_path)
        col_name = "code"
    else:
        print(f"Unsupported file type: {suffix}")
        return

    collection = get_collection(client, col_name)
    existing   = collection.get(where={"source": full_path})
    if existing["ids"]:
        collection.delete(ids=existing["ids"])

    ids = [f"{filepath.stem}_{i}" for i in range(len(chunks))]
    collection.add(documents=chunks, ids=ids, metadatas=metadatas)
    print(f"Reindexed {full_path} into '{col_name}': {len(chunks)} chunks")


# ── Querying ──────────────────────────────────────────────────────────────────

def _print_result(doc, meta, dist, prefix=""):
    source   = meta.get("source", "?")
    section  = meta.get("section", "")
    part     = meta.get("part", "")
    part_str = f" pt.{part}" if part else ""
    print(f"\n─── {prefix}[{source}] {section}{part_str}  (dist: {dist:.3f}) ───")
    print(doc[:600] + ("..." if len(doc) > 600 else ""))


def search(query, collection_name="docs", n=5):
    client = get_client()

    if collection_name == "all":
        rows = []
        for col_name in ALL_COLLECTIONS:
            try:
                collection = get_collection(client, col_name)
                results    = collection.query(query_texts=[query], n_results=n)
                for doc, meta, dist in zip(
                    results["documents"][0],
                    results["metadatas"][0],
                    results["distances"][0],
                ):
                    rows.append((dist, col_name, doc, meta))
            except Exception as e:
                print(f"  Warning: could not query {col_name}: {e}")

        rows.sort(key=lambda r: r[0])
        for dist, col_name, doc, meta in rows[:n]:
            _print_result(doc, meta, dist, prefix=f"[{col_name}] ")
        print()
        return

    try:
        collection = get_collection(client, collection_name)
        results    = collection.query(query_texts=[query], n_results=n)
    except Exception as e:
        print(f"Error querying '{collection_name}': {e}")
        return

    if not results["documents"][0]:
        print("No results found.")
        return

    for doc, meta, dist in zip(
        results["documents"][0],
        results["metadatas"][0],
        results["distances"][0],
    ):
        _print_result(doc, meta, dist)

    print()


# ── CLI ───────────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(description="Local RAG search tool")
    parser.add_argument("--index",      action="store_true", help="Build/rebuild the full database")
    parser.add_argument("--reindex",    metavar="FILE",      help="Reindex a single file (.md or .py)")
    parser.add_argument("--query",      metavar="TEXT",      help="Search query")
    parser.add_argument("--collection", default="docs",
                        choices=ALL_COLLECTIONS + ["all"],
                        help="Which collection to query (default: docs)")
    parser.add_argument("--results",    type=int, default=5, help="Number of results (default: 5)")

    args = parser.parse_args()

    if args.index:
        client = get_client()
        for fn in [index_docs, index_code]:
            try:
                fn(client)
            except Exception as e:
                print(f"{fn.__name__} failed: {e}")

    elif args.reindex:
        reindex_file(get_client(), args.reindex)

    elif args.query:
        search(args.query, collection_name=args.collection, n=args.results)

    else:
        parser.print_help()


if __name__ == "__main__":
    main()

---
name: rag-search
description: Search the project's local RAG database for context relevant to the current task. ONLY invoke when the user is clearly asking about domain content — documentation, design docs, existing written content, or codebase internals. Do NOT invoke for general coding questions, syntax help, or anything that doesn't require searching the project corpus.
---

# RAG Search

Use this skill to query the local ChromaDB database for project context.

## Running a search

```bash
path/to/my/executable --query "your query" --collection docs --results 5
```

## Important

NEVER RUN THIS COMMAND WITH `--index` or `--reindex`!

## Collections

| Flag | Contents |
|---|---|
| `docs` | Markdown documentation, wikis, design docs, style guides |
| `code` | Python files, chunked by class and function |
| `all` | Searches all collections, returns top N results ranked by distance |

`--results N` controls how many results to return (default: 5). With `--collection all`, results are interleaved by semantic distance — the N best matches across all collections combined, not N per collection.

## Query style

Use keyword and concept queries, not natural language questions. The search is semantic but works best with nouns and concrete phrases.

**Good:** `authentication flow`, `database connection`, `retry logic`

**Bad:** `how does authentication work?`, `where is the database connected?`

## When to search what

- **docs** — concepts, architecture decisions, design rationale, API descriptions, onboarding material
- **code** — existing implementations, class/method names, how something was previously done
- **all** — when you need full context, or aren't sure where to look

## Workflow

1. Run one or more queries depending on what's needed
2. Read the returned chunks for relevant context
3. Use that context to inform your suggestion, implementation, or review

## Notes

- Lower distance = more semantically similar to your query
- If results look off, try rephrasing — shorter and more concrete usually works better