feat: allow switching between collections and specifying initial coll…

…ection name

docs/advanced/configuration.mdx

@@ -24,6 +24,10 @@ config = InitConfig(ef=embedding_functions.OpenAIEmbeddingFunction(
  ))
 naval_chat_bot = App(config)
 
+# Example: specify a custom collection name
+config = InitConfig(collection_name="naval_chat_bot")
+naval_chat_bot = App(config)
+
 # Example: define your own chunker config for `youtube_video`
 youtube_add_config = {
  "chunker": {

docs/advanced/query_configuration.mdx

@@ -9,6 +9,7 @@ title: '🔍 Query configurations'
 | log_level | log level | string | WARNING |
 | ef | embedding function | chromadb.utils.embedding_functions | \{text-embedding-ada-002\} |
 | db | vector database (experimental) | BaseVectorDB | ChromaDB |
+| collection_name | initial collection name for the database | string | embedchain_store |
 
 
 ## AddConfig

embedchain/config/InitConfig.py

@@ -11,19 +11,21 @@ class InitConfig(BaseConfig):
  Config to initialize an embedchain `App` instance.
  """
 
- def __init__(self, log_level=None, ef=None, db=None, host=None, port=None, id=None):
+ def __init__(self, log_level=None, ef=None, db=None, collection_name=None, host=None, port=None, id=None):
  """
  :param log_level: Optional. (String) Debug level
  ['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'].
  :param ef: Optional. Embedding function to use.
  :param db: Optional. (Vector) database to use for embeddings.
+ :param collection_name: Optional. Collection name for the database.
  :param id: Optional. ID of the app. Document metadata will have this id.
  :param host: Optional. Hostname for the database server.
  :param port: Optional. Port for the database server.
  """
  self._setup_logging(log_level)
  self.ef = ef
  self.db = db
+ self.collection_name = collection_name
  self.host = host
  self.port = port
  self.id = id

embedchain/embedchain.py

@@ -33,7 +33,7 @@ def __init__(self, config: InitConfig):
 
  self.config = config
  self.db_client = self.config.db.client
- self.collection = self.config.db.collection
+ self.collection = self.config.db._get_or_create_collection(self.config.collection_name)
  self.user_asks = []
  self.is_code_docs_instance = False
  self.online = False
@@ -201,6 +201,7 @@ def get_answer_from_llm(self, prompt, config: ChatConfig):
 
  def access_search_and_get_results(self, input_query):
  from langchain.tools import DuckDuckGoSearchRun
+
  search = DuckDuckGoSearchRun()
  logging.info(f"Access search to get answers for {input_query}")
  return search.run(input_query)
@@ -314,6 +315,14 @@ def dry_run(self, input_query, config: QueryConfig = None):
  logging.info(f"Prompt: {prompt}")
  return prompt
 
+ def set_collection(self, collection_name):
+ """
+ Set the collection to use.
+
+ :param name: The name of the collection to use.
+ """
+ self.collection = self.config.db._get_or_create_collection(collection_name)
+
  def count(self):
  """
  Count the number of embeddings.

embedchain/vectordb/chroma_db.py

@@ -32,9 +32,11 @@ def _get_or_create_db(self):
  """Get or create the database."""
  return chromadb.Client(self.client_settings)
 
- def _get_or_create_collection(self):
+ def _get_or_create_collection(self, name=None):
  """Get or create the collection."""
+ if name is None:
+ name = "embedchain_store"
  return self.client.get_or_create_collection(
- "embedchain_store",
+ name=name,
  embedding_function=self.ef,
  )

tests/vectordb/test_chroma_db.py

@@ -71,3 +71,33 @@ def test_init_with_host_and_port(self, mock_client):
 
  self.assertEqual(mock_client.call_args[0][0].chroma_server_host, None)
  self.assertEqual(mock_client.call_args[0][0].chroma_server_http_port, None)
+
+
+class TestChromaDbDefaultCollection(unittest.TestCase):
+ def test_init_with_default_collection(self):
+ """
+ Test if the `App` instance is initialized with the correct default collection name.
+ """
+ app = App()
+
+ self.assertEqual(app.collection.name, "embedchain_store")
+
+ def test_init_with_custom_collection(self):
+ """
+ Test if the `App` instance is initialized with the correct custom collection name.
+ """
+ config = InitConfig(collection_name="test_collection")
+ app = App(config)
+
+ self.assertEqual(app.collection.name, "test_collection")
+
+
+class TestChromaDbSetCollection(unittest.TestCase):
+ def test_init_with_host_and_port(self):
+ """
+ Test if the `App` collection is correctly switched using the `set_collection` method.
+ """
+ app = App()
+ app.set_collection("test_collection")
+
+ self.assertEqual(app.collection.name, "test_collection")