#!/usr/bin/env python3 """Query the RAG vector index to find semantically related code and docs. This is the query-time half of the RAG pipeline (the counterpart to indexer.py, which builds the index). All operations are local — no network calls are made once the embedding model has been downloaded (see embeddings.py for details on model download and caching). Given a natural-language query, it embeds the query with the same MiniLM model used at index time, then finds the closest chunks in the local LanceDB vector database by cosine similarity. The core functions (get_db_and_embedder, search_table, format_results, search_related) are imported by the MCP server for tool calls. The script can also be run standalone from the command line. The "related file" mode works differently from a text query: it reads the target file's chunks from the index, combines them into a synthetic query vector, and searches for the nearest chunks from *other* files. This surfaces files that are semantically similar to the target file. Usage: openmc_search.py "query" # Search code (default) openmc_search.py "query" --docs # Search documentation openmc_search.py "query" --all # Search both code and docs openmc_search.py --related src/particle.cpp # Find related code openmc_search.py "query" --top-k 20 # Return more results """ import argparse import sys from pathlib import Path # Same sys.path setup as indexer.py — needed for standalone CLI use. TOOLS_DIR = Path(__file__).resolve().parent.parent sys.path.insert(0, str(TOOLS_DIR / "rag")) OPENMC_ROOT = Path(__file__).resolve().parents[3] CACHE_DIR = OPENMC_ROOT / ".claude" / "cache" INDEX_DIR = CACHE_DIR / "rag_index" def get_db_and_embedder(): """Load the LanceDB database and embedding provider.""" import lancedb from embeddings import EmbeddingProvider if not INDEX_DIR.exists(): raise FileNotFoundError( "No RAG index found. Call openmc_rag_rebuild() to build one." ) db = lancedb.connect(str(INDEX_DIR)) embedder = EmbeddingProvider() return db, embedder def _table_names(db): """Return table names as a list, compatible with multiple LanceDB versions.""" result = db.table_names() if hasattr(db, "table_names") else db.list_tables() return result.tables if hasattr(result, "tables") else list(result) def search_table(db, embedder, table_name, query, top_k): """Search a LanceDB table with a text query.""" if table_name not in _table_names(db): print(f"Table '{table_name}' not found in index.", file=sys.stderr) return [] table = db.open_table(table_name) query_vec = embedder.embed_query(query) results = table.search(query_vec).limit(top_k).to_list() return results def format_results(results, label=""): """Format search results for display.""" if not results: return "No results found.\n" output = [] if label: output.append(f"=== {label} ===\n") for i, r in enumerate(results, 1): filepath = r["filepath"] start = r["start_line"] end = r["end_line"] kind = r["kind"] dist = r.get("_distance", 0) header = f"[{i}] {filepath}:{start}-{end} ({kind}, dist={dist:.3f})" output.append(header) # Show text preview (first 500 chars) text = r["text"][:500] if len(r["text"]) > 500: text += "\n ..." # Indent the text for line in text.split("\n"): output.append(f" {line}") output.append("") return "\n".join(output) def search_related(db, embedder, filepath, top_k): """Find code related to a given file.""" if "code" not in _table_names(db): print("No 'code' table in index.", file=sys.stderr) return [] table = db.open_table("code") # Normalize filepath fp = filepath if Path(filepath).is_absolute(): try: fp = str(Path(filepath).relative_to(OPENMC_ROOT)) except ValueError: pass # Get chunks from target file try: safe_fp = fp.replace("'", "''") target_chunks = table.search().where( f"filepath = '{safe_fp}'" ).limit(50).to_list() except Exception: # LanceDB where clause might not work in all versions # Fall back to fetching all and filtering all_data = table.to_pandas() target_rows = all_data[all_data["filepath"] == fp] if target_rows.empty: print(f"No chunks found for '{fp}'", file=sys.stderr) return [] target_chunks = target_rows.head(50).to_dict("records") if not target_chunks: print(f"No chunks found for '{fp}'", file=sys.stderr) return [] # Combine top chunks as the query combined_text = " ".join(c["text"][:200] for c in target_chunks[:5]) query_vec = embedder.embed_query(combined_text) # Search excluding the source file results = table.search(query_vec).limit(top_k + 10).to_list() # Filter out same file results = [r for r in results if r["filepath"] != fp][:top_k] return results def main(): parser = argparse.ArgumentParser( description="Semantic search across OpenMC codebase and docs", formatter_class=argparse.RawDescriptionHelpFormatter, epilog="""examples: %(prog)s "particle random number seed initialization" %(prog)s "how to define tallies" --docs %(prog)s "weight window variance reduction" --all %(prog)s "where is cross section data loaded" --top-k 15 %(prog)s --related src/simulation.cpp %(prog)s --related src/particle_restart.cpp --top-k 5""", ) parser.add_argument("query", nargs="?", help="Search query") parser.add_argument("--docs", action="store_true", help="Search documentation instead of code") parser.add_argument("--all", action="store_true", help="Search both code and documentation") parser.add_argument("--related", metavar="FILE", help="Find code related to a given file") parser.add_argument("--top-k", type=int, default=10, help="Number of results (default: 10)") args = parser.parse_args() if not args.query and not args.related: parser.print_help() sys.exit(1) db, embedder = get_db_and_embedder() if args.related: results = search_related(db, embedder, args.related, args.top_k) print(format_results(results, f"Code related to {args.related}")) elif args.all: code_results = search_table( db, embedder, "code", args.query, args.top_k) doc_results = search_table( db, embedder, "docs", args.query, args.top_k) print(format_results(code_results, "Code")) print(format_results(doc_results, "Documentation")) elif args.docs: results = search_table(db, embedder, "docs", args.query, args.top_k) print(format_results(results, "Documentation")) else: results = search_table(db, embedder, "code", args.query, args.top_k) print(format_results(results, "Code")) if __name__ == "__main__": main()