InitConfig.py 1.2 KB

12345678910111213141516171819202122232425262728293031323334353637383940
  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, stream_response=False):
  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. if not isinstance(stream_response, bool):
  28. raise ValueError("`stream_respone` should be bool")
  29. self.stream_response = stream_response
  30. return
  31. def _set_embedding_function(self, ef):
  32. self.ef = ef
  33. return