Update rag/retriever.py

rag/retriever.py

@@ -12,11 +12,27 @@ def retrieve_context(query: str):
     distances, indices = index.search(query_embedding, TOP_K_RETRIEVAL)
 
     results = []
-    for idx in indices[0]:
-        if idx < len(documents):
-            results.append({
-                "text": documents[idx],
-                "source": metadata[idx]["source"]
-            })
+
+    for i in range(len(indices[0])):
+        idx = int(indices[0][i])  # 🔥 CRITICAL: convert np.int64 → int
+
+        # Case 1: metadata is a list
+        if isinstance(metadata, list):
+            if idx < len(metadata):
+                meta = metadata[idx]
+                results.append({
+                    "text": documents[idx],
+                    "source": meta.get("source", "unknown")
+                })
+
+        # Case 2: metadata is a dict
+        elif isinstance(metadata, dict):
+            meta = metadata.get(idx)
+            if meta:
+                results.append({
+                    "text": documents[idx],
+                    "source": meta.get("source", "unknown")
+                })
 
     return results
+