refactor

utils.py

@@ -0,0 +1,48 @@
+import logging
+import os
+import platform
+from typing import Literal
+
+import chromadb
+from dotenv import load_dotenv
+
+os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
+
+
+# Set up logging
+logging.basicConfig(
+    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
+)
+logger = logging.getLogger(__name__)
+
+load_dotenv()
+os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
+# Set up logging
+logging.basicConfig(
+    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
+)
+logger = logging.getLogger(__name__)
+
+load_dotenv()
+
+
+def get_save_path() -> Literal["chroma/"] | Literal["/data/chroma/"]:
+    path = "chroma/" if platform.system() == "Darwin" else "/data/chroma/"
+    logger.info(f"Using save path: {path}")
+    return path
+
+
+def get_chroma_client():
+    logger.info("Initializing Chroma client")
+    SAVE_PATH = get_save_path()
+    return chromadb.PersistentClient(path=SAVE_PATH)
+
+
+def get_collection(chroma_client, embedding_function, collection_name):
+    logger.info(f"Getting or creating collection: {collection_name}")
+    return chroma_client.create_collection(
+        name=collection_name,
+        get_or_create=True,
+        embedding_function=embedding_function,
+        metadata={"hnsw:space": "cosine"},
+    )