InitConfig.py 1.7 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152
  1. import os
  2. import logging
  3. from embedchain.config.BaseConfig import BaseConfig
  4. class InitConfig(BaseConfig):
  5. """
  6. Config to initialize an embedchain `App` instance.
  7. """
  8. def __init__(self, log_level=None, ef=None, db=None):
  9. """
  10. :param log_level: Optional. (String) Debug level ['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'].
  11. :param ef: Optional. Embedding function to use.
  12. :param db: Optional. (Vector) database to use for embeddings.
  13. """
  14. self._setup_logging(log_level)
  15. # Embedding Function
  16. if ef is None:
  17. from chromadb.utils import embedding_functions
  18. self.ef = embedding_functions.OpenAIEmbeddingFunction(
  19. api_key=os.getenv("OPENAI_API_KEY"),
  20. organization_id=os.getenv("OPENAI_ORGANIZATION"),
  21. model_name="text-embedding-ada-002"
  22. )
  23. else:
  24. self.ef = ef
  25. if db is None:
  26. from embedchain.vectordb.chroma_db import ChromaDB
  27. self.db = ChromaDB(ef=self.ef)
  28. else:
  29. self.db = db
  30. return
  31. def _set_embedding_function(self, ef):
  32. self.ef = ef
  33. return
  34. def _setup_logging(self, debug_level):
  35. level = logging.WARNING # Default level
  36. if debug_level is not None:
  37. level = getattr(logging, debug_level.upper(), None)
  38. if not isinstance(level, int):
  39. raise ValueError(f'Invalid log level: {debug_level}')
  40. logging.basicConfig(format="%(asctime)s [%(name)s] [%(levelname)s] %(message)s",
  41. level=level)
  42. self.logger = logging.getLogger(__name__)
  43. return