Bump version to 0.0.37 (#431)

Add an example of api server so that devs can quickly get up a bot running along with its api

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/data_types.mdx

@@ -54,6 +54,18 @@ To add any code documentation website as a loader, use the data_type as `docs_si
 app.add("docs_site", "https://docs.embedchain.ai/")
 ```
 
+### Notion
+To use notion you must install the extra dependencies with `pip install embedchain[notion]`.
+
+To load a notion page, use the data_type as `notion`.
+The next argument must **end** with the `notion page id`. The id is a 32-character string. Eg:
+
+```python
+app.add("notion", "cfbc134ca6464fc980d0391613959196")
+app.add("notion", "my-page-cfbc134ca6464fc980d0391613959196")
+app.add("notion", "https://www.notion.so/my-page-cfbc134ca6464fc980d0391613959196")
+```
+
 ### Text
 
 To supply your own text, use the data_type as `text` and enter a string. The text is not processed, this can be very versatile. Eg:

docs/advanced/query_configuration.mdx

@@ -4,11 +4,13 @@ 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 |
+| collect_metrics | collect anonymous telemetry data to improve embedchain | boolean | true |
 
 
 ## AddConfig
@@ -45,6 +47,7 @@ Default values of chunker config parameters for different `data_type`:
 |pdf_file|1000|0|len|
 |youtube_video|2000|0|len|
 |docs_site|500|50|len|
+|notion|300|0|len|
 
 ### LoaderConfig
 

docs/advanced/showcase.mdx

@@ -17,6 +17,10 @@ Embedchain community has been super active in creating demos on top of Embedchai
 
 - [Taobot.io](https://taobot.io) - chatbot & knowledgebase hybrid by [cachho](https://github.com/cachho)
 - [Create Instant ChatBot 🤖 using embedchain](https://databutton.com/v/h3e680h9) by Avra, ([Tweet](https://twitter.com/Avra_b/status/1674704745154641920/))
+- [JOBO 🤖 — The AI-driven sidekick to craft your resume](https://try-jobo.com/) by Enrico Willemse, ([LinkedIn Post](https://www.linkedin.com/posts/enrico-willemse_jobai-gptfun-embedchain-activity-7090340080879374336-ueLB/))
+- [Explore Your Knowledge Base: Interactive chats over various forms of documents](https://chatdocs.dkedar.com/) by Kedar Dabhadkar, ([LinkedIn Post](https://www.linkedin.com/posts/dkedar7_machinelearning-llmops-activity-7092524836639424513-2O3L/))
+- [Chatbot trained on 1000+ videos of Ester hicks the co-author behind the famous book Secret](https://ask-abraham.thoughtseed.repl.co) by Mohan Kumar
+
 
 ## Templates
 
@@ -33,6 +37,7 @@ Embedchain community has been super active in creating demos on top of Embedchai
 - [Chatbot in Under 5 mins using Embedchain](https://medium.com/@ayush.wattal/chatbot-in-under-5-mins-using-embedchain-a4f161fcf9c5) by Ayush Wattal
 - [Understanding what the LLM framework embedchain does](https://zenn.dev/hijikix/articles/4bc8d60156a436) by Daisuke Hashimoto
 - [In bed with GPT and Node.js](https://dev.to/worldlinetech/in-bed-with-gpt-and-nodejs-4kh2) by Raphaël Semeteys, ([LinkedIn Post](https://www.linkedin.com/posts/raphaelsemeteys_in-bed-with-gpt-and-nodejs-activity-7088113552326029313-nn87/))
+- [Using Embedchain — A powerful LangChain Python wrapper to build Chat Bots even faster!⚡](https://medium.com/@avra42/using-embedchain-a-powerful-langchain-python-wrapper-to-build-chat-bots-even-faster-35c12994a360) by Avra, ([Tweet](https://twitter.com/Avra_b/status/1686767751560310784/))
 
 ### LinkedIn
 
@@ -56,6 +61,7 @@ Embedchain community has been super active in creating demos on top of Embedchai
 - [Highlights of EmbedChain](https://twitter.com/carl_AIwarts/status/1673542204328120321/) by carl_AIwarts
 - [Build Llama-2 chatbot in less than 5 minutes](https://twitter.com/Saboo_Shubham_/status/1682168956918833152/) by Shubham Saboo
 - [All cool features of embedchain](https://twitter.com/DhravyaShah/status/1683497882438217728/) by Dhravya Shah, ([LinkedIn Post](https://www.linkedin.com/posts/dhravyashah_what-if-i-tell-you-that-you-can-make-an-ai-activity-7089459599287726080-ZIYm/))
+- [Read paid Medium articles for Free using embedchain](https://twitter.com/kumarkaushal_/status/1688952961622585344) by Kaushal Kumar
 
 ## Videos
 
@@ -69,6 +75,10 @@ Embedchain community has been super active in creating demos on top of Embedchai
 - [🤖CHAT with ANY ONLINE RESOURCES using EMBEDCHAIN - a LangChain wrapper, in few lines of code !](https://www.youtube.com/watch?v=Mp7zJe4TIdM) by Avra
 - [Building resource-driven LLM-powered bots with Embedchain](https://www.youtube.com/watch?v=IVfcAgxTO4I) by BugBytes
 - [embedchain-streamlit-demo](https://www.youtube.com/watch?v=yJAWB13FhYQ) by Amjad Raza
+- [Embedchain - create your own AI chatbots using open source models](https://www.youtube.com/shorts/O3rJWKwSrWE) by Dhravya Shah
+- [AI ChatBot in 5 lines Python Code](https://www.youtube.com/watch?v=zjWvLJLksv8) by Data Engineering
+- [Interview with Karl Marx](https://www.youtube.com/watch?v=5Y4Tscwj1xk) by Alexander Ray Williams
+- [Vlog where we try to build a bot based on our content on the internet](https://www.youtube.com/watch?v=I2w8CWM3bx4) by DV, ([Tweet](https://twitter.com/dvcoolster/status/1688387017544261632))
 
 ## Mentions
 

docs/advanced/vector_database.mdx

@@ -0,0 +1,34 @@
+---
+title: '💾 Vector Database'
+---
+
+We support `Chroma` and `Elasticsearch` as two vector database. 
+`Chroma` is used as a default database.
+
+### Elasticsearch
+In order to use `Elasticsearch` as vector database we need to use App type `CustomApp`. 
+```python
+import os
+from embedchain import CustomApp
+from embedchain.config import CustomAppConfig, ElasticsearchDBConfig
+from embedchain.models import Providers, EmbeddingFunctions, VectorDatabases
+
+os.environ["OPENAI_API_KEY"] = 'OPENAI_API_KEY'
+
+es_config = ElasticsearchDBConfig(
+	# elasticsearch url or list of nodes url with different hosts and ports.
+	es_url='http://localhost:9200',
+	# pass named parameters supported by Python Elasticsearch client
+	ca_certs="/path/to/http_ca.crt",
+	basic_auth=("username", "password")
+)
+config = CustomAppConfig(
+	embedding_fn=EmbeddingFunctions.OPENAI, 
+	provider=Providers.OPENAI, 
+	db_type=VectorDatabases.ELASTICSEARCH, 
+	es_config=es_config,
+)
+es_app = CustomApp(config)
+```
+- Set `db_type=VectorDatabases.ELASTICSEARCH` and `es_config=ElasticsearchDBConfig(es_url='')` in `CustomAppConfig`.
+- `ElasticsearchDBConfig` accepts `es_url` as elasticsearch url or as list of nodes url with different hosts and ports. Additionally we can pass named paramaters supported by Python Elasticsearch client.

docs/examples/api_server.mdx

@@ -0,0 +1,46 @@
+---
+title: '🌍 API Server'
+---
+
+### 🐳 Docker Setup
+
+- Open variables.env, and edit it to add your 🔑 `OPENAI_API_KEY`.
+- To setup your api server using docker, run the following command inside this folder using your terminal.
+
+```bash
+docker-compose up --build
+```
+
+📝 Note: The build command might take a while to install all the packages depending on your system resources.
+
+### 🚀 Usage Instructions
+
+- Your api server is running on [http://localhost:5000/](http://localhost:5000/)
+- To use the api server, make an api call to the endpoints `/add` and `/query` using the json formats discussed below.
+- To add data sources to the bot:
+```json
+// Request
+{
+  "data_type": "your_data_type_here",
+  "url_or_text": "your_url_or_text_here"
+}
+
+// Response
+{
+  "data": "Added data_type: url_or_text"
+}
+```
+- To ask questions from the bot:
+```json
+// Request
+{
+  "question": "your_question_here"
+}
+
+// Response
+{
+  "data": "your_answer_here"
+}
+```
+
+🎉 Happy Chatting! 🎉

docs/mint.json

@@ -32,11 +32,11 @@
     },
     {
       "group": "Advanced",
-      "pages": ["advanced/app_types", "advanced/interface_types", "advanced/adding_data","advanced/data_types", "advanced/query_configuration", "advanced/configuration", "advanced/testing", "advanced/showcase"]
+      "pages": ["advanced/app_types", "advanced/interface_types", "advanced/adding_data","advanced/data_types", "advanced/query_configuration", "advanced/configuration", "advanced/testing", "advanced/vector_database", "advanced/showcase"]
     },
     {
       "group": "Examples",
-      "pages": ["examples/full_stack"]
+      "pages": ["examples/full_stack", "examples/api_server"]
     },
     {
       "group": "Contribution Guidelines",

embedchain/chunkers/notion.py

@@ -0,0 +1,20 @@
+from typing import Optional
+
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+
+from embedchain.chunkers.base_chunker import BaseChunker
+from embedchain.config.AddConfig import ChunkerConfig
+
+
+class NotionChunker(BaseChunker):
+    """Chunker for notion."""
+
+    def __init__(self, config: Optional[ChunkerConfig] = None):
+        if config is None:
+            config = ChunkerConfig(chunk_size=300, chunk_overlap=0, length_function=len)
+        text_splitter = RecursiveCharacterTextSplitter(
+            chunk_size=config.chunk_size,
+            chunk_overlap=config.chunk_overlap,
+            length_function=config.length_function,
+        )
+        super().__init__(text_splitter)

embedchain/config/__init__.py

@@ -5,3 +5,5 @@ from .apps.OpenSourceAppConfig import OpenSourceAppConfig  # noqa: F401
 from .BaseConfig import BaseConfig  # noqa: F401
 from .ChatConfig import ChatConfig  # noqa: F401
 from .QueryConfig import QueryConfig  # noqa: F401
+from .vectordbs.ElasticsearchDBConfig import \
+    ElasticsearchDBConfig  # noqa: F401

embedchain/config/apps/AppConfig.py

@@ -1,4 +1,5 @@
 import os
+from typing import Optional
 
 try:
     from chromadb.utils import embedding_functions
@@ -16,16 +17,32 @@ 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,
+        host=None,
+        port=None,
+        id=None,
+        collection_name=None,
+        collect_metrics: Optional[bool] = None,
+    ):
         """
         :param log_level: Optional. (String) Debug level
         ['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'].
         :param host: Optional. Hostname for the database server.
         :param port: Optional. Port for the database server.
         :param id: Optional. ID of the app. Document metadata will have this id.
+        :param collection_name: Optional. Collection name for the database.
+        :param collect_metrics: Defaults to True. Send anonymous telemetry to improve embedchain.
         """
         super().__init__(
-            log_level=log_level, embedding_fn=AppConfig.default_embedding_function(), host=host, port=port, id=id
+            log_level=log_level,
+            embedding_fn=AppConfig.default_embedding_function(),
+            host=host,
+            port=port,
+            id=id,
+            collection_name=collection_name,
+            collect_metrics=collect_metrics,
         )
 
     @staticmethod

embedchain/config/apps/BaseAppConfig.py

@@ -1,6 +1,8 @@
 import logging
 
 from embedchain.config.BaseConfig import BaseConfig
+from embedchain.config.vectordbs import ElasticsearchDBConfig
+from embedchain.models import VectorDatabases, VectorDimensions
 
 
 class BaseAppConfig(BaseConfig):
@@ -8,35 +10,78 @@ 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,
+        host=None,
+        port=None,
+        id=None,
+        collection_name=None,
+        collect_metrics: bool = True,
+        db_type: VectorDatabases = None,
+        vector_dim: VectorDimensions = None,
+        es_config: ElasticsearchDBConfig = 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 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.
+        :param id: Optional. ID of the app. Document metadata will have this id.
+        :param collection_name: Optional. Collection name for the database.
+        :param collect_metrics: Defaults to True. Send anonymous telemetry to improve embedchain.
+        :param db_type: Optional. type of Vector database to use
+        :param vector_dim: Vector dimension generated by embedding fn
+        :param es_config: Optional. elasticsearch database config to be used for connection
         """
         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 if collection_name else "embedchain_store"
+        self.db = BaseAppConfig.get_db(
+            db=db,
+            embedding_fn=embedding_fn,
+            host=host,
+            port=port,
+            db_type=db_type,
+            vector_dim=vector_dim,
+            collection_name=self.collection_name,
+            es_config=es_config,
+        )
         self.id = id
+        self.collect_metrics = True if (collect_metrics is True or collect_metrics is None) else False
         return
 
     @staticmethod
-    def default_db(embedding_fn, host, port):
+    def get_db(db, embedding_fn, host, port, db_type, vector_dim, collection_name, es_config):
         """
-        Sets database to default (`ChromaDb`).
-
+        Get db based on db_type, db with default database (`ChromaDb`)
+        :param Optional. (Vector) database to use for embeddings.
         :param embedding_fn: Embedding function to use in database.
         :param host: Optional. Hostname for the database server.
         :param port: Optional. Port for the database server.
-        :returns: Default database
+        :param db_type: Optional. db type to use. Supported values (`es`, `chroma`)
+        :param vector_dim: Vector dimension generated by embedding fn
+        :param collection_name: Optional. Collection name for the database.
+        :param es_config: Optional. elasticsearch database config to be used for connection
         :raises ValueError: BaseAppConfig knows no default embedding function.
+        :returns: database instance
         """
+        if db:
+            return db
+
         if embedding_fn is None:
             raise ValueError("ChromaDb cannot be instantiated without an embedding function")
+
+        if db_type == VectorDatabases.ELASTICSEARCH:
+            from embedchain.vectordb.elasticsearch_db import ElasticsearchDB
+
+            return ElasticsearchDB(
+                embedding_fn=embedding_fn, vector_dim=vector_dim, collection_name=collection_name, es_config=es_config
+            )
+
         from embedchain.vectordb.chroma_db import ChromaDB
 
         return ChromaDB(embedding_fn=embedding_fn, host=host, port=port)

embedchain/config/apps/CustomAppConfig.py

@@ -1,9 +1,11 @@
-from typing import Any
+from typing import Any, Optional
 
 from chromadb.api.types import Documents, Embeddings
 from dotenv import load_dotenv
 
-from embedchain.models import EmbeddingFunctions, Providers
+from embedchain.config.vectordbs import ElasticsearchDBConfig
+from embedchain.models import (EmbeddingFunctions, Providers, VectorDatabases,
+                               VectorDimensions)
 
 from .BaseAppConfig import BaseAppConfig
 
@@ -24,10 +26,13 @@ class CustomAppConfig(BaseAppConfig):
         host=None,
         port=None,
         id=None,
+        collection_name=None,
         provider: Providers = None,
-        model=None,
         open_source_app_config=None,
         deployment_name=None,
+        collect_metrics: Optional[bool] = None,
+        db_type: VectorDatabases = None,
+        es_config: ElasticsearchDBConfig = None,
     ):
         """
         :param log_level: Optional. (String) Debug level
@@ -35,11 +40,15 @@ class CustomAppConfig(BaseAppConfig):
         :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 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.
+        :param id: Optional. ID of the app. Document metadata will have this id.
+        :param collection_name: Optional. Collection name for the database.
         :param provider: Optional. (Providers): LLM Provider to use.
         :param open_source_app_config: Optional. Config instance needed for open source apps.
+        :param collect_metrics: Defaults to True. Send anonymous telemetry to improve embedchain.
+        :param db_type: Optional. type of Vector database to use.
+        :param es_config: Optional. elasticsearch database config to be used for connection
         """
         if provider:
             self.provider = provider
@@ -51,13 +60,17 @@ class CustomAppConfig(BaseAppConfig):
         super().__init__(
             log_level=log_level,
             embedding_fn=CustomAppConfig.embedding_function(
-                embedding_function=embedding_fn, model=embedding_fn_model,
-                deployment_name=deployment_name
+                embedding_function=embedding_fn, model=embedding_fn_model, deployment_name=deployment_name
             ),
             db=db,
             host=host,
             port=port,
             id=id,
+            collection_name=collection_name,
+            collect_metrics=collect_metrics,
+            db_type=db_type,
+            vector_dim=CustomAppConfig.get_vector_dimension(embedding_function=embedding_fn),
+            es_config=es_config,
         )
 
     @staticmethod
@@ -107,3 +120,20 @@ class CustomAppConfig(BaseAppConfig):
             from chromadb.utils import embedding_functions
 
             return embedding_functions.SentenceTransformerEmbeddingFunction(model_name=model)
+
+    @staticmethod
+    def get_vector_dimension(embedding_function: EmbeddingFunctions):
+        if not isinstance(embedding_function, EmbeddingFunctions):
+            raise ValueError(f"Invalid option: '{embedding_function}'.")
+
+        if embedding_function == EmbeddingFunctions.OPENAI:
+            return VectorDimensions.OPENAI.value
+
+        elif embedding_function == EmbeddingFunctions.HUGGING_FACE:
+            return VectorDimensions.HUGGING_FACE.value
+
+        elif embedding_function == EmbeddingFunctions.VERTEX_AI:
+            return VectorDimensions.VERTEX_AI.value
+
+        elif embedding_function == EmbeddingFunctions.GPT4ALL:
+            return VectorDimensions.GPT4ALL.value

embedchain/config/apps/OpenSourceAppConfig.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 from chromadb.utils import embedding_functions
 
 from .BaseAppConfig import BaseAppConfig
@@ -8,13 +10,24 @@ 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,
+        host=None,
+        port=None,
+        id=None,
+        collection_name=None,
+        collect_metrics: Optional[bool] = 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 collect_metrics: Defaults to True. Send anonymous telemetry to improve embedchain.
         :param model: Optional. GPT4ALL uses the model to instantiate the class.
         So unlike `App`, it has to be provided before querying.
         """
@@ -26,6 +39,8 @@ class OpenSourceAppConfig(BaseAppConfig):
             host=host,
             port=port,
             id=id,
+            collection_name=collection_name,
+            collect_metrics=collect_metrics,
         )
 
     @staticmethod

embedchain/config/vectordbs/ElasticsearchDBConfig.py

@@ -0,0 +1,15 @@
+from typing import Dict, List, Union
+
+from embedchain.config.BaseConfig import BaseConfig
+
+
+class ElasticsearchDBConfig(BaseConfig):
+    """
+    Config to initialize an elasticsearch client.
+    :param es_url. elasticsearch url or list of nodes url to be used for connection
+    :param ES_EXTRA_PARAMS: extra params dict that can be passed to elasticsearch.
+    """
+
+    def __init__(self, es_url: Union[str, List[str]] = None, **ES_EXTRA_PARAMS: Dict[str, any]):
+        self.ES_URL = es_url
+        self.ES_EXTRA_PARAMS = ES_EXTRA_PARAMS

embedchain/data_formatter/data_formatter.py

@@ -1,5 +1,6 @@
 from embedchain.chunkers.docs_site import DocsSiteChunker
 from embedchain.chunkers.docx_file import DocxFileChunker
+from embedchain.chunkers.notion import NotionChunker
 from embedchain.chunkers.pdf_file import PdfFileChunker
 from embedchain.chunkers.qna_pair import QnaPairChunker
 from embedchain.chunkers.text import TextChunker
@@ -45,8 +46,16 @@ class DataFormatter:
             "sitemap": SitemapLoader(),
             "docs_site": DocsSiteLoader(),
         }
+        lazy_loaders = ("notion",)
         if data_type in loaders:
             return loaders[data_type]
+        elif data_type in lazy_loaders:
+            if data_type == "notion":
+                from embedchain.loaders.notion import NotionLoader
+
+                return NotionLoader()
+            else:
+                raise ValueError(f"Unsupported data type: {data_type}")
         else:
             raise ValueError(f"Unsupported data type: {data_type}")
 
@@ -67,6 +76,7 @@ class DataFormatter:
             "docx": DocxFileChunker,
             "sitemap": WebPageChunker,
             "docs_site": DocsSiteChunker,
+            "notion": NotionChunker,
         }
         if data_type in chunker_classes:
             chunker_class = chunker_classes[data_type]

embedchain/embedchain.py

@@ -1,10 +1,14 @@
+import importlib.metadata
 import logging
 import os
+import threading
+from typing import Optional
 
-from chromadb.errors import InvalidDimensionException
+import requests
 from dotenv import load_dotenv
 from langchain.docstore.document import Document
 from langchain.memory import ConversationBufferMemory
+from tenacity import retry, stop_after_attempt, wait_fixed
 
 from embedchain.chunkers.base_chunker import BaseChunker
 from embedchain.config import AddConfig, ChatConfig, QueryConfig
@@ -31,12 +35,16 @@ class EmbedChain:
         """
 
         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.db = self.config.db
         self.user_asks = []
         self.is_docs_site_instance = False
         self.online = False
 
+        # Send anonymous telemetry
+        thread_telemetry = threading.Thread(target=self._send_telemetry_event, args=("init",))
+        thread_telemetry.start()
+
     def add(self, data_type, url, metadata=None, config: AddConfig = None):
         """
         Adds the data from the given URL to the vector db.
@@ -54,10 +62,21 @@ class EmbedChain:
 
         data_formatter = DataFormatter(data_type, config)
         self.user_asks.append([data_type, url, metadata])
-        self.load_and_embed(data_formatter.loader, data_formatter.chunker, url, metadata)
+        documents, _metadatas, _ids, new_chunks = self.load_and_embed(
+            data_formatter.loader, data_formatter.chunker, url, metadata
+        )
         if data_type in ("docs_site",):
             self.is_docs_site_instance = True
 
+        # Send anonymous telemetry
+        if self.config.collect_metrics:
+            # it's quicker to check the variable twice than to count words when they won't be submitted.
+            word_count = sum([len(document.split(" ")) for document in documents])
+
+            extra_metadata = {"data_type": data_type, "word_count": word_count, "chunks_count": new_chunks}
+            thread_telemetry = threading.Thread(target=self._send_telemetry_event, args=("add", extra_metadata))
+            thread_telemetry.start()
+
     def add_local(self, data_type, content, metadata=None, config: AddConfig = None):
         """
         Adds the data you supply to the vector db.
@@ -91,6 +110,7 @@ class EmbedChain:
         :param src: The data to be handled by the loader. Can be a URL for
         remote sources or local content for local loaders.
         :param metadata: Optional. Metadata associated with the data source.
+        :return: (List) documents (embedded text), (List) metadata, (list) ids, (int) number of chunks
         """
         embeddings_data = chunker.create_chunks(loader, src)
         documents = embeddings_data["documents"]
@@ -99,11 +119,10 @@ class EmbedChain:
         # get existing ids, and discard doc if any common id exist.
         where = {"app_id": self.config.id} if self.config.id is not None else {}
         # where={"url": src}
-        existing_docs = self.collection.get(
+        existing_ids = self.db.get(
             ids=ids,
             where=where,  # optional filter
         )
-        existing_ids = set(existing_docs["ids"])
 
         if len(existing_ids):
             data_dict = {id: (doc, meta) for id, doc, meta in zip(ids, documents, metadatas)}
@@ -111,7 +130,8 @@ class EmbedChain:
 
             if not data_dict:
                 print(f"All data from {src} already exists in the database.")
-                return
+                # Make sure to return a matching return type
+                return [], [], [], 0
 
             ids = list(data_dict.keys())
             documents, metadatas = zip(*data_dict.values())
@@ -128,8 +148,10 @@ class EmbedChain:
         # Add metadata to each document
         metadatas_with_metadata = [{**meta, **metadata} for meta in metadatas]
 
-        self.collection.add(documents=documents, metadatas=list(metadatas_with_metadata), ids=ids)
-        print((f"Successfully saved {src}. New chunks count: " f"{self.count() - chunks_before_addition}"))
+        self.db.add(documents=documents, metadatas=metadatas_with_metadata, ids=ids)
+        count_new_chunks = self.count() - chunks_before_addition
+        print((f"Successfully saved {src}. New chunks count: {count_new_chunks}"))
+        return list(documents), metadatas_with_metadata, ids, count_new_chunks
 
     def _format_result(self, results):
         return [
@@ -156,23 +178,13 @@ class EmbedChain:
         :param config: The query configuration.
         :return: The content of the document that matched your query.
         """
-        try:
-            where = {"app_id": self.config.id} if self.config.id is not None else {}  # optional filter
-            result = self.collection.query(
-                query_texts=[
-                    input_query,
-                ],
-                n_results=config.number_documents,
-                where=where,
-            )
-        except InvalidDimensionException as e:
-            raise InvalidDimensionException(
-                e.message()
-                + ". This is commonly a side-effect when an embedding function, different from the one used to add the embeddings, is used to retrieve an embedding from the database."  # noqa E501
-            ) from None
+        where = {"app_id": self.config.id} if self.config.id is not None else {}  # optional filter
+        contents = self.db.query(
+            input_query=input_query,
+            n_results=config.number_documents,
+            where=where,
+        )
 
-        results_formatted = self._format_result(result)
-        contents = [result[0].page_content for result in results_formatted]
         return contents
 
     def _append_search_and_context(self, context, web_search_result):
@@ -252,6 +264,10 @@ class EmbedChain:
 
         answer = self.get_answer_from_llm(prompt, config)
 
+        # Send anonymous telemetry
+        thread_telemetry = threading.Thread(target=self._send_telemetry_event, args=("query",))
+        thread_telemetry.start()
+
         if isinstance(answer, str):
             logging.info(f"Answer: {answer}")
             return answer
@@ -309,6 +325,10 @@ class EmbedChain:
 
         memory.chat_memory.add_user_message(input_query)
 
+        # Send anonymous telemetry
+        thread_telemetry = threading.Thread(target=self._send_telemetry_event, args=("chat",))
+        thread_telemetry.start()
+
         if isinstance(answer, str):
             memory.chat_memory.add_ai_message(answer)
             logging.info(f"Answer: {answer}")
@@ -325,17 +345,48 @@ class EmbedChain:
         memory.chat_memory.add_ai_message(streamed_answer)
         logging.info(f"Answer: {streamed_answer}")
 
-    def count(self):
+    def set_collection(self, collection_name):
+        """
+        Set the collection to use.
+
+        :param collection_name: The name of the collection to use.
+        """
+        self.collection = self.config.db._get_or_create_collection(collection_name)
+
+    def count(self) -> int:
         """
         Count the number of embeddings.
 
         :return: The number of embeddings.
         """
-        return self.collection.count()
+        return self.db.count()
 
     def reset(self):
         """
         Resets the database. Deletes all embeddings irreversibly.
         `App` has to be reinitialized after using this method.
         """
-        self.db_client.reset()
+        # Send anonymous telemetry
+        thread_telemetry = threading.Thread(target=self._send_telemetry_event, args=("reset",))
+        thread_telemetry.start()
+
+        self.db.reset()
+
+    @retry(stop=stop_after_attempt(3), wait=wait_fixed(1))
+    def _send_telemetry_event(self, method: str, extra_metadata: Optional[dict] = None):
+        if not self.config.collect_metrics:
+            return
+
+        with threading.Lock():
+            url = "https://api.embedchain.ai/api/v1/telemetry/"
+            metadata = {
+                "app_id": self.config.id,
+                "version": importlib.metadata.version(__package__ or __name__),
+                "method": method,
+                "language": "py",
+            }
+            if extra_metadata:
+                metadata.update(extra_metadata)
+
+            response = requests.post(url, json={"metadata": metadata})
+            response.raise_for_status()

embedchain/loaders/notion.py

@@ -0,0 +1,41 @@
+import logging
+import os
+
+try:
+    from llama_index import download_loader
+except ImportError:
+    raise ImportError("Notion requires extra dependencies. Install with `pip install embedchain[notion]`") from None
+
+
+from embedchain.loaders.base_loader import BaseLoader
+from embedchain.utils import clean_string
+
+
+class NotionLoader(BaseLoader):
+    def load_data(self, source):
+        """Load data from a PDF file."""
+
+        NotionPageReader = download_loader("NotionPageReader")
+
+        # Reformat Id to match notion expectation
+        id = source[-32:]
+        formatted_id = f"{id[:8]}-{id[8:12]}-{id[12:16]}-{id[16:20]}-{id[20:]}"
+        logging.debug(f"Extracted notion page id as: {formatted_id}")
+
+        # Get page through the notion api
+        integration_token = os.getenv("NOTION_INTEGRATION_TOKEN")
+        reader = NotionPageReader(integration_token=integration_token)
+        documents = reader.load_data(page_ids=[formatted_id])
+
+        # Extract text
+        raw_text = documents[0].text
+
+        # Clean text
+        text = clean_string(raw_text)
+
+        return [
+            {
+                "content": text,
+                "meta_data": {"url": f"notion-{formatted_id}"},
+            }
+        ]

embedchain/models/VectorDatabases.py

@@ -0,0 +1,6 @@
+from enum import Enum
+
+
+class VectorDatabases(Enum):
+    CHROMADB = "CHROMADB"
+    ELASTICSEARCH = "ELASTICSEARCH"

embedchain/models/VectorDimensions.py

@@ -0,0 +1,9 @@
+from enum import Enum
+
+
+# vector length created by embedding fn
+class VectorDimensions(Enum):
+    GPT4ALL = 384
+    OPENAI = 1536
+    VERTEX_AI = 768
+    HUGGING_FACE = 384

embedchain/models/__init__.py

@@ -1,2 +1,4 @@
 from .EmbeddingFunctions import EmbeddingFunctions  # noqa: F401
 from .Providers import Providers  # noqa: F401
+from .VectorDatabases import VectorDatabases  # noqa: F401
+from .VectorDimensions import VectorDimensions  # noqa: F401

embedchain/utils.py

@@ -60,22 +60,32 @@ def use_pysqlite3():
     import sqlite3
 
     if platform.system() == "Linux" and sqlite3.sqlite_version_info < (3, 35, 0):
-        # According to the Chroma team, this patch only works on Linux
-        import datetime
-        import subprocess
-        import sys
+        try:
+            # According to the Chroma team, this patch only works on Linux
+            import datetime
+            import subprocess
+            import sys
 
-        subprocess.check_call(
-            [sys.executable, "-m", "pip", "install", "pysqlite3-binary", "--quiet", "--disable-pip-version-check"]
-        )
+            subprocess.check_call(
+                [sys.executable, "-m", "pip", "install", "pysqlite3-binary", "--quiet", "--disable-pip-version-check"]
+            )
 
-        __import__("pysqlite3")
-        sys.modules["sqlite3"] = sys.modules.pop("pysqlite3")
+            __import__("pysqlite3")
+            sys.modules["sqlite3"] = sys.modules.pop("pysqlite3")
 
-        # Let the user know what happened.
-        current_time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S,%f")[:-3]
-        print(
-            f"{current_time} [embedchain] [INFO]",
-            "Swapped std-lib sqlite3 with pysqlite3 for ChromaDb compatibility.",
-            f"Your original version was {sqlite3.sqlite_version}.",
-        )
+            # Let the user know what happened.
+            current_time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S,%f")[:-3]
+            print(
+                f"{current_time} [embedchain] [INFO]",
+                "Swapped std-lib sqlite3 with pysqlite3 for ChromaDb compatibility.",
+                f"Your original version was {sqlite3.sqlite_version}.",
+            )
+        except Exception as e:
+            # Escape all exceptions
+            current_time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S,%f")[:-3]
+            print(
+                f"{current_time} [embedchain] [ERROR]",
+                "Failed to swap std-lib sqlite3 with pysqlite3 for ChromaDb compatibility.",
+                "Error:",
+                e,
+            )

embedchain/vectordb/base_vector_db.py

@@ -3,7 +3,6 @@ class BaseVectorDB:
 
     def __init__(self):
         self.client = self._get_or_create_db()
-        self.collection = self._get_or_create_collection()
 
     def _get_or_create_db(self):
         """Get or create the database."""
@@ -11,3 +10,18 @@ class BaseVectorDB:
 
     def _get_or_create_collection(self):
         raise NotImplementedError
+
+    def get(self):
+        raise NotImplementedError
+
+    def add(self):
+        raise NotImplementedError
+
+    def query(self):
+        raise NotImplementedError
+
+    def count(self):
+        raise NotImplementedError
+
+    def reset(self):
+        raise NotImplementedError

embedchain/vectordb/chroma_db.py

@@ -1,4 +1,8 @@
 import logging
+from typing import Any, Dict, List
+
+from chromadb.errors import InvalidDimensionException
+from langchain.docstore.document import Document
 
 try:
     import chromadb
@@ -7,6 +11,7 @@ except RuntimeError:
 
     use_pysqlite3()
     import chromadb
+
 from chromadb.config import Settings
 
 from embedchain.vectordb.base_vector_db import BaseVectorDB
@@ -39,9 +44,75 @@ class ChromaDB(BaseVectorDB):
         """Get or create the database."""
         return self.client
 
-    def _get_or_create_collection(self):
+    def _get_or_create_collection(self, name):
         """Get or create the collection."""
-        return self.client.get_or_create_collection(
-            "embedchain_store",
+        self.collection = self.client.get_or_create_collection(
+            name=name,
             embedding_function=self.embedding_fn,
         )
+        return self.collection
+
+    def get(self, ids: List[str], where: Dict[str, any]) -> List[str]:
+        """
+        Get existing doc ids present in vector database
+        :param ids: list of doc ids to check for existance
+        :param where: Optional. to filter data
+        """
+        existing_docs = self.collection.get(
+            ids=ids,
+            where=where,  # optional filter
+        )
+
+        return set(existing_docs["ids"])
+
+    def add(self, documents: List[str], metadatas: List[object], ids: List[str]) -> Any:
+        """
+        add data in vector database
+        :param documents: list of texts to add
+        :param metadatas: list of metadata associated with docs
+        :param ids: ids of docs
+        """
+        self.collection.add(documents=documents, metadatas=metadatas, ids=ids)
+
+    def _format_result(self, results):
+        return [
+            (Document(page_content=result[0], metadata=result[1] or {}), result[2])
+            for result in zip(
+                results["documents"][0],
+                results["metadatas"][0],
+                results["distances"][0],
+            )
+        ]
+
+    def query(self, input_query: List[str], n_results: int, where: Dict[str, any]) -> List[str]:
+        """
+        query contents from vector data base based on vector similarity
+        :param input_query: list of query string
+        :param n_results: no of similar documents to fetch from database
+        :param where: Optional. to filter data
+        :return: The content of the document that matched your query.
+        """
+        try:
+            result = self.collection.query(
+                query_texts=[
+                    input_query,
+                ],
+                n_results=n_results,
+                where=where,
+            )
+        except InvalidDimensionException as e:
+            raise InvalidDimensionException(
+                e.message()
+                + ". This is commonly a side-effect when an embedding function, different from the one used to add the embeddings, is used to retrieve an embedding from the database."  # noqa E501
+            ) from None
+
+        results_formatted = self._format_result(result)
+        contents = [result[0].page_content for result in results_formatted]
+        return contents
+
+    def count(self) -> int:
+        return self.collection.count()
+
+    def reset(self):
+        # Delete all data from the database
+        self.client.reset()

embedchain/vectordb/elasticsearch_db.py

@@ -0,0 +1,136 @@
+from typing import Any, Callable, Dict, List
+
+try:
+    from elasticsearch import Elasticsearch
+    from elasticsearch.helpers import bulk
+except ImportError:
+    raise ImportError(
+        "Elasticsearch requires extra dependencies. Install with `pip install embedchain[elasticsearch]`"
+    ) from None
+
+from embedchain.config import ElasticsearchDBConfig
+from embedchain.models.VectorDimensions import VectorDimensions
+from embedchain.vectordb.base_vector_db import BaseVectorDB
+
+
+class ElasticsearchDB(BaseVectorDB):
+    def __init__(
+        self,
+        es_config: ElasticsearchDBConfig = None,
+        embedding_fn: Callable[[list[str]], list[str]] = None,
+        vector_dim: VectorDimensions = None,
+        collection_name: str = None,
+    ):
+        """
+        Elasticsearch as vector database
+        :param es_config. elasticsearch database config to be used for connection
+        :param embedding_fn: Function to generate embedding vectors.
+        :param vector_dim: Vector dimension generated by embedding fn
+        :param collection_name: Optional. Collection name for the database.
+        """
+        if not hasattr(embedding_fn, "__call__"):
+            raise ValueError("Embedding function is not a function")
+        if es_config is None:
+            raise ValueError("ElasticsearchDBConfig is required")
+        if vector_dim is None:
+            raise ValueError("Vector Dimension is required to refer correct index and mapping")
+        if collection_name is None:
+            raise ValueError("collection name is required. It cannot be empty")
+        self.embedding_fn = embedding_fn
+        self.client = Elasticsearch(es_config.ES_URL, **es_config.ES_EXTRA_PARAMS)
+        self.vector_dim = vector_dim
+        self.es_index = f"{collection_name}_{self.vector_dim}"
+        index_settings = {
+            "mappings": {
+                "properties": {
+                    "text": {"type": "text"},
+                    "embeddings": {"type": "dense_vector", "index": False, "dims": self.vector_dim},
+                }
+            }
+        }
+        if not self.client.indices.exists(index=self.es_index):
+            # create index if not exist
+            print("Creating index", self.es_index, index_settings)
+            self.client.indices.create(index=self.es_index, body=index_settings)
+        super().__init__()
+
+    def _get_or_create_db(self):
+        return self.client
+
+    def _get_or_create_collection(self, name):
+        """Note: nothing to return here. Discuss later"""
+
+    def get(self, ids: List[str], where: Dict[str, any]) -> List[str]:
+        """
+        Get existing doc ids present in vector database
+        :param ids: list of doc ids to check for existance
+        :param where: Optional. to filter data
+        """
+        query = {"bool": {"must": [{"ids": {"values": ids}}]}}
+        if "app_id" in where:
+            app_id = where["app_id"]
+            query["bool"]["must"].append({"term": {"metadata.app_id": app_id}})
+        response = self.client.search(index=self.es_index, query=query, _source=False)
+        docs = response["hits"]["hits"]
+        ids = [doc["_id"] for doc in docs]
+        return set(ids)
+
+    def add(self, documents: List[str], metadatas: List[object], ids: List[str]) -> Any:
+        """
+        add data in vector database
+        :param documents: list of texts to add
+        :param metadatas: list of metadata associated with docs
+        :param ids: ids of docs
+        """
+        docs = []
+        embeddings = self.embedding_fn(documents)
+        for id, text, metadata, embeddings in zip(ids, documents, metadatas, embeddings):
+            docs.append(
+                {
+                    "_index": self.es_index,
+                    "_id": id,
+                    "_source": {"text": text, "metadata": metadata, "embeddings": embeddings},
+                }
+            )
+        bulk(self.client, docs)
+        self.client.indices.refresh(index=self.es_index)
+        return
+
+    def query(self, input_query: List[str], n_results: int, where: Dict[str, any]) -> List[str]:
+        """
+        query contents from vector data base based on vector similarity
+        :param input_query: list of query string
+        :param n_results: no of similar documents to fetch from database
+        :param where: Optional. to filter data
+        """
+        input_query_vector = self.embedding_fn(input_query)
+        query_vector = input_query_vector[0]
+        query = {
+            "script_score": {
+                "query": {"bool": {"must": [{"exists": {"field": "text"}}]}},
+                "script": {
+                    "source": "cosineSimilarity(params.input_query_vector, 'embeddings') + 1.0",
+                    "params": {"input_query_vector": query_vector},
+                },
+            }
+        }
+        if "app_id" in where:
+            app_id = where["app_id"]
+            query["script_score"]["query"]["bool"]["must"] = [{"term": {"metadata.app_id": app_id}}]
+        _source = ["text"]
+        response = self.client.search(index=self.es_index, query=query, _source=_source, size=n_results)
+        docs = response["hits"]["hits"]
+        contents = [doc["_source"]["text"] for doc in docs]
+        return contents
+
+    def count(self) -> int:
+        query = {"match_all": {}}
+        response = self.client.count(index=self.es_index, query=query)
+        doc_count = response["count"]
+        return doc_count
+
+    def reset(self):
+        # Delete all data from the database
+        if self.client.indices.exists(index=self.es_index):
+            # delete index in Es
+            self.client.indices.delete(index=self.es_index)

examples/api_server/.dockerignore

@@ -0,0 +1,8 @@
+__pycache__/
+database
+db
+pyenv
+venv
+.env
+.git
+trash_files/

examples/api_server/.gitignore

@@ -0,0 +1,8 @@
+__pycache__
+db
+database
+pyenv
+venv
+.env
+trash_files/
+.ideas.md

examples/api_server/Dockerfile

@@ -0,0 +1,11 @@
+FROM python:3.11 AS backend
+
+WORKDIR /usr/src/api
+COPY requirements.txt .
+RUN pip install -r requirements.txt
+
+COPY . .
+
+EXPOSE 5000
+
+CMD ["python", "api_server.py"]

examples/api_server/api_server.py

@@ -0,0 +1,42 @@
+from flask import Flask, jsonify, request
+
+from embedchain import App
+
+app = Flask(__name__)
+
+
+def initialize_chat_bot():
+    global chat_bot
+    chat_bot = App()
+
+
+@app.route("/add", methods=["POST"])
+def add():
+    data = request.get_json()
+    data_type = data.get("data_type")
+    url_or_text = data.get("url_or_text")
+    if data_type and url_or_text:
+        try:
+            chat_bot.add(data_type, url_or_text)
+            return jsonify({"data": f"Added {data_type}: {url_or_text}"}), 200
+        except Exception:
+            return jsonify({"error": f"Failed to add {data_type}: {url_or_text}"}), 500
+    return jsonify({"error": "Invalid request. Please provide 'data_type' and 'url_or_text' in JSON format."}), 400
+
+
+@app.route("/query", methods=["POST"])
+def query():
+    data = request.get_json()
+    question = data.get("question")
+    if question:
+        try:
+            response = chat_bot.chat(question)
+            return jsonify({"data": response}), 200
+        except Exception:
+            return jsonify({"error": "An error occurred. Please try again!"}), 500
+    return jsonify({"error": "Invalid request. Please provide 'question' in JSON format."}), 400
+
+
+if __name__ == "__main__":
+    initialize_chat_bot()
+    app.run(host="0.0.0.0", port=5000, debug=False)

examples/api_server/docker-compose.yml

@@ -0,0 +1,13 @@
+version: "3.9"
+
+services:
+  backend:
+    container_name: embedchain_api
+    restart: unless-stopped
+    build:
+      context: .
+      dockerfile: Dockerfile
+    env_file:
+      - variables.env
+    ports:
+      - "5000:5000"

examples/api_server/requirements.txt

@@ -0,0 +1,2 @@
+flask==2.3.2
+embedchain==0.0.30

examples/api_server/variables.env

@@ -0,0 +1 @@
+OPENAI_API_KEY=""

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "embedchain"
-version = "0.0.23"
+version = "0.0.37"
 description = "embedchain is a framework to easily create LLM powered bots over any dataset"
 authors = ["Taranjeet Singh"]
 license = "Apache License"
@@ -90,6 +90,8 @@ youtube-transcript-api = "^0.6.1"
 beautifulsoup4 = "^4.12.2"
 pypdf = "^3.11.0"
 pytube = "^15.0.0"
+llama-index = { version = "^0.7.21", optional = true }
+elasticsearch = { version = "^8.9.0", optional = true }
 
 
 
@@ -105,7 +107,8 @@ isort = "^5.12.0"
 
 [tool.poetry.extras]
 streamlit = ["streamlit"]
-
+community = ["llama-index"]
+elasticsearch = ["elasticsearch"]
 
 [tool.poetry.group.docs.dependencies]
 

setup.py

@@ -5,7 +5,7 @@ with open("README.md", "r", encoding="utf-8") as fh:
 
 setuptools.setup(
     name="embedchain",
-    version="0.0.32",
+    version="0.0.37",
     author="Taranjeet Singh",
     author_email="reachtotj@gmail.com",
     description="embedchain is a framework to easily create LLM powered bots over any dataset",  # noqa:E501
@@ -37,5 +37,9 @@ setuptools.setup(
         "replicate==0.9.0",
         "duckduckgo-search==3.8.4",
     ],
-    extras_require={"dev": ["black", "ruff", "isort", "pytest"]},
+    extras_require={
+        "dev": ["black", "ruff", "isort", "pytest"],
+        "community": ["llama-index==0.7.21"],
+        "elasticsearch": ["elasticsearch>=8.9.0"],
+    },
 )

tests/embedchain/test_add.py

@@ -3,13 +3,14 @@ import unittest
 from unittest.mock import MagicMock, patch
 
 from embedchain import App
+from embedchain.config import AppConfig
 
 
 class TestApp(unittest.TestCase):
     os.environ["OPENAI_API_KEY"] = "test_key"
 
     def setUp(self):
-        self.app = App()
+        self.app = App(config=AppConfig(collect_metrics=False))
 
     @patch("chromadb.api.models.Collection.Collection.add", MagicMock)
     def test_add(self):

tests/embedchain/test_chat.py

@@ -3,13 +3,14 @@ import unittest
 from unittest.mock import patch
 
 from embedchain import App
+from embedchain.config import AppConfig
 
 
 class TestApp(unittest.TestCase):
     os.environ["OPENAI_API_KEY"] = "test_key"
 
     def setUp(self):
-        self.app = App()
+        self.app = App(config=AppConfig(collect_metrics=False))
 
     @patch("embedchain.embedchain.memory", autospec=True)
     @patch.object(App, "retrieve_from_database", return_value=["Test context"])

tests/embedchain/test_embedchain.py

@@ -25,7 +25,7 @@ class TestChromaDbHostsLoglevel(unittest.TestCase):
         """
         Test if the `App` instance is initialized without a config that does not contain default hosts and ports.
         """
-        config = AppConfig(log_level="DEBUG")
+        config = AppConfig(log_level="DEBUG", collect_metrics=False)
 
         app = App(config)
 

tests/embedchain/test_generate_prompt.py

@@ -2,12 +2,12 @@ import unittest
 from string import Template
 
 from embedchain import App
-from embedchain.embedchain import QueryConfig
+from embedchain.config import AppConfig, QueryConfig
 
 
 class TestGeneratePrompt(unittest.TestCase):
     def setUp(self):
-        self.app = App()
+        self.app = App(config=AppConfig(collect_metrics=False))
 
     def test_generate_prompt_with_template(self):
         """

tests/embedchain/test_query.py

@@ -3,14 +3,14 @@ import unittest
 from unittest.mock import MagicMock, patch
 
 from embedchain import App
-from embedchain.embedchain import QueryConfig
+from embedchain.config import AppConfig, QueryConfig
 
 
 class TestApp(unittest.TestCase):
     os.environ["OPENAI_API_KEY"] = "test_key"
 
     def setUp(self):
-        self.app = App()
+        self.app = App(config=AppConfig(collect_metrics=False))
 
     @patch("chromadb.api.models.Collection.Collection.add", MagicMock)
     def test_query(self):

tests/vectordb/test_chroma_db.py

@@ -39,7 +39,7 @@ class TestChromaDbHostsInit(unittest.TestCase):
         host = "test-host"
         port = "1234"
 
-        config = AppConfig(host=host, port=port)
+        config = AppConfig(host=host, port=port, collect_metrics=False)
 
         _app = App(config)
 
@@ -54,7 +54,7 @@ class TestChromaDbHostsNone(unittest.TestCase):
         Test if the `App` instance is initialized without default hosts and ports.
         """
 
-        _app = App()
+        _app = App(config=AppConfig(collect_metrics=False))
 
         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)
@@ -68,7 +68,191 @@ class TestChromaDbHostsLoglevel(unittest.TestCase):
         """
         config = AppConfig(log_level="DEBUG")
 
-        _app = App(config)
+        _app = App(config=AppConfig(collect_metrics=False))
 
         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 TestChromaDbDuplicateHandling:
+    def test_duplicates_throw_warning(self, caplog):
+        """
+        Test that add duplicates throws an error.
+        """
+        # Start with a clean app
+        App().reset()
+
+        app = App(config=AppConfig(collect_metrics=False))
+        app.collection.add(embeddings=[[0, 0, 0]], ids=["0"])
+        app.collection.add(embeddings=[[0, 0, 0]], ids=["0"])
+        assert "Insert of existing embedding ID: 0" in caplog.text
+        assert "Add of existing embedding ID: 0" in caplog.text
+
+    def test_duplicates_collections_no_warning(self, caplog):
+        """
+        Test that different collections can have duplicates.
+        """
+        # NOTE: Not part of the TestChromaDbCollection because `unittest.TestCase` doesn't have caplog.
+
+        # Start with a clean app
+        App().reset()
+
+        app = App(config=AppConfig(collect_metrics=False))
+        app.set_collection("test_collection_1")
+        app.collection.add(embeddings=[[0, 0, 0]], ids=["0"])
+        app.set_collection("test_collection_2")
+        app.collection.add(embeddings=[[0, 0, 0]], ids=["0"])
+        assert "Insert of existing embedding ID: 0" not in caplog.text  # not
+        assert "Add of existing embedding ID: 0" not in caplog.text  # not
+
+
+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(config=AppConfig(collect_metrics=False))
+
+        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", collect_metrics=False)
+        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(config=AppConfig(collect_metrics=False))
+        app.set_collection("test_collection")
+
+        self.assertEqual(app.collection.name, "test_collection")
+
+    def test_changes_encapsulated(self):
+        """
+        Test that changes to one collection do not affect the other collection
+        """
+        # Start with a clean app
+        App().reset()
+
+        app = App(config=AppConfig(collect_metrics=False))
+        app.set_collection("test_collection_1")
+        # Collection should be empty when created
+        self.assertEqual(app.count(), 0)
+
+        app.collection.add(embeddings=[0, 0, 0], ids=["0"])
+        # After adding, should contain one item
+        self.assertEqual(app.count(), 1)
+
+        app.set_collection("test_collection_2")
+        # New collection is empty
+        self.assertEqual(app.count(), 0)
+
+        # Adding to new collection should not effect existing collection
+        app.collection.add(embeddings=[0, 0, 0], ids=["0"])
+        app.set_collection("test_collection_1")
+        # Should still be 1, not 2.
+        self.assertEqual(app.count(), 1)
+
+    def test_collections_are_persistent(self):
+        """
+        Test that a collection can be picked up later.
+        """
+        # Start with a clean app
+        App().reset()
+
+        app = App(config=AppConfig(collect_metrics=False))
+        app.set_collection("test_collection_1")
+        app.collection.add(embeddings=[[0, 0, 0]], ids=["0"])
+        del app
+
+        app = App(config=AppConfig(collect_metrics=False))
+        app.set_collection("test_collection_1")
+        self.assertEqual(app.count(), 1)
+
+    def test_parallel_collections(self):
+        """
+        Test that two apps can have different collections open in parallel.
+        Switching the names will allow instant access to the collection of
+        the other app.
+        """
+        # Start clean
+        App().reset()
+
+        # Create two apps
+        app1 = App(AppConfig(collection_name="test_collection_1", collect_metrics=False))
+        app2 = App(AppConfig(collection_name="test_collection_2", collect_metrics=False))
+
+        # app2 has been created last, but adding to app1 will still write to collection 1.
+        app1.collection.add(embeddings=[0, 0, 0], ids=["0"])
+        self.assertEqual(app1.count(), 1)
+        self.assertEqual(app2.count(), 0)
+
+        # Add data
+        app1.collection.add(embeddings=[[0, 0, 0], [1, 1, 1]], ids=["1", "2"])
+        app2.collection.add(embeddings=[0, 0, 0], ids=["0"])
+
+        # Swap names and test
+        app1.set_collection("test_collection_2")
+        self.assertEqual(app1.count(), 1)
+        app2.set_collection("test_collection_1")
+        self.assertEqual(app2.count(), 3)
+
+    def test_ids_share_collections(self):
+        """
+        Different ids should still share collections.
+        """
+        # Start clean
+        App().reset()
+
+        # Create two apps
+        app1 = App(AppConfig(collection_name="one_collection", id="new_app_id_1", collect_metrics=False))
+        app2 = App(AppConfig(collection_name="one_collection", id="new_app_id_2", collect_metrics=False))
+
+        # Add data
+        app1.collection.add(embeddings=[[0, 0, 0], [1, 1, 1]], ids=["0", "1"])
+        app2.collection.add(embeddings=[0, 0, 0], ids=["2"])
+
+        # Both should have the same collection
+        self.assertEqual(app1.count(), 3)
+        self.assertEqual(app2.count(), 3)
+
+    def test_reset(self):
+        """
+        Resetting should hit all collections and ids.
+        """
+        # Start clean
+        App().reset()
+
+        # Create four apps.
+        # app1, which we are about to reset, shares an app with one, and an id with the other, none with the last.
+        app1 = App(AppConfig(collection_name="one_collection", id="new_app_id_1", collect_metrics=False))
+        app2 = App(AppConfig(collection_name="one_collection", id="new_app_id_2", collect_metrics=False))
+        app3 = App(AppConfig(collection_name="three_collection", id="new_app_id_1", collect_metrics=False))
+        app4 = App(AppConfig(collection_name="four_collection", id="new_app_id_4", collect_metrics=False))
+
+        # Each one of them get data
+        app1.collection.add(embeddings=[0, 0, 0], ids=["1"])
+        app2.collection.add(embeddings=[0, 0, 0], ids=["2"])
+        app3.collection.add(embeddings=[0, 0, 0], ids=["3"])
+        app4.collection.add(embeddings=[0, 0, 0], ids=["4"])
+
+        # Resetting the first one should reset them all.
+        app1.reset()
+
+        # Reinstantiate them
+        app1 = App(AppConfig(collection_name="one_collection", id="new_app_id_1", collect_metrics=False))
+        app2 = App(AppConfig(collection_name="one_collection", id="new_app_id_2", collect_metrics=False))
+        app3 = App(AppConfig(collection_name="three_collection", id="new_app_id_3", collect_metrics=False))
+        app4 = App(AppConfig(collection_name="four_collection", id="new_app_id_3", collect_metrics=False))
+
+        # All should be empty
+        self.assertEqual(app1.count(), 0)
+        self.assertEqual(app2.count(), 0)
+        self.assertEqual(app3.count(), 0)
+        self.assertEqual(app4.count(), 0)

tests/vectordb/test_elasticsearch_db.py

@@ -0,0 +1,33 @@
+import unittest
+from unittest.mock import Mock
+
+from embedchain.config import ElasticsearchDBConfig
+from embedchain.vectordb.elasticsearch_db import ElasticsearchDB
+
+
+class TestEsDB(unittest.TestCase):
+    def setUp(self):
+        self.es_config = ElasticsearchDBConfig()
+        self.vector_dim = 384
+
+    def test_init_with_invalid_embedding_fn(self):
+        # Test if an exception is raised when an invalid embedding_fn is provided
+        with self.assertRaises(ValueError):
+            ElasticsearchDB(embedding_fn=None)
+
+    def test_init_with_invalid_es_config(self):
+        # Test if an exception is raised when an invalid es_config is provided
+        with self.assertRaises(ValueError):
+            ElasticsearchDB(embedding_fn=Mock(), es_config=None)
+
+    def test_init_with_invalid_vector_dim(self):
+        # Test if an exception is raised when an invalid vector_dim is provided
+        with self.assertRaises(ValueError):
+            ElasticsearchDB(embedding_fn=Mock(), es_config=self.es_config, vector_dim=None)
+
+    def test_init_with_invalid_collection_name(self):
+        # Test if an exception is raised when an invalid collection_name is provided
+        with self.assertRaises(ValueError):
+            ElasticsearchDB(
+                embedding_fn=Mock(), es_config=self.es_config, vector_dim=self.vector_dim, collection_name=None
+            )