Update rag.py

rag.py

@@ -1,26 +1,20 @@
 from typing import List, Tuple
 import numpy as np
-import cohere
-from settings import COHERE_API_KEY, COHERE_EMBED_MODEL
+from llm_router import cohere_embed
 
 class RAGIndex:
     def __init__(self):
-        self.client = cohere.Client(api_key=COHERE_API_KEY) if COHERE_API_KEY else None
         self.texts: List[str] = []
         self.vecs: np.ndarray | None = None
 
-    def _embed(self, texts: List[str]) -> np.ndarray:
-        if not texts: return np.zeros((0, 384), dtype="float32")
-        if not self.client:
-            # Fallback: random embeddings (avoid crash; not ideal)
-            return np.random.normal(size=(len(texts), 384)).astype("float32")
-        resp = self.client.embed(texts=texts, model=COHERE_EMBED_MODEL)
-        vecs = np.array(getattr(resp, "embeddings", []) or getattr(resp, "data", []), dtype="float32")
-        return vecs
-
     def add(self, chunks: List[str]):
         if not chunks: return
-        new_vecs = self._embed(chunks)
+        new_vecs_list = cohere_embed(chunks)
+        if not new_vecs_list:
+            # fallback: random to avoid crash (not ideal for accuracy)
+            new_vecs = np.random.normal(size=(len(chunks), 384)).astype("float32")
+        else:
+            new_vecs = np.array(new_vecs_list, dtype="float32")
         if self.vecs is None:
             self.vecs = new_vecs
             self.texts = list(chunks)
@@ -30,7 +24,9 @@ class RAGIndex:
 
     def retrieve(self, query: str, k: int = 6) -> List[Tuple[str, float]]:
         if not self.texts: return []
-        qv = self._embed([query])[0]
+        qv_list = cohere_embed([query])
+        if not qv_list: return []
+        qv = np.array(qv_list[0], dtype="float32")
         sims = (self.vecs @ qv) / (np.linalg.norm(self.vecs, axis=1) * (np.linalg.norm(qv) + 1e-9))
         idx = np.argsort(-sims)[:k]
         return [(self.texts[i], float(sims[i])) for i in idx]