feat: collection name everywhere

docs/advanced/configuration.mdx

@@ -20,6 +20,10 @@ from chromadb.utils import embedding_functions
 config = AppConfig(log_level="DEBUG")
 naval_chat_bot = App(config)
 
+# Example: specify a custom collection name
+config = AppConfig(collection_name="naval_chat_bot")
+naval_chat_bot = App(config)
+
 # Example: define your own chunker config for `youtube_video`
 chunker_config = ChunkerConfig(chunk_size=1000, chunk_overlap=100, length_function=len)
 naval_chat_bot.add("youtube_video", "https://www.youtube.com/watch?v=3qHkcs3kG44", AddConfig(chunker=chunker_config))

docs/advanced/query_configuration.mdx

@@ -4,11 +4,12 @@ title: '🔍 Query configurations'
 
 ## AppConfig
 
-| option  | description | type | default |
-|-------------|-----------------------|---------------------------------|------------------------|
-| log_level  | log level | string | WARNING |
+| option | description | type | default |
+|-----------|-----------------------|---------------------------------|------------------------|
+| log_level | log level | string | WARNING |
 | embedding_fn| embedding function | chromadb.utils.embedding_functions | \{text-embedding-ada-002\} |
-| db | vector database (experimental) | BaseVectorDB | ChromaDB |
+| db | vector database (experimental) | BaseVectorDB | ChromaDB |
+| collection_name | initial collection name for the database | string | embedchain_store |
 
 
 ## AddConfig

embedchain/config/apps/AppConfig.py

@@ -10,16 +10,22 @@ class AppConfig(BaseAppConfig):
  Config to initialize an embedchain custom `App` instance, with extra config options.
  """
 
- def __init__(self, log_level=None, host=None, port=None, id=None):
+ def __init__(self, log_level=None, collection_name=None, host=None, port=None, id=None):
  """
  :param log_level: Optional. (String) Debug level
  ['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'].
  :param id: Optional. ID of the app. Document metadata will have this id.
+ :param collection_name: Optional. Collection name for the database.
  :param host: Optional. Hostname for the database server.
  :param port: Optional. Port for the database server.
  """
  super().__init__(
- log_level=log_level, embedding_fn=AppConfig.default_embedding_function(), host=host, port=port, id=id
+ log_level=log_level,
+ collection_name=collection_name,
+ embedding_fn=AppConfig.default_embedding_function(),
+ host=host,
+ port=port,
+ id=id,
  )
 
  @staticmethod

embedchain/config/apps/BaseAppConfig.py

@@ -8,19 +8,21 @@ class BaseAppConfig(BaseConfig):
  Parent config to initialize an instance of `App`, `OpenSourceApp` or `CustomApp`.
  """
 
- def __init__(self, log_level=None, embedding_fn=None, db=None, host=None, port=None, id=None):
+ def __init__(self, log_level=None, embedding_fn=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 embedding_fn: Embedding function to use.
  :param db: Optional. (Vector) database instance 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.db = db if db else BaseAppConfig.default_db(embedding_fn=embedding_fn, host=host, port=port)
+ self.collection_name = collection_name
  self.id = id
  return
 

embedchain/config/apps/CustomAppConfig.py

@@ -21,6 +21,7 @@ def __init__(
  embedding_fn: EmbeddingFunctions = None,
  embedding_fn_model=None,
  db=None,
+ collection_name=None,
  host=None,
  port=None,
  id=None,
@@ -34,6 +35,7 @@ def __init__(
  :param embedding_fn: Optional. Embedding function to use.
  :param embedding_fn_model: Optional. Model name to use for embedding function.
  :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.
@@ -51,6 +53,7 @@ def __init__(
  log_level=log_level,
  embedding_fn=CustomAppConfig.embedding_function(embedding_function=embedding_fn, model=embedding_fn_model),
  db=db,
+ collection_name=collection_name,
  host=host,
  port=port,
  id=id,

embedchain/config/apps/OpenSourceAppConfig.py

@@ -8,11 +8,12 @@ class OpenSourceAppConfig(BaseAppConfig):
  Config to initialize an embedchain custom `OpenSourceApp` instance, with extra config options.
  """
 
- def __init__(self, log_level=None, host=None, port=None, id=None, model=None):
+ def __init__(self, log_level=None, collection_name=None, host=None, port=None, id=None, model=None):
  """
  :param log_level: Optional. (String) Debug level
  ['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'].
  :param id: Optional. ID of the app. Document metadata will have this id.
+ :param collection_name: Optional. Collection name for the database.
  :param host: Optional. Hostname for the database server.
  :param port: Optional. Port for the database server.
  :param model: Optional. GPT4ALL uses the model to instantiate the class.
@@ -23,6 +24,7 @@ def __init__(self, log_level=None, host=None, port=None, id=None, model=None):
  super().__init__(
  log_level=log_level,
  embedding_fn=OpenSourceAppConfig.default_embedding_function(),
+ collection_name=collection_name,
  host=host,
  port=port,
  id=id,

embedchain/embedchain.py

@@ -30,7 +30,7 @@ def __init__(self, config: BaseAppConfig):
 
  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_docs_site_instance = False
  self.online = False
@@ -322,6 +322,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

@@ -35,9 +35,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.embedding_fn,
  )

tests/vectordb/test_chroma_db.py

@@ -3,7 +3,7 @@
 import unittest
 from unittest.mock import patch
 
-from embedchain import App
+from embedchain.apps.App import App
 from embedchain.config import AppConfig
 from embedchain.vectordb.chroma_db import ChromaDB, chromadb
 
@@ -71,3 +71,31 @@ 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 TestChromaDbCollection(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 = AppConfig(collection_name="test_collection")
+ app = App(config)
+
+ self.assertEqual(app.collection.name, "test_collection")
+
+ def test_set_collection(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")