Update rag_mini.py

rag_mini.py

@@ -1,55 +1,85 @@
 from __future__ import annotations
-import os, sys
+import os, math
 from pathlib import Path
 from typing import List, Tuple
 
+# ---- paths / constants ----
 ROOT_DIR   = Path(__file__).parent.resolve()
 MM_ROOT    = ROOT_DIR / "MaterialMind"
 DEFAULT_TOPK = 5
 
-INDEX_DS   = os.getenv("INDEX_DS", "").strip()
-INDEX_DIR_ENV = os.getenv("INDEX_DIR", "").strip()
+# ---- where the index lives ----
+INDEX_DS       = os.getenv("INDEX_DS", "").strip()
+INDEX_DIR_ENV  = os.getenv("INDEX_DIR", "").strip()
+INDEX_COLLECTION = os.getenv("INDEX_COLLECTION", "").strip()  # e.g., "materialmind"
 
-_EMB_FAST=None; _EMB_ST=None
-EMB_MODEL = "BAAI/bge-small-en-v1.5"
+# ---- embedding settings (match local!) ----
+# Use BGE-small (384-d) everywhere to avoid mismatch
+EMB_PROVIDER = os.getenv("EMB_PROVIDER", "hf").strip().lower()  # "hf" or "openai"
+EMB_MODEL    = os.getenv("EMB_MODEL", "BAAI/bge-small-en-v1.5").strip()
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")  # only used if EMB_PROVIDER=openai
+
+# backends
+_EMB_FAST = None
+_EMB_ST   = None
+_EMB_OAI  = None
+
+def _l2norm(vec: List[float]) -> List[float]:
+    s = math.sqrt(sum(x*x for x in vec)) or 1.0
+    return [x/s for x in vec]
 
 def _init_embedder():
-    global _EMB_FAST, _EMB_ST
-    if _EMB_FAST or _EMB_ST:
-        return
+    """Initialize exactly one embedding backend based on EMB_PROVIDER."""
+    global _EMB_FAST, _EMB_ST, _EMB_OAI
+    if EMB_PROVIDER in ("openai","oai"):
+        try:
+            from openai import OpenAI
+            _EMB_OAI = OpenAI(api_key=OPENAI_API_KEY)
+            print(f"[EMB] OpenAI embeddings ready: {EMB_MODEL}", flush=True)
+            return
+        except Exception as e:
+            print("[EMB] OpenAI embeddings unavailable:", e, flush=True)
+    # HF path (FastEmbed → SentenceTransformers fallback)
     try:
         from fastembed import TextEmbedding
-        _EMB_FAST = TextEmbedding(model_name=EMB_MODEL)
-        print("[EMB] FastEmbed ready:", EMB_MODEL, flush=True)
+        _EMB_FAST = TextEmbedding(model_name=EMB_MODEL)  # we’ll L2-normalize ourselves
+        print(f"[EMB] FastEmbed ready: {EMB_MODEL}", flush=True)
         return
     except Exception as e1:
         print("[EMB] FastEmbed unavailable:", e1, flush=True)
     try:
         from sentence_transformers import SentenceTransformer
         _EMB_ST = SentenceTransformer(EMB_MODEL)
-        print("[EMB] SentenceTransformers ready:", EMB_MODEL, flush=True)
+        print(f"[EMB] SentenceTransformers ready: {EMB_MODEL}", flush=True)
         return
     except Exception as e2:
         print("[EMB] SentenceTransformers unavailable:", e2, flush=True)
         print("[EMB] ERROR: No embedding backend available. Install 'fastembed' or 'sentence-transformers'.", flush=True)
 
-def _embed(texts:List[str])->List[List[float]]:
+def _embed(texts: List[str]) -> List[List[float]]:
     _init_embedder()
+    if _EMB_OAI is not None:
+        r = _EMB_OAI.embeddings.create(model=EMB_MODEL, input=texts)
+        return [_l2norm(d.embedding) for d in r.data]
     if _EMB_FAST is not None:
