Update rag_mini.py

rag_mini.py

@@ -10,17 +10,29 @@ from pathlib import Path
 from typing import List, Tuple
 
 # ---------- Paths ----------
-ROOT_DIR   = Path(__file__).parent.resolve()
-DATA_ROOT  = ROOT_DIR / "MaterialMind"
-DATA_DIR   = DATA_ROOT / "sources"               # kept for compatibility; not required
-DEFAULT_TOPK = 5
-DEFAULT_MODEL = ""                               # kept for compatibility with imports
+ROOT_DIR      = Path(__file__).parent.resolve()
+DATA_ROOT     = ROOT_DIR / "MaterialMind"
+DEFAULT_TOPK  = 5
+DEFAULT_MODEL = ""   # kept for compatibility with imports
+
+def _normalize_to_dir(p: Path) -> Path:
+    """If p is a file (e.g., .../chroma.sqlite3) return its parent; else p."""
+    try:
+        if p.is_file():
+            return p.parent
+        # If someone passed a string ending with .sqlite/.sqlite3, treat as file
+        if p.name.endswith(".sqlite3") or p.name.endswith(".sqlite"):
+            return p.parent
+    except Exception:
+        pass
+    return p
 
 # Base folder where the Chroma catalog (or its subfolder) lives
-_BASE_INDEX_DIR = Path(os.getenv("INDEX_DIR", str(DATA_ROOT / "index" / "chroma_v3"))).resolve()
+_base_from_env = Path(os.getenv("INDEX_DIR", DATA_ROOT / "index" / "chroma_v3"))
+_BASE_INDEX_DIR = _normalize_to_dir(_base_from_env).resolve()
 
 # ---------- Embedding ----------
-EMB_MODEL = "BAAI/bge-small-en-v1.5"
+EMB_MODEL  = "BAAI/bge-small-en-v1.5"
 _EMBED_FAST = None
 _EMBED_ST   = None
 
@@ -41,45 +53,44 @@ def _init_embedder():
 def _embed(texts: List[str]) -> List[List[float]]:
     _init_embedder()
     if _EMBED_FAST is not None:
-        # FastEmbed returns a generator of np arrays
-        return [v for v in _EMBED_FAST.embed(texts)]
+        # Ensure plain Python lists
+        return [vec.tolist() for vec in _EMBED_FAST.embed(texts)]
     return _EMBED_ST.encode(texts, normalize_embeddings=True).tolist()
 
 # ---------- Catalog discovery ----------
 def _has_catalog(dirpath: Path) -> bool:
     """Heuristic check for a Chroma catalog in dirpath."""
-    key_files = [
+    if not dirpath.exists() or not dirpath.is_dir():
+        return False
+    key_files = {
         "chroma.sqlite3",
+        "chroma.sqlite",
+        "chroma-collections.parquet",
         "index_metadata.pickle",
         "data_level0.bin",
         "length.bin",
         "link_lists.bin",
         "header.bin",
-        "chroma.sqlite",
-        "chroma-collections.parquet",
-    ]
+    }
     return any((dirpath / kf).exists() for kf in key_files)
 
 def _find_catalog_dir(base: Path) -> Path:
-    """
-    Return the directory that actually contains the Chroma catalog.
-    Handles a nested UUID subfolder (as in your screenshot).
-    """
-    # direct
+    """Return the directory that actually contains the Chroma catalog."""
+    base = _normalize_to_dir(base)
     if _has_catalog(base):
         return base
 
-    # look at one or two levels deep for chroma.sqlite3 (or other catalog markers)
-    hits = list(base.rglob("chroma.sqlite3"))
-    if hits:
-        return hits[0].parent
-
-    # fallback to any of the other known files
-    for pat in ["**/chroma-collections.parquet", "**/index_metadata.pickle",
-                "**/data_level0.bin", "**/chroma.sqlite"]:
-        alt = list(base.rglob(pat))
-        if alt:
-            return alt[0].parent
+    # look a bit deeper
+    for pat in (
+        "**/chroma.sqlite3",
+        "**/chroma.sqlite",
+        "**/chroma-collections.parquet",
+        "**/index_metadata.pickle",
+        "**/data_level0.bin",
+    ):
+        found = list(base.rglob(pat))
+        if found:
+            return found[0].parent
 
     return base  # may be empty; ensure_ready() will warn
 
@@ -89,7 +100,7 @@ INDEX_DIR = _find_catalog_dir(_BASE_INDEX_DIR)
 def _get_client_and_collection():
     """
     Open the persistent Chroma client on INDEX_DIR and return an existing collection.
-    If no collection is listed, fall back to 'materialmind'.
+    If none is listed, fall back to 'materialmind'.
     """
     import chromadb
     client = chromadb.PersistentClient(path=str(INDEX_DIR))
@@ -99,23 +110,20 @@ def _get_client_and_collection():
             # Use the first existing collection (the one you built locally)
             return client, client.get_collection(cols[0].name)
     except Exception as e:
-        print(f"[RAG] Warning: list_collections failed: {e}")
+        print(f"[RAG] list_collections failed: {e}")
 
     # Fallback: create/get 'materialmind'
     return client, client.get_or_create_collection(name="materialmind")
 
 # ---------- Public API ----------
 def ensure_ready() -> None:
-    """
-    Called by app.py at startup. Verifies the index path and prints a small stat.
-    """
+    """Called by app.py at startup. Verifies the index path and prints a small stat."""
     INDEX_DIR.mkdir(parents=True, exist_ok=True)
     if not _has_catalog(INDEX_DIR):
         print(f"[RAG] WARNING: No Chroma catalog found in: {INDEX_DIR}")
-        print("      If you uploaded a nested folder (UUID), set INDEX_DIR env/secret to that folder,")
-        print("      or keep the layout: MaterialMind/index/chroma_v3/<uuid>/chroma.sqlite3")
+        print("      If your files are inside a UUID folder, set INDEX_DIR to that folder,")
+        print("      e.g. MaterialMind/index/chroma_v3/22ddc0e5-...-d7a3fde16dd5")
         return
-
     try:
         stats = index_stats()
         print(f"[RAG] Index OK at {INDEX_DIR} — chunks={stats.get('count')}")
@@ -130,7 +138,6 @@ def search(query: str, k: int = DEFAULT_TOPK) -> List[Tuple[str, str]]:
     if not _has_catalog(INDEX_DIR):
         print("[RAG] search aborted: catalog not found")
         return []
-
     try:
         _, col = _get_client_and_collection()
         qvec = _embed([query])[0]
@@ -158,9 +165,7 @@ def search(query: str, k: int = DEFAULT_TOPK) -> List[Tuple[str, str]]:
     return hits
 
 def index_stats() -> dict:
-    """
-    Return simple stats {count: int}. Safe if catalog missing.
-    """
+    """Return simple stats {count: int}. Safe if catalog missing."""
     if not _has_catalog(INDEX_DIR):
         return {"count": 0, "note": f"no catalog at {INDEX_DIR}"}
     try: