import os import hashlib from langchain_community.vectorstores import Chroma from langchain_openai import OpenAIEmbeddings def get_base_directory() -> str: """ HF Spaces-safe data root: - Prefer explicit override via APP_DATA_DIR or PERSISTENT_STORAGE_PATH - Else use /data/ only if it exists and is writable - Else use $HOME/.cache/ - Else use /tmp/ """ override = (os.environ.get("APP_DATA_DIR") or os.environ.get("PERSISTENT_STORAGE_PATH") or "").strip() if override: os.makedirs(override, exist_ok=True) return override if os.path.isdir("/data") and os.access("/data", os.W_OK): base = "/data/pdf-rag" os.makedirs(base, exist_ok=True) return base home = os.environ.get("HOME") if home: base = os.path.join(home, ".cache", "pdf-rag") os.makedirs(base, exist_ok=True) return base base = "/tmp/pdf-rag" os.makedirs(base, exist_ok=True) return base # Define path relative to the base directory CHROMA_DB_BASE = os.path.join(get_base_directory(), "root databases") os.makedirs(CHROMA_DB_BASE, exist_ok=True) def compute_file_hash(content: bytes) -> str: return hashlib.md5(content).hexdigest() def get_vectorstore(category: str, embedding_model): path = os.path.join(CHROMA_DB_BASE, category.strip()) os.makedirs(path, exist_ok=True) return Chroma(persist_directory=path, embedding_function=embedding_model)