-        return [v for v in _EMB_FAST.embed(texts)]
+        return [_l2norm(v) for v in _EMB_FAST.embed(texts)]
     if _EMB_ST is not None:
-        return _EMB_ST.encode(texts, normalize_embeddings=True).tolist()
-    # Fallback: no embeddings – return zeros to avoid crashing
-    return [[0.0]*384 for _ in texts]  # length doesn’t matter; Chroma ignores if we don’t query
+        # ST can normalize internally, but we also L2-normalize for safety
+        from numpy import array
+        arr = _EMB_ST.encode(texts, normalize_embeddings=True)
+        return [_l2norm(list(v)) for v in array(arr).tolist()]
+    # last resort: zeros (prevents crashes; yields 0 hits)
+    return [[0.0]*384 for _ in texts]
 
-def _has_catalog(dirpath:Path)->bool:
+# ---- index discovery ----
+def _has_catalog(dirpath: Path) -> bool:
     for f in ["chroma.sqlite3","chroma.sqlite","chroma-collections.parquet",
               "index_metadata.pickle","data_level0.bin"]:
         if (dirpath/f).exists():
             return True
     return False
 
-def _locate_local_index()->Path:
+def _locate_local_index() -> Path:
     if INDEX_DIR_ENV:
         return (ROOT_DIR / INDEX_DIR_ENV).resolve()
     base = (MM_ROOT / "index" / "chroma_v3").resolve()
@@ -78,19 +108,30 @@ def ensure_ready():
     else:
         print(f"[RAG] Index OK at {local}", flush=True)
 
+# ---- Chroma access ----
 def _get_collection():
     import chromadb
     local = _locate_local_index()
     client = chromadb.PersistentClient(path=str(local))
+    if INDEX_COLLECTION:
+        try:
+            return client.get_collection(INDEX_COLLECTION)
+        except Exception:
+            # create with cosine metric to match unit-normalized embeddings
+            return client.get_or_create_collection(
+                name=INDEX_COLLECTION, metadata={"hnsw:space": "cosine"}
+            )
     try:
         cols = client.list_collections()
         if cols:
             return client.get_collection(cols[0].name)
     except Exception:
         pass
-    return client.get_or_create_collection(name="materialmind")
+    return client.get_or_create_collection(
+        name="materialmind", metadata={"hnsw:space": "cosine"}
+    )
 
-def search(query:str, k:int=DEFAULT_TOPK)->List[Tuple[str,str]]:
+def search(query: str, k: int = DEFAULT_TOPK) -> List[Tuple[str, str]]:
     local = _locate_local_index()
     if not _has_catalog(local):
         return []
@@ -104,12 +145,31 @@ def search(query:str, k:int=DEFAULT_TOPK)->List[Tuple[str,str]]:
         return []
     docs  = (res.get("documents") or [[]])[0]
     metas = (res.get("metadatas") or [[]])[0]
-    hits=[]
+    hits = []
     for d, m in zip(docs, metas):
         if not d:
             continue
         src = (m or {}).get("source") or (m or {}).get("path") or "unknown"
-        page= (m or {}).get("page")
+        page = (m or {}).get("page")
         cite = f"{src}" + (f":p.{page}" if page else "")
         hits.append((d, cite))
     return hits
+
+# ---- tiny debugger (optional) ----
+def rag_debug_info():
+    import chromadb
+    local = _locate_local_index()
+    client = chromadb.PersistentClient(path=str(local))
+    info = {"index_path": str(local), "collections": [], "emb": {
+        "provider": EMB_PROVIDER, "model": EMB_MODEL
+    }}
+    try:
+        for c in client.list_collections():
+            try:
+                cnt = c.count()
+            except Exception:
+                cnt = -1
+            info["collections"].append({"name": c.name, "count": cnt})
+    except Exception as e:
+        info["collections"].append({"error": str(e)})
+    return info