Fixing Storage

src/storage/chroma_store.py

@@ -1,13 +1,27 @@
 import os
 import chromadb
-from chromadb.config import Settings
 from src.storage.paths import nb_root
 
-def chroma_client(username: str, notebook_id: str):
-    persist_dir = os.path.join(nb_root(username, notebook_id), "chroma")
-    os.makedirs(persist_dir, exist_ok=True)
-    return chromadb.PersistentClient(path=persist_dir, settings=Settings(anonymized_telemetry=False))
 
-def get_collection(username: str, notebook_id: str):
+def chroma_client(username, notebook_id):
+
+    persist_dir = os.path.join(
+        nb_root(username, notebook_id),
+        "chroma"
+    )
+
+    return chromadb.Client(
+        chromadb.config.Settings(
+            persist_directory=persist_dir,
+            anonymized_telemetry=False
+        )
+    )
+
+
+def get_collection(username, notebook_id):
+
     client = chroma_client(username, notebook_id)
-    return client.get_or_create_collection(name="docs")
+
+    return client.get_or_create_collection(
+        name="notebook"
+    )

src/storage/paths.py

@@ -1,33 +1,18 @@
 import os
-from pathlib import Path
 
-# If DATA_ROOT env var is not set:
-# - Locally: write to ./data (project folder)
-# - On HF: you will set DATA_ROOT=/data in Space variables (or leave it as /data there)
-DEFAULT_LOCAL_DATA = str(Path(__file__).resolve().parents[2] / "data")
+DATA_ROOT = os.getenv("DATA_ROOT", "./data")
 
-DATA_ROOT = os.environ.get("DATA_ROOT", DEFAULT_LOCAL_DATA)
+def user_root(username: str):
+    return os.path.join(DATA_ROOT, username)
 
-def user_root(username: str) -> str:
-    return os.path.join(DATA_ROOT, "users", username, "notebooks")
-
-def index_path(username: str) -> str:
-    return os.path.join(user_root(username), "index.json")
-
-def nb_root(username: str, notebook_id: str) -> str:
+def nb_root(username: str, notebook_id: str):
     return os.path.join(user_root(username), notebook_id)
 
 def ensure_tree(username: str, notebook_id: str):
+
     base = nb_root(username, notebook_id)
-    paths = [
-        user_root(username),
-        os.path.join(base, "files_raw"),
-        os.path.join(base, "files_extracted"),
-        os.path.join(base, "chroma"),
-        os.path.join(base, "chat"),
-        os.path.join(base, "artifacts", "reports"),
-        os.path.join(base, "artifacts", "quizzes"),
-        os.path.join(base, "artifacts", "podcasts"),
-    ]
-    for p in paths:
-        os.makedirs(p, exist_ok=True)
+
+    os.makedirs(base, exist_ok=True)
+    os.makedirs(os.path.join(base, "files_raw"), exist_ok=True)
+    os.makedirs(os.path.join(base, "files_extracted"), exist_ok=True)
+    os.makedirs(os.path.join(base, "artifacts"), exist_ok=True)