Implement rag-search skill for semantic search
Add new skill for semantic search across personal state files and external documentation using ChromaDB and sentence-transformers. Components: - search.py: Main search interface (--index, --top-k flags) - index_personal.py: Index ~/.claude/state files - index_docs.py: Index external docs (git repos) - add_doc_source.py: Manage doc sources - test_rag.py: Test suite (5/5 passing) Features: - Two indexes: personal (116 chunks) and docs (k0s: 846 chunks) - all-MiniLM-L6-v2 embeddings (384 dimensions) - ChromaDB persistent storage - JSON output with ranked results and metadata Documentation: - Added to component-registry.json with triggers - Added /rag command alias - Updated skills/README.md - Resolved fc-013 (vector database for agent memory) 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
184
skills/rag-search/scripts/search.py
Executable file
184
skills/rag-search/scripts/search.py
Executable file
@@ -0,0 +1,184 @@
|
||||
#!/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()
|
||||
Reference in New Issue
Block a user