AppConfig.py 1.9 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152
  1. import os
  2. try:
  3. from chromadb.utils import embedding_functions
  4. except RuntimeError:
  5. from embedchain.utils import use_pysqlite3
  6. use_pysqlite3()
  7. from chromadb.utils import embedding_functions
  8. from .BaseAppConfig import BaseAppConfig
  9. class AppConfig(BaseAppConfig):
  10. """
  11. Config to initialize an embedchain custom `App` instance, with extra config options.
  12. """
  13. def __init__(self, log_level=None, host=None, port=None, id=None, collection_name=None):
  14. """
  15. :param log_level: Optional. (String) Debug level
  16. ['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'].
  17. :param host: Optional. Hostname for the database server.
  18. :param port: Optional. Port for the database server.
  19. :param id: Optional. ID of the app. Document metadata will have this id.
  20. :param collection_name: Optional. Collection name for the database.
  21. """
  22. super().__init__(
  23. log_level=log_level,
  24. embedding_fn=AppConfig.default_embedding_function(),
  25. host=host,
  26. port=port,
  27. id=id,
  28. collection_name=collection_name,
  29. )
  30. @staticmethod
  31. def default_embedding_function():
  32. """
  33. Sets embedding function to default (`text-embedding-ada-002`).
  34. :raises ValueError: If the template is not valid as template should contain
  35. $context and $query
  36. :returns: The default embedding function for the app class.
  37. """
  38. if os.getenv("OPENAI_API_KEY") is None and os.getenv("OPENAI_ORGANIZATION") is None:
  39. raise ValueError("OPENAI_API_KEY or OPENAI_ORGANIZATION environment variables not provided") # noqa:E501
  40. return embedding_functions.OpenAIEmbeddingFunction(
  41. api_key=os.getenv("OPENAI_API_KEY"),
  42. organization_id=os.getenv("OPENAI_ORGANIZATION"),
  43. model_name="text-embedding-ada-002",
  44. )