AppConfig.py 2.1 KB

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