utils.py

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"

from chromadb.config import Settings

# 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()

CHROMA_CLIENT = None


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():
    global CHROMA_CLIENT
    if CHROMA_CLIENT is None:
        SAVE_PATH = get_save_path()
        CHROMA_CLIENT = chromadb.PersistentClient(
            path=SAVE_PATH,
            settings=Settings(anonymized_telemetry=False, is_persistent=True),
        )
    return CHROMA_CLIENT


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"},
    )