[Data loader] Make json data loader work for other languages (#897)

Co-authored-by: Deven Patel <deven298@yahoo.com>

.gitignore

@@ -165,12 +165,14 @@ cython_debug/
 
 # Database
 db
+test-db
 
 .vscode
-/poetry.lock
 .idea/
 
 .DS_Store
 
 notebooks/*.yaml
-.ipynb_checkpoints/
+.ipynb_checkpoints/
+
+!configs/*.yaml

CONTRIBUTING.md

@@ -1,10 +1,10 @@
 # Contributing to embedchain
 
-Let us make contributing easy, collaborative and fun.
+Let us make contribution easy, collaborative and fun.
 
 ## Submit your Contribution through PR
 
-To make a contribution, follow the following steps:
+To make a contribution, follow these steps:
 
 1. Fork and clone this repository
 2. Do the changes on your fork with dedicated feature branch `feature/f1`
@@ -35,7 +35,7 @@ poetry shell
 
 ### 📌 Pre-commit
 
-To ensure our standards, make sure to install pre-commit before star to contribute.
+To ensure our standards, make sure to install pre-commit before starting to contribute.
 
 ```bash
 pre-commit install
@@ -51,7 +51,7 @@ make lint
 
 Make sure that the linter does not report any errors or warnings before submitting a pull request.
 
-### Code Format with `black`
+### Code Formatting with `black`
 
 We use `black` to reformat the code by running the following command:
 

README.md

@@ -1,5 +1,7 @@
 # embedchain
 
+<a href="https://runacap.com/ross-index/q3-2023/" target="_blank" rel="noopener"><img style="width: 260px; height: 56px" src="https://runacap.com/wp-content/uploads/2023/10/ROSS_badge_black_Q3_2023.svg" alt="ROSS Index - Fastest Growing Open-Source Startups in Q3 2023 | Runa Capital" width="260" height="56"/></a>
+
 [![PyPI](https://img.shields.io/pypi/v/embedchain)](https://pypi.org/project/embedchain/)
 [![Slack](https://img.shields.io/badge/slack-embedchain-brightgreen.svg?logo=slack)](https://join.slack.com/t/embedchain/shared_invite/zt-22uwz3c46-Zg7cIh5rOBteT_xe1jwLDw)
 [![Discord](https://dcbadge.vercel.app/api/server/6PzXDgEjG5?style=flat)](https://discord.gg/CUU9FPhRNt)
@@ -38,19 +40,23 @@ The documentation for embedchain can be found at [docs.embedchain.ai](https://do
 
 Embedchain empowers you to create ChatGPT like apps, on your own dynamic dataset.
 
-### Data Types Supported
+### Data types supported
 
 * Youtube video
 * PDF file
+* CSV file
 * Web page
+* MDX file
+* XML file
 * Sitemap
 * Doc file
-* JSON file
-* Code documentation website loader
 * Notion
+* JSON file
+* OpenAPI specs
+* Code docs website
 * Unstructured file loader and many more
 
-You can find the full list of data types on [our documentation](https://docs.embedchain.ai/data-sources/csv).
+You can find the full list of data types on [our documentation](https://docs.embedchain.ai/data-sources/).
 
 ### Queries
 
@@ -58,7 +64,7 @@ For example, you can use Embedchain to create an Elon Musk bot using the followi
 
 ```python
 import os
-from embedchain import App
+from embedchain import Pipeline as App
 
 # Create a bot instance
 os.environ["OPENAI_API_KEY"] = "YOUR API KEY"
@@ -72,6 +78,17 @@ elon_bot.add("https://www.youtube.com/watch?v=RcYjXbSJBN8")
 # Query the bot
 elon_bot.query("How many companies does Elon Musk run and name those?")
 # Answer: Elon Musk currently runs several companies. As of my knowledge, he is the CEO and lead designer of SpaceX, the CEO and product architect of Tesla, Inc., the CEO and founder of Neuralink, and the CEO and founder of The Boring Company. However, please note that this information may change over time, so it's always good to verify the latest updates.
+
+# (Optional): Deploy app to Embedchain Platform
+app.deploy()
+# 🔑 Enter your Embedchain API key. You can find the API key at https://app.embedchain.ai/settings/keys/
+# ec-xxxxxx
+
+# 🛠️ Creating pipeline on the platform...
+# 🎉🎉🎉 Pipeline created successfully! View your pipeline: https://app.embedchain.ai/pipelines/xxxxx
+
+# 🛠️ Adding data to your pipeline...
+# ✅ Data of type: web_page, value: https://www.forbes.com/profile/elon-musk added successfully.
 ```
 
 ## Examples

configs/gpt4all.yaml

@@ -1,7 +1,7 @@
 llm:
   provider: gpt4all
-  model: 'orca-mini-3b.ggmlv3.q4_0.bin'
   config:
+    model: 'orca-mini-3b.ggmlv3.q4_0.bin'
     temperature: 0.5
     max_tokens: 1000
     top_p: 1
@@ -9,5 +9,3 @@ llm:
 
 embedder:
   provider: gpt4all
-  config:
-    model: 'all-MiniLM-L6-v2'

configs/opensource.yaml

@@ -6,8 +6,8 @@ app:
 
 llm:
   provider: gpt4all
-  model: 'orca-mini-3b.ggmlv3.q4_0.bin'
   config:
+    model: 'orca-mini-3b.ggmlv3.q4_0.bin'
     temperature: 0.5
     max_tokens: 1000
     top_p: 1
@@ -23,5 +23,4 @@ vectordb:
 embedder:
   provider: gpt4all
   config:
-    model: 'all-MiniLM-L6-v2'
     deployment_name: null

configs/pipeline.yaml

@@ -0,0 +1,26 @@
+pipeline:
+  config:
+    name: Example pipeline
+    id: pipeline-1  # Make sure that id is different every time you create a new pipeline
+
+vectordb:
+  provider: chroma
+  config:
+    collection_name: pipeline-1
+    dir: db
+    allow_reset: true
+
+llm:
+  provider: gpt4all
+  config:
+    model: 'orca-mini-3b.ggmlv3.q4_0.bin'
+    temperature: 0.5
+    max_tokens: 1000
+    top_p: 1
+    stream: false
+
+embedding_model:
+  provider: gpt4all
+  config:
+    model: 'all-MiniLM-L6-v2'
+    deployment_name: null

docs/components/embedding-models.mdx

@@ -24,7 +24,7 @@ Once you have obtained the key, you can use it like this:
 
 ```python main.py
 import os
-from embedchain import App
+from embedchain import Pipeline as App
 
 os.environ['OPENAI_API_KEY'] = 'xxx'
 
@@ -52,7 +52,7 @@ To use Azure OpenAI embedding model, you have to set some of the azure openai re
 
 ```python main.py
 import os
-from embedchain import App
+from embedchain import Pipeline as App
 
 os.environ["OPENAI_API_TYPE"] = "azure"
 os.environ["OPENAI_API_BASE"] = "https://xxx.openai.azure.com/"
@@ -90,7 +90,7 @@ GPT4All supports generating high quality embeddings of arbitrary length document
 <CodeGroup>
 
 ```python main.py
-from embedchain import App
+from embedchain import Pipeline as App
 
 # load embedding model configuration from config.yaml file
 app = App.from_config(yaml_path="config.yaml")
@@ -108,8 +108,6 @@ llm:
 
 embedder:
   provider: gpt4all
-  config:
-    model: 'all-MiniLM-L6-v2'
 ```
 
 </CodeGroup>
@@ -121,7 +119,7 @@ Hugging Face supports generating embeddings of arbitrary length documents of tex
 <CodeGroup>
 
 ```python main.py
-from embedchain import App
+from embedchain import Pipeline as App
 
 # load embedding model configuration from config.yaml file
 app = App.from_config(yaml_path="config.yaml")
@@ -152,7 +150,7 @@ Embedchain supports Google's VertexAI embeddings model through a simple interfac
 <CodeGroup>
 
 ```python main.py
-from embedchain import App
+from embedchain import Pipeline as App
 
 # load embedding model configuration from config.yaml file
 app = App.from_config(yaml_path="config.yaml")

docs/components/llms.mdx

@@ -26,7 +26,7 @@ Once you have obtained the key, you can use it like this:
 
 ```python
 import os
-from embedchain import App
+from embedchain import Pipeline as App
 
 os.environ['OPENAI_API_KEY'] = 'xxx'
 
@@ -41,7 +41,7 @@ If you are looking to configure the different parameters of the LLM, you can do
 
 ```python main.py
 import os
-from embedchain import App
+from embedchain import Pipeline as App
 
 os.environ['OPENAI_API_KEY'] = 'xxx'
 
@@ -71,7 +71,7 @@ To use Azure OpenAI model, you have to set some of the azure openai related envi
 
 ```python main.py
 import os
-from embedchain import App
+from embedchain import Pipeline as App
 
 os.environ["OPENAI_API_TYPE"] = "azure"
 os.environ["OPENAI_API_BASE"] = "https://xxx.openai.azure.com/"
@@ -110,7 +110,7 @@ To use anthropic's model, please set the `ANTHROPIC_API_KEY` which you find on t
 
 ```python main.py
 import os
-from embedchain import App
+from embedchain import Pipeline as App
 
 os.environ["ANTHROPIC_API_KEY"] = "xxx"
 
@@ -147,7 +147,7 @@ Once you have the API key, you are all set to use it with Embedchain.
 
 ```python main.py
 import os
-from embedchain import App
+from embedchain import Pipeline as App
 
 os.environ["COHERE_API_KEY"] = "xxx"
 
@@ -180,7 +180,7 @@ GPT4all is a free-to-use, locally running, privacy-aware chatbot. No GPU or inte
 <CodeGroup>
 
 ```python main.py
-from embedchain import App
+from embedchain import Pipeline as App
 
 # load llm configuration from config.yaml file
 app = App.from_config(yaml_path="config.yaml")
@@ -198,8 +198,6 @@ llm:
 
 embedder:
   provider: gpt4all
-  config:
-    model: 'all-MiniLM-L6-v2'
 ```
 </CodeGroup>
 
@@ -214,7 +212,7 @@ Once you have the key, load the app using the config yaml file:
 
 ```python main.py
 import os
-from embedchain import App
+from embedchain import Pipeline as App
 
 os.environ["JINACHAT_API_KEY"] = "xxx"
 # load llm configuration from config.yaml file
@@ -250,7 +248,7 @@ Once you have the token, load the app using the config yaml file:
 
 ```python main.py
 import os
-from embedchain import App
+from embedchain import Pipeline as App
 
 os.environ["HUGGINGFACE_ACCESS_TOKEN"] = "xxx"
 
@@ -280,7 +278,7 @@ Once you have the token, load the app using the config yaml file:
 
 ```python main.py
 import os
-from embedchain import App
+from embedchain import Pipeline as App
 
 os.environ["REPLICATE_API_TOKEN"] = "xxx"
 
@@ -307,7 +305,7 @@ Setup Google Cloud Platform application credentials by following the instruction
 <CodeGroup>
 
 ```python main.py
-from embedchain import App
+from embedchain import Pipeline as App
 
 # load llm configuration from config.yaml file
 app = App.from_config(yaml_path="config.yaml")

docs/components/vector-databases.mdx

@@ -22,7 +22,7 @@ Utilizing a vector database alongside Embedchain is a seamless process. All you
 <CodeGroup>
 
 ```python main.py
-from embedchain import App
+from embedchain import Pipeline as App
 
 # load chroma configuration from yaml file
 app = App.from_config(yaml_path="config1.yaml")
@@ -61,7 +61,7 @@ pip install --upgrade 'embedchain[elasticsearch]'
 <CodeGroup>
 
 ```python main.py
-from embedchain import App
+from embedchain import Pipeline as App
 
 # load elasticsearch configuration from yaml file
 app = App.from_config(yaml_path="config.yaml")
@@ -89,7 +89,7 @@ pip install --upgrade 'embedchain[opensearch]'
 <CodeGroup>
 
 ```python main.py
-from embedchain import App
+from embedchain import Pipeline as App
 
 # load opensearch configuration from yaml file
 app = App.from_config(yaml_path="config.yaml")
@@ -125,7 +125,7 @@ Set the Zilliz environment variables `ZILLIZ_CLOUD_URI` and `ZILLIZ_CLOUD_TOKEN`
 
 ```python main.py
 import os
-from embedchain import App
+from embedchain import Pipeline as App
 
 os.environ['ZILLIZ_CLOUD_URI'] = 'https://xxx.zillizcloud.com'
 os.environ['ZILLIZ_CLOUD_TOKEN'] = 'xxx'
@@ -164,7 +164,7 @@ In order to use Pinecone as vector database, set the environment variables `PINE
 <CodeGroup>
 
 ```python main.py
-from embedchain import App
+from embedchain import Pipeline as App
 
 # load pinecone configuration from yaml file
 app = App.from_config(yaml_path="config.yaml")
@@ -187,7 +187,7 @@ In order to use Qdrant as a vector database, set the environment variables `QDRA
 
 <CodeGroup>
 ```python main.py
-from embedchain import App
+from embedchain import Pipeline as App
 
 # load qdrant configuration from yaml file
 app = App.from_config(yaml_path="config.yaml")
@@ -207,7 +207,7 @@ In order to use Weaviate as a vector database, set the environment variables `WE
 
 <CodeGroup>
 ```python main.py
-from embedchain import App
+from embedchain import Pipeline as App
 
 # load weaviate configuration from yaml file
 app = App.from_config(yaml_path="config.yaml")

docs/data-sources/csv.mdx

@@ -5,7 +5,7 @@ title: '📊 CSV'
 To add any csv file, use the data_type as `csv`. `csv` allows remote urls and conventional file paths. Headers are included for each line, so if you have an `age` column, `18` will be added as `age: 18`. Eg:
 
 ```python
-from embedchain import App
+from embedchain import Pipeline as App
 
 app = App()
 app.add('https://people.sc.fsu.edu/~jburkardt/data/csv/airtravel.csv', data_type="csv")

docs/data-sources/data-type-handling.mdx

@@ -35,7 +35,7 @@ Default behavior is to create a persistent vector db in the directory **./db**.
 Create a local index:
 
 ```python
-from embedchain import App
+from embedchain import Pipeline as App
 
 naval_chat_bot = App()
 naval_chat_bot.add("https://www.youtube.com/watch?v=3qHkcs3kG44")
@@ -45,7 +45,7 @@ naval_chat_bot.add("https://navalmanack.s3.amazonaws.com/Eric-Jorgenson_The-Alma
 You can reuse the local index with the same code, but without adding new documents:
 
 ```python
-from embedchain import App
+from embedchain import Pipeline as App
 
 naval_chat_bot = App()
 print(naval_chat_bot.query("What unique capacity does Naval argue humans possess when it comes to understanding explanations or concepts?"))

docs/data-sources/docs-site.mdx

@@ -5,7 +5,7 @@ title: '📚🌐 Code documentation'
 To add any code documentation website as a loader, use the data_type as `docs_site`. Eg:
 
 ```python
-from embedchain import App
+from embedchain import Pipeline as App
 
 app = App()
 app.add("https://docs.embedchain.ai/", data_type="docs_site")

docs/data-sources/docx.mdx

@@ -7,7 +7,7 @@ title: '📄 Docx file'
 To add any doc/docx file, use the data_type as `docx`. `docx` allows remote urls and conventional file paths. Eg:
 
 ```python
-from embedchain import App
+from embedchain import Pipeline as App
 
 app = App()
 app.add('https://example.com/content/intro.docx', data_type="docx")

docs/data-sources/gmail.mdx

@@ -0,0 +1,35 @@
+---
+title: '📬 Gmail'
+---
+
+To use GmailLoader you must install the extra dependencies with `pip install --upgrade embedchain[gmail]`.
+
+The `source` must be a valid Gmail search query, you can refer `https://support.google.com/mail/answer/7190?hl=en` to build a query.
+
+To load Gmail messages, you MUST use the data_type as `gmail`. Otherwise the source will be detected as simple `text`.
+
+To use this you need to save `credentials.json` in the directory from where you will run the loader. Follow these steps to get the credentials
+
+1. Go to the [Google Cloud Console](https://console.cloud.google.com/apis/credentials).
+2. Create a project if you don't have one already.
+3. Create an `OAuth Consent Screen` in the project. You may need to select the `external` option.
+4. Make sure the consent screen is published.
+5. Enable the [Gmail API](https://console.cloud.google.com/apis/api/gmail.googleapis.com)
+6. Create credentials from the `Credentials` tab.
+7. Select the type `OAuth Client ID`.
+8. Choose the application type `Web application`. As a name you can choose `embedchain` or any other name as per your use case.
+9. Add an authorized redirect URI for `http://localhost:8080/`.
+10. You can leave everything else at default, finish the creation.
+11. When you are done, a modal opens where you can download the details in `json` format.
+12. Put the `.json` file in your current directory and rename it to `credentials.json`
+
+```python
+import os
+from embedchain.apps.app import App
+from embedchain.models.data_type import DataType
+app = App()
+
+query = "to: me label:inbox"
+app.add(query, data_type=DataType.GMAIL)
+app.query("Summarize my email conversations")
+```

docs/data-sources/json.mdx

@@ -0,0 +1,36 @@
+---
+title: '📃 JSON'
+---
+
+To add any json file, use the data_type as `json`. `json` allows remote urls and conventional file paths. Headers are included for each line, so if you have an `age` column, `18` will be added as `age: 18`. Eg:
+
+```python
+import os
+
+from embedchain.apps.app import App
+
+os.environ["OPENAI_API_KEY"] = "openai_api_key"
+
+app = App()
+
+response = app.query("What is the net worth of Elon Musk as of October 2023?")
+
+print(response)
+"I'm sorry, but I don't have access to real-time information or future predictions. Therefore, I don't know the net worth of Elon Musk as of October 2023."
+
+source_id = app.add("temp.json")
+
+response = app.query("What is the net worth of Elon Musk as of October 2023?")
+
+print(response)
+"As of October 2023, Elon Musk's net worth is $255.2 billion."
+```
+
+```temp.json
+{
+    "question": "What is your net worth, Elon Musk?",
+    "answer": "As of October 2023, Elon Musk's net worth is $255.2 billion, making him one of the wealthiest individuals in the world."
+}
+```
+
+

docs/data-sources/mdx.mdx

@@ -5,7 +5,7 @@ title: '📝 Mdx file'
 To add any `.mdx` file to your app, use the data_type (first argument to `.add()` method) as `mdx`. Note that this supports support mdx file present on machine, so this should be a file path. Eg:
 
 ```python
-from embedchain import App
+from embedchain import Pipeline as App
 
 app = App()
 app.add('path/to/file.mdx', data_type='mdx')

docs/data-sources/notion.mdx

@@ -2,13 +2,13 @@
 title: '📓 Notion'
 ---
 
-To use notion you must install the extra dependencies with `pip install --upgrade embedchain[notion]`.
+To use notion you must install the extra dependencies with `pip install --upgrade embedchain[community]`.
 
 To load a notion page, use the data_type as `notion`. Since it is hard to automatically detect, it is advised to specify the `data_type` when adding a notion document.
 The next argument must **end** with the `notion page id`. The id is a 32-character string. Eg:
 
 ```python
-from embedchain import App
+from embedchain import Pipeline as App
 
 app = App()
 

docs/data-sources/openapi.mdx

@@ -0,0 +1,23 @@
+---
+title: 🙌 OpenAPI
+---
+
+To add any OpenAPI spec yaml file (currently the json file will be detected as JSON data type), use the data_type as 'openapi'. 'openapi' allows remote urls and conventional file paths. Headers are included for each line, so if you have an `age` column, `18` will be added as `age: 18`. Eg:
+
+```python
+from embedchain.apps.app import App
+import os
+
+os.environ["OPENAI_API_KEY"] = "sk-xxx"
+
+app = App()
+
+app.add("https://github.com/openai/openai-openapi/blob/master/openapi.yaml", data_type="openapi")
+# Or add using the local file path
+# app.add("configs/openai_openapi.yaml", data_type="openapi")
+
+response = app.query("What can OpenAI API endpoint do? Can you list the things it can learn from?")
+# Answer: The OpenAI API endpoint allows users to interact with OpenAI's models and perform various tasks such as generating text, answering questions, summarizing documents, translating languages, and more. The specific capabilities and tasks that the API can learn from may vary depending on the models and features provided by OpenAI. For more detailed information, it is recommended to refer to the OpenAI API documentation at https://platform.openai.com/docs/api-reference.
+```
+
+NOTE: The yaml file added to the App must have the required OpenAPI fields otherwise the adding OpenAPI spec will fail. Please refer to [OpenAPI Spec Doc](https://spec.openapis.org/oas/v3.1.0)

docs/data-sources/overview.mdx

@@ -6,6 +6,7 @@ Embedchain comes with built-in support for various data sources. We handle the c
 
 <CardGroup cols={4}>
   <Card title="📊 csv" href="/data-sources/csv"></Card>
+  <Card title="📃 JSON" href="/data-sources/json"></Card>
   <Card title="📚🌐 docs site" href="/data-sources/docs-site"></Card>
   <Card title="📄 docx" href="/data-sources/docx"></Card>
   <Card title="📝 mdx" href="/data-sources/mdx"></Card>
@@ -16,7 +17,9 @@ Embedchain comes with built-in support for various data sources. We handle the c
   <Card title="📝 text" href="/data-sources/text"></Card>
   <Card title="🌐📄 web page" href="/data-sources/web-page"></Card>
   <Card title="🧾 xml" href="/data-sources/xml"></Card>
+  <Card title="🙌 OpenAPI" href="/data-sources/openapi"></Card>
   <Card title="🎥📺 youtube video" href="/data-sources/youtube-video"></Card>
+  <Card title="📬 Gmail" href="/data-sources/gmail"></Card>
 </CardGroup>
 
 <br/ >

docs/data-sources/pdf-file.mdx

@@ -5,7 +5,7 @@ title: '📰 PDF file'
 To add any pdf file, use the data_type as `pdf_file`. Eg:
 
 ```python
-from embedchain import App
+from embedchain import Pipeline as App
 
 app = App()
 

docs/data-sources/qna.mdx

@@ -5,7 +5,7 @@ title: '❓💬 Queston and answer pair'
 QnA pair is a local data type. To supply your own QnA pair, use the data_type as `qna_pair` and enter a tuple. Eg:
 
 ```python
-from embedchain import App
+from embedchain import Pipeline as App
 
 app = App()
 

docs/data-sources/sitemap.mdx

@@ -5,7 +5,7 @@ title: '🗺️ Sitemap'
 Add all web pages from an xml-sitemap. Filters non-text files. Use the data_type as `sitemap`. Eg:
 
 ```python
-from embedchain import App
+from embedchain import Pipeline as App
 
 app = App()
 

docs/data-sources/text.mdx

@@ -7,7 +7,7 @@ title: '📝 Text'
 Text is a local data type. 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:
 
 ```python
-from embedchain import App
+from embedchain import Pipeline as App
 
 app = App()
 

docs/data-sources/web-page.mdx

@@ -5,7 +5,7 @@ title: '🌐📄 Web page'
 To add any web page, use the data_type as `web_page`. Eg:
 
 ```python
-from embedchain import App
+from embedchain import Pipeline as App
 
 app = App()
 

docs/data-sources/xml.mdx

@@ -7,7 +7,7 @@ title: '🧾 XML file'
 To add any xml file, use the data_type as `xml`. Eg:
 
 ```python
-from embedchain import App
+from embedchain import Pipeline as App
 
 app = App()
 

docs/data-sources/youtube-video.mdx

@@ -6,7 +6,7 @@ title: '🎥📺 Youtube video'
 To add any youtube video to your app, use the data_type (first argument to `.add()` method) as `youtube_video`. Eg:
 
 ```python
-from embedchain import App
+from embedchain import Pipeline as App
 
 app = App()
 app.add('a_valid_youtube_url_here', data_type='youtube_video')

docs/get-started/faq.mdx

@@ -9,7 +9,7 @@ description: 'Collections of all the frequently asked questions'
 
 ```python main.py
 import os
-from embedchain import App
+from embedchain import Pipeline as App
 
 os.environ['OPENAI_API_KEY'] = 'xxx'
 
@@ -36,7 +36,7 @@ llm:
 
 ```python main.py
 import os
-from embedchain import App
+from embedchain import Pipeline as App
 
 os.environ['OPENAI_API_KEY'] = 'xxx'
 

docs/get-started/introduction.mdx

@@ -3,30 +3,106 @@ title: 📚 Introduction
 description: '📝 Embedchain is a Data Platform for LLMs - load, index, retrieve, and sync any unstructured data'
 ---
 
-## 🤔 What is Embedchain?
+## 🌐 What is Embedchain?
 
-Embedchain abstracts the entire process of loading data, chunking it, creating embeddings, and storing it in a vector database.
+Embedchain simplifies data handling by automatically processing unstructured data, breaking it into chunks, generating embeddings, and storing it in a vector database.
 
-You can add data from different data sources using the `.add()` method. Then, simply use the `.query()` method to find answers from the added datasets.
+Through various APIs, you can obtain contextual information for queries, find answers to specific questions, and engage in chat conversations using your data.
+## 🔍 Search
 
-If you want to create a Naval Ravikant bot with a YouTube video, a book in PDF format, two blog posts, and a question and answer pair, all you need to do is add the respective links. Embedchain will take care of the rest, creating a bot for you.
+Embedchain lets you get most relevant context by doing semantic search over your data sources for a provided query. See the example below:
 
 ```python
-from embedchain import App
+from embedchain import Pipeline as App
 
-naval_bot = App()
-# Add online data
-naval_bot.add("https://www.youtube.com/watch?v=3qHkcs3kG44")
-naval_bot.add("https://navalmanack.s3.amazonaws.com/Eric-Jorgenson_The-Almanack-of-Naval-Ravikant_Final.pdf")
-naval_bot.add("https://nav.al/feedback")
-naval_bot.add("https://nav.al/agi")
-naval_bot.add("The Meanings of Life", 'text', metadata={'chapter': 'philosphy'})
+# Initialize app
+app = App()
 
-# Add local resources
-naval_bot.add(("Who is Naval Ravikant?", "Naval Ravikant is an Indian-American entrepreneur and investor."))
+# Add data source
+app.add("https://www.forbes.com/profile/elon-musk")
 
-naval_bot.query("What unique capacity does Naval argue humans possess when it comes to understanding explanations or concepts?")
-# Answer: Naval argues that humans possess the unique capacity to understand explanations or concepts to the maximum extent possible in this physical reality.
+# Get relevant context using semantic search
+context = app.search("What is the net worth of Elon?", num_documents=2)
+print(context)
+# Context:
+# [
+#     {
+#         'context': 'Elon Musk PROFILEElon MuskCEO, Tesla$221.9BReal Time Net Worthas of 10/29/23Reflects change since 5 pm ET of prior trading day. 1 in the world todayPhoto by Martin Schoeller for ForbesAbout Elon MuskElon Musk cofounded six companies, including electric car maker Tesla, rocket producer SpaceX and tunneling startup Boring Company.He owns about 21% of Tesla between stock and options, but has pledged more than half his shares as collateral for personal loans of up to $3.5 billion.SpaceX, founded in',
+#         'source': 'https://www.forbes.com/profile/elon-musk',
+#         'document_id': 'some_document_id'
+#     },
+#     {
+#         'context': 'company, which is now called X.Wealth HistoryHOVER TO REVEAL NET WORTH BY YEARForbes Lists 1Forbes 400 (2023)The Richest Person In Every State (2023) 2Billionaires (2023) 1Innovative Leaders (2019) 25Powerful People (2018) 12Richest In Tech (2017)Global Game Changers (2016)More ListsPersonal StatsAge52Source of WealthTesla, SpaceX, Self MadeSelf-Made Score8Philanthropy Score1ResidenceAustin, TexasCitizenshipUnited StatesMarital StatusSingleChildren11EducationBachelor of Arts/Science, University',
+#         'source': 'https://www.forbes.com/profile/elon-musk',
+#         'document_id': 'some_document_id'
+#     }
+# ]
+```
+
+## ❓Query
+
+Embedchain empowers developers to ask questions and receive relevant answers through a user-friendly query API. Refer to the following example to learn how to utilize the query API:
+
+```python
+from embedchain import Pipeline as App
+
+# Initialize app
+app = App()
+
+# Add data source
+app.add("https://www.forbes.com/profile/elon-musk")
+
+# Get relevant answer for your query
+answer = app.query("What is the net worth of Elon?")
+print(answer)
+# Answer: The net worth of Elon Musk is $221.9 billion.
+```
+
+## 💬 Chat
+
+Embedchain allows easy chatting over your data sources using a user-friendly chat API. Check out the example below to understand how to use the chat API:
+
+```python
+from embedchain import Pipeline as App
+
+# Initialize app
+app = App()
+
+# Add data source
+app.add("https://www.forbes.com/profile/elon-musk")
+
+# Chat on your data using `.chat()`
+answer = app.chat("How much did Elon pay for Twitter?")
+print(answer)
+# Answer: Elon Musk paid $44 billion for Twitter.
+```
+
+## 🚀 Deploy
+
+Embedchain enables developers to deploy their LLM-powered apps in production using the Embedchain platform. The platform offers free access to context on your data through its REST API. Once the pipeline is deployed, you can update your data sources anytime after deployment.
+
+See the example below on how to use the deploy API:
+
+```python
+from embedchain import Pipeline as App
+
+# Initialize app
+app = App()
+
+# Add data source
+app.add("https://www.forbes.com/profile/elon-musk")
+
+# Deploy your pipeline to Embedchain Platform
+app.deploy()
+
+# 🔑 Enter your Embedchain API key. You can find the API key at https://app.embedchain.ai/settings/keys/
+# ec-xxxxxx
+
+# 🛠️ Creating pipeline on the platform...
+# 🎉🎉🎉 Pipeline created successfully! View your pipeline: https://app.embedchain.ai/pipelines/xxxxx
+
+# 🛠️ Adding data to your pipeline...
+# ✅ Data of type: web_page, value: https://www.forbes.com/profile/elon-musk added successfully.
 ```
 
 ## 🚀 How it works?

docs/get-started/quickstart.mdx

@@ -16,23 +16,36 @@ Creating an app involves 3 steps:
 <Steps>
   <Step title="⚙️ Import app instance">
 ```python
-from embedchain import App
+from embedchain import Pipeline as App
 app = App()
 ```
   </Step>
   <Step title="🗃️ Add data sources">
 ```python
 # Add different data sources
-elon_bot.add("https://en.wikipedia.org/wiki/Elon_Musk")
-elon_bot.add("https://www.forbes.com/profile/elon-musk")
+app.add("https://en.wikipedia.org/wiki/Elon_Musk")
+app.add("https://www.forbes.com/profile/elon-musk")
 # You can also add local data sources such as pdf, csv files etc.
-# elon_bot.add("/path/to/file.pdf")
+# app.add("/path/to/file.pdf")
 ```
   </Step>
-  <Step title="💬 Query or chat on your data and get answers">
+  <Step title="💬 Query or chat or search context on your data">
 ```python
-elon_bot.query("What is the net worth of Elon Musk today?")
+app.query("What is the net worth of Elon Musk today?")
 # Answer: The net worth of Elon Musk today is $258.7 billion.
+```
+  </Step>
+  <Step title="🚀 (Optional) Deploy your pipeline to Embedchain Platform">
+```python
+app.deploy()
+# 🔑 Enter your Embedchain API key. You can find the API key at https://app.embedchain.ai/settings/keys/
+# ec-xxxxxx
+
+# 🛠️ Creating pipeline on the platform...
+# 🎉🎉🎉 Pipeline created successfully! View your pipeline: https://app.embedchain.ai/pipelines/xxxxx
+
+# 🛠️ Adding data to your pipeline...
+# ✅ Data of type: web_page, value: https://www.forbes.com/profile/elon-musk added successfully.
 ```
   </Step>
 </Steps>
@@ -41,18 +54,28 @@ Putting it together, you can run your first app using the following code. Make s
 
 ```python
 import os
-from embedchain import App
+from embedchain import Pipeline as App
 
 os.environ["OPENAI_API_KEY"] = "xxx"
-elon_bot = App()
+app = App()
 
 # Add different data sources
-elon_bot.add("https://en.wikipedia.org/wiki/Elon_Musk")
-elon_bot.add("https://www.forbes.com/profile/elon-musk")
+app.add("https://en.wikipedia.org/wiki/Elon_Musk")
+app.add("https://www.forbes.com/profile/elon-musk")
 # You can also add local data sources such as pdf, csv files etc.
-# elon_bot.add("/path/to/file.pdf")
+# app.add("/path/to/file.pdf")
 
-response = elon_bot.query("What is the net worth of Elon Musk today?")
+response = app.query("What is the net worth of Elon Musk today?")
 print(response)
 # Answer: The net worth of Elon Musk today is $258.7 billion.
+
+app.deploy()
+# 🔑 Enter your Embedchain API key. You can find the API key at https://app.embedchain.ai/settings/keys/
+# ec-xxxxxx
+
+# 🛠️ Creating pipeline on the platform...
+# 🎉🎉🎉 Pipeline created successfully! View your pipeline: https://app.embedchain.ai/pipelines/xxxxx
+
+# 🛠️ Adding data to your pipeline...
+# ✅ Data of type: web_page, value: https://www.forbes.com/profile/elon-musk added successfully.
 ```

docs/integration/langsmith.mdx

@@ -39,7 +39,7 @@ os.environ['LANGCHAIN_PROJECT] = <your-project>
 
 
 ```python
-from embedchain import App
+from embedchain import Pipeline as App
 
 app = App()
 app.add("https://en.wikipedia.org/wiki/Elon_Musk")

docs/mint.json

@@ -7,9 +7,20 @@
   },
   "favicon": "/favicon.png",
   "colors": {
-    "primary": "#12A7D3",
-    "light": "#81D7F7",
-    "dark": "#004E7A"
+    "primary": "#2B48EE",
+    "light": "#2B48EE",
+    "dark": "#2B48EE",
+    "background": {
+      "dark": "#020415"
+    }
+  },
+  "metadata": {
+    "og:image": "/images/og.png",
+    "twitter:site": "@embedchain"
+  },
+  "topAnchor": {
+    "name": "Documentation",
+    "icon": "book-open"
   },
   "topbarLinks": [
     {
@@ -26,8 +37,11 @@
     }
   ],
   "topbarCtaButton": {
-    "name": "GitHub",
-    "url": "https://embedchain.ai"
+    "name": "Get started",
+    "url": "https://app.embedchain.ai"
+  },
+  "primaryTab": {
+    "name": "Docs"
   },
   "navigation": [
     {
@@ -46,6 +60,7 @@
           "group": "Supported data sources",
           "pages": [
             "data-sources/csv",
+            "data-sources/json",
             "data-sources/docs-site",
             "data-sources/docx",
             "data-sources/mdx",
@@ -55,6 +70,7 @@
             "data-sources/sitemap",
             "data-sources/text",
             "data-sources/web-page",
+            "data-sources/openapi",
             "data-sources/youtube-video"
           ]
         },
@@ -98,7 +114,6 @@
     }
 
   ],
-
   "footerSocials": {
     "website": "https://embedchain.ai",
     "github": "https://github.com/embedchain/embedchain",
@@ -107,7 +122,19 @@
     "twitter": "https://twitter.com/embedchain",
     "linkedin": "https://www.linkedin.com/company/embedchain"
   },
-  "backgroundImage": "/background.png",
   "isWhiteLabeled": true,
-  "feedback.thumbsRating": true
+  "analytics": {
+    "posthog": {
+      "apiKey": "phc_PHQDA5KwztijnSojsxJ2c1DuJd52QCzJzT2xnSGvjN2",
+      "apiHost": "https://app.embedchain.ai/ingest"
+    }
+  },
+  "feedback": {
+    "suggestEdit": true,
+    "raiseIssue": true,
+    "thumbsRating": true
+  },
+  "search": {
+    "prompt": "✨ Search embedchain docs..."
+  }
 }

embedchain/__init__.py

@@ -3,5 +3,6 @@ import importlib.metadata
 __version__ = importlib.metadata.version(__package__ or __name__)
 
 from embedchain.apps.app import App  # noqa: F401
+from embedchain.client import Client  # noqa: F401
 from embedchain.pipeline import Pipeline  # noqa: F401
 from embedchain.vectordb.chroma import ChromaDB  # noqa: F401

embedchain/chunkers/gmail.py

@@ -0,0 +1,22 @@
+from typing import Optional
+
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+
+from embedchain.chunkers.base_chunker import BaseChunker
+from embedchain.config.add_config import ChunkerConfig
+from embedchain.helper.json_serializable import register_deserializable
+
+
+@register_deserializable
+class GmailChunker(BaseChunker):
+    """Chunker for gmail."""
+
+    def __init__(self, config: Optional[ChunkerConfig] = None):
+        if config is None:
+            config = ChunkerConfig(chunk_size=1000, 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/chunkers/openapi.py

@@ -0,0 +1,18 @@
+from typing import Optional
+
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+
+from embedchain.chunkers.base_chunker import BaseChunker
+from embedchain.config.add_config import ChunkerConfig
+
+
+class OpenAPIChunker(BaseChunker):
+    def __init__(self, config: Optional[ChunkerConfig] = None):
+        if config is None:
+            config = ChunkerConfig(chunk_size=1000, 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/client.py

@@ -0,0 +1,102 @@
+import json
+import logging
+import os
+import uuid
+
+import requests
+
+from embedchain.embedchain import CONFIG_DIR, CONFIG_FILE
+
+
+class Client:
+    def __init__(self, api_key=None, host="https://apiv2.embedchain.ai"):
+        self.config_data = self.load_config()
+        self.host = host
+
+        if api_key:
+            if self.check(api_key):
+                self.api_key = api_key
+                self.save()
+            else:
+                raise ValueError(
+                    "Invalid API key provided. You can find your API key on https://app.embedchain.ai/settings/keys."
+                )
+        else:
+            if "api_key" in self.config_data:
+                self.api_key = self.config_data["api_key"]
+                logging.info("API key loaded successfully!")
+            else:
+                raise ValueError(
+                    "You are not logged in. Please obtain an API key from https://app.embedchain.ai/settings/keys/"
+                )
+
+    @classmethod
+    def setup_dir(self):
+        """
+        Loads the user id from the config file if it exists, otherwise generates a new
+        one and saves it to the config file.
+
+        :return: user id
+        :rtype: str
+        """
+        if not os.path.exists(CONFIG_DIR):
+            os.makedirs(CONFIG_DIR)
+
+        if os.path.exists(CONFIG_FILE):
+            with open(CONFIG_FILE, "r") as f:
+                data = json.load(f)
+                if "user_id" in data:
+                    return data["user_id"]
+
+        u_id = str(uuid.uuid4())
+        with open(CONFIG_FILE, "w") as f:
+            json.dump({"user_id": u_id}, f)
+
+    @classmethod
+    def load_config(cls):
+        if not os.path.exists(CONFIG_FILE):
+            cls.setup_dir()
+
+        with open(CONFIG_FILE, "r") as config_file:
+            return json.load(config_file)
+
+    def save(self):
+        self.config_data["api_key"] = self.api_key
+        with open(CONFIG_FILE, "w") as config_file:
+            json.dump(self.config_data, config_file, indent=4)
+
+        logging.info("API key saved successfully!")
+
+    def clear(self):
+        if "api_key" in self.config_data:
+            del self.config_data["api_key"]
+            with open(CONFIG_FILE, "w") as config_file:
+                json.dump(self.config_data, config_file, indent=4)
+            self.api_key = None
+            logging.info("API key deleted successfully!")
+        else:
+            logging.warning("API key not found in the configuration file.")
+
+    def update(self, api_key):
+        if self.check(api_key):
+            self.api_key = api_key
+            self.save()
+            logging.info("API key updated successfully!")
+        else:
+            logging.warning("Invalid API key provided. API key not updated.")
+
+    def check(self, api_key):
+        validation_url = f"{self.host}/api/v1/accounts/api_keys/validate/"
+        response = requests.post(validation_url, headers={"Authorization": f"Token {api_key}"})
+        if response.status_code == 200:
+            return True
+        else:
+            logging.warning(f"Response from API: {response.text}")
+            logging.warning("Invalid API key. Unable to validate.")
+            return False
+
+    def get(self):
+        return self.api_key
+
+    def __str__(self):
+        return self.api_key

embedchain/config/__init__.py

@@ -2,7 +2,6 @@
 
 from .add_config import AddConfig, ChunkerConfig
 from .apps.app_config import AppConfig
-from .pipeline_config import PipelineConfig
 from .base_config import BaseConfig
 from .embedder.base import BaseEmbedderConfig
 from .embedder.base import BaseEmbedderConfig as EmbedderConfig

embedchain/config/apps/app_config.py

@@ -15,7 +15,7 @@ class AppConfig(BaseAppConfig):
         self,
         log_level: str = "WARNING",
         id: Optional[str] = None,
-        collect_metrics: Optional[bool] = None,
+        collect_metrics: Optional[bool] = True,
         collection_name: Optional[str] = None,
     ):
         """

embedchain/config/pipeline_config.py

@@ -16,7 +16,7 @@ class PipelineConfig(BaseAppConfig):
         log_level: str = "WARNING",
         id: Optional[str] = None,
         name: Optional[str] = None,
-        collect_metrics: Optional[bool] = False,
+        collect_metrics: Optional[bool] = True,
     ):
         """
         Initializes a configuration class instance for an App. This is the simplest form of an embedchain app.

embedchain/data_formatter/data_formatter.py

@@ -1,37 +1,10 @@
+from importlib import import_module
+
 from embedchain.chunkers.base_chunker import BaseChunker
-from embedchain.chunkers.docs_site import DocsSiteChunker
-from embedchain.chunkers.docx_file import DocxFileChunker
-from embedchain.chunkers.images import ImagesChunker
-from embedchain.chunkers.json import JSONChunker
-from embedchain.chunkers.mdx import MdxChunker
-from embedchain.chunkers.notion import NotionChunker
-from embedchain.chunkers.pdf_file import PdfFileChunker
-from embedchain.chunkers.qna_pair import QnaPairChunker
-from embedchain.chunkers.sitemap import SitemapChunker
-from embedchain.chunkers.table import TableChunker
-from embedchain.chunkers.text import TextChunker
-from embedchain.chunkers.unstructured_file import UnstructuredFileChunker
-from embedchain.chunkers.web_page import WebPageChunker
-from embedchain.chunkers.xml import XmlChunker
-from embedchain.chunkers.youtube_video import YoutubeVideoChunker
 from embedchain.config import AddConfig
 from embedchain.config.add_config import ChunkerConfig, LoaderConfig
 from embedchain.helper.json_serializable import JSONSerializable
 from embedchain.loaders.base_loader import BaseLoader
-from embedchain.loaders.csv import CsvLoader
-from embedchain.loaders.docs_site_loader import DocsSiteLoader
-from embedchain.loaders.docx_file import DocxFileLoader
-from embedchain.loaders.images import ImagesLoader
-from embedchain.loaders.json import JSONLoader
-from embedchain.loaders.local_qna_pair import LocalQnaPairLoader
-from embedchain.loaders.local_text import LocalTextLoader
-from embedchain.loaders.mdx import MdxLoader
-from embedchain.loaders.pdf_file import PdfFileLoader
-from embedchain.loaders.sitemap import SitemapLoader
-from embedchain.loaders.unstructured_file import UnstructuredLoader
-from embedchain.loaders.web_page import WebPageLoader
-from embedchain.loaders.xml import XmlLoader
-from embedchain.loaders.youtube_video import YoutubeVideoLoader
 from embedchain.models.data_type import DataType
 
 
@@ -54,6 +27,11 @@ class DataFormatter(JSONSerializable):
         self.loader = self._get_loader(data_type=data_type, config=config.loader)
         self.chunker = self._get_chunker(data_type=data_type, config=config.chunker)
 
+    def _lazy_load(self, module_path: str):
+        module_path, class_name = module_path.rsplit(".", 1)
+        module = import_module(module_path)
+        return getattr(module, class_name)
+
     def _get_loader(self, data_type: DataType, config: LoaderConfig) -> BaseLoader:
         """
         Returns the appropriate data loader for the given data type.
@@ -67,67 +45,55 @@ class DataFormatter(JSONSerializable):
         :rtype: BaseLoader
         """
         loaders = {
-            DataType.YOUTUBE_VIDEO: YoutubeVideoLoader,
-            DataType.PDF_FILE: PdfFileLoader,
-            DataType.WEB_PAGE: WebPageLoader,
-            DataType.QNA_PAIR: LocalQnaPairLoader,
-            DataType.TEXT: LocalTextLoader,
-            DataType.DOCX: DocxFileLoader,
-            DataType.SITEMAP: SitemapLoader,
-            DataType.XML: XmlLoader,
-            DataType.DOCS_SITE: DocsSiteLoader,
-            DataType.CSV: CsvLoader,
-            DataType.MDX: MdxLoader,
-            DataType.IMAGES: ImagesLoader,
-            DataType.UNSTRUCTURED: UnstructuredLoader,
-            DataType.JSON: JSONLoader,
+            DataType.YOUTUBE_VIDEO: "embedchain.loaders.youtube_video.YoutubeVideoLoader",
+            DataType.PDF_FILE: "embedchain.loaders.pdf_file.PdfFileLoader",
+            DataType.WEB_PAGE: "embedchain.loaders.web_page.WebPageLoader",
+            DataType.QNA_PAIR: "embedchain.loaders.local_qna_pair.LocalQnaPairLoader",
+            DataType.TEXT: "embedchain.loaders.local_text.LocalTextLoader",
+            DataType.DOCX: "embedchain.loaders.docx_file.DocxFileLoader",
+            DataType.SITEMAP: "embedchain.loaders.sitemap.SitemapLoader",
+            DataType.XML: "embedchain.loaders.xml.XmlLoader",
+            DataType.DOCS_SITE: "embedchain.loaders.docs_site_loader.DocsSiteLoader",
+            DataType.CSV: "embedchain.loaders.csv.CsvLoader",
+            DataType.MDX: "embedchain.loaders.mdx.MdxLoader",
+            DataType.IMAGES: "embedchain.loaders.images.ImagesLoader",
+            DataType.UNSTRUCTURED: "embedchain.loaders.unstructured_file.UnstructuredLoader",
+            DataType.JSON: "embedchain.loaders.json.JSONLoader",
+            DataType.OPENAPI: "embedchain.loaders.openapi.OpenAPILoader",
+            DataType.GMAIL: "embedchain.loaders.gmail.GmailLoader",
+            DataType.NOTION: "embedchain.loaders.notion.NotionLoader",
         }
-        lazy_loaders = {DataType.NOTION}
         if data_type in loaders:
-            loader_class: type = loaders[data_type]
-            loader: BaseLoader = loader_class()
-            return loader
-        elif data_type in lazy_loaders:
-            if data_type == DataType.NOTION:
-                from embedchain.loaders.notion import NotionLoader
-
-                return NotionLoader()
-            else:
-                raise ValueError(f"Unsupported data type: {data_type}")
+            loader_class: type = self._lazy_load(loaders[data_type])
+            return loader_class()
         else:
             raise ValueError(f"Unsupported data type: {data_type}")
 
     def _get_chunker(self, data_type: DataType, config: ChunkerConfig) -> BaseChunker:
-        """Returns the appropriate chunker for the given data type.
-
-        :param data_type: The type of the data to chunk.
-        :type data_type: DataType
-        :param config: Config to initialize the chunker with.
-        :type config: ChunkerConfig
-        :raises ValueError: If an unsupported data type is provided.
-        :return: The chunker for the given data type.
-        :rtype: BaseChunker
-        """
+        """Returns the appropriate chunker for the given data type (updated for lazy loading)."""
         chunker_classes = {
-            DataType.YOUTUBE_VIDEO: YoutubeVideoChunker,
-            DataType.PDF_FILE: PdfFileChunker,
-            DataType.WEB_PAGE: WebPageChunker,
-            DataType.QNA_PAIR: QnaPairChunker,
-            DataType.TEXT: TextChunker,
-            DataType.DOCX: DocxFileChunker,
-            DataType.DOCS_SITE: DocsSiteChunker,
-            DataType.SITEMAP: SitemapChunker,
-            DataType.NOTION: NotionChunker,
-            DataType.CSV: TableChunker,
-            DataType.MDX: MdxChunker,
-            DataType.IMAGES: ImagesChunker,
-            DataType.XML: XmlChunker,
-            DataType.UNSTRUCTURED: UnstructuredFileChunker,
-            DataType.JSON: JSONChunker,
+            DataType.YOUTUBE_VIDEO: "embedchain.chunkers.youtube_video.YoutubeVideoChunker",
+            DataType.PDF_FILE: "embedchain.chunkers.pdf_file.PdfFileChunker",
+            DataType.WEB_PAGE: "embedchain.chunkers.web_page.WebPageChunker",
+            DataType.QNA_PAIR: "embedchain.chunkers.qna_pair.QnaPairChunker",
+            DataType.TEXT: "embedchain.chunkers.text.TextChunker",
+            DataType.DOCX: "embedchain.chunkers.docx_file.DocxFileChunker",
+            DataType.SITEMAP: "embedchain.chunkers.sitemap.SitemapChunker",
+            DataType.XML: "embedchain.chunkers.xml.XmlChunker",
+            DataType.DOCS_SITE: "embedchain.chunkers.docs_site.DocsSiteChunker",
+            DataType.CSV: "embedchain.chunkers.table.TableChunker",
+            DataType.MDX: "embedchain.chunkers.mdx.MdxChunker",
+            DataType.IMAGES: "embedchain.chunkers.images.ImagesChunker",
+            DataType.UNSTRUCTURED: "embedchain.chunkers.unstructured_file.UnstructuredFileChunker",
+            DataType.JSON: "embedchain.chunkers.json.JSONChunker",
+            DataType.OPENAPI: "embedchain.chunkers.openapi.OpenAPIChunker",
+            DataType.GMAIL: "embedchain.chunkers.gmail.GmailChunker",
+            DataType.NOTION: "embedchain.chunkers.notion.NotionChunker",
         }
+
         if data_type in chunker_classes:
-            chunker_class: type = chunker_classes[data_type]
-            chunker: BaseChunker = chunker_class(config)
+            chunker_class = self._lazy_load(chunker_classes[data_type])
+            chunker = chunker_class(config)
             chunker.set_data_type(data_type)
             return chunker
         else:

embedchain/embedchain.py

@@ -1,17 +1,13 @@
 import hashlib
-import importlib.metadata
 import json
 import logging
 import os
-import threading
-import uuid
+import sqlite3
 from pathlib import Path
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Tuple, Union
 
-import requests
 from dotenv import load_dotenv
 from langchain.docstore.document import Document
-from tenacity import retry, stop_after_attempt, wait_fixed
 
 from embedchain.chunkers.base_chunker import BaseChunker
 from embedchain.config import AddConfig, BaseLlmConfig
@@ -23,6 +19,7 @@ from embedchain.llm.base import BaseLlm
 from embedchain.loaders.base_loader import BaseLoader
 from embedchain.models.data_type import (DataType, DirectDataType,
                                          IndirectDataType, SpecialDataType)
+from embedchain.telemetry.posthog import AnonymousTelemetry
 from embedchain.utils import detect_datatype
 from embedchain.vectordb.base import BaseVectorDB
 
@@ -32,6 +29,7 @@ ABS_PATH = os.getcwd()
 HOME_DIR = str(Path.home())
 CONFIG_DIR = os.path.join(HOME_DIR, ".embedchain")
 CONFIG_FILE = os.path.join(CONFIG_DIR, "config.json")
+SQLITE_PATH = os.path.join(CONFIG_DIR, "embedchain.db")
 
 
 class EmbedChain(JSONSerializable):
@@ -87,13 +85,29 @@ class EmbedChain(JSONSerializable):
         self.user_asks = []
 
         # Send anonymous telemetry
-        self.s_id = self.config.id if self.config.id else str(uuid.uuid4())
-        self.u_id = self._load_or_generate_user_id()
-        # NOTE: Uncomment the next two lines when running tests to see if any test fires a telemetry event.
-        # if (self.config.collect_metrics):
-        #     raise ConnectionRefusedError("Collection of metrics should not be allowed.")
-        thread_telemetry = threading.Thread(target=self._send_telemetry_event, args=("init",))
-        thread_telemetry.start()
+        self._telemetry_props = {"class": self.__class__.__name__}
+        self.telemetry = AnonymousTelemetry(enabled=self.config.collect_metrics)
+        # Establish a connection to the SQLite database
+        self.connection = sqlite3.connect(SQLITE_PATH)
+        self.cursor = self.connection.cursor()
+
+        # Create the 'data_sources' table if it doesn't exist
+        self.cursor.execute(
+            """
+            CREATE TABLE IF NOT EXISTS data_sources (
+                pipeline_id TEXT,
+                hash TEXT,
+                type TEXT,
+                value TEXT,
+                metadata TEXT,
+                is_uploaded INTEGER DEFAULT 0,
+                PRIMARY KEY (pipeline_id, hash)
+            )
+        """
+        )
+        self.connection.commit()
+        # Send anonymous telemetry
+        self.telemetry.capture(event_name="init", properties=self._telemetry_props)
 
     @property
     def collect_metrics(self):
@@ -115,29 +129,6 @@ class EmbedChain(JSONSerializable):
             raise ValueError(f"Boolean value expected but got {type(value)}.")
         self.llm.online = value
 
-    def _load_or_generate_user_id(self) -> str:
-        """
-        Loads the user id from the config file if it exists, otherwise generates a new
-        one and saves it to the config file.
-
-        :return: user id
-        :rtype: str
-        """
-        if not os.path.exists(CONFIG_DIR):
-            os.makedirs(CONFIG_DIR)
-
-        if os.path.exists(CONFIG_FILE):
-            with open(CONFIG_FILE, "r") as f:
-                data = json.load(f)
-                if "user_id" in data:
-                    return data["user_id"]
-
-        u_id = str(uuid.uuid4())
-        with open(CONFIG_FILE, "w") as f:
-            json.dump({"user_id": u_id}, f)
-
-        return u_id
-
     def add(
         self,
         source: Any,
@@ -163,7 +154,7 @@ class EmbedChain(JSONSerializable):
         :raises ValueError: Invalid data type
         :param dry_run: Optional. A dry run displays the chunks to ensure that the loader and chunker work as intended.
         deafaults to False
-        :return: source_id, a md5-hash of the source, in hexadecimal representation.
+        :return: source_hash, a md5-hash of the source, in hexadecimal representation.
         :rtype: str
         """
         if config is None:
@@ -192,18 +183,40 @@ class EmbedChain(JSONSerializable):
         if not data_type:
             data_type = detect_datatype(source)
 
-        # `source_id` is the hash of the source argument
+        # `source_hash` is the md5 hash of the source argument
         hash_object = hashlib.md5(str(source).encode("utf-8"))
-        source_id = hash_object.hexdigest()
+        source_hash = hash_object.hexdigest()
+
+        # Check if the data hash already exists, if so, skip the addition
+        self.cursor.execute(
+            "SELECT 1 FROM data_sources WHERE hash = ? AND pipeline_id = ?", (source_hash, self.config.id)
+        )
+        existing_data = self.cursor.fetchone()
+
+        if existing_data:
+            print(f"Data with hash {source_hash} already exists. Skipping addition.")
+            return source_hash
 
         data_formatter = DataFormatter(data_type, config)
         self.user_asks.append([source, data_type.value, metadata])
         documents, metadatas, _ids, new_chunks = self.load_and_embed(
-            data_formatter.loader, data_formatter.chunker, source, metadata, source_id, dry_run
+            data_formatter.loader, data_formatter.chunker, source, metadata, source_hash, dry_run
         )
         if data_type in {DataType.DOCS_SITE}:
             self.is_docs_site_instance = True
 
+        # Insert the data into the 'data' table
+        self.cursor.execute(
+            """
+            INSERT INTO data_sources (hash, pipeline_id, type, value, metadata)
+            VALUES (?, ?, ?, ?, ?)
+        """,
+            (source_hash, self.config.id, data_type.value, str(source), json.dumps(metadata)),
+        )
+
+        # Commit the transaction
+        self.connection.commit()
+
         if dry_run:
             data_chunks_info = {"chunks": documents, "metadata": metadatas, "count": len(documents), "type": data_type}
             logging.debug(f"Dry run info : {data_chunks_info}")
@@ -214,11 +227,16 @@ class EmbedChain(JSONSerializable):
             # it's quicker to check the variable twice than to count words when they won't be submitted.
             word_count = data_formatter.chunker.get_word_count(documents)
 
-            extra_metadata = {"data_type": data_type.value, "word_count": word_count, "chunks_count": new_chunks}
-            thread_telemetry = threading.Thread(target=self._send_telemetry_event, args=("add", extra_metadata))
-            thread_telemetry.start()
+            # Send anonymous telemetry
+            event_properties = {
+                **self._telemetry_props,
+                "data_type": data_type.value,
+                "word_count": word_count,
+                "chunks_count": new_chunks,
+            }
+            self.telemetry.capture(event_name="add", properties=event_properties)
 
-        return source_id
+        return source_hash
 
     def add_local(
         self,
@@ -245,7 +263,7 @@ class EmbedChain(JSONSerializable):
         :param config: The `AddConfig` instance to use as configuration options., defaults to None
         :type config: Optional[AddConfig], optional
         :raises ValueError: Invalid data type
-        :return: source_id, a md5-hash of the source, in hexadecimal representation.
+        :return: source_hash, a md5-hash of the source, in hexadecimal representation.
         :rtype: str
         """
         logging.warning(
@@ -313,7 +331,7 @@ class EmbedChain(JSONSerializable):
         chunker: BaseChunker,
         src: Any,
         metadata: Optional[Dict[str, Any]] = None,
-        source_id: Optional[str] = None,
+        source_hash: Optional[str] = None,
         dry_run=False,
     ):
         """
@@ -324,7 +342,7 @@ class EmbedChain(JSONSerializable):
         :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.
-        :param source_id: Hexadecimal hash of the source.
+        :param source_hash: Hexadecimal hash of the source.
         :param dry_run: Optional. A dry run returns chunks and doesn't update DB.
         :type dry_run: bool, defaults to False
         :return: (List) documents (embedded text), (List) metadata, (list) ids, (int) number of chunks
@@ -382,7 +400,7 @@ class EmbedChain(JSONSerializable):
                 m["app_id"] = self.config.id
 
             # Add hashed source
-            m["hash"] = source_id
+            m["hash"] = source_hash
 
             # Note: Metadata is the function argument
             if metadata:
@@ -420,7 +438,9 @@ class EmbedChain(JSONSerializable):
             )
         ]
 
-    def retrieve_from_database(self, input_query: str, config: Optional[BaseLlmConfig] = None, where=None) -> List[str]:
+    def retrieve_from_database(
+        self, input_query: str, config: Optional[BaseLlmConfig] = None, where=None, citations: bool = False
+    ) -> Union[List[Tuple[str, str, str]], List[str]]:
         """
         Queries the vector database based on the given input query.
         Gets relevant doc based on the query
@@ -431,6 +451,8 @@ class EmbedChain(JSONSerializable):
         :type config: Optional[BaseLlmConfig], optional
         :param where: A dictionary of key-value pairs to filter the database results, defaults to None
         :type where: _type_, optional
+        :param citations: A boolean to indicate if db should fetch citation source
+        :type citations: bool
         :return: List of contents of the document that matched your query
         :rtype: List[str]
         """
@@ -455,15 +477,24 @@ class EmbedChain(JSONSerializable):
 
             db_query = ClipProcessor.get_text_features(query=input_query)
 
-        contents = self.db.query(
+        contexts = self.db.query(
             input_query=db_query,
             n_results=query_config.number_documents,
             where=where,
             skip_embedding=(hasattr(config, "query_type") and config.query_type == "Images"),
+            citations=citations,
         )
-        return contents
 
-    def query(self, input_query: str, config: BaseLlmConfig = None, dry_run=False, where: Optional[Dict] = None) -> str:
+        return contexts
+
+    def query(
+        self,
+        input_query: str,
+        config: BaseLlmConfig = None,
+        dry_run=False,
+        where: Optional[Dict] = None,
+        **kwargs: Dict[str, Any],
+    ) -> Union[Tuple[str, List[Tuple[str, str, str]]], str]:
         """
         Queries the vector database based on the given input query.
         Gets relevant doc based on the query and then passes it to an
@@ -479,17 +510,31 @@ class EmbedChain(JSONSerializable):
         :type dry_run: bool, optional
         :param where: A dictionary of key-value pairs to filter the database results., defaults to None
         :type where: Optional[Dict[str, str]], optional
-        :return: The answer to the query or the dry run result
-        :rtype: str
+        :param kwargs: To read more params for the query function. Ex. we use citations boolean
+        param to return context along with the answer
+        :type kwargs: Dict[str, Any]
+        :return: The answer to the query, with citations if the citation flag is True
+        or the dry run result
+        :rtype: str, if citations is False, otherwise Tuple[str,List[Tuple[str,str,str]]]
         """
-        contexts = self.retrieve_from_database(input_query=input_query, config=config, where=where)
-        answer = self.llm.query(input_query=input_query, contexts=contexts, config=config, dry_run=dry_run)
+        citations = kwargs.get("citations", False)
+        contexts = self.retrieve_from_database(input_query=input_query, config=config, where=where, citations=citations)
+        if citations and len(contexts) > 0 and isinstance(contexts[0], tuple):
+            contexts_data_for_llm_query = list(map(lambda x: x[0], contexts))
+        else:
+            contexts_data_for_llm_query = contexts
+
+        answer = self.llm.query(
+            input_query=input_query, contexts=contexts_data_for_llm_query, config=config, dry_run=dry_run
+        )
 
         # Send anonymous telemetry
-        thread_telemetry = threading.Thread(target=self._send_telemetry_event, args=("query",))
-        thread_telemetry.start()
+        self.telemetry.capture(event_name="query", properties=self._telemetry_props)
 
-        return answer
+        if citations:
+            return answer, contexts
+        else:
+            return answer
 
     def chat(
         self,
@@ -497,6 +542,7 @@ class EmbedChain(JSONSerializable):
         config: Optional[BaseLlmConfig] = None,
         dry_run=False,
         where: Optional[Dict[str, str]] = None,
+        **kwargs: Dict[str, Any],
     ) -> str:
         """
         Queries the vector database on the given input query.
@@ -515,17 +561,31 @@ class EmbedChain(JSONSerializable):
         :type dry_run: bool, optional
         :param where: A dictionary of key-value pairs to filter the database results., defaults to None
         :type where: Optional[Dict[str, str]], optional
-        :return: The answer to the query or the dry run result
-        :rtype: str
+        :param kwargs: To read more params for the query function. Ex. we use citations boolean
+        param to return context along with the answer
+        :type kwargs: Dict[str, Any]
+        :return: The answer to the query, with citations if the citation flag is True
+        or the dry run result
+        :rtype: str, if citations is False, otherwise Tuple[str,List[Tuple[str,str,str]]]
         """
-        contexts = self.retrieve_from_database(input_query=input_query, config=config, where=where)
-        answer = self.llm.chat(input_query=input_query, contexts=contexts, config=config, dry_run=dry_run)
+        citations = kwargs.get("citations", False)
+        contexts = self.retrieve_from_database(input_query=input_query, config=config, where=where, citations=citations)
+        if citations and len(contexts) > 0 and isinstance(contexts[0], tuple):
+            contexts_data_for_llm_query = list(map(lambda x: x[0], contexts))
+        else:
+            contexts_data_for_llm_query = contexts
+
+        answer = self.llm.chat(
+            input_query=input_query, contexts=contexts_data_for_llm_query, config=config, dry_run=dry_run
+        )
 
         # Send anonymous telemetry
-        thread_telemetry = threading.Thread(target=self._send_telemetry_event, args=("chat",))
-        thread_telemetry.start()
+        self.telemetry.capture(event_name="chat", properties=self._telemetry_props)
 
-        return answer
+        if citations:
+            return answer, contexts
+        else:
+            return answer
 
     def set_collection_name(self, name: str):
         """
@@ -558,36 +618,9 @@ class EmbedChain(JSONSerializable):
         """
         Resets the database. Deletes all embeddings irreversibly.
         `App` does not have to be reinitialized after using this method.
-
-        DEPRECATED IN FAVOR OF `db.reset()`
         """
-        # Send anonymous telemetry
-        thread_telemetry = threading.Thread(target=self._send_telemetry_event, args=("reset",))
-        thread_telemetry.start()
-
-        logging.warning("DEPRECATION WARNING: Please use `app.db.reset()` instead of `App.reset()`.")
         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):
-        """
-        Send telemetry event to the embedchain server. This is anonymous. It can be toggled off in `AppConfig`.
-        """
-        if not self.config.collect_metrics:
-            return
-
-        with threading.Lock():
-            url = "https://api.embedchain.ai/api/v1/telemetry/"
-            metadata = {
-                "s_id": self.s_id,
-                "version": importlib.metadata.version(__package__ or __name__),
-                "method": method,
-                "language": "py",
-                "u_id": self.u_id,
-            }
-            if extra_metadata:
-                metadata.update(extra_metadata)
-
-            response = requests.post(url, json={"metadata": metadata})
-            if response.status_code != 200:
-                logging.warning(f"Telemetry event failed with status code {response.status_code}")
+        self.cursor.execute("DELETE FROM data_sources WHERE pipeline_id = ?", (self.config.id,))
+        self.connection.commit()
+        # Send anonymous telemetry
+        self.telemetry.capture(event_name="reset", properties=self._telemetry_props)

embedchain/embedder/gpt4all.py

@@ -1,7 +1,5 @@
 from typing import Optional
 
-from chromadb.utils import embedding_functions
-
 from embedchain.config import BaseEmbedderConfig
 from embedchain.embedder.base import BaseEmbedder
 from embedchain.models import VectorDimensions
@@ -9,12 +7,13 @@ from embedchain.models import VectorDimensions
 
 class GPT4AllEmbedder(BaseEmbedder):
     def __init__(self, config: Optional[BaseEmbedderConfig] = None):
-        # Note: We could use langchains GPT4ALL embedding, but it's not available in all versions.
         super().__init__(config=config)
-        if self.config.model is None:
-            self.config.model = "all-MiniLM-L6-v2"
 
-        embedding_fn = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=self.config.model)
+        from langchain.embeddings import \
+            GPT4AllEmbeddings as LangchainGPT4AllEmbeddings
+
+        embeddings = LangchainGPT4AllEmbeddings()
+        embedding_fn = BaseEmbedder._langchain_default_concept(embeddings)
         self.set_embedding_fn(embedding_fn=embedding_fn)
 
         vector_dimension = VectorDimensions.GPT4ALL.value

embedchain/factory.py

@@ -41,15 +41,16 @@ class LlmFactory:
 
 class EmbedderFactory:
     provider_to_class = {
+        "azure_openai": "embedchain.embedder.openai.OpenAIEmbedder",
         "gpt4all": "embedchain.embedder.gpt4all.GPT4AllEmbedder",
         "huggingface": "embedchain.embedder.huggingface.HuggingFaceEmbedder",
-        "vertexai": "embedchain.embedder.vertexai.VertexAIEmbedder",
-        "azure_openai": "embedchain.embedder.openai.OpenAIEmbedder",
         "openai": "embedchain.embedder.openai.OpenAIEmbedder",
+        "vertexai": "embedchain.embedder.vertexai.VertexAIEmbedder",
     }
     provider_to_config_class = {
-        "openai": "embedchain.config.embedder.base.BaseEmbedderConfig",
         "azure_openai": "embedchain.config.embedder.base.BaseEmbedderConfig",
+        "openai": "embedchain.config.embedder.base.BaseEmbedderConfig",
+        "gpt4all": "embedchain.config.embedder.base.BaseEmbedderConfig",
     }
 
     @classmethod
@@ -72,16 +73,18 @@ class VectorDBFactory:
         "elasticsearch": "embedchain.vectordb.elasticsearch.ElasticsearchDB",
         "opensearch": "embedchain.vectordb.opensearch.OpenSearchDB",
         "pinecone": "embedchain.vectordb.pinecone.PineconeDB",
-        "weaviate": "embedchain.vectordb.weaviate.WeaviateDB",
         "qdrant": "embedchain.vectordb.qdrant.QdrantDB",
+        "weaviate": "embedchain.vectordb.weaviate.WeaviateDB",
+        "zilliz": "embedchain.vectordb.zilliz.ZillizVectorDB",
     }
     provider_to_config_class = {
         "chroma": "embedchain.config.vectordb.chroma.ChromaDbConfig",
         "elasticsearch": "embedchain.config.vectordb.elasticsearch.ElasticsearchDBConfig",
         "opensearch": "embedchain.config.vectordb.opensearch.OpenSearchDBConfig",
         "pinecone": "embedchain.config.vectordb.pinecone.PineconeDBConfig",
-        "weaviate": "embedchain.config.vectordb.weaviate.WeaviateDBConfig",
         "qdrant": "embedchain.config.vectordb.qdrant.QdrantDBConfig",
+        "weaviate": "embedchain.config.vectordb.weaviate.WeaviateDBConfig",
+        "zilliz": "embedchain.config.vectordb.zilliz.ZillizDBConfig",
     }
 
     @classmethod

embedchain/llm/gpt4all.py

@@ -1,5 +1,8 @@
 from typing import Iterable, Optional, Union
 
+from langchain.callbacks.stdout import StdOutCallbackHandler
+from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
+
 from embedchain.config import BaseLlmConfig
 from embedchain.helper.json_serializable import register_deserializable
 from embedchain.llm.base import BaseLlm
@@ -12,6 +15,7 @@ class GPT4ALLLlm(BaseLlm):
         if self.config.model is None:
             self.config.model = "orca-mini-3b.ggmlv3.q4_0.bin"
         self.instance = GPT4ALLLlm._get_instance(self.config.model)
+        self.instance.streaming = self.config.stream
 
     def get_llm_model_answer(self, prompt):
         return self._get_answer(prompt=prompt, config=self.config)
@@ -19,13 +23,13 @@ class GPT4ALLLlm(BaseLlm):
     @staticmethod
     def _get_instance(model):
         try:
-            from gpt4all import GPT4All
+            from langchain.llms.gpt4all import GPT4All as LangchainGPT4All
         except ModuleNotFoundError:
             raise ModuleNotFoundError(
                 "The GPT4All python package is not installed. Please install it with `pip install --upgrade embedchain[opensource]`"  # noqa E501
             ) from None
 
-        return GPT4All(model_name=model)
+        return LangchainGPT4All(model=model, allow_download=True)
 
     def _get_answer(self, prompt: str, config: BaseLlmConfig) -> Union[str, Iterable]:
         if config.model and config.model != self.config.model:
@@ -33,14 +37,21 @@ class GPT4ALLLlm(BaseLlm):
                 "GPT4ALLLlm does not support switching models at runtime. Please create a new app instance."
             )
 
+        messages = []
         if config.system_prompt:
-            raise ValueError("GPT4ALLLlm does not support `system_prompt`")
+            messages.append(config.system_prompt)
+        messages.append(prompt)
+        kwargs = {
+            "temp": config.temperature,
+            "max_tokens": config.max_tokens,
+        }
+        if config.top_p:
+            kwargs["top_p"] = config.top_p
 
-        response = self.instance.generate(
-            prompt=prompt,
-            streaming=config.stream,
-            top_p=config.top_p,
-            max_tokens=config.max_tokens,
-            temp=config.temperature,
-        )
-        return response
+        callbacks = [StreamingStdOutCallbackHandler()] if config.stream else [StdOutCallbackHandler()]
+
+        response = self.instance.generate(prompts=messages, callbacks=callbacks, **kwargs)
+        answer = ""
+        for generations in response.generations:
+            answer += " ".join(map(lambda generation: generation.text, generations))
+        return answer

embedchain/llm/llama2.py

@@ -8,18 +8,17 @@ from embedchain.config import BaseLlmConfig
 from embedchain.helper.json_serializable import register_deserializable
 from embedchain.llm.base import BaseLlm
 
-try:
-    importlib.import_module("replicate")
-except ModuleNotFoundError:
-    raise ModuleNotFoundError(
-        "The required dependencies for Llama2 are not installed."
-        'Please install with `pip install --upgrade "embedchain[llama2]"`'
-    ) from None
-
 
 @register_deserializable
 class Llama2Llm(BaseLlm):
     def __init__(self, config: Optional[BaseLlmConfig] = None):
+        try:
+            importlib.import_module("replicate")
+        except ModuleNotFoundError:
+            raise ModuleNotFoundError(
+                "The required dependencies for Llama2 are not installed."
+                'Please install with `pip install --upgrade "embedchain[llama2]"`'
+            ) from None
         if "REPLICATE_API_TOKEN" not in os.environ:
             raise ValueError("Please set the REPLICATE_API_TOKEN environment variable.")
 

embedchain/llm/vertex_ai.py

@@ -6,18 +6,17 @@ from embedchain.config import BaseLlmConfig
 from embedchain.helper.json_serializable import register_deserializable
 from embedchain.llm.base import BaseLlm
 
-try:
-    importlib.import_module("vertexai")
-except ModuleNotFoundError:
-    raise ModuleNotFoundError(
-        "The required dependencies for VertexAI are not installed."
-        'Please install with `pip install --upgrade "embedchain[vertexai]"`'
-    ) from None
-
 
 @register_deserializable
 class VertexAILlm(BaseLlm):
     def __init__(self, config: Optional[BaseLlmConfig] = None):
+        try:
+            importlib.import_module("vertexai")
+        except ModuleNotFoundError:
+            raise ModuleNotFoundError(
+                "The required dependencies for VertexAI are not installed."
+                'Please install with `pip install --upgrade "embedchain[vertexai]"`'
+            ) from None
         super().__init__(config=config)
 
     def get_llm_model_answer(self, prompt):

embedchain/loaders/gmail.py

@@ -0,0 +1,123 @@
+import hashlib
+import logging
+import os
+import quopri
+from textwrap import dedent
+
+from bs4 import BeautifulSoup
+
+try:
+    from llama_hub.gmail.base import GmailReader
+except ImportError:
+    raise ImportError("Gmail requires extra dependencies. Install with `pip install embedchain[gmail]`") from None
+
+from embedchain.loaders.base_loader import BaseLoader
+from embedchain.utils import clean_string
+
+
+def get_header(text: str, header: str) -> str:
+    start_string_position = text.find(header)
+    pos_start = text.find(":", start_string_position) + 1
+    pos_end = text.find("\n", pos_start)
+    header = text[pos_start:pos_end]
+    return header.strip()
+
+
+class GmailLoader(BaseLoader):
+    def load_data(self, query):
+        """Load data from gmail."""
+        if not os.path.isfile("credentials.json"):
+            raise FileNotFoundError(
+                "You must download the valid credentials file from your google \
+                dev account. Refer this `https://cloud.google.com/docs/authentication/api-keys`"
+            )
+
+        loader = GmailReader(query=query, service=None, results_per_page=20)
+        documents = loader.load_data()
+        logging.info(f"Gmail Loader: {len(documents)} mails found for query- {query}")
+
+        data = []
+        data_contents = []
+        logging.info(f"Gmail Loader: {len(documents)} mails found")
+        for document in documents:
+            original_size = len(document.text)
+
+            snippet = document.metadata.get("snippet")
+            meta_data = {
+                "url": document.metadata.get("id"),
+                "date": get_header(document.text, "Date"),
+                "subject": get_header(document.text, "Subject"),
+                "from": get_header(document.text, "From"),
+                "to": get_header(document.text, "To"),
+                "search_query": query,
+            }
+
+            # Decode
+            decoded_bytes = quopri.decodestring(document.text)
+            decoded_str = decoded_bytes.decode("utf-8", errors="replace")
+
+            # Slice
+            mail_start = decoded_str.find("<!DOCTYPE")
+            email_data = decoded_str[mail_start:]
+
+            # Web Page HTML Processing
+            soup = BeautifulSoup(email_data, "html.parser")
+
+            tags_to_exclude = [
+                "nav",
+                "aside",
+                "form",
+                "header",
+                "noscript",
+                "svg",
+                "canvas",
+                "footer",
+                "script",
+                "style",
+            ]
+
+            for tag in soup(tags_to_exclude):
+                tag.decompose()
+
+            ids_to_exclude = ["sidebar", "main-navigation", "menu-main-menu"]
+            for id in ids_to_exclude:
+                tags = soup.find_all(id=id)
+                for tag in tags:
+                    tag.decompose()
+
+            classes_to_exclude = [
+                "elementor-location-header",
+                "navbar-header",
+                "nav",
+                "header-sidebar-wrapper",
+                "blog-sidebar-wrapper",
+                "related-posts",
+            ]
+
+            for class_name in classes_to_exclude:
+                tags = soup.find_all(class_=class_name)
+                for tag in tags:
+                    tag.decompose()
+
+            content = soup.get_text()
+            content = clean_string(content)
+
+            cleaned_size = len(content)
+            if original_size != 0:
+                logging.info(
+                    f"[{id}] Cleaned page size: {cleaned_size} characters, down from {original_size} (shrunk: {original_size-cleaned_size} chars, {round((1-(cleaned_size/original_size)) * 100, 2)}%)"  # noqa:E501
+                )
+
+            result = f"""
+            email from '{meta_data.get('from')}' to '{meta_data.get('to')}'
+            subject: {meta_data.get('subject')}
+            date: {meta_data.get('date')}
+            preview: {snippet}
+            content: f{content}
+            """
+            data_content = dedent(result)
+            data.append({"content": data_content, "meta_data": meta_data})
+            data_contents.append(data_content)
+        doc_id = hashlib.sha256((query + ", ".join(data_contents)).encode()).hexdigest()
+        response_data = {"doc_id": doc_id, "data": data}
+        return response_data

embedchain/loaders/json.py

@@ -1,24 +1,57 @@
 import hashlib
+import json
+import os
+import re
 
-from langchain.document_loaders.json_loader import \
-    JSONLoader as LangchainJSONLoader
+import requests
 
 from embedchain.loaders.base_loader import BaseLoader
+from embedchain.utils import clean_string
 
-langchain_json_jq_schema = 'to_entries | map("\(.key): \(.value|tostring)") | .[]'
+VALID_URL_PATTERN = "^https:\/\/[0-9A-z.]+.[0-9A-z.]+.[a-z]+\/.*\.json$"
 
 
 class JSONLoader(BaseLoader):
     @staticmethod
     def load_data(content):
         """Load a json file. Each data point is a key value pair."""
+        try:
+            from llama_hub.jsondata.base import \
+                JSONDataReader as LLHBUBJSONLoader
+        except ImportError:
+            raise Exception(
+                f"Couldn't import the required packages to load {content}, \
+                Do `pip install --upgrade 'embedchain[json]`"
+            )
+
+        loader = LLHBUBJSONLoader()
+
+        if not isinstance(content, str):
+            print(f"Invaid content input. Provide the correct path to the json file saved locally in {content}")
+
         data = []
         data_content = []
-        loader = LangchainJSONLoader(content, text_content=False, jq_schema=langchain_json_jq_schema)
-        docs = loader.load()
+
+        # Load json data from various sources. TODO: add support for dictionary
+        if os.path.isfile(content):
+            with open(content, "r", encoding="utf-8") as json_file:
+                json_data = json.load(json_file)
+        elif re.match(VALID_URL_PATTERN, content):
+            response = requests.get(content)
+            if response.status_code == 200:
+                json_data = response.json()
+            else:
+                raise ValueError(
+                    f"Loading data from the given url: {content} failed. \
+                    Make sure the url is working."
+                )
+        else:
+            raise ValueError(f"Invalid content to load json data from: {content}")
+
+        docs = loader.load_data(json_data)
         for doc in docs:
-            meta_data = doc.metadata
-            data.append({"content": doc.page_content, "meta_data": {"url": content, "row": meta_data["seq_num"]}})
-            data_content.append(doc.page_content)
+            doc_content = clean_string(doc.text)
+            data.append({"content": doc_content, "meta_data": {"url": content}})
+            data_content.append(doc_content)
         doc_id = hashlib.sha256((content + ", ".join(data_content)).encode()).hexdigest()
         return {"doc_id": doc_id, "data": data}

embedchain/loaders/openapi.py

@@ -0,0 +1,42 @@
+import hashlib
+from io import StringIO
+from urllib.parse import urlparse
+
+import requests
+import yaml
+
+from embedchain.loaders.base_loader import BaseLoader
+
+
+class OpenAPILoader(BaseLoader):
+    @staticmethod
+    def _get_file_content(content):
+        url = urlparse(content)
+        if all([url.scheme, url.netloc]) and url.scheme not in ["file", "http", "https"]:
+            raise ValueError("Not a valid URL.")
+
+        if url.scheme in ["http", "https"]:
+            response = requests.get(content)
+            response.raise_for_status()
+            return StringIO(response.text)
+        elif url.scheme == "file":
+            path = url.path
+            return open(path)
+        else:
+            return open(content)
+
+    @staticmethod
+    def load_data(content):
+        """Load yaml file of openapi. Each pair is a document."""
+        data = []
+        file_path = content
+        data_content = []
+        with OpenAPILoader._get_file_content(content=content) as file:
+            yaml_data = yaml.load(file, Loader=yaml.Loader)
+            for i, (key, value) in enumerate(yaml_data.items()):
+                string_data = f"{key}: {value}"
+                meta_data = {"url": file_path, "row": i + 1}
+                data.append({"content": string_data, "meta_data": meta_data})
+                data_content.append(string_data)
+        doc_id = hashlib.sha256((content + ", ".join(data_content)).encode()).hexdigest()
+        return {"doc_id": doc_id, "data": data}

embedchain/models/data_type.py

@@ -27,6 +27,8 @@ class IndirectDataType(Enum):
     IMAGES = "images"
     UNSTRUCTURED = "unstructured"
     JSON = "json"
+    OPENAPI = "openapi"
+    GMAIL = "gmail"
 
 
 class SpecialDataType(Enum):
@@ -53,3 +55,5 @@ class DataType(Enum):
     IMAGES = IndirectDataType.IMAGES.value
     UNSTRUCTURED = IndirectDataType.UNSTRUCTURED.value
     JSON = IndirectDataType.JSON.value
+    OPENAPI = IndirectDataType.OPENAPI.value
+    GMAIL = IndirectDataType.GMAIL.value

embedchain/pipeline.py

@@ -1,17 +1,29 @@
-import threading
+import ast
+import json
+import logging
+import os
+import sqlite3
 import uuid
 
+import requests
 import yaml
+from fastapi import FastAPI, HTTPException
 
+from embedchain import Client
 from embedchain.config import PipelineConfig
-from embedchain.embedchain import EmbedChain
+from embedchain.embedchain import CONFIG_DIR, EmbedChain
 from embedchain.embedder.base import BaseEmbedder
 from embedchain.embedder.openai import OpenAIEmbedder
-from embedchain.factory import EmbedderFactory, VectorDBFactory
+from embedchain.factory import EmbedderFactory, LlmFactory, VectorDBFactory
 from embedchain.helper.json_serializable import register_deserializable
+from embedchain.llm.base import BaseLlm
+from embedchain.llm.openai import OpenAILlm
+from embedchain.telemetry.posthog import AnonymousTelemetry
 from embedchain.vectordb.base import BaseVectorDB
 from embedchain.vectordb.chroma import ChromaDB
 
+SQLITE_PATH = os.path.join(CONFIG_DIR, "embedchain.db")
+
 
 @register_deserializable
 class Pipeline(EmbedChain):
@@ -21,7 +33,18 @@ class Pipeline(EmbedChain):
     and vector database.
     """
 
-    def __init__(self, config: PipelineConfig = None, db: BaseVectorDB = None, embedding_model: BaseEmbedder = None):
+    def __init__(
+        self,
+        id: str = None,
+        name: str = None,
+        config: PipelineConfig = None,
+        db: BaseVectorDB = None,
+        embedding_model: BaseEmbedder = None,
+        llm: BaseLlm = None,
+        yaml_path: str = None,
+        log_level=logging.INFO,
+        auto_deploy: bool = False,
+    ):
         """
         Initialize a new `App` instance.
 
@@ -31,68 +54,360 @@ class Pipeline(EmbedChain):
         :type db: BaseVectorDB, optional
         :param embedding_model: The embedding model used to calculate embeddings, defaults to None
         :type embedding_model: BaseEmbedder, optional
+        :param llm: The LLM model used to calculate embeddings, defaults to None
+        :type llm: BaseLlm, optional
+        :param yaml_path: Path to the YAML configuration file, defaults to None
+        :type yaml_path: str, optional
+        :param log_level: Log level to use, defaults to logging.INFO
+        :type log_level: int, optional
+        :param auto_deploy: Whether to deploy the pipeline automatically, defaults to False
+        :type auto_deploy: bool, optional
+        :raises Exception: If an error occurs while creating the pipeline
         """
-        super().__init__()
+        if id and yaml_path:
+            raise Exception("Cannot provide both id and config. Please provide only one of them.")
+
+        if id and name:
+            raise Exception("Cannot provide both id and name. Please provide only one of them.")
+
+        if name and config:
+            raise Exception("Cannot provide both name and config. Please provide only one of them.")
+
+        logging.basicConfig(level=log_level, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+        self.logger = logging.getLogger(__name__)
+
+        self.auto_deploy = auto_deploy
+
+        # Store the yaml config as an attribute to be able to send it
+        self.yaml_config = None
+        self.client = None
+        # pipeline_id from the backend
+        self.id = None
+
         self.config = config or PipelineConfig()
         self.name = self.config.name
-        self.id = self.config.id or str(uuid.uuid4())
+
+        self.config.id = self.local_id = str(uuid.uuid4()) if self.config.id is None else self.config.id
+
+        if yaml_path:
+            with open(yaml_path, "r") as file:
+                config_data = yaml.safe_load(file)
+                self.yaml_config = config_data
+
+        if id is not None:
+            # Init client first since user is trying to fetch the pipeline
+            # details from the platform
+            self._init_client()
+            pipeline_details = self._get_pipeline(id)
+            self.config.id = self.local_id = pipeline_details["metadata"]["local_id"]
+            self.id = id
+
+        if name is not None:
+            self.name = name
 
         self.embedding_model = embedding_model or OpenAIEmbedder()
         self.db = db or ChromaDB()
-        self._initialize_db()
+        self.llm = llm or OpenAILlm()
+        self._init_db()
 
-        self.user_asks = []  # legacy defaults
+        # Send anonymous telemetry
+        self._telemetry_props = {"class": self.__class__.__name__}
+        self.telemetry = AnonymousTelemetry(enabled=self.config.collect_metrics)
 
-        self.s_id = self.config.id or str(uuid.uuid4())
-        self.u_id = self._load_or_generate_user_id()
+        # Establish a connection to the SQLite database
+        self.connection = sqlite3.connect(SQLITE_PATH)
+        self.cursor = self.connection.cursor()
 
-        thread_telemetry = threading.Thread(target=self._send_telemetry_event, args=("pipeline_init",))
-        thread_telemetry.start()
+        # Create the 'data_sources' table if it doesn't exist
+        self.cursor.execute(
+            """
+            CREATE TABLE IF NOT EXISTS data_sources (
+                pipeline_id TEXT,
+                hash TEXT,
+                type TEXT,
+                value TEXT,
+                metadata TEXT,
+                is_uploaded INTEGER DEFAULT 0,
+                PRIMARY KEY (pipeline_id, hash)
+            )
+        """
+        )
+        self.connection.commit()
+        # Send anonymous telemetry
+        self.telemetry.capture(event_name="init", properties=self._telemetry_props)
 
-    def _initialize_db(self):
+        self.user_asks = []
+        if self.auto_deploy:
+            self.deploy()
+
+    def _init_db(self):
         """
         Initialize the database.
         """
         self.db._set_embedder(self.embedding_model)
         self.db._initialize()
-        self.db.set_collection_name(self.name)
+        self.db.set_collection_name(self.db.config.collection_name)
+
+    def _init_client(self):
+        """
+        Initialize the client.
+        """
+        config = Client.load_config()
+        if config.get("api_key"):
+            self.client = Client()
+        else:
+            api_key = input(
+                "🔑 Enter your Embedchain API key. You can find the API key at https://app.embedchain.ai/settings/keys/ \n"  # noqa: E501
+            )
+            self.client = Client(api_key=api_key)
+
+    def _get_pipeline(self, id):
+        """
+        Get existing pipeline
+        """
+        print("🛠️ Fetching pipeline details from the platform...")
+        url = f"{self.client.host}/api/v1/pipelines/{id}/cli/"
+        r = requests.get(
+            url,
+            headers={"Authorization": f"Token {self.client.api_key}"},
+        )
+        if r.status_code == 404:
+            raise Exception(f"❌ Pipeline with id {id} not found!")
+
+        print(
+            f"🎉 Pipeline loaded successfully! Pipeline url: https://app.embedchain.ai/pipelines/{r.json()['id']}\n"  # noqa: E501
+        )
+        return r.json()
+
+    def _create_pipeline(self):
+        """
+        Create a pipeline on the platform.
+        """
+        print("🛠️ Creating pipeline on the platform...")
+        # self.yaml_config is a dict. Pass it inside the key 'yaml_config' to the backend
+        payload = {
+            "yaml_config": json.dumps(self.yaml_config),
+            "name": self.name,
+            "local_id": self.local_id,
+        }
+        url = f"{self.client.host}/api/v1/pipelines/cli/create/"
+        r = requests.post(
+            url,
+            json=payload,
+            headers={"Authorization": f"Token {self.client.api_key}"},
+        )
+        if r.status_code not in [200, 201]:
+            raise Exception(f"❌ Error occurred while creating pipeline. API response: {r.text}")
+
+        if r.status_code == 200:
+            print(
+                f"🎉🎉🎉 Existing pipeline found! View your pipeline: https://app.embedchain.ai/pipelines/{r.json()['id']}\n"  # noqa: E501
+            )  # noqa: E501
+        elif r.status_code == 201:
+            print(
+                f"🎉🎉🎉 Pipeline created successfully! View your pipeline: https://app.embedchain.ai/pipelines/{r.json()['id']}\n"  # noqa: E501
+            )
+        return r.json()
+
+    def _get_presigned_url(self, data_type, data_value):
+        payload = {"data_type": data_type, "data_value": data_value}
+        r = requests.post(
+            f"{self.client.host}/api/v1/pipelines/{self.id}/cli/presigned_url/",
+            json=payload,
+            headers={"Authorization": f"Token {self.client.api_key}"},
+        )
+        r.raise_for_status()
+        return r.json()
 
     def search(self, query, num_documents=3):
         """
         Search for similar documents related to the query in the vector database.
         """
-        where = {"app_id": self.id}
-        return self.db.query(
-            query,
-            n_results=num_documents,
-            where=where,
-            skip_embedding=False,
+        # Send anonymous telemetry
+        self.telemetry.capture(event_name="search", properties=self._telemetry_props)
+
+        # TODO: Search will call the endpoint rather than fetching the data from the db itself when deploy=True.
+        if self.id is None:
+            where = {"app_id": self.local_id}
+            context = self.db.query(
+                query,
+                n_results=num_documents,
+                where=where,
+                skip_embedding=False,
+                citations=True,
+            )
+            result = []
+            for c in context:
+                result.append(
+                    {
+                        "context": c[0],
+                        "source": c[1],
+                        "document_id": c[2],
+                    }
+                )
+            return result
+        else:
+            # Make API call to the backend to get the results
+            NotImplementedError("Search is not implemented yet for the prod mode.")
+
+    def _upload_file_to_presigned_url(self, presigned_url, file_path):
+        try:
+            with open(file_path, "rb") as file:
+                response = requests.put(presigned_url, data=file)
+                response.raise_for_status()
+                return response.status_code == 200
+        except Exception as e:
+            self.logger.exception(f"Error occurred during file upload: {str(e)}")
+            print("❌ Error occurred during file upload!")
+            return False
+
+    def _upload_data_to_pipeline(self, data_type, data_value, metadata=None):
+        payload = {
+            "data_type": data_type,
+            "data_value": data_value,
+            "metadata": metadata,
+        }
+        try:
+            self._send_api_request(f"/api/v1/pipelines/{self.id}/cli/add/", payload)
+            # print the local file path if user tries to upload a local file
+            printed_value = metadata.get("file_path") if metadata.get("file_path") else data_value
+            print(f"✅ Data of type: {data_type}, value: {printed_value} added successfully.")
+        except Exception as e:
+            print(f"❌ Error occurred during data upload for type {data_type}!. Error: {str(e)}")
+
+    def _send_api_request(self, endpoint, payload):
+        url = f"{self.client.host}{endpoint}"
+        headers = {"Authorization": f"Token {self.client.api_key}"}
+        response = requests.post(url, json=payload, headers=headers)
+        response.raise_for_status()
+        return response
+
+    def _process_and_upload_data(self, data_hash, data_type, data_value):
+        if os.path.isabs(data_value):
+            presigned_url_data = self._get_presigned_url(data_type, data_value)
+            presigned_url = presigned_url_data["presigned_url"]
+            s3_key = presigned_url_data["s3_key"]
+            if self._upload_file_to_presigned_url(presigned_url, file_path=data_value):
+                metadata = {"file_path": data_value, "s3_key": s3_key}
+                data_value = presigned_url
+            else:
+                self.logger.error(f"File upload failed for hash: {data_hash}")
+                return False
+        else:
+            if data_type == "qna_pair":
+                data_value = list(ast.literal_eval(data_value))
+            metadata = {}
+
+        try:
+            self._upload_data_to_pipeline(data_type, data_value, metadata)
+            self._mark_data_as_uploaded(data_hash)
+            return True
+        except Exception:
+            print(f"❌ Error occurred during data upload for hash {data_hash}!")
+            return False
+
+    def _mark_data_as_uploaded(self, data_hash):
+        self.cursor.execute(
+            "UPDATE data_sources SET is_uploaded = 1 WHERE hash = ? AND pipeline_id = ?",
+            (data_hash, self.local_id),
         )
+        self.connection.commit()
+
+    def deploy(self):
+        if self.client is None:
+            self._init_client()
+
+        pipeline_data = self._create_pipeline()
+        self.id = pipeline_data["id"]
+
+        results = self.cursor.execute(
+            "SELECT * FROM data_sources WHERE pipeline_id = ? AND is_uploaded = 0", (self.local_id,)  # noqa:E501
+        ).fetchall()
+
+        if len(results) > 0:
+            print("🛠️ Adding data to your pipeline...")
+        for result in results:
+            data_hash, data_type, data_value = result[1], result[2], result[3]
+            self._process_and_upload_data(data_hash, data_type, data_value)
+
+        # Send anonymous telemetry
+        self.telemetry.capture(event_name="deploy", properties=self._telemetry_props)
 
     @classmethod
-    def from_config(cls, yaml_path: str):
+    def from_config(cls, yaml_path: str, auto_deploy: bool = False):
         """
         Instantiate a Pipeline object from a YAML configuration file.
 
         :param yaml_path: Path to the YAML configuration file.
         :type yaml_path: str
+        :param auto_deploy: Whether to deploy the pipeline automatically, defaults to False
+        :type auto_deploy: bool, optional
         :return: An instance of the Pipeline class.
         :rtype: Pipeline
         """
         with open(yaml_path, "r") as file:
             config_data = yaml.safe_load(file)
 
-        pipeline_config_data = config_data.get("pipeline", {})
+        pipeline_config_data = config_data.get("pipeline", {}).get("config", {})
         db_config_data = config_data.get("vectordb", {})
         embedding_model_config_data = config_data.get("embedding_model", {})
+        llm_config_data = config_data.get("llm", {})
 
         pipeline_config = PipelineConfig(**pipeline_config_data)
 
         db_provider = db_config_data.get("provider", "chroma")
         db = VectorDBFactory.create(db_provider, db_config_data.get("config", {}))
 
+        if llm_config_data:
+            llm_provider = llm_config_data.get("provider", "openai")
+            llm = LlmFactory.create(llm_provider, llm_config_data.get("config", {}))
+        else:
+            llm = None
+
         embedding_model_provider = embedding_model_config_data.get("provider", "openai")
         embedding_model = EmbedderFactory.create(
             embedding_model_provider, embedding_model_config_data.get("config", {})
         )
-        return cls(config=pipeline_config, db=db, embedding_model=embedding_model)
+
+        # Send anonymous telemetry
+        event_properties = {"init_type": "yaml_config"}
+        AnonymousTelemetry().capture(event_name="init", properties=event_properties)
+
+        return cls(
+            config=pipeline_config,
+            llm=llm,
+            db=db,
+            embedding_model=embedding_model,
+            yaml_path=yaml_path,
+            auto_deploy=auto_deploy,
+        )
+
+    def start(self, host="0.0.0.0", port=8000):
+        app = FastAPI()
+
+        @app.post("/add")
+        async def add_document(data_value: str, data_type: str = None):
+            """
+            Add a document to the pipeline.
+            """
+            try:
+                document = {"data_value": data_value, "data_type": data_type}
+                self.add(document)
+                return {"message": "Document added successfully"}
+            except Exception as e:
+                raise HTTPException(status_code=500, detail=str(e))
+
+        @app.post("/query")
+        async def query_documents(query: str, num_documents: int = 3):
+            """
+            Query for similar documents in the pipeline.
+            """
+            try:
+                results = self.search(query, num_documents)
+                return results
+            except Exception as e:
+                raise HTTPException(status_code=500, detail=str(e))
+
+        import uvicorn
+
+        uvicorn.run(app, host=host, port=port)

embedchain/telemetry/posthog.py

@@ -0,0 +1,67 @@
+import json
+import logging
+import os
+import uuid
+from pathlib import Path
+
+from posthog import Posthog
+
+import embedchain
+
+HOME_DIR = str(Path.home())
+CONFIG_DIR = os.path.join(HOME_DIR, ".embedchain")
+CONFIG_FILE = os.path.join(CONFIG_DIR, "config.json")
+
+logger = logging.getLogger(__name__)
+
+
+class AnonymousTelemetry:
+    def __init__(self, host="https://app.posthog.com", enabled=True):
+        self.project_api_key = "phc_PHQDA5KwztijnSojsxJ2c1DuJd52QCzJzT2xnSGvjN2"
+        self.host = host
+        self.posthog = Posthog(project_api_key=self.project_api_key, host=self.host)
+        self.user_id = self.get_user_id()
+        self.enabled = enabled
+
+        # Check if telemetry tracking is disabled via environment variable
+        if "EC_TELEMETRY" in os.environ and os.environ["EC_TELEMETRY"].lower() not in [
+            "1",
+            "true",
+            "yes",
+        ]:
+            self.enabled = False
+
+        if not self.enabled:
+            self.posthog.disabled = True
+
+        # Silence posthog logging
+        posthog_logger = logging.getLogger("posthog")
+        posthog_logger.disabled = True
+
+    def get_user_id(self):
+        if not os.path.exists(CONFIG_DIR):
+            os.makedirs(CONFIG_DIR)
+
+        if os.path.exists(CONFIG_FILE):
+            with open(CONFIG_FILE, "r") as f:
+                data = json.load(f)
+                if "user_id" in data:
+                    return data["user_id"]
+
+        user_id = str(uuid.uuid4())
+        with open(CONFIG_FILE, "w") as f:
+            json.dump({"user_id": user_id}, f)
+        return user_id
+
+    def capture(self, event_name, properties=None):
+        default_properties = {
+            "version": embedchain.__version__,
+            "language": "python",
+            "pid": os.getpid(),
+        }
+        properties.update(default_properties)
+
+        try:
+            self.posthog.capture(self.user_id, event_name, properties)
+        except Exception:
+            logger.exception(f"Failed to send telemetry {event_name=}")

embedchain/utils.py

@@ -115,6 +115,13 @@ def detect_datatype(source: Any) -> DataType:
     """
     from urllib.parse import urlparse
 
+    import requests
+    import yaml
+
+    def is_openapi_yaml(yaml_content):
+        # currently the following two fields are required in openapi spec yaml config
+        return "openapi" in yaml_content and "info" in yaml_content
+
     try:
         if not isinstance(source, str):
             raise ValueError("Source is not a string and thus cannot be a URL.")
@@ -155,6 +162,31 @@ def detect_datatype(source: Any) -> DataType:
             logging.debug(f"Source of `{formatted_source}` detected as `docx`.")
             return DataType.DOCX
 
+        if url.path.endswith(".yaml"):
+            try:
+                response = requests.get(source)
+                response.raise_for_status()
+                try:
+                    yaml_content = yaml.safe_load(response.text)
+                except yaml.YAMLError as exc:
+                    logging.error(f"Error parsing YAML: {exc}")
+                    raise TypeError(f"Not a valid data type. Error loading YAML: {exc}")
+
+                if is_openapi_yaml(yaml_content):
+                    logging.debug(f"Source of `{formatted_source}` detected as `openapi`.")
+                    return DataType.OPENAPI
+                else:
+                    logging.error(
+                        f"Source of `{formatted_source}` does not contain all the required \
+                        fields of OpenAPI yaml. Check 'https://spec.openapis.org/oas/v3.1.0'"
+                    )
+                    raise TypeError(
+                        "Not a valid data type. Check 'https://spec.openapis.org/oas/v3.1.0', \
+                        make sure you have all the required fields in YAML config data"
+                    )
+            except requests.exceptions.RequestException as e:
+                logging.error(f"Error fetching URL {formatted_source}: {e}")
+
         if url.path.endswith(".json"):
             logging.debug(f"Source of `{formatted_source}` detected as `json_file`.")
             return DataType.JSON
@@ -198,6 +230,22 @@ def detect_datatype(source: Any) -> DataType:
             logging.debug(f"Source of `{formatted_source}` detected as `xml`.")
             return DataType.XML
 
+        if source.endswith(".yaml"):
+            with open(source, "r") as file:
+                yaml_content = yaml.safe_load(file)
+                if is_openapi_yaml(yaml_content):
+                    logging.debug(f"Source of `{formatted_source}` detected as `openapi`.")
+                    return DataType.OPENAPI
+                else:
+                    logging.error(
+                        f"Source of `{formatted_source}` does not contain all the required \
+                                  fields of OpenAPI yaml. Check 'https://spec.openapis.org/oas/v3.1.0'"
+                    )
+                    raise ValueError(
+                        "Invalid YAML data. Check 'https://spec.openapis.org/oas/v3.1.0', \
+                        make sure to add all the required params"
+                    )
+
         if source.endswith(".json"):
             logging.debug(f"Source of `{formatted_source}` detected as `json`.")
             return DataType.JSON
@@ -211,6 +259,8 @@ def detect_datatype(source: Any) -> DataType:
     else:
         # Source is not a URL.
 
+        # TODO: check if source is gmail query
+
         # Use text as final fallback.
         logging.debug(f"Source of `{formatted_source}` detected as `text`.")
         return DataType.TEXT

embedchain/vectordb/chroma.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 from chromadb import Collection, QueryResult
 from langchain.docstore.document import Document
@@ -38,7 +38,7 @@ class ChromaDB(BaseVectorDB):
         else:
             self.config = ChromaDbConfig()
 
-        self.settings = Settings()
+        self.settings = Settings(anonymized_telemetry=False)
         self.settings.allow_reset = self.config.allow_reset if hasattr(self.config, "allow_reset") else False
         if self.config.chroma_settings:
             for key, value in self.config.chroma_settings.items():
@@ -191,7 +191,14 @@ class ChromaDB(BaseVectorDB):
             )
         ]
 
-    def query(self, input_query: List[str], n_results: int, where: Dict[str, any], skip_embedding: bool) -> List[str]:
+    def query(
+        self,
+        input_query: List[str],
+        n_results: int,
+        where: Dict[str, any],
+        skip_embedding: bool,
+        citations: bool = False,
+    ) -> Union[List[Tuple[str, str, str]], List[str]]:
         """
         Query contents from vector database based on vector similarity
 
@@ -203,9 +210,12 @@ class ChromaDB(BaseVectorDB):
         :type where: Dict[str, Any]
         :param skip_embedding: Optional. If True, then the input_query is assumed to be already embedded.
         :type skip_embedding: bool
+        :param citations: we use citations boolean param to return context along with the answer.
+        :type citations: bool, default is False.
         :raises InvalidDimensionException: Dimensions do not match.
-        :return: The content of the document that matched your query.
-        :rtype: List[str]
+        :return: The content of the document that matched your query,
+        along with url of the source and doc_id (if citations flag is true)
+        :rtype: List[str], if citations=False, otherwise List[Tuple[str, str, str]]
         """
         try:
             if skip_embedding:
@@ -231,8 +241,17 @@ class ChromaDB(BaseVectorDB):
                 " embeddings, is used to retrieve an embedding from the database."
             ) from None
         results_formatted = self._format_result(result)
-        contents = [result[0].page_content for result in results_formatted]
-        return contents
+        contexts = []
+        for result in results_formatted:
+            context = result[0].page_content
+            if citations:
+                metadata = result[0].metadata
+                source = metadata["url"]
+                doc_id = metadata["doc_id"]
+                contexts.append((context, source, doc_id))
+            else:
+                contexts.append(context)
+        return contexts
 
     def set_collection_name(self, name: str):
         """
@@ -262,9 +281,9 @@ class ChromaDB(BaseVectorDB):
         """
         Resets the database. Deletes all embeddings irreversibly.
         """
-        # Delete all data from the database
+        # Delete all data from the collection
         try:
-            self.client.reset()
+            self.client.delete_collection(self.config.collection_name)
         except ValueError:
             raise ValueError(
                 "For safety reasons, resetting is disabled. "

embedchain/vectordb/elasticsearch.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 try:
     from elasticsearch import Elasticsearch
@@ -135,7 +135,14 @@ class ElasticsearchDB(BaseVectorDB):
         bulk(self.client, docs)
         self.client.indices.refresh(index=self._get_index())
 
-    def query(self, input_query: List[str], n_results: int, where: Dict[str, any], skip_embedding: bool) -> List[str]:
+    def query(
+        self,
+        input_query: List[str],
+        n_results: int,
+        where: Dict[str, any],
+        skip_embedding: bool,
+        citations: bool = False,
+    ) -> Union[List[Tuple[str, str, str]], List[str]]:
         """
         query contents from vector data base based on vector similarity
 
@@ -147,8 +154,12 @@ class ElasticsearchDB(BaseVectorDB):
         :type where: Dict[str, any]
         :param skip_embedding: Optional. If True, then the input_query is assumed to be already embedded.
         :type skip_embedding: bool
-        :return: Database contents that are the result of the query
-        :rtype: List[str]
+        :return: The context of the document that matched your query, url of the source, doc_id
+        :param citations: we use citations boolean param to return context along with the answer.
+        :type citations: bool, default is False.
+        :return: The content of the document that matched your query,
+        along with url of the source and doc_id (if citations flag is true)
+        :rtype: List[str], if citations=False, otherwise List[Tuple[str, str, str]]
         """
         if skip_embedding:
             query_vector = input_query
@@ -156,6 +167,7 @@ class ElasticsearchDB(BaseVectorDB):
             input_query_vector = self.embedder.embedding_fn(input_query)
             query_vector = input_query_vector[0]
 
+        # `https://www.elastic.co/guide/en/elasticsearch/reference/7.17/query-dsl-script-score-query.html`
         query = {
             "script_score": {
                 "query": {"bool": {"must": [{"exists": {"field": "text"}}]}},
@@ -167,12 +179,21 @@ class ElasticsearchDB(BaseVectorDB):
         }
         if "app_id" in where:
             app_id = where["app_id"]
-            query["script_score"]["query"]["bool"]["must"] = [{"term": {"metadata.app_id": app_id}}]
-        _source = ["text"]
+            query["script_score"]["query"] = {"match": {"metadata.app_id": app_id}}
+        _source = ["text", "metadata.url", "metadata.doc_id"]
         response = self.client.search(index=self._get_index(), query=query, _source=_source, size=n_results)
         docs = response["hits"]["hits"]
-        contents = [doc["_source"]["text"] for doc in docs]
-        return contents
+        contexts = []
+        for doc in docs:
+            context = doc["_source"]["text"]
+            if citations:
+                metadata = doc["_source"]["metadata"]
+                source = metadata["url"]
+                doc_id = metadata["doc_id"]
+                contexts.append(tuple((context, source, doc_id)))
+            else:
+                contexts.append(context)
+        return contexts
 
     def set_collection_name(self, name: str):
         """

embedchain/vectordb/opensearch.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Dict, List, Optional, Set
+from typing import Dict, List, Optional, Set, Tuple, Union
 
 try:
     from opensearchpy import OpenSearch
@@ -145,7 +145,14 @@ class OpenSearchDB(BaseVectorDB):
         bulk(self.client, docs)
         self.client.indices.refresh(index=self._get_index())
 
-    def query(self, input_query: List[str], n_results: int, where: Dict[str, any], skip_embedding: bool) -> List[str]:
+    def query(
+        self,
+        input_query: List[str],
+        n_results: int,
+        where: Dict[str, any],
+        skip_embedding: bool,
+        citations: bool = False,
+    ) -> Union[List[Tuple[str, str, str]], List[str]]:
         """
         query contents from vector data base based on vector similarity
 
@@ -157,8 +164,11 @@ class OpenSearchDB(BaseVectorDB):
         :type where: Dict[str, any]
         :param skip_embedding: Optional. If True, then the input_query is assumed to be already embedded.
         :type skip_embedding: bool
-        :return: Database contents that are the result of the query
-        :rtype: List[str]
+        :param citations: we use citations boolean param to return context along with the answer.
+        :type citations: bool, default is False.
+        :return: The content of the document that matched your query,
+        along with url of the source and doc_id (if citations flag is true)
+        :rtype: List[str], if citations=False, otherwise List[Tuple[str, str, str]]
         """
         # TODO(rupeshbansal, deshraj): Add support for skip embeddings here if already exists
         embeddings = OpenAIEmbeddings()
@@ -185,8 +195,17 @@ class OpenSearchDB(BaseVectorDB):
             pre_filter=pre_filter,
             k=n_results,
         )
-        contents = [doc.page_content for doc in docs]
-        return contents
+
+        contexts = []
+        for doc in docs:
+            context = doc.page_content
+            if citations:
+                source = doc.metadata["url"]
+                doc_id = doc.metadata["doc_id"]
+                contexts.append(tuple((context, source, doc_id)))
+            else:
+                contexts.append(context)
+        return contexts
 
     def set_collection_name(self, name: str):
         """

embedchain/vectordb/pinecone.py

@@ -1,5 +1,5 @@
 import os
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Tuple, Union
 
 try:
     import pinecone
@@ -118,7 +118,14 @@ class PineconeDB(BaseVectorDB):
         for i in range(0, len(docs), self.BATCH_SIZE):
             self.client.upsert(docs[i : i + self.BATCH_SIZE])
 
-    def query(self, input_query: List[str], n_results: int, where: Dict[str, any], skip_embedding: bool) -> List[str]:
+    def query(
+        self,
+        input_query: List[str],
+        n_results: int,
+        where: Dict[str, any],
+        skip_embedding: bool,
+        citations: bool = False,
+    ) -> Union[List[Tuple[str, str, str]], List[str]]:
         """
         query contents from vector database based on vector similarity
         :param input_query: list of query string
@@ -129,16 +136,28 @@ class PineconeDB(BaseVectorDB):
         :type where: Dict[str, any]
         :param skip_embedding: Optional. if True, input_query is already embedded
         :type skip_embedding: bool
-        :return: Database contents that are the result of the query
-        :rtype: List[str]
+        :param citations: we use citations boolean param to return context along with the answer.
+        :type citations: bool, default is False.
+        :return: The content of the document that matched your query,
+        along with url of the source and doc_id (if citations flag is true)
+        :rtype: List[str], if citations=False, otherwise List[Tuple[str, str, str]]
         """
         if not skip_embedding:
             query_vector = self.embedder.embedding_fn([input_query])[0]
         else:
             query_vector = input_query
-        contents = self.client.query(vector=query_vector, filter=where, top_k=n_results, include_metadata=True)
-        embeddings = list(map(lambda content: content["metadata"]["text"], contents["matches"]))
-        return embeddings
+        data = self.client.query(vector=query_vector, filter=where, top_k=n_results, include_metadata=True)
+        contexts = []
+        for doc in data["matches"]:
+            metadata = doc["metadata"]
+            context = metadata["text"]
+            if citations:
+                source = metadata["url"]
+                doc_id = metadata["doc_id"]
+                contexts.append(tuple((context, source, doc_id)))
+            else:
+                contexts.append(context)
+        return contexts
 
     def set_collection_name(self, name: str):
         """

embedchain/vectordb/qdrant.py

@@ -1,7 +1,7 @@
 import copy
 import os
 import uuid
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Tuple, Union
 
 try:
     from qdrant_client import QdrantClient
@@ -160,7 +160,14 @@ class QdrantDB(BaseVectorDB):
                 ),
             )
 
-    def query(self, input_query: List[str], n_results: int, where: Dict[str, any], skip_embedding: bool) -> List[str]:
+    def query(
+        self,
+        input_query: List[str],
+        n_results: int,
+        where: Dict[str, any],
+        skip_embedding: bool,
+        citations: bool = False,
+    ) -> Union[List[Tuple[str, str, str]], List[str]]:
         """
         query contents from vector database based on vector similarity
         :param input_query: list of query string
@@ -172,8 +179,11 @@ class QdrantDB(BaseVectorDB):
         :param skip_embedding: A boolean flag indicating if the embedding for the documents to be added is to be
         generated or not
         :type skip_embedding: bool
-        :return: Database contents that are the result of the query
-        :rtype: List[str]
+        :param citations: we use citations boolean param to return context along with the answer.
+        :type citations: bool, default is False.
+        :return: The content of the document that matched your query,
+        along with url of the source and doc_id (if citations flag is true)
+        :rtype: List[str], if citations=False, otherwise List[Tuple[str, str, str]]
         """
         if not skip_embedding:
             query_vector = self.embedder.embedding_fn([input_query])[0]
@@ -199,10 +209,18 @@ class QdrantDB(BaseVectorDB):
             query_vector=query_vector,
             limit=n_results,
         )
-        response = []
+
+        contexts = []
         for result in results:
-            response.append(result.payload.get("text", ""))
-        return response
+            context = result.payload["text"]
+            if citations:
+                metadata = result.payload["metadata"]
+                source = metadata["url"]
+                doc_id = metadata["doc_id"]
+                contexts.append(tuple((context, source, doc_id)))
+            else:
+                contexts.append(context)
+        return contexts
 
     def count(self) -> int:
         response = self.client.get_collection(collection_name=self.collection_name)
@@ -211,3 +229,15 @@ class QdrantDB(BaseVectorDB):
     def reset(self):
         self.client.delete_collection(collection_name=self.collection_name)
         self._initialize()
+
+    def set_collection_name(self, name: str):
+        """
+        Set the name of the collection. A collection is an isolated space for vectors.
+
+        :param name: Name of the collection.
+        :type name: str
+        """
+        if not isinstance(name, str):
+            raise TypeError("Collection name must be a string")
+        self.config.collection_name = name
+        self.collection_name = self._get_or_create_collection()

embedchain/vectordb/weaviate.py

@@ -1,6 +1,6 @@
 import copy
 import os
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Tuple, Union
 
 try:
     import weaviate
@@ -58,10 +58,14 @@ class WeaviateDB(BaseVectorDB):
             raise ValueError("Embedder not set. Please set an embedder with `set_embedder` before initialization.")
 
         self.index_name = self._get_index_name()
-        self.metadata_keys = {"data_type", "doc_id", "url", "hash", "app_id", "text"}
+        self.metadata_keys = {"data_type", "doc_id", "url", "hash", "app_id"}
         if not self.client.schema.exists(self.index_name):
             # id is a reserved field in Weaviate, hence we had to change the name of the id field to identifier
             # The none vectorizer is crucial as we have our own custom embedding function
+            """
+            TODO: wait for weaviate to add indexing on `object[]` data-type so that we can add filter while querying.
+            Once that is done, change `dataType` of "metadata" field to `object[]` and update the query below.
+            """
             class_obj = {
                 "classes": [
                     {
@@ -106,10 +110,6 @@ class WeaviateDB(BaseVectorDB):
                                 "name": "app_id",
                                 "dataType": ["text"],
                             },
-                            {
-                                "name": "text",
-                                "dataType": ["text"],
-                            },
                         ],
                     },
                 ]
@@ -194,7 +194,14 @@ class WeaviateDB(BaseVectorDB):
                 )
                 batch.add_reference(obj_uuid, self.index_name, "metadata", metadata_uuid, self.index_name + "_metadata")
 
-    def query(self, input_query: List[str], n_results: int, where: Dict[str, any], skip_embedding: bool) -> List[str]:
+    def query(
+        self,
+        input_query: List[str],
+        n_results: int,
+        where: Dict[str, any],
+        skip_embedding: bool,
+        citations: bool = False,
+    ) -> Union[List[Tuple[str, str, str]], List[str]]:
         """
         query contents from vector database based on vector similarity
         :param input_query: list of query string
@@ -206,14 +213,23 @@ class WeaviateDB(BaseVectorDB):
         :param skip_embedding: A boolean flag indicating if the embedding for the documents to be added is to be
         generated or not
         :type skip_embedding: bool
-        :return: Database contents that are the result of the query
-        :rtype: List[str]
+        :param citations: we use citations boolean param to return context along with the answer.
+        :type citations: bool, default is False.
+        :return: The content of the document that matched your query,
+        along with url of the source and doc_id (if citations flag is true)
+        :rtype: List[str], if citations=False, otherwise List[Tuple[str, str, str]]
         """
         if not skip_embedding:
             query_vector = self.embedder.embedding_fn([input_query])[0]
         else:
             query_vector = input_query
+
         keys = set(where.keys() if where is not None else set())
+        data_fields = ["text"]
+
+        if citations:
+            data_fields.append(weaviate.LinkTo("metadata", self.index_name + "_metadata", list(self.metadata_keys)))
+
         if len(keys.intersection(self.metadata_keys)) != 0:
             weaviate_where_operands = []
             for key in keys:
@@ -231,7 +247,7 @@ class WeaviateDB(BaseVectorDB):
                 weaviate_where_clause = {"operator": "And", "operands": weaviate_where_operands}
 
             results = (
-                self.client.query.get(self.index_name, ["text"])
+                self.client.query.get(self.index_name, data_fields)
                 .with_where(weaviate_where_clause)
                 .with_near_vector({"vector": query_vector})
                 .with_limit(n_results)
@@ -239,16 +255,24 @@ class WeaviateDB(BaseVectorDB):
             )
         else:
             results = (
-                self.client.query.get(self.index_name, ["text"])
+                self.client.query.get(self.index_name, data_fields)
                 .with_near_vector({"vector": query_vector})
                 .with_limit(n_results)
                 .do()
             )
-        matched_tokens = []
-        for result in results["data"]["Get"].get(self.index_name):
-            matched_tokens.append(result["text"])
 
-        return matched_tokens
+        docs = results["data"]["Get"].get(self.index_name)
+        contexts = []
+        for doc in docs:
+            context = doc["text"]
+            if citations:
+                metadata = doc["metadata"][0]
+                source = metadata["url"]
+                doc_id = metadata["doc_id"]
+                contexts.append((context, source, doc_id))
+            else:
+                contexts.append(context)
+        return contexts
 
     def set_collection_name(self, name: str):
         """

embedchain/vectordb/zilliz.py

@@ -1,4 +1,5 @@
-from typing import Dict, List, Optional
+import logging
+from typing import Dict, List, Optional, Tuple, Union
 
 from embedchain.config import ZillizDBConfig
 from embedchain.helper.json_serializable import register_deserializable
@@ -61,6 +62,7 @@ class ZillizVectorDB(BaseVectorDB):
         :type name: str
         """
         if utility.has_collection(name):
+            logging.info(f"[ZillizDB]: found an existing collection {name}, make sure the auto-id is disabled.")
             self.collection = Collection(name)
         else:
             fields = [
@@ -124,7 +126,14 @@ class ZillizVectorDB(BaseVectorDB):
         self.collection.flush()
         self.client.flush(self.config.collection_name)
 
-    def query(self, input_query: List[str], n_results: int, where: Dict[str, any], skip_embedding: bool) -> List[str]:
+    def query(
+        self,
+        input_query: List[str],
+        n_results: int,
+        where: Dict[str, any],
+        skip_embedding: bool,
+        citations: bool = False,
+    ) -> Union[List[Tuple[str, str, str]], List[str]]:
         """
         Query contents from vector data base based on vector similarity
 
@@ -135,8 +144,11 @@ class ZillizVectorDB(BaseVectorDB):
         :param where: to filter data
         :type where: str
         :raises InvalidDimensionException: Dimensions do not match.
-        :return: The content of the document that matched your query.
-        :rtype: List[str]
+        :param citations: we use citations boolean param to return context along with the answer.
+        :type citations: bool, default is False.
+        :return: The content of the document that matched your query,
+        along with url of the source and doc_id (if citations flag is true)
+        :rtype: List[str], if citations=False, otherwise List[Tuple[str, str, str]]
         """
 
         if self.collection.is_empty:
@@ -145,13 +157,14 @@ class ZillizVectorDB(BaseVectorDB):
         if not isinstance(where, str):
             where = None
 
+        output_fields = ["text", "url", "doc_id"]
         if skip_embedding:
             query_vector = input_query
             query_result = self.client.search(
                 collection_name=self.config.collection_name,
                 data=query_vector,
                 limit=n_results,
-                output_fields=["text"],
+                output_fields=output_fields,
             )
 
         else:
@@ -162,14 +175,20 @@ class ZillizVectorDB(BaseVectorDB):
                 collection_name=self.config.collection_name,
                 data=[query_vector],
                 limit=n_results,
-                output_fields=["text"],
+                output_fields=output_fields,
             )
 
-        doc_list = []
+        contexts = []
         for query in query_result:
-            doc_list.append(query[0]["entity"]["text"])
-
-        return doc_list
+            data = query[0]["entity"]
+            context = data["text"]
+            if citations:
+                source = data["url"]
+                doc_id = data["doc_id"]
+                contexts.append(tuple((context, source, doc_id)))
+            else:
+                contexts.append(context)
+        return contexts
 
     def count(self) -> int:
         """

examples/api_server/Dockerfile

@@ -8,4 +8,9 @@ COPY . .
 
 EXPOSE 5000
 
-CMD ["python", "api_server.py"]
+ENV FLASK_APP=api_server.py
+
+ENV FLASK_RUN_EXTRA_FILES=/usr/src/api/*
+ENV FLASK_ENV=development
+
+CMD ["flask", "run", "--host=0.0.0.0", "--reload"]

examples/api_server/api_server.py

@@ -1,3 +1,5 @@
+import logging
+
 from flask import Flask, jsonify, request
 
 from embedchain import App
@@ -5,11 +7,6 @@ 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()
@@ -17,9 +14,10 @@ def add():
     url_or_text = data.get("url_or_text")
     if data_type and url_or_text:
         try:
-            chat_bot.add(data_type, url_or_text)
+            App().add(url_or_text, data_type=data_type)
             return jsonify({"data": f"Added {data_type}: {url_or_text}"}), 200
         except Exception:
+            logging.exception(f"Failed to add {data_type=}: {url_or_text=}")
             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
 
@@ -30,9 +28,10 @@ def query():
     question = data.get("question")
     if question:
         try:
-            response = chat_bot.query(question)
+            response = App().query(question)
             return jsonify({"data": response}), 200
         except Exception:
+            logging.exception(f"Failed to query {question=}")
             return jsonify({"error": "An error occurred. Please try again!"}), 500
     return jsonify({"error": "Invalid request. Please provide 'question' in JSON format."}), 400
 
@@ -43,13 +42,13 @@ def chat():
     question = data.get("question")
     if question:
         try:
-            response = chat_bot.chat(question)
+            response = App().chat(question)
             return jsonify({"data": response}), 200
         except Exception:
+            logging.exception(f"Failed to chat {question=}")
             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

@@ -10,4 +10,6 @@ services:
     env_file:
       - variables.env
     ports:
-      - "5000:5000"
+      - "5000:5000"
+    volumes:
+      - .:/usr/src/api

examples/api_server/requirements.txt

@@ -1,2 +1,2 @@
 flask==2.3.2
-embedchain==0.0.58
+embedchain[dataloaders]==0.0.78

examples/full_stack/frontend/public/icons/twitter.svg

@@ -1,24 +1,10 @@
-<svg
-  fill="currentColor"
-  version="1.1"
-  xmlns="http://www.w3.org/2000/svg"
-  xmlns:xlink="http://www.w3.org/1999/xlink"
-  viewBox="0 0 512 512"
-  xml:space="preserve"
->
-  <g id="SVGRepo_bgCarrier" stroke-width="0"></g>
-  <g
-    id="SVGRepo_tracerCarrier"
-    stroke-linecap="round"
-    stroke-linejoin="round"
-  ></g>
-  <g id="SVGRepo_iconCarrier">
-    <g id="7935ec95c421cee6d86eb22ecd12f847">
-      <path
-        style="display: inline;"
-        d="M459.186,151.787c0.203,4.501,0.305,9.023,0.305,13.565 c0,138.542-105.461,298.285-298.274,298.285c-59.209,0-114.322-17.357-160.716-47.104c8.212,0.973,16.546,1.47,25.012,1.47 c49.121,0,94.318-16.759,130.209-44.884c-45.887-0.841-84.596-31.154-97.938-72.804c6.408,1.227,12.968,1.886,19.73,1.886 c9.55,0,18.816-1.287,27.617-3.68c-47.955-9.633-84.1-52.001-84.1-102.795c0-0.446,0-0.882,0.011-1.318 c14.133,7.847,30.294,12.562,47.488,13.109c-28.134-18.796-46.637-50.885-46.637-87.262c0-19.212,5.16-37.218,14.193-52.7 c51.707,63.426,128.941,105.156,216.072,109.536c-1.784-7.675-2.718-15.674-2.718-23.896c0-57.891,46.941-104.832,104.832-104.832 c30.173,0,57.404,12.734,76.525,33.102c23.887-4.694,46.313-13.423,66.569-25.438c-7.827,24.485-24.434,45.025-46.089,58.002 c21.209-2.535,41.426-8.171,60.222-16.505C497.448,118.542,479.666,137.004,459.186,151.787z"
-      > 
-      </path>
-    </g>
-  </g>
-</svg>
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Generator: Adobe Illustrator 27.5.0, SVG Export Plug-In . SVG Version: 6.00 Build 0)  -->
+<svg version="1.1" id="svg5" xmlns:svg="http://www.w3.org/2000/svg"
+	 xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" x="0px" y="0px" viewBox="0 0 1668.56 1221.19"
+	 style="enable-background:new 0 0 1668.56 1221.19;" xml:space="preserve">
+<g id="layer1" transform="translate(52.390088,-25.058597)">
+	<path id="path1009" d="M283.94,167.31l386.39,516.64L281.5,1104h87.51l340.42-367.76L984.48,1104h297.8L874.15,558.3l361.92-390.99
+		h-87.51l-313.51,338.7l-253.31-338.7H283.94z M412.63,231.77h136.81l604.13,807.76h-136.81L412.63,231.77z"/>
+</g>
+</svg>

notebooks/anthropic.ipynb

@@ -34,6 +34,15 @@
         "!pip install embedchain"
       ]
     },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "!pip install embedchain[dataloaders]"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -54,7 +63,7 @@
       "outputs": [],
       "source": [
         "import os\n",
-        "from embedchain import App\n",
+        "from embedchain import Pipeline as App\n",
         "\n",
         "os.environ[\"OPENAI_API_KEY\"] = \"sk-xxx\"\n",
         "os.environ[\"ANTHROPIC_API_KEY\"] = \"xxx\""

notebooks/azure-openai.ipynb

@@ -26,6 +26,16 @@
     "!pip install embedchain"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "692ff37b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!pip install embedchain[dataloaders]"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "ac982a56",
@@ -44,7 +54,7 @@
    "outputs": [],
    "source": [
     "import os\n",
-    "from embedchain import App\n",
+    "from embedchain import Pipeline as App\n",
     "\n",
     "os.environ[\"OPENAI_API_TYPE\"] = \"azure\"\n",
     "os.environ[\"OPENAI_API_BASE\"] = \"https://xxx.openai.azure.com/\"\n",

notebooks/chromadb.ipynb

@@ -1,84 +1,84 @@
 {
-  "nbformat": 4,
-  "nbformat_minor": 0,
-  "metadata": {
-    "colab": {
-      "provenance": []
-    },
-    "kernelspec": {
-      "name": "python3",
-      "display_name": "Python 3"
-    },
-    "language_info": {
-      "name": "python"
-    }
-  },
   "cells": [
     {
       "cell_type": "markdown",
-      "source": [
-        "## Cookbook for using ChromaDB with Embedchain"
-      ],
       "metadata": {
         "id": "b02n_zJ_hl3d"
-      }
+      },
+      "source": [
+        "## Cookbook for using ChromaDB with Embedchain"
+      ]
     },
     {
       "cell_type": "markdown",
-      "source": [
-        "### Step-1: Install embedchain package"
-      ],
       "metadata": {
         "id": "gyJ6ui2vhtMY"
-      }
+      },
+      "source": [
+        "### Step-1: Install embedchain package"
+      ]
     },
     {
       "cell_type": "code",
-      "source": [
-        "!pip install embedchain"
-      ],
+      "execution_count": null,
       "metadata": {
         "id": "-NbXjAdlh0vJ"
       },
+      "outputs": [],
+      "source": [
+        "!pip install embedchain"
+      ]
+    },
+    {
+      "cell_type": "code",
       "execution_count": null,
-      "outputs": []
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "!pip install embedchain[dataloaders]"
+      ]
     },
     {
       "cell_type": "markdown",
+      "metadata": {
+        "id": "nGnpSYAAh2bQ"
+      },
       "source": [
         "### Step-2: Set OpenAI environment variables\n",
         "\n",
         "You can find this env variable on your [OpenAI dashboard](https://platform.openai.com/account/api-keys)."
-      ],
-      "metadata": {
-        "id": "nGnpSYAAh2bQ"
-      }
+      ]
     },
     {
       "cell_type": "code",
-      "source": [
-        "import os\n",
-        "from embedchain import App\n",
-        "\n",
-        "os.environ[\"OPENAI_API_KEY\"] = \"sk-xxx\""
-      ],
+      "execution_count": null,
       "metadata": {
         "id": "0fBdQ9GAiRvK"
       },
-      "execution_count": null,
-      "outputs": []
+      "outputs": [],
+      "source": [
+        "import os\n",
+        "from embedchain import Pipeline as App\n",
+        "\n",
+        "os.environ[\"OPENAI_API_KEY\"] = \"sk-xxx\""
+      ]
     },
     {
       "cell_type": "markdown",
-      "source": [
-        "### Step-3: Define your Vector Database config"
-      ],
       "metadata": {
         "id": "Ns6RhPfbiitr"
-      }
+      },
+      "source": [
+        "### Step-3: Define your Vector Database config"
+      ]
     },
     {
       "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "S9CkxVjriotB"
+      },
+      "outputs": [],
       "source": [
         "config = \"\"\"\n",
         "vectordb:\n",
@@ -95,64 +95,64 @@
         "# Write the multi-line string to a YAML file\n",
         "with open('chromadb.yaml', 'w') as file:\n",
         "    file.write(config)"
-      ],
-      "metadata": {
-        "id": "S9CkxVjriotB"
-      },
-      "execution_count": null,
-      "outputs": []
+      ]
     },
     {
       "cell_type": "markdown",
-      "source": [
-        "### Step-4 Create embedchain app based on the config"
-      ],
       "metadata": {
         "id": "PGt6uPLIi1CS"
-      }
+      },
+      "source": [
+        "### Step-4 Create embedchain app based on the config"
+      ]
     },
     {
       "cell_type": "code",
-      "source": [
-        "app = App.from_config(yaml_path=\"chromadb.yaml\")"
-      ],
+      "execution_count": null,
       "metadata": {
         "id": "Amzxk3m-i3tD"
       },
-      "execution_count": null,
-      "outputs": []
+      "outputs": [],
+      "source": [
+        "app = App.from_config(yaml_path=\"chromadb.yaml\")"
+      ]
     },
     {
       "cell_type": "markdown",
-      "source": [
-        "### Step-5: Add data sources to your app"
-      ],
       "metadata": {
         "id": "XNXv4yZwi7ef"
-      }
+      },
+      "source": [
+        "### Step-5: Add data sources to your app"
+      ]
     },
     {
       "cell_type": "code",
-      "source": [
-        "app.add(\"https://www.forbes.com/profile/elon-musk\")"
-      ],
+      "execution_count": null,
       "metadata": {
         "id": "Sn_0rx9QjIY9"
       },
-      "execution_count": null,
-      "outputs": []
+      "outputs": [],
+      "source": [
+        "app.add(\"https://www.forbes.com/profile/elon-musk\")"
+      ]
     },
     {
       "cell_type": "markdown",
-      "source": [
-        "### Step-6: All set. Now start asking questions related to your data"
-      ],
       "metadata": {
         "id": "_7W6fDeAjMAP"
-      }
+      },
+      "source": [
+        "### Step-6: All set. Now start asking questions related to your data"
+      ]
     },
     {
       "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "cvIK7dWRjN_f"
+      },
+      "outputs": [],
       "source": [
         "while(True):\n",
         "    question = input(\"Enter question: \")\n",
@@ -160,12 +160,21 @@
         "        break\n",
         "    answer = app.query(question)\n",
         "    print(answer)"
-      ],
-      "metadata": {
-        "id": "cvIK7dWRjN_f"
-      },
-      "execution_count": null,
-      "outputs": []
+      ]
     }
-  ]
-}
+  ],
+  "metadata": {
+    "colab": {
+      "provenance": []
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/cohere.ipynb

@@ -33,6 +33,15 @@
         "!pip install embedchain"
       ]
     },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "!pip install embedchain[dataloaders]"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -69,7 +78,7 @@
       "outputs": [],
       "source": [
         "import os\n",
-        "from embedchain import App\n",
+        "from embedchain import Pipeline as App\n",
         "\n",
         "os.environ[\"OPENAI_API_KEY\"] = \"sk-xxx\"\n",
         "os.environ[\"COHERE_API_KEY\"] = \"xxx\""

notebooks/elasticsearch.ipynb

@@ -1,95 +1,95 @@
 {
-  "nbformat": 4,
-  "nbformat_minor": 0,
-  "metadata": {
-    "colab": {
-      "provenance": []
-    },
-    "kernelspec": {
-      "name": "python3",
-      "display_name": "Python 3"
-    },
-    "language_info": {
-      "name": "python"
-    }
-  },
   "cells": [
     {
       "cell_type": "markdown",
-      "source": [
-        "## Cookbook for using ElasticSearchDB with Embedchain"
-      ],
       "metadata": {
         "id": "b02n_zJ_hl3d"
-      }
+      },
+      "source": [
+        "## Cookbook for using ElasticSearchDB with Embedchain"
+      ]
     },
     {
       "cell_type": "markdown",
-      "source": [
-        "### Step-1: Install embedchain package"
-      ],
       "metadata": {
         "id": "gyJ6ui2vhtMY"
-      }
+      },
+      "source": [
+        "### Step-1: Install embedchain package"
+      ]
     },
     {
       "cell_type": "code",
-      "source": [
-        "!pip install embedchain"
-      ],
+      "execution_count": null,
       "metadata": {
         "id": "-NbXjAdlh0vJ"
       },
+      "outputs": [],
+      "source": [
+        "!pip install embedchain"
+      ]
+    },
+    {
+      "cell_type": "code",
       "execution_count": null,
-      "outputs": []
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "!pip install embedchain[dataloaders]"
+      ]
     },
     {
       "cell_type": "markdown",
+      "metadata": {
+        "id": "nGnpSYAAh2bQ"
+      },
       "source": [
         "### Step-2: Set OpenAI environment variables and install the dependencies.\n",
         "\n",
         "You can find this env variable on your [OpenAI dashboard](https://platform.openai.com/account/api-keys). Now lets install the dependencies needed for Elasticsearch."
-      ],
-      "metadata": {
-        "id": "nGnpSYAAh2bQ"
-      }
+      ]
     },
     {
       "cell_type": "code",
-      "source": [
-        "!pip install --upgrade 'embedchain[elasticsearch]'"
-      ],
+      "execution_count": null,
       "metadata": {
         "id": "-MUFRfxV7Jk7"
       },
-      "execution_count": null,
-      "outputs": []
+      "outputs": [],
+      "source": [
+        "!pip install --upgrade 'embedchain[elasticsearch]'"
+      ]
     },
     {
       "cell_type": "code",
-      "source": [
-        "import os\n",
-        "from embedchain import App\n",
-        "\n",
-        "os.environ[\"OPENAI_API_KEY\"] = \"sk-xxx\""
-      ],
+      "execution_count": null,
       "metadata": {
         "id": "0fBdQ9GAiRvK"
       },
-      "execution_count": null,
-      "outputs": []
+      "outputs": [],
+      "source": [
+        "import os\n",
+        "from embedchain import Pipeline as App\n",
+        "\n",
+        "os.environ[\"OPENAI_API_KEY\"] = \"sk-xxx\""
+      ]
     },
     {
       "cell_type": "markdown",
-      "source": [
-        "### Step-3: Define your Vector Database config"
-      ],
       "metadata": {
         "id": "Ns6RhPfbiitr"
-      }
+      },
+      "source": [
+        "### Step-3: Define your Vector Database config"
+      ]
     },
     {
       "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "S9CkxVjriotB"
+      },
+      "outputs": [],
       "source": [
         "config = \"\"\"\n",
         "vectordb:\n",
@@ -104,64 +104,64 @@
         "# Write the multi-line string to a YAML file\n",
         "with open('elasticsearch.yaml', 'w') as file:\n",
         "    file.write(config)"
-      ],
-      "metadata": {
-        "id": "S9CkxVjriotB"
-      },
-      "execution_count": null,
-      "outputs": []
+      ]
     },
     {
       "cell_type": "markdown",
-      "source": [
-        "### Step-4 Create embedchain app based on the config"
-      ],
       "metadata": {
         "id": "PGt6uPLIi1CS"
-      }
+      },
+      "source": [
+        "### Step-4 Create embedchain app based on the config"
+      ]
     },
     {
       "cell_type": "code",
-      "source": [
-        "app = App.from_config(yaml_path=\"elasticsearch.yaml\")"
-      ],
+      "execution_count": null,
       "metadata": {
         "id": "Amzxk3m-i3tD"
       },
-      "execution_count": null,
-      "outputs": []
+      "outputs": [],
+      "source": [
+        "app = App.from_config(yaml_path=\"elasticsearch.yaml\")"
+      ]
     },
     {
       "cell_type": "markdown",
-      "source": [
-        "### Step-5: Add data sources to your app"
-      ],
       "metadata": {
         "id": "XNXv4yZwi7ef"
-      }
+      },
+      "source": [
+        "### Step-5: Add data sources to your app"
+      ]
     },
     {
       "cell_type": "code",
-      "source": [
-        "app.add(\"https://www.forbes.com/profile/elon-musk\")"
-      ],
+      "execution_count": null,
       "metadata": {
         "id": "Sn_0rx9QjIY9"
       },
-      "execution_count": null,
-      "outputs": []
+      "outputs": [],
+      "source": [
+        "app.add(\"https://www.forbes.com/profile/elon-musk\")"
+      ]
     },
     {
       "cell_type": "markdown",
-      "source": [
-        "### Step-6: All set. Now start asking questions related to your data"
-      ],
       "metadata": {
         "id": "_7W6fDeAjMAP"
-      }
+      },
+      "source": [
+        "### Step-6: All set. Now start asking questions related to your data"
+      ]
     },
     {
       "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "cvIK7dWRjN_f"
+      },
+      "outputs": [],
       "source": [
         "while(True):\n",
         "    question = input(\"Enter question: \")\n",
@@ -169,12 +169,21 @@
         "        break\n",
         "    answer = app.query(question)\n",
         "    print(answer)"
-      ],
-      "metadata": {
-        "id": "cvIK7dWRjN_f"
-      },
-      "execution_count": null,
-      "outputs": []
+      ]
     }
-  ]
-}
+  ],
+  "metadata": {
+    "colab": {
+      "provenance": []
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/embedchain-chromadb-server.ipynb

@@ -33,7 +33,7 @@
    "outputs": [],
    "source": [
     "import os\n",
-    "from embedchain import App\n",
+    "from embedchain import Pipeline as App\n",
     "from embedchain.config import AppConfig\n",
     "\n",
     "\n",

notebooks/embedchain-docs-site-example.ipynb

@@ -7,7 +7,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from embedchain import App\n",
+    "from embedchain import Pipeline as App\n",
     "\n",
     "embedchain_docs_bot = App()"
    ]

notebooks/gpt4all.ipynb

@@ -33,6 +33,15 @@
         "!pip install embedchain"
       ]
     },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "!pip install embedchain[dataloaders]"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -67,7 +76,7 @@
       },
       "outputs": [],
       "source": [
-        "from embedchain import App"
+        "from embedchain import Pipeline as App"
       ]
     },
     {

notebooks/hugging_face_hub.ipynb

@@ -34,6 +34,15 @@
         "!pip install embedchain"
       ]
     },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "!pip install embedchain[dataloaders]"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -84,7 +93,7 @@
       "outputs": [],
       "source": [
         "import os\n",
-        "from embedchain import App\n",
+        "from embedchain import Pipeline as App\n",
         "\n",
         "os.environ[\"HUGGINGFACE_ACCESS_TOKEN\"] = \"hf_xxx\""
       ]

notebooks/jina.ipynb

@@ -34,6 +34,15 @@
         "!pip install embedchain"
       ]
     },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "!pip install embedchain[dataloaders]"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -54,7 +63,7 @@
       "outputs": [],
       "source": [
         "import os\n",
-        "from embedchain import App\n",
+        "from embedchain import Pipeline as App\n",
         "\n",
         "os.environ[\"OPENAI_API_KEY\"] = \"sk-xxx\"\n",
         "os.environ[\"JINACHAT_API_KEY\"] = \"xxx\""

notebooks/llama2.ipynb

@@ -33,6 +33,15 @@
         "!pip install embedchain"
       ]
     },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "!pip install embedchain[dataloaders]"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -64,7 +73,7 @@
       "outputs": [],
       "source": [
         "import os\n",
-        "from embedchain import App\n",
+        "from embedchain import Pipeline as App\n",
         "\n",
         "os.environ[\"OPENAI_API_KEY\"] = \"sk-xxx\"\n",
         "os.environ[\"REPLICATE_API_TOKEN\"] = \"xxx\""

notebooks/openai.ipynb

@@ -34,6 +34,15 @@
         "!pip install embedchain"
       ]
     },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "!pip install embedchain[dataloaders]"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -54,7 +63,7 @@
       "outputs": [],
       "source": [
         "import os\n",
-        "from embedchain import App\n",
+        "from embedchain import Pipeline as App\n",
         "\n",
         "os.environ[\"OPENAI_API_KEY\"] = \"sk-xxx\""
       ]

notebooks/opensearch.ipynb

@@ -1,95 +1,95 @@
 {
-  "nbformat": 4,
-  "nbformat_minor": 0,
-  "metadata": {
-    "colab": {
-      "provenance": []
-    },
-    "kernelspec": {
-      "name": "python3",
-      "display_name": "Python 3"
-    },
-    "language_info": {
-      "name": "python"
-    }
-  },
   "cells": [
     {
       "cell_type": "markdown",
-      "source": [
-        "## Cookbook for using OpenSearchDB with Embedchain"
-      ],
       "metadata": {
         "id": "b02n_zJ_hl3d"
-      }
+      },
+      "source": [
+        "## Cookbook for using OpenSearchDB with Embedchain"
+      ]
     },
     {
       "cell_type": "markdown",
-      "source": [
-        "### Step-1: Install embedchain package"
-      ],
       "metadata": {
         "id": "gyJ6ui2vhtMY"
-      }
+      },
+      "source": [
+        "### Step-1: Install embedchain package"
+      ]
     },
     {
       "cell_type": "code",
-      "source": [
-        "!pip install embedchain"
-      ],
+      "execution_count": null,
       "metadata": {
         "id": "-NbXjAdlh0vJ"
       },
+      "outputs": [],
+      "source": [
+        "!pip install embedchain"
+      ]
+    },
+    {
+      "cell_type": "code",
       "execution_count": null,
-      "outputs": []
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "!pip install embedchain[dataloaders]"
+      ]
     },
     {
       "cell_type": "markdown",
+      "metadata": {
+        "id": "nGnpSYAAh2bQ"
+      },
       "source": [
         "### Step-2: Set OpenAI environment variables and install the dependencies.\n",
         "\n",
         "You can find this env variable on your [OpenAI dashboard](https://platform.openai.com/account/api-keys). Now lets install the dependencies needed for Opensearch."
-      ],
-      "metadata": {
-        "id": "nGnpSYAAh2bQ"
-      }
+      ]
     },
     {
       "cell_type": "code",
-      "source": [
-        "!pip install --upgrade 'embedchain[opensearch]'"
-      ],
+      "execution_count": null,
       "metadata": {
         "id": "-MUFRfxV7Jk7"
       },
-      "execution_count": null,
-      "outputs": []
+      "outputs": [],
+      "source": [
+        "!pip install --upgrade 'embedchain[opensearch]'"
+      ]
     },
     {
       "cell_type": "code",
-      "source": [
-        "import os\n",
-        "from embedchain import App\n",
-        "\n",
-        "os.environ[\"OPENAI_API_KEY\"] = \"sk-xxx\""
-      ],
+      "execution_count": null,
       "metadata": {
         "id": "0fBdQ9GAiRvK"
       },
-      "execution_count": null,
-      "outputs": []
+      "outputs": [],
+      "source": [
+        "import os\n",
+        "from embedchain import Pipeline as App\n",
+        "\n",
+        "os.environ[\"OPENAI_API_KEY\"] = \"sk-xxx\""
+      ]
     },
     {
       "cell_type": "markdown",
-      "source": [
-        "### Step-3: Define your Vector Database config"
-      ],
       "metadata": {
         "id": "Ns6RhPfbiitr"
-      }
+      },
+      "source": [
+        "### Step-3: Define your Vector Database config"
+      ]
     },
     {
       "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "S9CkxVjriotB"
+      },
+      "outputs": [],
       "source": [
         "config = \"\"\"\n",
         "vectordb:\n",
@@ -108,64 +108,64 @@
         "# Write the multi-line string to a YAML file\n",
         "with open('opensearch.yaml', 'w') as file:\n",
         "    file.write(config)"
-      ],
-      "metadata": {
-        "id": "S9CkxVjriotB"
-      },
-      "execution_count": null,
-      "outputs": []
+      ]
     },
     {
       "cell_type": "markdown",
-      "source": [
-        "### Step-4 Create embedchain app based on the config"
-      ],
       "metadata": {
         "id": "PGt6uPLIi1CS"
-      }
+      },
+      "source": [
+        "### Step-4 Create embedchain app based on the config"
+      ]
     },
     {
       "cell_type": "code",
-      "source": [
-        "app = App.from_config(yaml_path=\"opensearch.yaml\")"
-      ],
+      "execution_count": null,
       "metadata": {
         "id": "Amzxk3m-i3tD"
       },
-      "execution_count": null,
-      "outputs": []
+      "outputs": [],
+      "source": [
+        "app = App.from_config(yaml_path=\"opensearch.yaml\")"
+      ]
     },
     {
       "cell_type": "markdown",
-      "source": [
-        "### Step-5: Add data sources to your app"
-      ],
       "metadata": {
         "id": "XNXv4yZwi7ef"
-      }
+      },
+      "source": [
+        "### Step-5: Add data sources to your app"
+      ]
     },
     {
       "cell_type": "code",
-      "source": [
-        "app.add(\"https://www.forbes.com/profile/elon-musk\")"
-      ],
+      "execution_count": null,
       "metadata": {
         "id": "Sn_0rx9QjIY9"
       },
-      "execution_count": null,
-      "outputs": []
+      "outputs": [],
+      "source": [
+        "app.add(\"https://www.forbes.com/profile/elon-musk\")"
+      ]
     },
     {
       "cell_type": "markdown",
-      "source": [
-        "### Step-6: All set. Now start asking questions related to your data"
-      ],
       "metadata": {
         "id": "_7W6fDeAjMAP"
-      }
+      },
+      "source": [
+        "### Step-6: All set. Now start asking questions related to your data"
+      ]
     },
     {
       "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "cvIK7dWRjN_f"
+      },
+      "outputs": [],
       "source": [
         "while(True):\n",
         "    question = input(\"Enter question: \")\n",
@@ -173,12 +173,21 @@
         "        break\n",
         "    answer = app.query(question)\n",
         "    print(answer)"
-      ],
-      "metadata": {
-        "id": "cvIK7dWRjN_f"
-      },
-      "execution_count": null,
-      "outputs": []
+      ]
     }
-  ]
-}
+  ],
+  "metadata": {
+    "colab": {
+      "provenance": []
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/pinecone.ipynb

@@ -29,6 +29,15 @@
         "!pip install embedchain"
       ]
     },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "!pip install embedchain[dataloaders]"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -60,7 +69,7 @@
       "outputs": [],
       "source": [
         "import os\n",
-        "from embedchain import App\n",
+        "from embedchain import Pipeline as App\n",
         "\n",
         "os.environ[\"OPENAI_API_KEY\"] = \"sk-xxx\"\n",
         "os.environ[\"PINECONE_API_KEY\"] = \"xxx\"\n",

notebooks/vertex_ai.ipynb

@@ -33,6 +33,15 @@
         "!pip install embedchain"
       ]
     },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "!pip install embedchain[dataloaders]"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -64,7 +73,7 @@
       "outputs": [],
       "source": [
         "import os\n",
-        "from embedchain import App\n",
+        "from embedchain import Pipeline as App\n",
         "\n",
         "os.environ[\"OPENAI_API_KEY\"] = \"sk-xxx\""
       ]

pyproject.toml

@@ -1,8 +1,11 @@
 [tool.poetry]
 name = "embedchain"
-version = "0.0.75"
+version = "0.0.90"
 description = "Data platform for LLMs - Load, index, retrieve and sync any unstructured data"
-authors = ["Taranjeet Singh, Deshraj Yadav"]
+authors = [
+    "Taranjeet Singh <taranjeet@embedchain.ai>",
+    "Deshraj Yadav <deshraj@embedchain.ai>",
+]
 license = "Apache License"
 readme = "README.md"
 exclude = [
@@ -87,16 +90,17 @@ color = true
 [tool.poetry.dependencies]
 python = ">=3.9,<3.13"
 python-dotenv = "^1.0.0"
-langchain = "^0.0.279"
+langchain = "^0.0.303"
 requests = "^2.31.0"
-openai = "^0.27.5"
-tiktoken = { version="^0.4.0", optional=true }
-chromadb ="^0.4.8"
-youtube-transcript-api = { version="^0.6.1", optional=true }
-beautifulsoup4 = { version="^4.12.2", optional=true }
-pypdf = { version="^3.11.0", optional=true }
-pytube = { version="^15.0.0", optional=true }
-duckduckgo-search = { version="^3.8.5", optional=true }
+openai = ">=0.28.0"
+chromadb = "^0.4.8"
+posthog = "^3.0.2"
+tiktoken = { version = "^0.4.0", optional = true }
+youtube-transcript-api = { version = "^0.6.1", optional = true }
+beautifulsoup4 = { version = "^4.12.2", optional = true }
+pypdf = { version = "^3.11.0", optional = true }
+pytube = { version = "^15.0.0", optional = true }
+duckduckgo-search = { version = "^3.8.5", optional = true }
 llama-hub = { version = "^0.0.29", optional = true }
 sentence-transformers = { version = "^2.2.2", optional = true }
 torch = { version = "2.0.0", optional = true }
@@ -110,12 +114,12 @@ twilio = { version = "^8.5.0", optional = true }
 fastapi-poe = { version = "0.0.16", optional = true }
 discord = { version = "^2.3.2", optional = true }
 slack-sdk = { version = "3.21.3", optional = true }
-cohere = { version = "^4.27", optional= true }
-weaviate-client = { version = "^3.24.1", optional= true }
-docx2txt = { version="^0.8", optional=true }
+cohere = { version = "^4.27", optional = true }
+weaviate-client = { version = "^3.24.1", optional = true }
+docx2txt = { version = "^0.8", optional = true }
 pinecone-client = { version = "^2.2.4", optional = true }
 qdrant-client = { version = "1.6.3", optional = true }
-unstructured = {extras = ["local-inference"], version = "^0.10.18", optional=true}
+unstructured = {extras = ["local-inference"], version = "^0.10.18", optional = true}
 pillow = { version = "10.0.1", optional = true }
 torchvision = { version = ">=0.15.1, !=0.15.2", optional = true }
 ftfy = { version = "6.1.1", optional = true }
@@ -124,7 +128,6 @@ huggingface_hub = { version = "^0.17.3", optional = true }
 pymilvus = { version = "2.3.1", optional = true }
 google-cloud-aiplatform = { version = "^1.26.1", optional = true }
 replicate = { version = "^0.15.4", optional = true }
-jq = { version=">=1.6.0", optional = true}
 
 [tool.poetry.group.dev.dependencies]
 black = "^23.3.0"
@@ -160,16 +163,25 @@ milvus = ["pymilvus"]
 dataloaders=[
     "youtube-transcripts-api",
     "beautifulsoup4",
+    "docx2txt",
+    "duckduckgo-search",
     "pypdf",
     "pytube",
-    "duckduckgo-search",
-    "docx2txt",
-    "unstructured",
     "sentence-transformers",
-    "jq",
+    "unstructured",
 ]
 vertexai = ["google-cloud-aiplatform"]
 llama2 = ["replicate"]
+gmail = [
+    "llama-hub",
+    "requests",
+    "google-api-python-client",
+    "google-auth",
+    "google-auth-oauthlib",
+    "google-auth-httplib2",
+    "google-api-core",
+]
+json = ["llama-hub"]
 
 [tool.poetry.group.docs.dependencies]
 

tests/apps/test_apps.py

@@ -135,6 +135,7 @@ class TestAppFromConfig:
 
         # Validate the LLM config values
         llm_config = config_data["llm"]["config"]
+        assert app.llm.config.model == llm_config["model"]
         assert app.llm.config.temperature == llm_config["temperature"]
         assert app.llm.config.max_tokens == llm_config["max_tokens"]
         assert app.llm.config.top_p == llm_config["top_p"]
@@ -148,5 +149,4 @@ class TestAppFromConfig:
 
         # Validate the Embedder config values
         embedder_config = config_data["embedder"]["config"]
-        assert app.embedder.config.model == embedder_config["model"]
         assert app.embedder.config.deployment_name == embedder_config["deployment_name"]

tests/chunkers/test_chunkers.py

@@ -1,8 +1,10 @@
 from embedchain.chunkers.docs_site import DocsSiteChunker
 from embedchain.chunkers.docx_file import DocxFileChunker
+from embedchain.chunkers.gmail import GmailChunker
 from embedchain.chunkers.json import JSONChunker
 from embedchain.chunkers.mdx import MdxChunker
 from embedchain.chunkers.notion import NotionChunker
+from embedchain.chunkers.openapi import OpenAPIChunker
 from embedchain.chunkers.pdf_file import PdfFileChunker
 from embedchain.chunkers.qna_pair import QnaPairChunker
 from embedchain.chunkers.sitemap import SitemapChunker
@@ -29,6 +31,8 @@ chunker_common_config = {
     XmlChunker: {"chunk_size": 500, "chunk_overlap": 50, "length_function": len},
     YoutubeVideoChunker: {"chunk_size": 2000, "chunk_overlap": 0, "length_function": len},
     JSONChunker: {"chunk_size": 1000, "chunk_overlap": 0, "length_function": len},
+    OpenAPIChunker: {"chunk_size": 1000, "chunk_overlap": 0, "length_function": len},
+    GmailChunker: {"chunk_size": 1000, "chunk_overlap": 0, "length_function": len},
 }
 
 

tests/conftest.py

@@ -0,0 +1,23 @@
+import os
+
+import pytest
+
+
+def clean_db():
+    db_path = os.path.expanduser("~/.embedchain/embedchain.db")
+    if os.path.exists(db_path):
+        os.remove(db_path)
+
+
+@pytest.fixture
+def setup():
+    clean_db()
+    yield
+    clean_db()
+
+
+@pytest.fixture(autouse=True)
+def disable_telemetry():
+    os.environ["EC_TELEMETRY"] = "false"
+    yield
+    del os.environ["EC_TELEMETRY"]

tests/embedchain/test_add.py

@@ -16,19 +16,20 @@ def app(mocker):
 
 
 def test_add(app):
-    app.add("https://example.com", metadata={"meta": "meta-data"})
-    assert app.user_asks == [["https://example.com", "web_page", {"meta": "meta-data"}]]
+    app.add("https://example.com", metadata={"foo": "bar"})
+    assert app.user_asks == [["https://example.com", "web_page", {"foo": "bar"}]]
 
 
-def test_add_sitemap(app):
-    app.add("https://www.google.com/sitemap.xml", metadata={"meta": "meta-data"})
-    assert app.user_asks == [["https://www.google.com/sitemap.xml", "sitemap", {"meta": "meta-data"}]]
+# TODO: Make this test faster by generating a sitemap locally rather than using a remote one
+# def test_add_sitemap(app):
+#     app.add("https://www.google.com/sitemap.xml", metadata={"foo": "bar"})
+#     assert app.user_asks == [["https://www.google.com/sitemap.xml", "sitemap", {"foo": "bar"}]]
 
 
 def test_add_forced_type(app):
     data_type = "text"
-    app.add("https://example.com", data_type=data_type, metadata={"meta": "meta-data"})
-    assert app.user_asks == [["https://example.com", data_type, {"meta": "meta-data"}]]
+    app.add("https://example.com", data_type=data_type, metadata={"foo": "bar"})
+    assert app.user_asks == [["https://example.com", data_type, {"foo": "bar"}]]
 
 
 def test_dry_run(app):

tests/embedchain/test_utils.py

@@ -39,6 +39,12 @@ class TestApp(unittest.TestCase):
     def test_detect_datatype_local_docx(self):
         self.assertEqual(detect_datatype("file:///home/user/document.docx"), DataType.DOCX)
 
+    def test_detect_data_type_json(self):
+        self.assertEqual(detect_datatype("https://www.example.com/data.json"), DataType.JSON)
+
+    def test_detect_data_type_local_json(self):
+        self.assertEqual(detect_datatype("file:///home/user/data.json"), DataType.JSON)
+
     @patch("os.path.isfile")
     def test_detect_datatype_regular_filesystem_docx(self, mock_isfile):
         with tempfile.NamedTemporaryFile(suffix=".docx", delete=True) as tmp:

tests/llm/test_gpt4all.py

@@ -0,0 +1,60 @@
+import pytest
+from langchain.llms.gpt4all import GPT4All as LangchainGPT4All
+
+from embedchain.config import BaseLlmConfig
+from embedchain.llm.gpt4all import GPT4ALLLlm
+
+
+@pytest.fixture
+def config():
+    config = BaseLlmConfig(
+        temperature=0.7,
+        max_tokens=50,
+        top_p=0.8,
+        stream=False,
+        system_prompt="System prompt",
+        model="orca-mini-3b.ggmlv3.q4_0.bin",
+    )
+    yield config
+
+
+@pytest.fixture
+def gpt4all_with_config(config):
+    return GPT4ALLLlm(config=config)
+
+
+@pytest.fixture
+def gpt4all_without_config():
+    return GPT4ALLLlm()
+
+
+def test_gpt4all_init_with_config(config, gpt4all_with_config):
+    assert gpt4all_with_config.config.temperature == config.temperature
+    assert gpt4all_with_config.config.max_tokens == config.max_tokens
+    assert gpt4all_with_config.config.top_p == config.top_p
+    assert gpt4all_with_config.config.stream == config.stream
+    assert gpt4all_with_config.config.system_prompt == config.system_prompt
+    assert gpt4all_with_config.config.model == config.model
+
+    assert isinstance(gpt4all_with_config.instance, LangchainGPT4All)
+
+
+def test_gpt4all_init_without_config(gpt4all_without_config):
+    assert gpt4all_without_config.config.model == "orca-mini-3b.ggmlv3.q4_0.bin"
+    assert isinstance(gpt4all_without_config.instance, LangchainGPT4All)
+
+
+def test_get_llm_model_answer(mocker, gpt4all_with_config):
+    test_query = "Test query"
+    test_answer = "Test answer"
+
+    mocked_get_answer = mocker.patch("embedchain.llm.gpt4all.GPT4ALLLlm._get_answer", return_value=test_answer)
+    answer = gpt4all_with_config.get_llm_model_answer(test_query)
+
+    assert answer == test_answer
+    mocked_get_answer.assert_called_once_with(prompt=test_query, config=gpt4all_with_config.config)
+
+
+def test_gpt4all_model_switching(gpt4all_with_config):
+    with pytest.raises(RuntimeError, match="GPT4ALLLlm does not support switching models at runtime."):
+        gpt4all_with_config._get_answer("Test prompt", BaseLlmConfig(model="new_model"))

tests/loaders/test_gmail.py

@@ -0,0 +1,44 @@
+import pytest
+from llama_hub.readwise.base import Document
+
+from embedchain.loaders.gmail import GmailLoader
+
+
+@pytest.fixture
+def mock_quopri(mocker):
+    return mocker.patch("embedchain.loaders.gmail.quopri.decodestring", return_value=b"your_test_decoded_string")
+
+
+@pytest.fixture
+def mock_beautifulsoup(mocker):
+    return mocker.patch("embedchain.loaders.gmail.BeautifulSoup", return_value=mocker.MagicMock())
+
+
+@pytest.fixture
+def gmail_loader(mock_quopri, mock_beautifulsoup):
+    return GmailLoader()
+
+
+def test_load_data_file_not_found(gmail_loader, mocker):
+    with pytest.raises(FileNotFoundError):
+        with mocker.patch("os.path.isfile", return_value=False):
+            gmail_loader.load_data("your_query")
+
+
+@pytest.mark.skip(reason="TODO: Fix this test. Failing due to some googleapiclient import issue.")
+def test_load_data(gmail_loader, mocker):
+    mock_gmail_reader_instance = mocker.MagicMock()
+    text = "your_test_email_text"
+    metadata = {
+        "id": "your_test_id",
+        "snippet": "your_test_snippet",
+    }
+    mock_gmail_reader_instance.load_data.return_value = [Document(text=text, extra_info=metadata)]
+
+    with mocker.patch("os.path.isfile", return_value=True):
+        response_data = gmail_loader.load_data("your_query")
+
+    assert "doc_id" in response_data
+    assert "data" in response_data
+    assert isinstance(response_data["doc_id"], str)
+    assert isinstance(response_data["data"], list)

tests/loaders/test_json.py

@@ -1,32 +1,91 @@
 import hashlib
-from unittest.mock import patch
 
-from langchain.docstore.document import Document
-from langchain.document_loaders.json_loader import \
-    JSONLoader as LangchainJSONLoader
+import pytest
+from llama_index.readers.schema.base import Document
 
 from embedchain.loaders.json import JSONLoader
 
 
-def test_load_data():
-    mock_document = [
-        Document(page_content="content1", metadata={"seq_num": 1}),
-        Document(page_content="content2", metadata={"seq_num": 2}),
+def test_load_data(mocker):
+    content = "temp.json"
+
+    mock_document = {
+        "doc_id": hashlib.sha256((content + ", ".join(["content1", "content2"])).encode()).hexdigest(),
+        "data": [
+            {"content": "content1", "meta_data": {"url": content}},
+            {"content": "content2", "meta_data": {"url": content}},
+        ],
+    }
+
+    mocker.patch("embedchain.loaders.json.JSONLoader.load_data", return_value=mock_document)
+
+    json_loader = JSONLoader()
+
+    result = json_loader.load_data(content)
+
+    assert "doc_id" in result
+    assert "data" in result
+
+    expected_data = [
+        {"content": "content1", "meta_data": {"url": content}},
+        {"content": "content2", "meta_data": {"url": content}},
     ]
-    with patch.object(LangchainJSONLoader, "load", return_value=mock_document):
-        content = "temp.json"
 
-        result = JSONLoader.load_data(content)
+    assert result["data"] == expected_data
 
-        assert "doc_id" in result
-        assert "data" in result
+    expected_doc_id = hashlib.sha256((content + ", ".join(["content1", "content2"])).encode()).hexdigest()
+    assert result["doc_id"] == expected_doc_id
 
-        expected_data = [
-            {"content": "content1", "meta_data": {"url": content, "row": 1}},
-            {"content": "content2", "meta_data": {"url": content, "row": 2}},
-        ]
 
-        assert result["data"] == expected_data
+def test_load_data_url(mocker):
+    content = "https://example.com/posts.json"
 
-        expected_doc_id = hashlib.sha256((content + ", ".join(["content1", "content2"])).encode()).hexdigest()
-        assert result["doc_id"] == expected_doc_id
+    mocker.patch("os.path.isfile", return_value=False)  # Mocking os.path.isfile to simulate a URL case
+    mocker.patch(
+        "llama_hub.jsondata.base.JSONDataReader.load_data",
+        return_value=[Document(text="content1"), Document(text="content2")],
+    )
+
+    mock_response = mocker.Mock()
+    mock_response.status_code = 200
+    mock_response.json.return_value = {"document1": "content1", "document2": "content2"}
+
+    mocker.patch("requests.get", return_value=mock_response)
+
+    result = JSONLoader.load_data(content)
+
+    assert "doc_id" in result
+    assert "data" in result
+
+    expected_data = [
+        {"content": "content1", "meta_data": {"url": content}},
+        {"content": "content2", "meta_data": {"url": content}},
+    ]
+
+    assert result["data"] == expected_data
+
+    expected_doc_id = hashlib.sha256((content + ", ".join(["content1", "content2"])).encode()).hexdigest()
+    assert result["doc_id"] == expected_doc_id
+
+
+def test_load_data_invalid_content(mocker):
+    mocker.patch("os.path.isfile", return_value=False)
+    mocker.patch("requests.get")
+
+    content = "123"
+
+    with pytest.raises(ValueError, match="Invalid content to load json data from"):
+        JSONLoader.load_data(content)
+
+
+def test_load_data_invalid_url(mocker):
+    mocker.patch("os.path.isfile", return_value=False)
+
+    mock_response = mocker.Mock()
+    mock_response.status_code = 404
+    mocker.patch("requests.get", return_value=mock_response)
+
+    content = "http://invalid-url.com/"
+
+    with pytest.raises(ValueError, match=f"Invalid content to load json data from: {content}"):
+        JSONLoader.load_data(content)

tests/loaders/test_openapi.py

@@ -0,0 +1,26 @@
+import pytest
+
+from embedchain.loaders.openapi import OpenAPILoader
+
+
+@pytest.fixture
+def openapi_loader():
+    return OpenAPILoader()
+
+
+def test_load_data(openapi_loader, mocker):
+    mocker.patch("builtins.open", mocker.mock_open(read_data="key1: value1\nkey2: value2"))
+
+    mocker.patch("hashlib.sha256", return_value=mocker.Mock(hexdigest=lambda: "mock_hash"))
+
+    file_path = "configs/openai_openapi.yaml"
+    result = openapi_loader.load_data(file_path)
+
+    expected_doc_id = "mock_hash"
+    expected_data = [
+        {"content": "key1: value1", "meta_data": {"url": file_path, "row": 1}},
+        {"content": "key2: value2", "meta_data": {"url": file_path, "row": 2}},
+    ]
+
+    assert result["doc_id"] == expected_doc_id
+    assert result["data"] == expected_data

tests/telemetry/test_posthog.py

@@ -0,0 +1,63 @@
+import logging
+import os
+
+from embedchain.telemetry.posthog import AnonymousTelemetry
+
+
+class TestAnonymousTelemetry:
+    def test_init(self, mocker):
+        # Enable telemetry specifically for this test
+        os.environ["EC_TELEMETRY"] = "true"
+        mock_posthog = mocker.patch("embedchain.telemetry.posthog.Posthog")
+        telemetry = AnonymousTelemetry()
+        assert telemetry.project_api_key == "phc_PHQDA5KwztijnSojsxJ2c1DuJd52QCzJzT2xnSGvjN2"
+        assert telemetry.host == "https://app.posthog.com"
+        assert telemetry.enabled is True
+        assert telemetry.user_id
+        mock_posthog.assert_called_once_with(project_api_key=telemetry.project_api_key, host=telemetry.host)
+
+    def test_init_with_disabled_telemetry(self, mocker, monkeypatch):
+        mocker.patch("embedchain.telemetry.posthog.Posthog")
+        telemetry = AnonymousTelemetry()
+        assert telemetry.enabled is False
+        assert telemetry.posthog.disabled is True
+
+    def test_get_user_id(self, mocker, tmpdir):
+        mock_uuid = mocker.patch("embedchain.telemetry.posthog.uuid.uuid4")
+        mock_uuid.return_value = "unique_user_id"
+        config_file = tmpdir.join("config.json")
+        mocker.patch("embedchain.telemetry.posthog.CONFIG_FILE", str(config_file))
+        telemetry = AnonymousTelemetry()
+
+        user_id = telemetry.get_user_id()
+        assert user_id == "unique_user_id"
+        assert config_file.read() == '{"user_id": "unique_user_id"}'
+
+    def test_capture(self, mocker):
+        # Enable telemetry specifically for this test
+        os.environ["EC_TELEMETRY"] = "true"
+        mock_posthog = mocker.patch("embedchain.telemetry.posthog.Posthog")
+        telemetry = AnonymousTelemetry()
+        event_name = "test_event"
+        properties = {"key": "value"}
+        telemetry.capture(event_name, properties)
+
+        mock_posthog.assert_called_once_with(
+            project_api_key=telemetry.project_api_key,
+            host=telemetry.host,
+        )
+        mock_posthog.return_value.capture.assert_called_once_with(
+            telemetry.user_id,
+            event_name,
+            properties,
+        )
+
+    def test_capture_with_exception(self, mocker, caplog):
+        mock_posthog = mocker.patch("embedchain.telemetry.posthog.Posthog")
+        mock_posthog.return_value.capture.side_effect = Exception("Test Exception")
+        telemetry = AnonymousTelemetry()
+        event_name = "test_event"
+        properties = {"key": "value"}
+        with caplog.at_level(logging.ERROR):
+            telemetry.capture(event_name, properties)
+        assert "Failed to send telemetry event" in caplog.text

tests/test_client.py

@@ -0,0 +1,53 @@
+import pytest
+
+from embedchain import Client
+
+
+class TestClient:
+    @pytest.fixture
+    def mock_requests_post(self, mocker):
+        return mocker.patch("embedchain.client.requests.post")
+
+    def test_valid_api_key(self, mock_requests_post):
+        mock_requests_post.return_value.status_code = 200
+        client = Client(api_key="valid_api_key")
+        assert client.check("valid_api_key") is True
+
+    def test_invalid_api_key(self, mock_requests_post):
+        mock_requests_post.return_value.status_code = 401
+        with pytest.raises(ValueError):
+            Client(api_key="invalid_api_key")
+
+    def test_update_valid_api_key(self, mock_requests_post):
+        mock_requests_post.return_value.status_code = 200
+        client = Client(api_key="valid_api_key")
+        client.update("new_valid_api_key")
+        assert client.get() == "new_valid_api_key"
+
+    def test_clear_api_key(self, mock_requests_post):
+        mock_requests_post.return_value.status_code = 200
+        client = Client(api_key="valid_api_key")
+        client.clear()
+        assert client.get() is None
+
+    def test_save_api_key(self, mock_requests_post):
+        mock_requests_post.return_value.status_code = 200
+        api_key_to_save = "valid_api_key"
+        client = Client(api_key=api_key_to_save)
+        client.save()
+        assert client.get() == api_key_to_save
+
+    def test_load_api_key_from_config(self, mocker):
+        mocker.patch("embedchain.Client.load_config", return_value={"api_key": "test_api_key"})
+        client = Client()
+        assert client.get() == "test_api_key"
+
+    def test_load_invalid_api_key_from_config(self, mocker):
+        mocker.patch("embedchain.Client.load_config", return_value={})
+        with pytest.raises(ValueError):
+            Client()
+
+    def test_load_missing_api_key_from_config(self, mocker):
+        mocker.patch("embedchain.Client.load_config", return_value={})
+        with pytest.raises(ValueError):
+            Client()

tests/vectordb/test_chroma_db.py

@@ -1,366 +1,343 @@
-# ruff: noqa: E501
-
-import unittest
+import os
+import shutil
 from unittest.mock import patch
 
+import pytest
 from chromadb.config import Settings
 
 from embedchain import App
 from embedchain.config import AppConfig, ChromaDbConfig
 from embedchain.vectordb.chroma import ChromaDB
 
-
-class TestChromaDbHosts(unittest.TestCase):
-    def test_init_with_host_and_port(self):
-        """
-        Test if the `ChromaDB` instance is initialized with the correct host and port values.
-        """
-        host = "test-host"
-        port = "1234"
-        config = ChromaDbConfig(host=host, port=port)
-
-        db = ChromaDB(config=config)
-        settings = db.client.get_settings()
-        self.assertEqual(settings.chroma_server_host, host)
-        self.assertEqual(settings.chroma_server_http_port, port)
-
-    def test_init_with_basic_auth(self):
-        host = "test-host"
-        port = "1234"
-
-        chroma_config = {
-            "host": host,
-            "port": port,
-            "chroma_settings": {
-                "chroma_client_auth_provider": "chromadb.auth.basic.BasicAuthClientProvider",
-                "chroma_client_auth_credentials": "admin:admin",
-            },
-        }
-
-        config = ChromaDbConfig(**chroma_config)
-        db = ChromaDB(config=config)
-        settings = db.client.get_settings()
-        self.assertEqual(settings.chroma_server_host, host)
-        self.assertEqual(settings.chroma_server_http_port, port)
-        self.assertEqual(
-            settings.chroma_client_auth_provider, chroma_config["chroma_settings"]["chroma_client_auth_provider"]
-        )
-        self.assertEqual(
-            settings.chroma_client_auth_credentials, chroma_config["chroma_settings"]["chroma_client_auth_credentials"]
-        )
+os.environ["OPENAI_API_KEY"] = "test-api-key"
 
 
-# Review this test
-class TestChromaDbHostsInit(unittest.TestCase):
-    @patch("embedchain.vectordb.chroma.chromadb.Client")
-    def test_app_init_with_host_and_port(self, mock_client):
-        """
-        Test if the `App` instance is initialized with the correct host and port values.
-        """
-        host = "test-host"
-        port = "1234"
-
-        config = AppConfig(collect_metrics=False)
-        db_config = ChromaDbConfig(host=host, port=port)
-
-        _app = App(config, db_config=db_config)
-
-        called_settings: Settings = mock_client.call_args[0][0]
-
-        self.assertEqual(called_settings.chroma_server_host, host)
-        self.assertEqual(called_settings.chroma_server_http_port, port)
+@pytest.fixture
+def chroma_db():
+    return ChromaDB(config=ChromaDbConfig(host="test-host", port="1234"))
 
 
-class TestChromaDbHostsNone(unittest.TestCase):
-    @patch("embedchain.vectordb.chroma.chromadb.Client")
-    def test_init_with_host_and_port_none(self, mock_client):
-        """
-        Test if the `App` instance is initialized without default hosts and ports.
-        """
-
-        _app = App(config=AppConfig(collect_metrics=False))
-
-        called_settings: Settings = mock_client.call_args[0][0]
-        self.assertEqual(called_settings.chroma_server_host, None)
-        self.assertEqual(called_settings.chroma_server_http_port, None)
+@pytest.fixture
+def app_with_settings():
+    chroma_config = ChromaDbConfig(allow_reset=True, dir="test-db")
+    app_config = AppConfig(collect_metrics=False)
+    return App(config=app_config, db_config=chroma_config)
 
 
-class TestChromaDbHostsLoglevel(unittest.TestCase):
-    @patch("embedchain.vectordb.chroma.chromadb.Client")
-    def test_init_with_host_and_port_log_level(self, mock_client):
-        """
-        Test if the `App` instance is initialized without a config that does not contain default hosts and ports.
-        """
-
-        _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)
+@pytest.fixture(scope="session", autouse=True)
+def cleanup_db():
+    yield
+    try:
+        shutil.rmtree("test-db")
+    except OSError as e:
+        print("Error: %s - %s." % (e.filename, e.strerror))
 
 
-class TestChromaDbDuplicateHandling:
-    chroma_config = ChromaDbConfig(allow_reset=True)
-    app_config = AppConfig(collection_name=False, collect_metrics=False)
-    app_with_settings = App(config=app_config, db_config=chroma_config)
-
-    def test_duplicates_throw_warning(self, caplog):
-        """
-        Test that add duplicates throws an error.
-        """
-        # Start with a clean app
-        self.app_with_settings.reset()
-
-        app = App(config=AppConfig(collect_metrics=False))
-        app.db.collection.add(embeddings=[[0, 0, 0]], ids=["0"])
-        app.db.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
-        self.app_with_settings.reset()
-
-        app = App(config=AppConfig(collect_metrics=False))
-        app.set_collection_name("test_collection_1")
-        app.db.collection.add(embeddings=[[0, 0, 0]], ids=["0"])
-        app.set_collection_name("test_collection_2")
-        app.db.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
+def test_chroma_db_init_with_host_and_port(chroma_db):
+    settings = chroma_db.client.get_settings()
+    assert settings.chroma_server_host == "test-host"
+    assert settings.chroma_server_http_port == "1234"
 
 
-class TestChromaDbCollection(unittest.TestCase):
-    chroma_config = ChromaDbConfig(allow_reset=True)
-    app_config = AppConfig(collection_name=False, collect_metrics=False)
-    app_with_settings = App(config=app_config, db_config=chroma_config)
+def test_chroma_db_init_with_basic_auth():
+    chroma_config = {
+        "host": "test-host",
+        "port": "1234",
+        "chroma_settings": {
+            "chroma_client_auth_provider": "chromadb.auth.basic.BasicAuthClientProvider",
+            "chroma_client_auth_credentials": "admin:admin",
+        },
+    }
 
-    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))
+    db = ChromaDB(config=ChromaDbConfig(**chroma_config))
+    settings = db.client.get_settings()
+    assert settings.chroma_server_host == "test-host"
+    assert settings.chroma_server_http_port == "1234"
+    assert settings.chroma_client_auth_provider == chroma_config["chroma_settings"]["chroma_client_auth_provider"]
+    assert settings.chroma_client_auth_credentials == chroma_config["chroma_settings"]["chroma_client_auth_credentials"]
 
-        self.assertEqual(app.db.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(collect_metrics=False)
-        app = App(config=config)
-        app.set_collection_name(name="test_collection")
+@patch("embedchain.vectordb.chroma.chromadb.Client")
+def test_app_init_with_host_and_port(mock_client):
+    host = "test-host"
+    port = "1234"
+    config = AppConfig(collect_metrics=False)
+    db_config = ChromaDbConfig(host=host, port=port)
+    _app = App(config, db_config=db_config)
 
-        self.assertEqual(app.db.collection.name, "test_collection")
+    called_settings: Settings = mock_client.call_args[0][0]
+    assert called_settings.chroma_server_host == host
+    assert called_settings.chroma_server_http_port == port
 
-    def test_set_collection_name(self):
-        """
-        Test if the `App` collection is correctly switched using the `set_collection_name` method.
-        """
-        app = App(config=AppConfig(collect_metrics=False))
-        app.set_collection_name("test_collection")
 
-        self.assertEqual(app.db.collection.name, "test_collection")
+@patch("embedchain.vectordb.chroma.chromadb.Client")
+def test_app_init_with_host_and_port_none(mock_client):
+    _app = App(config=AppConfig(collect_metrics=False), db_config=ChromaDbConfig(allow_reset=True, dir="test-db"))
 
-    def test_changes_encapsulated(self):
-        """
-        Test that changes to one collection do not affect the other collection
-        """
-        # Start with a clean app
-        self.app_with_settings.reset()
+    called_settings: Settings = mock_client.call_args[0][0]
+    assert called_settings.chroma_server_host is None
+    assert called_settings.chroma_server_http_port is None
 
-        app = App(config=AppConfig(collect_metrics=False))
-        app.set_collection_name("test_collection_1")
-        # Collection should be empty when created
-        self.assertEqual(app.db.count(), 0)
 
-        app.db.collection.add(embeddings=[0, 0, 0], ids=["0"])
-        # After adding, should contain one item
-        self.assertEqual(app.db.count(), 1)
+def test_chroma_db_duplicates_throw_warning(caplog):
+    app = App(config=AppConfig(collect_metrics=False), db_config=ChromaDbConfig(allow_reset=True, dir="test-db"))
+    app.db.collection.add(embeddings=[[0, 0, 0]], ids=["0"])
+    app.db.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
+    app.db.reset()
 
-        app.set_collection_name("test_collection_2")
-        # New collection is empty
-        self.assertEqual(app.db.count(), 0)
 
-        # Adding to new collection should not effect existing collection
-        app.db.collection.add(embeddings=[0, 0, 0], ids=["0"])
-        app.set_collection_name("test_collection_1")
-        # Should still be 1, not 2.
-        self.assertEqual(app.db.count(), 1)
+def test_chroma_db_duplicates_collections_no_warning(caplog):
+    app = App(config=AppConfig(collect_metrics=False), db_config=ChromaDbConfig(allow_reset=True, dir="test-db"))
+    app.set_collection_name("test_collection_1")
+    app.db.collection.add(embeddings=[[0, 0, 0]], ids=["0"])
+    app.set_collection_name("test_collection_2")
+    app.db.collection.add(embeddings=[[0, 0, 0]], ids=["0"])
+    assert "Insert of existing embedding ID: 0" not in caplog.text
+    assert "Add of existing embedding ID: 0" not in caplog.text
+    app.db.reset()
+    app.set_collection_name("test_collection_1")
+    app.db.reset()
 
-    def test_add_with_skip_embedding(self):
-        """
-        Test that changes to one collection do not affect the other collection
-        """
-        # Start with a clean app
-        self.app_with_settings.reset()
-        # app = App(config=AppConfig(collect_metrics=False), db=db)
 
-        # Collection should be empty when created
-        self.assertEqual(self.app_with_settings.db.count(), 0)
+def test_chroma_db_collection_init_with_default_collection():
+    app = App(config=AppConfig(collect_metrics=False), db_config=ChromaDbConfig(allow_reset=True, dir="test-db"))
+    assert app.db.collection.name == "embedchain_store"
 
-        self.app_with_settings.db.add(
+
+def test_chroma_db_collection_init_with_custom_collection():
+    app = App(config=AppConfig(collect_metrics=False), db_config=ChromaDbConfig(allow_reset=True, dir="test-db"))
+    app.set_collection_name(name="test_collection")
+    assert app.db.collection.name == "test_collection"
+
+
+def test_chroma_db_collection_set_collection_name():
+    app = App(config=AppConfig(collect_metrics=False), db_config=ChromaDbConfig(allow_reset=True, dir="test-db"))
+    app.set_collection_name("test_collection")
+    assert app.db.collection.name == "test_collection"
+
+
+def test_chroma_db_collection_changes_encapsulated():
+    app = App(config=AppConfig(collect_metrics=False), db_config=ChromaDbConfig(allow_reset=True, dir="test-db"))
+    app.set_collection_name("test_collection_1")
+    assert app.db.count() == 0
+
+    app.db.collection.add(embeddings=[0, 0, 0], ids=["0"])
+    assert app.db.count() == 1
+
+    app.set_collection_name("test_collection_2")
+    assert app.db.count() == 0
+
+    app.db.collection.add(embeddings=[0, 0, 0], ids=["0"])
+    app.set_collection_name("test_collection_1")
+    assert app.db.count() == 1
+    app.db.reset()
+    app.set_collection_name("test_collection_2")
+    app.db.reset()
+
+
+def test_chroma_db_collection_add_with_skip_embedding(app_with_settings):
+    # Start with a clean app
+    app_with_settings.db.reset()
+
+    assert app_with_settings.db.count() == 0
+
+    app_with_settings.db.add(
+        embeddings=[[0, 0, 0]],
+        documents=["document"],
+        metadatas=[{"url": "url_1", "doc_id": "doc_id_1"}],
+        ids=["id"],
+        skip_embedding=True,
+    )
+
+    assert app_with_settings.db.count() == 1
+
+    data = app_with_settings.db.get(["id"], limit=1)
+    expected_value = {
+        "documents": ["document"],
+        "embeddings": None,
+        "ids": ["id"],
+        "metadatas": [{"url": "url_1", "doc_id": "doc_id_1"}],
+    }
+
+    assert data == expected_value
+
+    data_without_citations = app_with_settings.db.query(
+        input_query=[0, 0, 0], where={}, n_results=1, skip_embedding=True
+    )
+    expected_value_without_citations = ["document"]
+    assert data_without_citations == expected_value_without_citations
+
+    app_with_settings.db.reset()
+
+
+def test_chroma_db_collection_add_with_invalid_inputs(app_with_settings):
+    # Start with a clean app
+    app_with_settings.db.reset()
+
+    assert app_with_settings.db.count() == 0
+
+    with pytest.raises(ValueError):
+        app_with_settings.db.add(
             embeddings=[[0, 0, 0]],
-            documents=["document"],
+            documents=["document", "document2"],
             metadatas=[{"value": "somevalue"}],
             ids=["id"],
             skip_embedding=True,
         )
-        # After adding, should contain one item
-        self.assertEqual(self.app_with_settings.db.count(), 1)
 
-        # Validate if the get utility of the database is working as expected
-        data = self.app_with_settings.db.get(["id"], limit=1)
-        expected_value = {
-            "documents": ["document"],
-            "embeddings": None,
-            "ids": ["id"],
-            "metadatas": [{"value": "somevalue"}],
-        }
-        self.assertEqual(data, expected_value)
+    assert app_with_settings.db.count() == 0
 
-        # Validate if the query utility of the database is working as expected
-        data = self.app_with_settings.db.query(input_query=[0, 0, 0], where={}, n_results=1, skip_embedding=True)
-        expected_value = ["document"]
-        self.assertEqual(data, expected_value)
+    with pytest.raises(ValueError):
+        app_with_settings.db.add(
+            embeddings=None,
+            documents=["document", "document2"],
+            metadatas=[{"value": "somevalue"}],
+            ids=["id"],
+            skip_embedding=True,
+        )
 
-    def test_add_with_invalid_inputs(self):
-        """
-        Test add fails with invalid inputs
-        """
-        # Start with a clean app
-        self.app_with_settings.reset()
-        # app = App(config=AppConfig(collect_metrics=False), db=db)
+    assert app_with_settings.db.count() == 0
+    app_with_settings.db.reset()
 
-        # Collection should be empty when created
-        self.assertEqual(self.app_with_settings.db.count(), 0)
 
-        with self.assertRaises(ValueError):
-            self.app_with_settings.db.add(
-                embeddings=[[0, 0, 0]],
-                documents=["document", "document2"],
-                metadatas=[{"value": "somevalue"}],
-                ids=["id"],
-                skip_embedding=True,
-            )
-        # After adding, should contain no item
-        self.assertEqual(self.app_with_settings.db.count(), 0)
+def test_chroma_db_collection_collections_are_persistent():
+    app = App(config=AppConfig(collect_metrics=False), db_config=ChromaDbConfig(allow_reset=True, dir="test-db"))
+    app.set_collection_name("test_collection_1")
+    app.db.collection.add(embeddings=[[0, 0, 0]], ids=["0"])
+    del app
 
-        with self.assertRaises(ValueError):
-            self.app_with_settings.db.add(
-                embeddings=None,
-                documents=["document", "document2"],
-                metadatas=[{"value": "somevalue"}],
-                ids=["id"],
-                skip_embedding=True,
-            )
+    app = App(config=AppConfig(collect_metrics=False), db_config=ChromaDbConfig(allow_reset=True, dir="test-db"))
+    app.set_collection_name("test_collection_1")
+    assert app.db.count() == 1
 
-        # After adding, should contain no item
-        self.assertEqual(self.app_with_settings.db.count(), 0)
+    app.db.reset()
 
-    def test_collections_are_persistent(self):
-        """
-        Test that a collection can be picked up later.
-        """
-        # Start with a clean app
-        self.app_with_settings.reset()
 
-        app = App(config=AppConfig(collect_metrics=False))
-        app.set_collection_name("test_collection_1")
-        app.db.collection.add(embeddings=[[0, 0, 0]], ids=["0"])
-        del app
+def test_chroma_db_collection_parallel_collections():
+    app1 = App(
+        AppConfig(collection_name="test_collection_1", collect_metrics=False),
+        db_config=ChromaDbConfig(allow_reset=True, dir="test-db"),
+    )
+    app2 = App(
+        AppConfig(collection_name="test_collection_2", collect_metrics=False),
+        db_config=ChromaDbConfig(allow_reset=True, dir="test-db"),
+    )
 
-        app = App(config=AppConfig(collect_metrics=False))
-        app.set_collection_name("test_collection_1")
-        self.assertEqual(app.db.count(), 1)
+    # cleanup if any previous tests failed or were interrupted
+    app1.db.reset()
+    app2.db.reset()
 
-    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
-        self.app_with_settings.reset()
+    app1.db.collection.add(embeddings=[0, 0, 0], ids=["0"])
+    assert app1.db.count() == 1
+    assert app2.db.count() == 0
 
-        # 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))
+    app1.db.collection.add(embeddings=[[0, 0, 0], [1, 1, 1]], ids=["1", "2"])
+    app2.db.collection.add(embeddings=[0, 0, 0], ids=["0"])
 
-        # app2 has been created last, but adding to app1 will still write to collection 1.
-        app1.db.collection.add(embeddings=[0, 0, 0], ids=["0"])
-        self.assertEqual(app1.db.count(), 1)
-        self.assertEqual(app2.db.count(), 0)
+    app1.set_collection_name("test_collection_2")
+    assert app1.db.count() == 1
+    app2.set_collection_name("test_collection_1")
+    assert app2.db.count() == 3
 
-        # Add data
-        app1.db.collection.add(embeddings=[[0, 0, 0], [1, 1, 1]], ids=["1", "2"])
-        app2.db.collection.add(embeddings=[0, 0, 0], ids=["0"])
+    # cleanup
+    app1.db.reset()
+    app2.db.reset()
 
-        # Swap names and test
-        app1.set_collection_name("test_collection_2")
-        self.assertEqual(app1.count(), 1)
-        app2.set_collection_name("test_collection_1")
-        self.assertEqual(app2.count(), 3)
 
-    def test_ids_share_collections(self):
-        """
-        Different ids should still share collections.
-        """
-        # Start clean
-        self.app_with_settings.reset()
+def test_chroma_db_collection_ids_share_collections():
+    app1 = App(
+        AppConfig(id="new_app_id_1", collect_metrics=False), db_config=ChromaDbConfig(allow_reset=True, dir="test-db")
+    )
+    app1.set_collection_name("one_collection")
+    app2 = App(
+        AppConfig(id="new_app_id_2", collect_metrics=False), db_config=ChromaDbConfig(allow_reset=True, dir="test-db")
+    )
+    app2.set_collection_name("one_collection")
 
-        # Create two apps
-        app1 = App(AppConfig(id="new_app_id_1", collect_metrics=False))
-        app1.set_collection_name("one_collection")
-        app2 = App(AppConfig(id="new_app_id_2", collect_metrics=False))
-        app2.set_collection_name("one_collection")
+    app1.db.collection.add(embeddings=[[0, 0, 0], [1, 1, 1]], ids=["0", "1"])
+    app2.db.collection.add(embeddings=[0, 0, 0], ids=["2"])
 
-        # Add data
-        app1.db.collection.add(embeddings=[[0, 0, 0], [1, 1, 1]], ids=["0", "1"])
-        app2.db.collection.add(embeddings=[0, 0, 0], ids=["2"])
+    assert app1.db.count() == 3
+    assert app2.db.count() == 3
 
-        # Both should have the same collection
-        self.assertEqual(app1.count(), 3)
-        self.assertEqual(app2.count(), 3)
+    # cleanup
+    app1.db.reset()
+    app2.db.reset()
 
-    def test_reset(self):
-        """
-        Resetting should hit all collections and ids.
-        """
-        # Start clean
-        self.app_with_settings.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(id="new_app_id_1", collect_metrics=False), db_config=self.chroma_config)
-        app1.set_collection_name("one_collection")
-        app2 = App(AppConfig(id="new_app_id_2", collect_metrics=False))
-        app2.set_collection_name("one_collection")
-        app3 = App(AppConfig(id="new_app_id_1", collect_metrics=False))
-        app3.set_collection_name("three_collection")
-        app4 = App(AppConfig(id="new_app_id_4", collect_metrics=False))
-        app4.set_collection_name("four_collection")
+def test_chroma_db_collection_reset():
+    app1 = App(
+        AppConfig(id="new_app_id_1", collect_metrics=False), db_config=ChromaDbConfig(allow_reset=True, dir="test-db")
+    )
+    app1.set_collection_name("one_collection")
+    app2 = App(
+        AppConfig(id="new_app_id_2", collect_metrics=False), db_config=ChromaDbConfig(allow_reset=True, dir="test-db")
+    )
+    app2.set_collection_name("two_collection")
+    app3 = App(
+        AppConfig(id="new_app_id_1", collect_metrics=False), db_config=ChromaDbConfig(allow_reset=True, dir="test-db")
+    )
+    app3.set_collection_name("three_collection")
+    app4 = App(
+        AppConfig(id="new_app_id_4", collect_metrics=False), db_config=ChromaDbConfig(allow_reset=True, dir="test-db")
+    )
+    app4.set_collection_name("four_collection")
 
-        # Each one of them get data
-        app1.db.collection.add(embeddings=[0, 0, 0], ids=["1"])
-        app2.db.collection.add(embeddings=[0, 0, 0], ids=["2"])
-        app3.db.collection.add(embeddings=[0, 0, 0], ids=["3"])
-        app4.db.collection.add(embeddings=[0, 0, 0], ids=["4"])
+    app1.db.collection.add(embeddings=[0, 0, 0], ids=["1"])
+    app2.db.collection.add(embeddings=[0, 0, 0], ids=["2"])
+    app3.db.collection.add(embeddings=[0, 0, 0], ids=["3"])
+    app4.db.collection.add(embeddings=[0, 0, 0], ids=["4"])
 
-        # Resetting the first one should reset them all.
-        app1.reset()
+    app1.db.reset()
 
-        # Reinstantiate app2-4, app1 doesn't have to be reinstantiated (PR #319)
-        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))
+    assert app1.db.count() == 0
+    assert app2.db.count() == 1
+    assert app3.db.count() == 1
+    assert app4.db.count() == 1
 
-        # All should be empty
-        self.assertEqual(app1.count(), 0)
-        self.assertEqual(app2.count(), 0)
-        self.assertEqual(app3.count(), 0)
-        self.assertEqual(app4.count(), 0)
+    # cleanup
+    app2.db.reset()
+    app3.db.reset()
+    app4.db.reset()
+
+
+def test_chroma_db_collection_query(app_with_settings):
+    app_with_settings.db.reset()
+
+    assert app_with_settings.db.count() == 0
+
+    app_with_settings.db.add(
+        embeddings=[[0, 0, 0]],
+        documents=["document"],
+        metadatas=[{"url": "url_1", "doc_id": "doc_id_1"}],
+        ids=["id"],
+        skip_embedding=True,
+    )
+
+    assert app_with_settings.db.count() == 1
+
+    app_with_settings.db.add(
+        embeddings=[[0, 1, 0]],
+        documents=["document2"],
+        metadatas=[{"url": "url_2", "doc_id": "doc_id_2"}],
+        ids=["id2"],
+        skip_embedding=True,
+    )
+
+    assert app_with_settings.db.count() == 2
+
+    data_without_citations = app_with_settings.db.query(
+        input_query=[0, 0, 0], where={}, n_results=2, skip_embedding=True
+    )
+    expected_value_without_citations = ["document", "document2"]
+    assert data_without_citations == expected_value_without_citations
+
+    data_with_citations = app_with_settings.db.query(
+        input_query=[0, 0, 0], where={}, n_results=2, skip_embedding=True, citations=True
+    )
+    expected_value_with_citations = [("document", "url_1", "doc_id_1"), ("document2", "url_2", "doc_id_2")]
+    assert data_with_citations == expected_value_with_citations
+
+    app_with_settings.db.reset()

tests/vectordb/test_elasticsearch_db.py

@@ -31,7 +31,7 @@ class TestEsDB(unittest.TestCase):
         # Create some dummy data.
         embeddings = [[1, 2, 3], [4, 5, 6]]
         documents = ["This is a document.", "This is another document."]
-        metadatas = [{}, {}]
+        metadatas = [{"url": "url_1", "doc_id": "doc_id_1"}, {"url": "url_2", "doc_id": "doc_id_2"}]
         ids = ["doc_1", "doc_2"]
 
         # Add the data to the database.
@@ -40,8 +40,17 @@ class TestEsDB(unittest.TestCase):
         search_response = {
             "hits": {
                 "hits": [
-                    {"_source": {"text": "This is a document."}, "_score": 0.9},
-                    {"_source": {"text": "This is another document."}, "_score": 0.8},
+                    {
+                        "_source": {"text": "This is a document.", "metadata": {"url": "url_1", "doc_id": "doc_id_1"}},
+                        "_score": 0.9,
+                    },
+                    {
+                        "_source": {
+                            "text": "This is another document.",
+                            "metadata": {"url": "url_2", "doc_id": "doc_id_2"},
+                        },
+                        "_score": 0.8,
+                    },
                 ]
             }
         }
@@ -51,10 +60,16 @@ class TestEsDB(unittest.TestCase):
 
         # Query the database for the documents that are most similar to the query "This is a document".
         query = ["This is a document"]
-        results = self.db.query(query, n_results=2, where={}, skip_embedding=False)
+        results_without_citations = self.db.query(query, n_results=2, where={}, skip_embedding=False)
+        expected_results_without_citations = ["This is a document.", "This is another document."]
+        self.assertEqual(results_without_citations, expected_results_without_citations)
 
-        # Assert that the results are correct.
-        self.assertEqual(results, ["This is a document.", "This is another document."])
+        results_with_citations = self.db.query(query, n_results=2, where={}, skip_embedding=False, citations=True)
+        expected_results_with_citations = [
+            ("This is a document.", "url_1", "doc_id_1"),
+            ("This is another document.", "url_2", "doc_id_2"),
+        ]
+        self.assertEqual(results_with_citations, expected_results_with_citations)
 
     @patch("embedchain.vectordb.elasticsearch.Elasticsearch")
     def test_query_with_skip_embedding(self, mock_client):
@@ -68,7 +83,7 @@ class TestEsDB(unittest.TestCase):
         # Create some dummy data.
         embeddings = [[1, 2, 3], [4, 5, 6]]
         documents = ["This is a document.", "This is another document."]
-        metadatas = [{}, {}]
+        metadatas = [{"url": "url_1", "doc_id": "doc_id_1"}, {"url": "url_2", "doc_id": "doc_id_2"}]
         ids = ["doc_1", "doc_2"]
 
         # Add the data to the database.
@@ -77,8 +92,17 @@ class TestEsDB(unittest.TestCase):
         search_response = {
             "hits": {
                 "hits": [
-                    {"_source": {"text": "This is a document."}, "_score": 0.9},
-                    {"_source": {"text": "This is another document."}, "_score": 0.8},
+                    {
+                        "_source": {"text": "This is a document.", "metadata": {"url": "url_1", "doc_id": "doc_id_1"}},
+                        "_score": 0.9,
+                    },
+                    {
+                        "_source": {
+                            "text": "This is another document.",
+                            "metadata": {"url": "url_2", "doc_id": "doc_id_2"},
+                        },
+                        "_score": 0.8,
+                    },
                 ]
             }
         }