chroma.py 1.6 KB

12345678910111213141516171819202122232425262728293031323334353637
  1. from typing import Optional
  2. from embedchain.config.vectordb.base import BaseVectorDbConfig
  3. from embedchain.helpers.json_serializable import register_deserializable
  4. @register_deserializable
  5. class ChromaDbConfig(BaseVectorDbConfig):
  6. def __init__(
  7. self,
  8. collection_name: Optional[str] = None,
  9. dir: Optional[str] = None,
  10. host: Optional[str] = None,
  11. port: Optional[str] = None,
  12. allow_reset=False,
  13. chroma_settings: Optional[dict] = None,
  14. ):
  15. """
  16. Initializes a configuration class instance for ChromaDB.
  17. :param collection_name: Default name for the collection, defaults to None
  18. :type collection_name: Optional[str], optional
  19. :param dir: Path to the database directory, where the database is stored, defaults to None
  20. :type dir: Optional[str], optional
  21. :param host: Database connection remote host. Use this if you run Embedchain as a client, defaults to None
  22. :type host: Optional[str], optional
  23. :param port: Database connection remote port. Use this if you run Embedchain as a client, defaults to None
  24. :type port: Optional[str], optional
  25. :param allow_reset: Resets the database. defaults to False
  26. :type allow_reset: bool
  27. :param chroma_settings: Chroma settings dict, defaults to None
  28. :type chroma_settings: Optional[dict], optional
  29. """
  30. self.chroma_settings = chroma_settings
  31. self.allow_reset = allow_reset
  32. super().__init__(collection_name=collection_name, dir=dir, host=host, port=port)