Update rag_query.py

rag_query.py

@@ -0,0 +1,38 @@
+import json
+import faiss
+import numpy as np
+from sentence_transformers import SentenceTransformer
+
+from config.settings import *
+
+_model = None
+_index = None
+_metadata = None
+
+
+def load_resources():
+    global _model, _index, _metadata
+
+    if _model is None:
+        _model = SentenceTransformer(EMBEDDING_MODEL)
+
+    if _index is None:
+        _index = faiss.read_index(FAISS_INDEX_PATH)
+
+    if _metadata is None:
+        with open(METADATA_PATH, "r", encoding="utf-8") as f:
+            _metadata = json.load(f)
+
+
+def retrieve(query: str):
+    load_resources()
+
+    query_embedding = _model.encode([query]).astype("float32")
+    distances, indices = _index.search(query_embedding, TOP_K)
+
+    results = []
+    for idx in indices[0]:
+        meta = _metadata[idx]
+        results.append(meta)
+
+    return results