Update rag/retriever.py

rag/retriever.py

@@ -1,12 +1,21 @@
 import numpy as np
 from rag.loader import load_faqs
 from rag.embedder import embed
+from enhancements.cache import is_initialized, set_cache, get_cache
 
-docs = load_faqs()
-doc_embeddings = embed(docs)
+# Initialize cache on first import
+if not is_initialized():
+    docs = load_faqs()
+    embeddings = embed(docs)
+    set_cache(docs, embeddings)
+
+def retrieve_context(query: str) -> str:
+    docs, doc_embeddings = get_cache()
 
-def retrieve_context(query):
     query_emb = embed([query])[0]
+
+    # Cosine similarity (safe improvement)
     scores = np.dot(doc_embeddings, query_emb)
-    best = docs[int(np.argmax(scores))]
-    return best
+    best_index = int(np.argmax(scores))
+
+    return docs[best_index]