chroma_db.py 1.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142
  1. import logging
  2. import os
  3. import chromadb
  4. from chromadb.utils import embedding_functions
  5. from embedchain.vectordb.base_vector_db import BaseVectorDB
  6. class ChromaDB(BaseVectorDB):
  7. """Vector database using ChromaDB."""
  8. def __init__(self, db_dir=None, ef=None, host=None, port=None):
  9. self.ef = ef
  10. if host and port:
  11. logging.info(f"Connecting to ChromaDB server: {host}:{port}")
  12. self.client_settings = chromadb.config.Settings(
  13. chroma_api_impl="rest",
  14. chroma_server_host=host,
  15. chroma_server_http_port=port,
  16. )
  17. else:
  18. if db_dir is None:
  19. db_dir = "db"
  20. self.client_settings = chromadb.config.Settings(
  21. chroma_db_impl="duckdb+parquet",
  22. persist_directory=db_dir,
  23. anonymized_telemetry=False,
  24. )
  25. super().__init__()
  26. def _get_or_create_db(self):
  27. """Get or create the database."""
  28. return chromadb.Client(self.client_settings)
  29. def _get_or_create_collection(self):
  30. """Get or create the collection."""
  31. return self.client.get_or_create_collection(
  32. "embedchain_store",
  33. embedding_function=self.ef,
  34. )