claude-code/skills/rag-search/scripts/search.py

#!/usr/bin/env python3
"""
RAG Search - Main search entry point

Searches personal and/or docs indexes for semantically similar content.
"""

import argparse
import json
import sys
from pathlib import Path
from typing import Optional

# Add venv site-packages to path
VENV_PATH = Path(__file__).parent.parent / "venv" / "lib" / "python3.13" / "site-packages"
if str(VENV_PATH) not in sys.path:
    sys.path.insert(0, str(VENV_PATH))

import chromadb
from sentence_transformers import SentenceTransformer

# Constants
DATA_DIR = Path.home() / ".claude" / "data" / "rag-search"
CHROMA_DIR = DATA_DIR / "chroma"
MODEL_NAME = "all-MiniLM-L6-v2"
DEFAULT_TOP_K = 5

# Lazy-loaded globals
_model: Optional[SentenceTransformer] = None
_client: Optional[chromadb.PersistentClient] = None


def get_model() -> SentenceTransformer:
    """Lazy-load the embedding model."""
    global _model
    if _model is None:
        _model = SentenceTransformer(MODEL_NAME)
    return _model


def get_client() -> chromadb.PersistentClient:
    """Lazy-load the ChromaDB client."""
    global _client
    if _client is None:
        CHROMA_DIR.mkdir(parents=True, exist_ok=True)
        _client = chromadb.PersistentClient(path=str(CHROMA_DIR))
    return _client


def search(
    query: str,
    index: Optional[str] = None,
    top_k: int = DEFAULT_TOP_K,
) -> dict:
    """
    Search for semantically similar content.

    Args:
        query: The search query
        index: Which index to search ("personal", "docs", or None for both)
        top_k: Number of results to return per collection

    Returns:
        dict with query, results, and metadata
    """
    client = get_client()
    model = get_model()

    # Embed the query
    query_embedding = model.encode(query).tolist()

    # Determine which collections to search
    collections_to_search = []
    if index is None or index == "personal":
        try:
            collections_to_search.append(("personal", client.get_collection("personal")))
        except Exception:
            pass  # Collection doesn't exist
    if index is None or index == "docs":
        try:
            collections_to_search.append(("docs", client.get_collection("docs")))
        except Exception:
            pass  # Collection doesn't exist

    if not collections_to_search:
        return {
            "query": query,
            "results": [],
            "searched_collections": [],
            "total_chunks_searched": 0,
            "error": f"No collections found for index: {index or 'any'}"
        }

    # Search each collection
    all_results = []
    total_chunks = 0
    searched_collections = []

    for coll_name, collection in collections_to_search:
        searched_collections.append(coll_name)
        count = collection.count()
        total_chunks += count

        if count == 0:
            continue

        results = collection.query(
            query_embeddings=[query_embedding],
            n_results=min(top_k, count),
            include=["documents", "metadatas", "distances"]
        )

        # Process results
        if results["documents"] and results["documents"][0]:
            for i, (doc, metadata, distance) in enumerate(zip(
                results["documents"][0],
                results["metadatas"][0],
                results["distances"][0]
            )):
                # Convert distance to similarity score (cosine distance to similarity)
                score = 1 - (distance / 2)  # Normalized for cosine distance
                all_results.append({
                    "source": coll_name,
                    "file": metadata.get("file", "unknown"),
                    "chunk": doc,
                    "score": round(score, 3),
                    "metadata": {k: v for k, v in metadata.items() if k != "file"}
                })

    # Sort by score and add ranks
    all_results.sort(key=lambda x: x["score"], reverse=True)
    for i, result in enumerate(all_results[:top_k]):
        result["rank"] = i + 1

    return {
        "query": query,
        "results": all_results[:top_k],
        "searched_collections": searched_collections,
        "total_chunks_searched": total_chunks
    }


def main():
    parser = argparse.ArgumentParser(
        description="Search the RAG index for relevant content",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  %(prog)s "how did I configure ArgoCD sync?"
  %(prog)s --index personal "past decisions about caching"
  %(prog)s --index docs "k0s node maintenance"
  %(prog)s --top-k 10 "prometheus alerting rules"
"""
    )
    parser.add_argument("query", help="Search query")
    parser.add_argument(
        "--index", "-i",
        choices=["personal", "docs"],
        help="Search only this index (default: both)"
    )
    parser.add_argument(
        "--top-k", "-k",
        type=int,
        default=DEFAULT_TOP_K,
        help=f"Number of results to return (default: {DEFAULT_TOP_K})"
    )
    parser.add_argument(
        "--raw",
        action="store_true",
        help="Output raw JSON (default: formatted)"
    )

    args = parser.parse_args()

    results = search(args.query, args.index, args.top_k)

    if args.raw:
        print(json.dumps(results))
    else:
        print(json.dumps(results, indent=2))


if __name__ == "__main__":
    main()