InitConfig.py 1.0 KB

123456789101112131415161718192021222324252627282930313233343536
  1. import os
  2. from embedchain.config.BaseConfig import BaseConfig
  3. class InitConfig(BaseConfig):
  4. """
  5. Config to initialize an embedchain `App` instance.
  6. """
  7. def __init__(self, ef=None, db=None):
  8. """
  9. :param ef: Optional. Embedding function to use.
  10. :param db: Optional. (Vector) database to use for embeddings.
  11. """
  12. # Embedding Function
  13. if ef is None:
  14. from chromadb.utils import embedding_functions
  15. self.ef = embedding_functions.OpenAIEmbeddingFunction(
  16. api_key=os.getenv("OPENAI_API_KEY"),
  17. organization_id=os.getenv("OPENAI_ORGANIZATION"),
  18. model_name="text-embedding-ada-002"
  19. )
  20. else:
  21. self.ef = ef
  22. if db is None:
  23. from embedchain.vectordb.chroma_db import ChromaDB
  24. self.db = ChromaDB(ef=self.ef)
  25. else:
  26. self.db = db
  27. return
  28. def _set_embedding_function(self, ef):
  29. self.ef = ef
  30. return