utils.py

import chromadb
from chromadb.config import Settings
import chromadb.utils.embedding_functions as embedding_functions
from dotenv import load_dotenv
import os

load_dotenv()
openai_key = os.getenv("OPENAI_API_KEY")


def get_chroma_client(
    host: str = "chroma.brianknows.org",
    port: int = 443,
) -> chromadb.HttpClient:
    chroma_client = chromadb.HttpClient(
        host=host,
        port=443,
        ssl=port,
        settings=Settings(
            allow_reset=True,
        ),
    )

    return chroma_client


def get_embedding_function(model_name="text-embedding-ada-002"):
    openai_ef = embedding_functions.OpenAIEmbeddingFunction(
        api_key=openai_key, model_name=model_name
    )
    return openai_ef