Update retriever.py

retriever.py

@@ -1,49 +1,45 @@
-\
-import os, json
-from typing import List
-import faiss
-import numpy as np
+import logging
+
+try:
+    import faiss
+    _HAS_FAISS = True
+except ImportError:
+    logging.warning("FAISS not installed — retrieval will be disabled. Install faiss-cpu or faiss-gpu for full functionality.")
+    _HAS_FAISS = False
+
 from sentence_transformers import SentenceTransformer
 
-class Retriever:
-    def __init__(self, index_path: str, meta_path: str):
-        if not (os.path.exists(index_path) and os.path.exists(meta_path)):
-            self._ready = False
-            self._err = f"Missing index or meta at {index_path} / {meta_path}"
-            return
-        self.index = faiss.read_index(index_path)
-        meta = json.loads(open(meta_path, "r", encoding="utf-8").read())
-        self.docs = meta["docs"]
-        self.model_name = meta["model"]
-        self.embed = SentenceTransformer(self.model_name)
-        self._ready = True
-        self._err = None
-
-    def ready(self) -> bool:
-        return self._ready
-
-    def reason(self) -> str:
-        return self._err or ""
-
-    def retrieve(self, query: str, k: int = 6) -> List[str]:
-        if not self._ready: return []
-        q = self.embed.encode([query], convert_to_numpy=True, normalize_embeddings=True)
-        D, I = self.index.search(q.astype(np.float32), k)
-        chunks = []
-        for idx in I[0]:
-            if 0 <= idx < len(self.docs):
-                chunks.append(self.docs[idx]["text"])
-        return chunks
-
-_retriever = None
-def init_retriever(index_path="rag_store/index.faiss", meta_path="rag_store/meta.json"):
-    global _retriever
-    if _retriever is None:
-        _retriever = Retriever(index_path, meta_path)
-    return _retriever
-
-def retrieve_context(query: str, k: int = 6) -> str:
-    r = init_retriever()
-    if not r.ready():
-        return "(No policy index found. Run build_policy_index.py to enable RAG.)"
-    return "\n---\n".join(r.retrieve(query, k=k))
+# load embedding model (still works even if FAISS missing)
+_model = SentenceTransformer("all-MiniLM-L6-v2")
+
+_index = None
+_docs = []
+
+def init_retriever(docs=None):
+    """
+    Initialize FAISS index if FAISS is available.
+    docs: list[str] to index
+    """
+    global _index, _docs
+    if not _HAS_FAISS:
+        _docs = docs or []
+        return
+
+    if docs:
+        _docs = docs
+        embeddings = _model.encode(docs, convert_to_numpy=True)
+        d = embeddings.shape[1]
+        _index = faiss.IndexFlatL2(d)
+        _index.add(embeddings)
+
+def retrieve_context(query: str, k: int = 5):
+    """
+    Retrieve top-k docs matching query.
+    Falls back to empty list if FAISS unavailable.
+    """
+    if not _HAS_FAISS or _index is None or not _docs:
+        return []
+
+    q_emb = _model.encode([query], convert_to_numpy=True)
+    D, I = _index.search(q_emb, k)
+    return [_docs[i] for i in I[0] if i < len(_docs)]