Update src/rag_index.py

src/rag_index.py

@@ -93,22 +93,42 @@ def _format_e5_query(text: str) -> str:
     return f"query: {text.strip()}"
 
 def search_index(query_text: str, *args, top_k: int = 5, **kwargs) -> List[Dict[str, Any]]:
-    """Flexible signature: supports legacy (query, top_k) or explicit embedder/bundle."""
+    """
+    Flexible signature:
+      - search_index(query, top_k=5)
+      - search_index(query, embedder, bundle, top_k=5)
+    """
+    # --- Determine how the function was called ---
     if len(args) >= 2:
         embedder = args[0]
         bundle = args[1]
+    elif len(args) == 1:
+        # Called as search_index(query, embedder, top_k=5)
+        embedder = args[0]
+        _, bundle = load_index_bundle()
     else:
         embedder, bundle = load_index_bundle()
 
-    index: faiss.Index = bundle["index"]
-    chunks: List[Dict[str, Any]] = bundle["chunks"]
+    # --- Ensure top_k is an integer ---
+    try:
+        top_k = int(top_k)
+    except Exception:
+        raise TypeError(f"top_k must be int, got {type(top_k)}")
+
+    index = bundle.get("index")
+    chunks: List[Dict[str, Any]] = bundle.get("chunks", [])
 
+    if index is None or not hasattr(index, "search"):
+        raise ValueError("Invalid FAISS index bundle.")
+
+    # --- Embed and search ---
     q = _format_e5_query(query_text)
     qv = np.asarray(embedder.encode([q], normalize_embeddings=True), dtype="float32")
 
     scores, idxs = index.search(qv, top_k)
     idxs, scores = idxs[0], scores[0]
 
+    # --- Collect results ---
     results: List[Dict[str, Any]] = []
     for rank, (i, s) in enumerate(zip(idxs, scores)):
         if i < 0 or i >= len(chunks):
@@ -123,6 +143,7 @@ def search_index(query_text: str, *args, top_k: int = 5, **kwargs) -> List[Dict[
         })
     return results
 
+
 # ------------------------------ stats helper ---------------------------------
 
 def index_stats(idx_dir: Path | None = None) -> Dict[str, Any]: