fix

ragg/ingest_all.py

@@ -1,9 +1,10 @@
 # ragg/ingest_all.py
 import os
 import sys
+from pathlib import Path
 from langchain_community.vectorstores import Chroma
 
-# Allow running both as a module and as a script
+# Support both module and direct runs
 if __package__ in (None, ""):
     sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
     from ragg.rag_backend import ingest_pdfs_from_folder, get_embeddings
@@ -12,58 +13,67 @@ else:
     from ragg.rag_backend import ingest_pdfs_from_folder, get_embeddings
     import ragg.rag_backend as rag_backend
 
-IS_PROD = os.getenv("ENV", "dev").lower() == "prod"
-HERE = os.path.dirname(__file__)
+# Detect environment
+IS_HF = bool(os.getenv("HF_HOME") or os.getenv("SPACE_ID"))
+HERE = Path(__file__).resolve().parent
 
-# ✅ Point to the actual PDF location inside the Space
-PDFS_ROOT = os.path.abspath(os.path.join(HERE, "pdfs"))  # /app/ragg/pdfs (fallback)
-if not os.path.isdir(PDFS_ROOT):
-    # Hugging Face uses /app as working dir → /app/pdfs/*
-    PDFS_ROOT = os.path.abspath(os.path.join(HERE, "..", "pdfs"))
+# PDF root auto-detect
+PDFS_ROOT = (HERE / "pdfs")
+if not PDFS_ROOT.is_dir():
+    PDFS_ROOT = (HERE.parent / "pdfs")  # Works for /app/pdfs/*
 
-# ✅ Always persist Chroma under /data/chroma (mounted persistent dir)
-CHROMA_BASE = os.getenv("CHROMA_ROOT", "/data/chroma")
+# Chroma root auto-detect
+CHROMA_BASE = Path(os.getenv("CHROMA_ROOT") or ("/data/chroma" if IS_HF else "./chroma"))
+CHROMA_BASE.mkdir(parents=True, exist_ok=True)
 
 def ingest_all_levels():
-    """
-    Ingest PDFs from ./pdfs/{low,mid,high} into separate Chroma DBs.
-    Works both locally and on Hugging Face.
-    """
-    print("\n🚀 Starting ingestion\n")
+    print("\n🚀 Starting ingestion")
     print(f"📂 PDFs root: {PDFS_ROOT}")
-    print(f"💾 Chroma base: {CHROMA_BASE} (ENV={'prod' if IS_PROD else 'dev'})\n")
+    print(f"💾 Chroma base: {CHROMA_BASE}\n")
 
-    levels = ["low", "mid", "high"]
-
-    for name in levels:
-        folder_path = os.path.join(PDFS_ROOT, name)
-        if not os.path.isdir(folder_path):
-            print(f"⚠️ Skip {name}: {folder_path} not found")
+    for level in ["low", "mid", "high"]:
+        folder_path = PDFS_ROOT / level
+        if not folder_path.is_dir():
+            print(f"⚠️ Skip {level}: {folder_path} not found")
             continue
 
-        chroma_dir = os.path.join(CHROMA_BASE, name)
-        os.makedirs(chroma_dir, exist_ok=True)
+        chroma_dir = CHROMA_BASE / level
+        chroma_dir.mkdir(parents=True, exist_ok=True)
+        collection_name = f"pylearn_{level}"
 
-        # Rebind vectorstore factory used by your backend for this run
+        # Monkey-patch vectorstore factory for this run
         def get_vectorstore_for_level():
             print(f"🔹 Using Chroma at: {chroma_dir}")
-            return Chroma(
-                persist_directory=chroma_dir,
-                embedding_function=get_embeddings()
+            print(f"🔹 Collection: {collection_name}")
+            vs = Chroma(
+                collection_name=collection_name,
+                persist_directory=str(chroma_dir),
+                embedding_function=get_embeddings(),
             )
+            return vs
 
         rag_backend.get_vectorstore = get_vectorstore_for_level
 
         try:
             print(f"➡️ Ingesting {folder_path}")
             res = ingest_pdfs_from_folder(
-                folder_path, subject="English", grade="5", chapter=name
+                folder_path=str(folder_path),
+                subject="English",
+                grade="5",
+                chapter=level,
             )
-            print(f"✅ {name}: {res}")
+            # Persist the vectorstore to disk
+            vs = rag_backend.get_vectorstore()
+            if hasattr(vs, "persist"):
+                vs.persist()
+                print("📝 Called persist()")
+
+            print(f"✅ {level}: {res}")
             print(f"📦 Stored in: {chroma_dir}\n")
+
         except Exception as e:
             import traceback
-            print(f"❌ {name}: {e}")
+            print(f"❌ {level}: {e}")
             print(traceback.format_exc())
 
     print("🎯 Ingestion complete.\n")

ragg/rag_backend.py

@@ -83,15 +83,47 @@ def _vs_count_safe(vs) -> Optional[int]:
         except Exception:
             return None
 
+# def get_vectorstore():
+#     global _vectorstore
+#     if _vectorstore is None:
+#         os.makedirs(CHROMA_DIR, exist_ok=True)
+#         print(f"🔹 Loading Chroma vectorstore at: {CHROMA_DIR}")
+#         _vectorstore = Chroma(
+#             persist_directory=CHROMA_DIR,
+#             embedding_function=get_embeddings()
+#         )
+#     cnt = _vs_count_safe(_vectorstore)
+#     if cnt is not None:
+#         print(f"📦 Vectorstore currently has ~{cnt} chunks.")
+#     else:
+#         print("📦 Vectorstore count not available (skipping).")
+#     return _vectorstore
+
 def get_vectorstore():
+    """
+    Returns a Chroma vectorstore that works in both local and Hugging Face environments.
+    - Uses CHROMA_DIR if defined (e.g., /data/chroma/low)
+    - Defaults to ./chroma when running locally
+    - Monkey-patching from ingest_all.py overrides this function
+    """
     global _vectorstore
-    if _vectorstore is None:
-        os.makedirs(CHROMA_DIR, exist_ok=True)
-        print(f"🔹 Loading Chroma vectorstore at: {CHROMA_DIR}")
-        _vectorstore = Chroma(
-            persist_directory=CHROMA_DIR,
-            embedding_function=get_embeddings()
-        )
+    if _vectorstore is not None:
+        return _vectorstore
+
+    # Detect environment (Hugging Face vs Local)
+    hf_env = bool(os.getenv("HF_HOME") or os.getenv("SPACE_ID"))
+    default_chroma = "/data/chroma" if hf_env else "./chroma"
+
+    # Final resolved path (can be overridden via env variable)
+    chroma_dir = os.getenv("CHROMA_DIR", default_chroma)
+    os.makedirs(chroma_dir, exist_ok=True)
+
+    print(f"🔹 Loading Chroma vectorstore at: {chroma_dir}")
+    _vectorstore = Chroma(
+        persist_directory=chroma_dir,
+        embedding_function=get_embeddings()
+    )
+
     cnt = _vs_count_safe(_vectorstore)
     if cnt is not None:
         print(f"📦 Vectorstore currently has ~{cnt} chunks.")
@@ -99,6 +131,7 @@ def get_vectorstore():
         print("📦 Vectorstore count not available (skipping).")
     return _vectorstore
 
+
 # ---------------- Text Splitter ---------------- #
 def chunk_docs(docs: List[Document], chunk_size=1200, chunk_overlap=150) -> List[Document]:
     splitter = RecursiveCharacterTextSplitter(