Bump version to 0.0.68 (#785)

.env.example

@@ -1 +1 @@
-OPENAI_API_KEY=
+OPENAI_API_KEY="your-openai-api-key"

README.md

@@ -36,7 +36,7 @@ The documentation for embedchain can be found at [docs.embedchain.ai](https://do
 
 ## 💻 Usage
 
-Embedchain empowers you to create chatbot models similar to ChatGPT, using your own evolving dataset.
+Embedchain empowers you to create ChatGPT like apps, on your own dynamic dataset.
 
 ### Data Types Supported
 
@@ -46,7 +46,9 @@ Embedchain empowers you to create chatbot models similar to ChatGPT, using your
 * Sitemap
 * Doc file
 * Code documentation website loader
-* Notion
+* Notion and many more.
+
+You can find the full list of data types on [our documentation](https://docs.embedchain.ai/data-sources/csv).
 
 ### Queries
 
@@ -87,7 +89,7 @@ If you utilize this repository, please consider citing it with:
 
 ```
 @misc{embedchain,
-  author = {Taranjeet Singh},
+  author = {Taranjeet Singh, Deshraj Yadav},
   title = {Embedchain: Framework to easily create LLM powered bots over any dataset},
   year = {2023},
   publisher = {GitHub},

docs/advanced/app_types.mdx

@@ -27,6 +27,8 @@ The following LLM providers are supported by Embedchain:
 - GPT4ALL
 - AZURE_OPENAI
 - LLAMA2
+- JINA
+- COHERE
 
 You can choose one by importing it from `embedchain.llm`. E.g.:
 

docs/data-sources/xml.mdx

@@ -0,0 +1,13 @@
+---
+title: 'XML File'
+---
+
+### XML file
+
+To add any xml file, use the data_type as `xml`. Eg:
+
+```python
+app.add('content/data.xml')
+```
+
+Note: Only the text content of the xml file will be added to the app. The tags will be ignored.

embedchain/__init__.py

@@ -6,6 +6,6 @@ from embedchain.apps.app import App  # noqa: F401
 from embedchain.apps.custom_app import CustomApp  # noqa: F401
 from embedchain.apps.Llama2App import Llama2App  # noqa: F401
 from embedchain.apps.open_source_app import OpenSourceApp  # noqa: F401
-from embedchain.apps.PersonApp import (PersonApp,  # noqa: F401
+from embedchain.apps.person_app import (PersonApp,  # noqa: F401
                                        PersonOpenSourceApp)
 from embedchain.vectordb.chroma import ChromaDB  # noqa: F401

embedchain/chunkers/sitemap.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 SitemapChunker(BaseChunker):
+    """Chunker for sitemap."""
+
+    def __init__(self, config: Optional[ChunkerConfig] = None):
+        if config is None:
+            config = ChunkerConfig(chunk_size=500, 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/xml.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 XmlChunker(BaseChunker):
+    """Chunker for XML files."""
+
+    def __init__(self, config: Optional[ChunkerConfig] = None):
+        if config is None:
+            config = ChunkerConfig(chunk_size=500, chunk_overlap=50, 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/data_formatter/data_formatter.py

@@ -6,9 +6,11 @@ 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.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
@@ -24,6 +26,7 @@ from embedchain.loaders.mdx import MdxLoader
 from embedchain.loaders.pdf_file import PdfFileLoader
 from embedchain.loaders.sitemap import SitemapLoader
 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
 
@@ -67,6 +70,7 @@ class DataFormatter(JSONSerializable):
             DataType.TEXT: LocalTextLoader,
             DataType.DOCX: DocxFileLoader,
             DataType.SITEMAP: SitemapLoader,
+            DataType.XML: XmlLoader,
             DataType.DOCS_SITE: DocsSiteLoader,
             DataType.CSV: CsvLoader,
             DataType.MDX: MdxLoader,
@@ -106,10 +110,12 @@ class DataFormatter(JSONSerializable):
             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,
         }
         if data_type in chunker_classes:
             chunker_class: type = chunker_classes[data_type]

embedchain/embedchain.py

@@ -352,7 +352,7 @@ class EmbedChain(JSONSerializable):
         # get existing ids, and discard doc if any common id exist.
         where = {"url": src}
         if self.config.id is not None:
-            where.update({"metadata.app_id": self.config.id})
+            where["app_id"] = self.config.id
 
         db_result = self.db.get(ids=ids, where=where)  # optional filter
         existing_ids = set(db_result["ids"])

embedchain/llm/cohere.py

@@ -0,0 +1,43 @@
+import importlib
+import os
+from typing import Optional
+
+from langchain.llms import Cohere
+
+from embedchain.config import BaseLlmConfig
+from embedchain.helper.json_serializable import register_deserializable
+from embedchain.llm.base import BaseLlm
+
+
+@register_deserializable
+class CohereLlm(BaseLlm):
+    def __init__(self, config: Optional[BaseLlmConfig] = None):
+        if "COHERE_API_KEY" not in os.environ:
+            raise ValueError("Please set the COHERE_API_KEY environment variable.")
+
+        try:
+            importlib.import_module("cohere")
+        except ModuleNotFoundError:
+            raise ModuleNotFoundError(
+                "The required dependencies for Cohere are not installed."
+                'Please install with `pip install --upgrade "embedchain[cohere]"`'
+            ) from None
+
+        super().__init__(config=config)
+
+    def get_llm_model_answer(self, prompt):
+        if self.config.system_prompt:
+            raise ValueError("CohereLlm does not support `system_prompt`")
+        return CohereLlm._get_answer(prompt=prompt, config=self.config)
+
+    @staticmethod
+    def _get_answer(prompt: str, config: BaseLlmConfig) -> str:
+        llm = Cohere(
+            cohere_api_key=os.environ["COHERE_API_KEY"],
+            model=config.model,
+            max_tokens=config.max_tokens,
+            temperature=config.temperature,
+            p=config.top_p,
+        )
+
+        return llm(prompt)

embedchain/llm/hugging_face_hub.py

@@ -0,0 +1,51 @@
+import importlib
+import os
+from typing import Optional
+
+from langchain.llms import HuggingFaceHub
+
+from embedchain.config import BaseLlmConfig
+from embedchain.helper.json_serializable import register_deserializable
+from embedchain.llm.base import BaseLlm
+
+
+@register_deserializable
+class HuggingFaceHubLlm(BaseLlm):
+    def __init__(self, config: Optional[BaseLlmConfig] = None):
+        if "HUGGINGFACEHUB_ACCESS_TOKEN" not in os.environ:
+            raise ValueError("Please set the HUGGINGFACEHUB_ACCESS_TOKEN environment variable.")
+
+        try:
+            importlib.import_module("huggingface_hub")
+        except ModuleNotFoundError:
+            raise ModuleNotFoundError(
+                "The required dependencies for HuggingFaceHub are not installed."
+                'Please install with `pip install --upgrade "embedchain[huggingface_hub]"`'
+            ) from None
+
+        super().__init__(config=config)
+
+    def get_llm_model_answer(self, prompt):
+        if self.config.system_prompt:
+            raise ValueError("HuggingFaceHubLlm does not support `system_prompt`")
+        return HuggingFaceHubLlm._get_answer(prompt=prompt, config=self.config)
+
+    @staticmethod
+    def _get_answer(prompt: str, config: BaseLlmConfig) -> str:
+        model_kwargs = {
+            "temperature": config.temperature or 0.1,
+            "max_new_tokens": config.max_tokens,
+        }
+
+        if config.top_p > 0.0 and config.top_p < 1.0:
+            model_kwargs["top_p"] = config.top_p
+        else:
+            raise ValueError("`top_p` must be > 0.0 and < 1.0")
+
+        llm = HuggingFaceHub(
+            huggingfacehub_api_token=os.environ["HUGGINGFACEHUB_ACCESS_TOKEN"],
+            repo_id=config.model or "google/flan-t5-xxl",
+            model_kwargs=model_kwargs,
+        )
+
+        return llm(prompt)

embedchain/llm/jina.py

@@ -0,0 +1,42 @@
+import os
+from typing import Optional
+
+from langchain.chat_models import JinaChat
+from langchain.schema import HumanMessage, SystemMessage
+
+from embedchain.config import BaseLlmConfig
+from embedchain.helper.json_serializable import register_deserializable
+from embedchain.llm.base import BaseLlm
+
+
+@register_deserializable
+class JinaLlm(BaseLlm):
+    def __init__(self, config: Optional[BaseLlmConfig] = None):
+        if "JINACHAT_API_KEY" not in os.environ:
+            raise ValueError("Please set the JINACHAT_API_KEY environment variable.")
+        super().__init__(config=config)
+
+    def get_llm_model_answer(self, prompt):
+        response = JinaLlm._get_answer(prompt, self.config)
+        return response
+
+    @staticmethod
+    def _get_answer(prompt: str, config: BaseLlmConfig) -> str:
+        messages = []
+        if config.system_prompt:
+            messages.append(SystemMessage(content=config.system_prompt))
+        messages.append(HumanMessage(content=prompt))
+        kwargs = {
+            "temperature": config.temperature,
+            "max_tokens": config.max_tokens,
+            "model_kwargs": {},
+        }
+        if config.top_p:
+            kwargs["model_kwargs"]["top_p"] = config.top_p
+        if config.stream:
+            from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
+
+            chat = JinaChat(**kwargs, streaming=config.stream, callbacks=[StreamingStdOutCallbackHandler()])
+        else:
+            chat = JinaChat(**kwargs)
+        return chat(messages).content

embedchain/llm/vertex_ai.py

@@ -12,10 +12,10 @@ class VertexAiLlm(BaseLlm):
         super().__init__(config=config)
 
     def get_llm_model_answer(self, prompt):
-        return VertexAiLlm._get_athrophic_answer(prompt=prompt, config=self.config)
+        return VertexAiLlm._get_answer(prompt=prompt, config=self.config)
 
     @staticmethod
-    def _get_athrophic_answer(prompt: str, config: BaseLlmConfig) -> str:
+    def _get_answer(prompt: str, config: BaseLlmConfig) -> str:
         from langchain.chat_models import ChatVertexAI
 
         chat = ChatVertexAI(temperature=config.temperature, model=config.model)

embedchain/loaders/images.py

@@ -16,15 +16,15 @@ class ImagesLoader(BaseLoader):
         # load model and image preprocessing
         from embedchain.models.clip_processor import ClipProcessor
 
-        model, preprocess = ClipProcessor.load_model()
+        model = ClipProcessor.load_model()
         if os.path.isfile(image_url):
-            data = [ClipProcessor.get_image_features(image_url, model, preprocess)]
+            data = [ClipProcessor.get_image_features(image_url, model)]
         else:
             data = []
             for filename in os.listdir(image_url):
                 filepath = os.path.join(image_url, filename)
                 try:
-                    data.append(ClipProcessor.get_image_features(filepath, model, preprocess))
+                    data.append(ClipProcessor.get_image_features(filepath, model))
                 except Exception as e:
                     # Log the file that was not loaded
                     logging.exception("Failed to load the file {}. Exception {}".format(filepath, e))

embedchain/loaders/sitemap.py

@@ -36,9 +36,8 @@ class SitemapLoader(BaseLoader):
         for link in links:
             try:
                 each_load_data = web_page_loader.load_data(link)
-
-                if is_readable(each_load_data[0].get("content")):
-                    output.append(each_load_data)
+                if is_readable(each_load_data.get("data")[0].get("content")):
+                    output.append(each_load_data.get("data"))
                 else:
                     logging.warning(f"Page is not readable (too many invalid characters): {link}")
             except ParserRejectedMarkup as e:

embedchain/loaders/web_page.py

@@ -15,7 +15,25 @@ class WebPageLoader(BaseLoader):
         """Load data from a web page."""
         response = requests.get(url)
         data = response.content
-        soup = BeautifulSoup(data, "html.parser")
+        content = self._get_clean_content(data, url)
+
+        meta_data = {
+            "url": url,
+        }
+
+        doc_id = hashlib.sha256((content + url).encode()).hexdigest()
+        return {
+            "doc_id": doc_id,
+            "data": [
+                {
+                    "content": content,
+                    "meta_data": meta_data,
+                }
+            ],
+        }
+
+    def _get_clean_content(self, html, url) -> str:
+        soup = BeautifulSoup(html, "html.parser")
         original_size = len(str(soup.get_text()))
 
         tags_to_exclude = [
@@ -61,17 +79,4 @@ class WebPageLoader(BaseLoader):
                 f"[{url}] 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
             )
 
-        meta_data = {
-            "url": url,
-        }
-        content = content
-        doc_id = hashlib.sha256((content + url).encode()).hexdigest()
-        return {
-            "doc_id": doc_id,
-            "data": [
-                {
-                    "content": content,
-                    "meta_data": meta_data,
-                }
-            ],
-        }
+        return content

embedchain/loaders/xml.py

@@ -0,0 +1,26 @@
+import hashlib
+
+from langchain.document_loaders import UnstructuredXMLLoader
+
+from embedchain.helper.json_serializable import register_deserializable
+from embedchain.loaders.base_loader import BaseLoader
+from embedchain.utils import clean_string
+
+
+@register_deserializable
+class XmlLoader(BaseLoader):
+    def load_data(self, xml_url):
+        """Load data from a XML file."""
+        loader = UnstructuredXMLLoader(xml_url)
+        data = loader.load()
+        content = data[0].page_content
+        content = clean_string(content)
+        meta_data = data[0].metadata
+        meta_data["url"] = meta_data["source"]
+        del meta_data["source"]
+        output = [{"content": content, "meta_data": meta_data}]
+        doc_id = hashlib.sha256((content + xml_url).encode()).hexdigest()
+        return {
+            "doc_id": doc_id,
+            "data": output,
+        }

embedchain/models/clip_processor.py

@@ -1,31 +1,27 @@
 try:
-    import clip
-    import torch
     from PIL import Image, UnidentifiedImageError
+    from sentence_transformers import SentenceTransformer
 except ImportError:
     raise ImportError(
-        "Images requires extra dependencies. Install with `pip install 'embedchain[images]' git+https://github.com/openai/CLIP.git#a1d0717`"  # noqa: E501
+        "Images requires extra dependencies. Install with `pip install 'embedchain[images]'"
     ) from None
 
-MODEL_NAME = "ViT-B/32"
+MODEL_NAME = "clip-ViT-B-32"
 
 
 class ClipProcessor:
     @staticmethod
     def load_model():
         """Load data from a director of images."""
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-
         # load model and image preprocessing
-        model, preprocess = clip.load(MODEL_NAME, device=device, jit=False)
-        return model, preprocess
+        model = SentenceTransformer(MODEL_NAME)
+        return model
 
     @staticmethod
-    def get_image_features(image_url, model, preprocess):
+    def get_image_features(image_url, model):
         """
         Applies the CLIP model to evaluate the vector representation of the supplied image
         """
-        device = "cuda" if torch.cuda.is_available() else "cpu"
         try:
             # load image
             image = Image.open(image_url)
@@ -34,27 +30,15 @@ class ClipProcessor:
         except UnidentifiedImageError:
             raise UnidentifiedImageError("The supplied file is not an image`")
 
-        # pre-process image
-        processed_image = preprocess(image).unsqueeze(0).to(device)
-        with torch.no_grad():
-            image_features = model.encode_image(processed_image)
-            image_features /= image_features.norm(dim=-1, keepdim=True)
-
-        image_features = image_features.cpu().detach().numpy().tolist()[0]
+        image_features = model.encode(image)
         meta_data = {"url": image_url}
-        return {"content": image_url, "embedding": image_features, "meta_data": meta_data}
+        return {"content": image_url, "embedding": image_features.tolist(), "meta_data": meta_data}
 
     @staticmethod
     def get_text_features(query):
         """
         Applies the CLIP model to evaluate the vector representation of the supplied text
         """
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-
-        model, preprocess = ClipProcessor.load_model()
-        text = clip.tokenize(query).to(device)
-        with torch.no_grad():
-            text_features = model.encode_text(text)
-            text_features /= text_features.norm(dim=-1, keepdim=True)
-
-        return text_features.cpu().numpy().tolist()[0]
+        model = ClipProcessor.load_model()
+        text_features = model.encode(query)
+        return text_features.tolist()

embedchain/models/data_type.py

@@ -18,6 +18,7 @@ class IndirectDataType(Enum):
     PDF_FILE = "pdf_file"
     WEB_PAGE = "web_page"
     SITEMAP = "sitemap"
+    XML = "xml"
     DOCX = "docx"
     DOCS_SITE = "docs_site"
     NOTION = "notion"
@@ -40,6 +41,7 @@ class DataType(Enum):
     PDF_FILE = IndirectDataType.PDF_FILE.value
     WEB_PAGE = IndirectDataType.WEB_PAGE.value
     SITEMAP = IndirectDataType.SITEMAP.value
+    XML = IndirectDataType.XML.value
     DOCX = IndirectDataType.DOCX.value
     DOCS_SITE = IndirectDataType.DOCS_SITE.value
     NOTION = IndirectDataType.NOTION.value

embedchain/utils.py

@@ -190,6 +190,10 @@ def detect_datatype(source: Any) -> DataType:
             logging.debug(f"Source of `{formatted_source}` detected as `csv`.")
             return DataType.CSV
 
+        if source.endswith(".xml"):
+            logging.debug(f"Source of `{formatted_source}` detected as `xml`.")
+            return DataType.XML
+
         # If the source is a valid file, that's not detectable as a type, an error is raised.
         # It does not fallback to text.
         raise ValueError(

embedchain/vectordb/chroma.py

@@ -126,12 +126,16 @@ class ChromaDB(BaseVectorDB):
         """
         Add vectors to chroma database
 
+        :param embeddings: list of embeddings to add
+        :type embeddings: List[List[str]]
         :param documents: Documents
         :type documents: List[str]
         :param metadatas: Metadatas
         :type metadatas: List[object]
         :param ids: ids
         :type ids: List[str]
+        :param skip_embedding: Optional. If True, then the embeddings are assumed to be already generated.
+        :type skip_embedding: bool
         """
         if skip_embedding:
             self.collection.add(embeddings=embeddings, documents=documents, metadatas=metadatas, ids=ids)
@@ -158,7 +162,7 @@ class ChromaDB(BaseVectorDB):
 
     def query(self, input_query: List[str], n_results: int, where: Dict[str, any], skip_embedding: bool) -> List[str]:
         """
-        Query contents from vector data base based on vector similarity
+        Query contents from vector database based on vector similarity
 
         :param input_query: list of query string
         :type input_query: List[str]
@@ -166,6 +170,8 @@ class ChromaDB(BaseVectorDB):
         :type n_results: int
         :param where: to filter data
         :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
         :raises InvalidDimensionException: Dimensions do not match.
         :return: The content of the document that matched your query.
         :rtype: List[str]

embedchain/vectordb/elasticsearch.py

@@ -110,12 +110,16 @@ class ElasticsearchDB(BaseVectorDB):
     ) -> Any:
         """
         add data in vector database
+        :param embeddings: list of embeddings to add
+        :type embeddings: List[List[str]]
         :param documents: list of texts to add
         :type documents: List[str]
         :param metadatas: list of metadata associated with docs
         :type metadatas: List[object]
         :param ids: ids of docs
         :type ids: List[str]
+        :param skip_embedding: Optional. If True, then the input_query is assumed to be already embedded.
+        :type skip_embedding: bool
         """
 
         docs = []
@@ -143,6 +147,8 @@ class ElasticsearchDB(BaseVectorDB):
         :type n_results: int
         :param where: Optional. to filter data
         :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]
         """

embedchain/vectordb/opensearch.py

@@ -93,7 +93,7 @@ class OpenSearchDB(BaseVectorDB):
 
         if "app_id" in where:
             app_id = where["app_id"]
-            query["query"]["bool"]["must"].append({"term": {"metadata.app_id": app_id}})
+            query["query"]["bool"]["must"].append({"term": {"metadata.app_id.keyword": app_id}})
 
         # OpenSearch syntax is different from Elasticsearch
         response = self.client.search(index=self._get_index(), body=query, _source=True, size=limit)
@@ -110,12 +110,12 @@ class OpenSearchDB(BaseVectorDB):
         return result
 
     def add(
-        self, embeddings: List[str], documents: List[str], metadatas: List[object], ids: List[str], skip_embedding: bool
-    ):
+        self, embeddings: List[List[str]], documents: List[str], metadatas: List[object], ids: List[str],
+            skip_embedding: bool):
         """add data in vector database
 
         :param embeddings: list of embeddings to add
-        :type embeddings: List[str]
+        :type embeddings: List[List[str]]
         :param documents: list of texts to add
         :type documents: List[str]
         :param metadatas: list of metadata associated with docs
@@ -127,8 +127,8 @@ class OpenSearchDB(BaseVectorDB):
         """
 
         docs = []
-        # TODO(rupeshbansal, deshraj): Add support for skip embeddings here if already exists
-        embeddings = self.embedder.embedding_fn(documents)
+        if not skip_embedding:
+            embeddings = self.embedder.embedding_fn(documents)
         for id, text, metadata, embeddings in zip(ids, documents, metadatas, embeddings):
             docs.append(
                 {
@@ -168,7 +168,7 @@ class OpenSearchDB(BaseVectorDB):
         pre_filter = {"match_all": {}}  # default
         if "app_id" in where:
             app_id = where["app_id"]
-            pre_filter = {"bool": {"must": [{"term": {"metadata.app_id": app_id}}]}}
+            pre_filter = {"bool": {"must": [{"term": {"metadata.app_id.keyword": app_id}}]}}
         docs = docsearch.similarity_search(
             input_query,
             search_type="script_scoring",

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "embedchain"
-version = "0.0.66"
+version = "0.0.68"
 description = "embedchain is a framework to easily create LLM powered bots over any dataset"
 authors = ["Taranjeet Singh"]
 license = "Apache License"
@@ -105,11 +105,14 @@ 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 }
 docx2txt = "^0.8"
+unstructured = {extras = ["local-inference"], version = "^0.10.18"}
 pillow = { version = "10.0.1", optional = true }
 torchvision = { version = ">=0.15.1, !=0.15.2", optional = true }
 ftfy = { version = "6.1.1", optional = true }
 regex = { version = "2023.8.8", optional = true }
+huggingface_hub = { version = "^0.17.3", optional = true }
 
 [tool.poetry.group.dev.dependencies]
 black = "^23.3.0"
@@ -121,6 +124,7 @@ pytest-env = "^0.8.1"
 click = "^8.1.3"
 isort = "^5.12.0"
 pytest-cov = "^4.1.0"
+responses = "^0.23.3"
 
 [tool.poetry.extras]
 streamlit = ["streamlit"]
@@ -133,6 +137,8 @@ discord = ["discord"]
 slack = ["slack-sdk", "flask"]
 whatsapp = ["twilio", "flask"]
 images = ["torch", "ftfy", "regex", "pillow", "torchvision"]
+huggingface_hub=["huggingface_hub"]
+cohere = ["cohere"]
 
 [tool.poetry.group.docs.dependencies]
 

tests/chunkers/test_base_chunker.py

@@ -0,0 +1,84 @@
+import hashlib
+import pytest
+from unittest.mock import MagicMock
+from embedchain.chunkers.base_chunker import BaseChunker
+from embedchain.models.data_type import DataType
+
+
+@pytest.fixture
+def text_splitter_mock():
+    return MagicMock()
+
+
+@pytest.fixture
+def loader_mock():
+    return MagicMock()
+
+
+@pytest.fixture
+def app_id():
+    return "test_app"
+
+
+@pytest.fixture
+def data_type():
+    return DataType.TEXT
+
+
+@pytest.fixture
+def chunker(text_splitter_mock, data_type):
+    text_splitter = text_splitter_mock
+    chunker = BaseChunker(text_splitter)
+    chunker.set_data_type(data_type)
+    return chunker
+
+
+def test_create_chunks(chunker, text_splitter_mock, loader_mock, app_id, data_type):
+    text_splitter_mock.split_text.return_value = ["Chunk 1", "Chunk 2"]
+    loader_mock.load_data.return_value = {
+        "data": [{"content": "Content 1", "meta_data": {"url": "URL 1"}}],
+        "doc_id": "DocID",
+    }
+
+    result = chunker.create_chunks(loader_mock, "test_src", app_id)
+    expected_ids = [
+        hashlib.sha256(("Chunk 1" + "URL 1").encode()).hexdigest(),
+        hashlib.sha256(("Chunk 2" + "URL 1").encode()).hexdigest(),
+    ]
+
+    assert result["documents"] == ["Chunk 1", "Chunk 2"]
+    assert result["ids"] == expected_ids
+    assert result["metadatas"] == [
+        {
+            "url": "URL 1",
+            "data_type": data_type.value,
+            "doc_id": f"{app_id}--DocID",
+        },
+        {
+            "url": "URL 1",
+            "data_type": data_type.value,
+            "doc_id": f"{app_id}--DocID",
+        },
+    ]
+    assert result["doc_id"] == f"{app_id}--DocID"
+
+
+def test_get_chunks(chunker, text_splitter_mock):
+    text_splitter_mock.split_text.return_value = ["Chunk 1", "Chunk 2"]
+
+    content = "This is a test content."
+    result = chunker.get_chunks(content)
+
+    assert len(result) == 2
+    assert result == ["Chunk 1", "Chunk 2"]
+
+
+def test_set_data_type(chunker):
+    chunker.set_data_type(DataType.MDX)
+    assert chunker.data_type == DataType.MDX
+
+
+def test_get_word_count(chunker):
+    documents = ["This is a test.", "Another test."]
+    result = chunker.get_word_count(documents)
+    assert result == 6

tests/chunkers/test_chunkers.py

@@ -0,0 +1,46 @@
+from embedchain.chunkers.docs_site import DocsSiteChunker
+from embedchain.chunkers.docx_file import DocxFileChunker
+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.web_page import WebPageChunker
+from embedchain.chunkers.xml import XmlChunker
+from embedchain.chunkers.youtube_video import YoutubeVideoChunker
+from embedchain.config.add_config import ChunkerConfig
+
+chunker_config = ChunkerConfig(chunk_size=500, chunk_overlap=0, length_function=len)
+
+chunker_common_config = {
+    DocsSiteChunker: {"chunk_size": 500, "chunk_overlap": 50, "length_function": len},
+    DocxFileChunker: {"chunk_size": 1000, "chunk_overlap": 0, "length_function": len},
+    PdfFileChunker: {"chunk_size": 1000, "chunk_overlap": 0, "length_function": len},
+    TextChunker: {"chunk_size": 300, "chunk_overlap": 0, "length_function": len},
+    MdxChunker: {"chunk_size": 1000, "chunk_overlap": 0, "length_function": len},
+    NotionChunker: {"chunk_size": 300, "chunk_overlap": 0, "length_function": len},
+    QnaPairChunker: {"chunk_size": 300, "chunk_overlap": 0, "length_function": len},
+    TableChunker: {"chunk_size": 300, "chunk_overlap": 0, "length_function": len},
+    SitemapChunker: {"chunk_size": 500, "chunk_overlap": 0, "length_function": len},
+    WebPageChunker: {"chunk_size": 500, "chunk_overlap": 0, "length_function": len},
+    XmlChunker: {"chunk_size": 500, "chunk_overlap": 50, "length_function": len},
+    YoutubeVideoChunker: {"chunk_size": 2000, "chunk_overlap": 0, "length_function": len},
+}
+
+
+def test_default_config_values():
+    for chunker_class, config in chunker_common_config.items():
+        chunker = chunker_class()
+        assert chunker.text_splitter._chunk_size == config["chunk_size"]
+        assert chunker.text_splitter._chunk_overlap == config["chunk_overlap"]
+        assert chunker.text_splitter._length_function == config["length_function"]
+
+
+def test_custom_config_values():
+    for chunker_class, _ in chunker_common_config.items():
+        chunker = chunker_class(config=chunker_config)
+        assert chunker.text_splitter._chunk_size == 500
+        assert chunker.text_splitter._chunk_overlap == 0
+        assert chunker.text_splitter._length_function == len

tests/embedchain/test_add.py

@@ -27,6 +27,14 @@ class TestApp(unittest.TestCase):
         self.app.add("https://example.com", metadata={"meta": "meta-data"})
         self.assertEqual(self.app.user_asks, [["https://example.com", "web_page", {"meta": "meta-data"}]])
 
+    @patch("chromadb.api.models.Collection.Collection.add", MagicMock)
+    def test_add_sitemap(self):
+        """
+        In addition to the test_add function, this test checks that sitemaps can be added with the correct data type.
+        """
+        self.app.add("https://www.google.com/sitemap.xml", metadata={"meta": "meta-data"})
+        self.assertEqual(self.app.user_asks, [["https://www.google.com/sitemap.xml", "sitemap", {"meta": "meta-data"}]])
+
     @patch("chromadb.api.models.Collection.Collection.add", MagicMock)
     def test_add_forced_type(self):
         """

tests/llm/test_cohere.py

@@ -0,0 +1,33 @@
+import os
+import unittest
+from unittest.mock import patch
+
+from embedchain.config import BaseLlmConfig
+from embedchain.llm.cohere import CohereLlm
+
+
+class TestCohereLlm(unittest.TestCase):
+    def setUp(self):
+        os.environ["COHERE_API_KEY"] = "test_api_key"
+        self.config = BaseLlmConfig(model="gptd-instruct-tft", max_tokens=50, temperature=0.7, top_p=0.8)
+
+    def test_init_raises_value_error_without_api_key(self):
+        os.environ.pop("COHERE_API_KEY")
+        with self.assertRaises(ValueError):
+            CohereLlm()
+
+    def test_get_llm_model_answer_raises_value_error_for_system_prompt(self):
+        llm = CohereLlm(self.config)
+        llm.config.system_prompt = "system_prompt"
+        with self.assertRaises(ValueError):
+            llm.get_llm_model_answer("prompt")
+
+    @patch("embedchain.llm.cohere.CohereLlm._get_answer")
+    def test_get_llm_model_answer(self, mock_get_answer):
+        mock_get_answer.return_value = "Test answer"
+
+        llm = CohereLlm(self.config)
+        answer = llm.get_llm_model_answer("Test query")
+
+        self.assertEqual(answer, "Test answer")
+        mock_get_answer.assert_called_once()

tests/llm/test_hugging_face_hub.py

@@ -0,0 +1,64 @@
+import importlib
+import os
+import unittest
+from unittest.mock import patch, MagicMock
+
+from embedchain.config import BaseLlmConfig
+from embedchain.llm.hugging_face_hub import HuggingFaceHubLlm
+
+
+class TestHuggingFaceHubLlm(unittest.TestCase):
+    def setUp(self):
+        os.environ["HUGGINGFACEHUB_ACCESS_TOKEN"] = "test_access_token"
+        self.config = BaseLlmConfig(model="google/flan-t5-xxl", max_tokens=50, temperature=0.7, top_p=0.8)
+
+    def test_init_raises_value_error_without_api_key(self):
+        os.environ.pop("HUGGINGFACEHUB_ACCESS_TOKEN")
+        with self.assertRaises(ValueError):
+            HuggingFaceHubLlm()
+
+    def test_get_llm_model_answer_raises_value_error_for_system_prompt(self):
+        llm = HuggingFaceHubLlm(self.config)
+        llm.config.system_prompt = "system_prompt"
+        with self.assertRaises(ValueError):
+            llm.get_llm_model_answer("prompt")
+
+    def test_top_p_value_within_range(self):
+        config = BaseLlmConfig(top_p=1.0)
+        with self.assertRaises(ValueError):
+            HuggingFaceHubLlm._get_answer("test_prompt", config)
+
+    def test_dependency_is_imported(self):
+        importlib_installed = True
+        try:
+            importlib.import_module("huggingface_hub")
+        except ImportError:
+            importlib_installed = False
+        self.assertTrue(importlib_installed)
+
+    @patch("embedchain.llm.hugging_face_hub.HuggingFaceHubLlm._get_answer")
+    def test_get_llm_model_answer(self, mock_get_answer):
+        mock_get_answer.return_value = "Test answer"
+
+        llm = HuggingFaceHubLlm(self.config)
+        answer = llm.get_llm_model_answer("Test query")
+
+        self.assertEqual(answer, "Test answer")
+        mock_get_answer.assert_called_once()
+
+    @patch("embedchain.llm.hugging_face_hub.HuggingFaceHub")
+    def test_hugging_face_mock(self, mock_hugging_face_hub):
+        mock_llm_instance = MagicMock()
+        mock_llm_instance.return_value = "Test answer"
+        mock_hugging_face_hub.return_value = mock_llm_instance
+
+        llm = HuggingFaceHubLlm(self.config)
+        answer = llm.get_llm_model_answer("Test query")
+
+        self.assertEqual(answer, "Test answer")
+        mock_hugging_face_hub.assert_called_once_with(
+            huggingfacehub_api_token="test_access_token",
+            repo_id="google/flan-t5-xxl",
+            model_kwargs={"temperature": 0.7, "max_new_tokens": 50, "top_p": 0.8},
+        )
+        mock_llm_instance.assert_called_once_with("Test query")

tests/llm/test_jina.py

@@ -0,0 +1,40 @@
+import os
+import unittest
+from unittest.mock import patch
+
+from embedchain.config import BaseLlmConfig
+from embedchain.llm.jina import JinaLlm
+
+
+class TestJinaLlm(unittest.TestCase):
+    def setUp(self):
+        os.environ["JINACHAT_API_KEY"] = "test_api_key"
+        self.config = BaseLlmConfig(
+            temperature=0.7, max_tokens=50, top_p=0.8, stream=False, system_prompt="System prompt"
+        )
+
+    def test_init_raises_value_error_without_api_key(self):
+        os.environ.pop("JINACHAT_API_KEY")
+        with self.assertRaises(ValueError):
+            JinaLlm()
+
+    @patch("embedchain.llm.jina.JinaLlm._get_answer")
+    def test_get_llm_model_answer(self, mock_get_answer):
+        mock_get_answer.return_value = "Test answer"
+
+        llm = JinaLlm(self.config)
+        answer = llm.get_llm_model_answer("Test query")
+
+        self.assertEqual(answer, "Test answer")
+        mock_get_answer.assert_called_once()
+
+    @patch("embedchain.llm.jina.JinaLlm._get_answer")
+    def test_get_llm_model_answer_with_system_prompt(self, mock_get_answer):
+        self.config.system_prompt = "Custom system prompt"
+        mock_get_answer.return_value = "Test answer"
+
+        llm = JinaLlm(self.config)
+        answer = llm.get_llm_model_answer("Test query")
+
+        self.assertEqual(answer, "Test answer")
+        mock_get_answer.assert_called_once()

tests/loaders/test_csv.py

@@ -2,6 +2,7 @@ import csv
 import os
 import pathlib
 import tempfile
+from unittest.mock import MagicMock, patch
 
 import pytest
 
@@ -84,3 +85,29 @@ def test_load_data_with_file_uri(delimiter):
 
         # Cleaning up the temporary file
         os.unlink(tmpfile.name)
+
+
+@pytest.mark.parametrize("content", ["ftp://example.com", "sftp://example.com", "mailto://example.com"])
+def test_get_file_content(content):
+    with pytest.raises(ValueError):
+        loader = CsvLoader()
+        loader._get_file_content(content)
+
+
+@pytest.mark.parametrize("content", ["http://example.com", "https://example.com"])
+def test_get_file_content_http(content):
+    """
+    Test _get_file_content method of CsvLoader for http and https URLs
+    """
+
+    with patch("requests.get") as mock_get:
+        mock_response = MagicMock()
+        mock_response.text = "Name,Age,Occupation\nAlice,28,Engineer\nBob,35,Doctor\nCharlie,22,Student"
+        mock_get.return_value = mock_response
+
+        loader = CsvLoader()
+        file_content = loader._get_file_content(content)
+
+        mock_get.assert_called_once_with(content)
+        mock_response.raise_for_status.assert_called_once()
+        assert file_content.read() == mock_response.text

tests/loaders/test_docs_site.py

@@ -0,0 +1,128 @@
+import hashlib
+import pytest
+from unittest.mock import Mock, patch
+from requests import Response
+from embedchain.loaders.docs_site_loader import DocsSiteLoader
+
+
+@pytest.fixture
+def mock_requests_get():
+    with patch("requests.get") as mock_get:
+        yield mock_get
+
+
+@pytest.fixture
+def docs_site_loader():
+    return DocsSiteLoader()
+
+
+def test_get_child_links_recursive(mock_requests_get, docs_site_loader):
+    mock_response = Mock()
+    mock_response.status_code = 200
+    mock_response.text = """
+        <html>
+            <a href="/page1">Page 1</a>
+            <a href="/page2">Page 2</a>
+        </html>
+    """
+    mock_requests_get.return_value = mock_response
+
+    docs_site_loader._get_child_links_recursive("https://example.com")
+
+    assert len(docs_site_loader.visited_links) == 2
+    assert "https://example.com/page1" in docs_site_loader.visited_links
+    assert "https://example.com/page2" in docs_site_loader.visited_links
+
+
+def test_get_child_links_recursive_status_not_200(mock_requests_get, docs_site_loader):
+    mock_response = Mock()
+    mock_response.status_code = 404
+    mock_requests_get.return_value = mock_response
+
+    docs_site_loader._get_child_links_recursive("https://example.com")
+
+    assert len(docs_site_loader.visited_links) == 0
+
+
+def test_get_all_urls(mock_requests_get, docs_site_loader):
+    mock_response = Mock()
+    mock_response.status_code = 200
+    mock_response.text = """
+        <html>
+            <a href="/page1">Page 1</a>
+            <a href="/page2">Page 2</a>
+            <a href="https://example.com/external">External</a>
+        </html>
+    """
+    mock_requests_get.return_value = mock_response
+
+    all_urls = docs_site_loader._get_all_urls("https://example.com")
+
+    assert len(all_urls) == 3
+    assert "https://example.com/page1" in all_urls
+    assert "https://example.com/page2" in all_urls
+    assert "https://example.com/external" in all_urls
+
+
+def test_load_data_from_url(mock_requests_get, docs_site_loader):
+    mock_response = Mock()
+    mock_response.status_code = 200
+    mock_response.content = """
+        <html>
+            <nav>
+                <h1>Navigation</h1>
+            </nav>
+            <article class="bd-article">
+                <p>Article Content</p>
+            </article>
+        </html>
+    """.encode()
+    mock_requests_get.return_value = mock_response
+
+    data = docs_site_loader._load_data_from_url("https://example.com/page1")
+
+    assert len(data) == 1
+    assert data[0]["content"] == "Article Content"
+    assert data[0]["meta_data"]["url"] == "https://example.com/page1"
+
+
+def test_load_data_from_url_status_not_200(mock_requests_get, docs_site_loader):
+    mock_response = Mock()
+    mock_response.status_code = 404
+    mock_requests_get.return_value = mock_response
+
+    data = docs_site_loader._load_data_from_url("https://example.com/page1")
+
+    assert data == []
+    assert len(data) == 0
+
+
+def test_load_data(mock_requests_get, docs_site_loader):
+    mock_response = Response()
+    mock_response.status_code = 200
+    mock_response._content = """
+        <html>
+            <a href="/page1">Page 1</a>
+            <a href="/page2">Page 2</a>
+        """.encode()
+    mock_requests_get.return_value = mock_response
+
+    url = "https://example.com"
+    data = docs_site_loader.load_data(url)
+    expected_doc_id = hashlib.sha256((" ".join(docs_site_loader.visited_links) + url).encode()).hexdigest()
+
+    assert len(data["data"]) == 2
+    assert data["doc_id"] == expected_doc_id
+
+
+def test_if_response_status_not_200(mock_requests_get, docs_site_loader):
+    mock_response = Response()
+    mock_response.status_code = 404
+    mock_requests_get.return_value = mock_response
+
+    url = "https://example.com"
+    data = docs_site_loader.load_data(url)
+    expected_doc_id = hashlib.sha256((" ".join(docs_site_loader.visited_links) + url).encode()).hexdigest()
+
+    assert len(data["data"]) == 0
+    assert data["doc_id"] == expected_doc_id

tests/loaders/test_docs_site_loader.py

@@ -0,0 +1,218 @@
+import pytest
+import responses
+from bs4 import BeautifulSoup
+
+
+@pytest.mark.parametrize(
+    "ignored_tag",
+    [
+        "<nav>This is a navigation bar.</nav>",
+        "<aside>This is an aside.</aside>",
+        "<form>This is a form.</form>",
+        "<header>This is a header.</header>",
+        "<noscript>This is a noscript.</noscript>",
+        "<svg>This is an SVG.</svg>",
+        "<canvas>This is a canvas.</canvas>",
+        "<footer>This is a footer.</footer>",
+        "<script>This is a script.</script>",
+        "<style>This is a style.</style>",
+    ],
+    ids=["nav", "aside", "form", "header", "noscript", "svg", "canvas", "footer", "script", "style"],
+)
+@pytest.mark.parametrize(
+    "selectee",
+    [
+        """
+<article class="bd-article">
+    <h2>Article Title</h2>
+    <p>Article content goes here.</p>
+    {ignored_tag}
+</article>
+""",
+        """
+<article role="main">
+    <h2>Main Article Title</h2>
+    <p>Main article content goes here.</p>
+    {ignored_tag}
+</article>
+""",
+        """
+<div class="md-content">
+    <h2>Markdown Content</h2>
+    <p>Markdown content goes here.</p>
+    {ignored_tag}
+</div>
+""",
+        """
+<div role="main">
+    <h2>Main Content</h2>
+    <p>Main content goes here.</p>
+    {ignored_tag}
+</div>
+""",
+        """
+<div class="container">
+    <h2>Container</h2>
+    <p>Container content goes here.</p>
+    {ignored_tag}
+</div>
+        """,
+        """
+<div class="section">
+    <h2>Section</h2>
+    <p>Section content goes here.</p>
+    {ignored_tag}
+</div>
+        """,
+        """
+<article>
+    <h2>Generic Article</h2>
+    <p>Generic article content goes here.</p>
+    {ignored_tag}
+</article>
+        """,
+        """
+<main>
+    <h2>Main Content</h2>
+    <p>Main content goes here.</p>
+    {ignored_tag}
+</main>
+""",
+    ],
+    ids=[
+        "article.bd-article",
+        'article[role="main"]',
+        "div.md-content",
+        'div[role="main"]',
+        "div.container",
+        "div.section",
+        "article",
+        "main",
+    ],
+)
+def test_load_data_gets_by_selectors_and_ignored_tags(selectee, ignored_tag, loader, mocked_responses, mocker):
+    child_url = "https://docs.embedchain.ai/quickstart"
+    selectee = selectee.format(ignored_tag=ignored_tag)
+    html_body = """
+<!DOCTYPE html>
+<html lang="en">
+<body>
+    {selectee}
+</body>
+</html>
+"""
+    html_body = html_body.format(selectee=selectee)
+    mocked_responses.get(child_url, body=html_body, status=200, content_type="text/html")
+
+    url = "https://docs.embedchain.ai/"
+    html_body = """
+<!DOCTYPE html>
+<html lang="en">
+<body>
+    <li><a href="/quickstart">Quickstart</a></li>
+</body>
+</html>
+"""
+    mocked_responses.get(url, body=html_body, status=200, content_type="text/html")
+
+    mock_sha256 = mocker.patch("embedchain.loaders.docs_site_loader.hashlib.sha256")
+    doc_id = "mocked_hash"
+    mock_sha256.return_value.hexdigest.return_value = doc_id
+
+    result = loader.load_data(url)
+    selector_soup = BeautifulSoup(selectee, "html.parser")
+    expected_content = " ".join((selector_soup.select_one("h2").get_text(), selector_soup.select_one("p").get_text()))
+    assert result["doc_id"] == doc_id
+    assert result["data"] == [
+        {
+            "content": expected_content,
+            "meta_data": {"url": "https://docs.embedchain.ai/quickstart"},
+        }
+    ]
+
+
+def test_load_data_gets_child_links_recursively(loader, mocked_responses, mocker):
+    child_url = "https://docs.embedchain.ai/quickstart"
+    html_body = """
+<!DOCTYPE html>
+<html lang="en">
+<body>
+    <li><a href="/">..</a></li>
+    <li><a href="/quickstart">.</a></li>
+</body>
+</html>
+"""
+    mocked_responses.get(child_url, body=html_body, status=200, content_type="text/html")
+
+    child_url = "https://docs.embedchain.ai/introduction"
+    html_body = """
+<!DOCTYPE html>
+<html lang="en">
+<body>
+    <li><a href="/">..</a></li>
+    <li><a href="/introduction">.</a></li>
+</body>
+</html>
+"""
+    mocked_responses.get(child_url, body=html_body, status=200, content_type="text/html")
+
+    url = "https://docs.embedchain.ai/"
+    html_body = """
+<!DOCTYPE html>
+<html lang="en">
+<body>
+    <li><a href="/quickstart">Quickstart</a></li>
+    <li><a href="/introduction">Introduction</a></li>
+</body>
+</html>
+"""
+    mocked_responses.get(url, body=html_body, status=200, content_type="text/html")
+
+    mock_sha256 = mocker.patch("embedchain.loaders.docs_site_loader.hashlib.sha256")
+    doc_id = "mocked_hash"
+    mock_sha256.return_value.hexdigest.return_value = doc_id
+
+    result = loader.load_data(url)
+    assert result["doc_id"] == doc_id
+    expected_data = [
+        {"content": "..\n.", "meta_data": {"url": "https://docs.embedchain.ai/quickstart"}},
+        {"content": "..\n.", "meta_data": {"url": "https://docs.embedchain.ai/introduction"}},
+    ]
+    assert all(item in expected_data for item in result["data"])
+
+
+def test_load_data_fails_to_fetch_website(loader, mocked_responses, mocker):
+    child_url = "https://docs.embedchain.ai/introduction"
+    mocked_responses.get(child_url, status=404)
+
+    url = "https://docs.embedchain.ai/"
+    html_body = """
+<!DOCTYPE html>
+<html lang="en">
+<body>
+    <li><a href="/introduction">Introduction</a></li>
+</body>
+</html>
+"""
+    mocked_responses.get(url, body=html_body, status=200, content_type="text/html")
+
+    mock_sha256 = mocker.patch("embedchain.loaders.docs_site_loader.hashlib.sha256")
+    doc_id = "mocked_hash"
+    mock_sha256.return_value.hexdigest.return_value = doc_id
+
+    result = loader.load_data(url)
+    assert result["doc_id"] is doc_id
+    assert result["data"] == []
+
+
+@pytest.fixture
+def loader():
+    from embedchain.loaders.docs_site_loader import DocsSiteLoader
+
+    return DocsSiteLoader()
+
+
+@pytest.fixture
+def mocked_responses():
+    with responses.RequestsMock() as rsps:
+        yield rsps

tests/loaders/test_docx_file.py

@@ -0,0 +1,37 @@
+import hashlib
+import pytest
+from unittest.mock import MagicMock, patch
+from embedchain.loaders.docx_file import DocxFileLoader
+
+
+@pytest.fixture
+def mock_docx2txt_loader():
+    with patch("embedchain.loaders.docx_file.Docx2txtLoader") as mock_loader:
+        yield mock_loader
+
+
+@pytest.fixture
+def docx_file_loader():
+    return DocxFileLoader()
+
+
+def test_load_data(mock_docx2txt_loader, docx_file_loader):
+    mock_url = "mock_docx_file.docx"
+
+    mock_loader = MagicMock()
+    mock_loader.load.return_value = [MagicMock(page_content="Sample Docx Content", metadata={"url": "local"})]
+
+    mock_docx2txt_loader.return_value = mock_loader
+
+    result = docx_file_loader.load_data(mock_url)
+
+    assert "doc_id" in result
+    assert "data" in result
+
+    expected_content = "Sample Docx Content"
+    assert result["data"][0]["content"] == expected_content
+
+    assert result["data"][0]["meta_data"]["url"] == "local"
+
+    expected_doc_id = hashlib.sha256((expected_content + mock_url).encode()).hexdigest()
+    assert result["doc_id"] == expected_doc_id

tests/loaders/test_local_qna_pair.py

@@ -0,0 +1,30 @@
+import hashlib
+import pytest
+from embedchain.loaders.local_qna_pair import LocalQnaPairLoader
+
+
+@pytest.fixture
+def qna_pair_loader():
+    return LocalQnaPairLoader()
+
+
+def test_load_data(qna_pair_loader):
+    question = "What is the capital of France?"
+    answer = "The capital of France is Paris."
+
+    content = (question, answer)
+    result = qna_pair_loader.load_data(content)
+
+    assert "doc_id" in result
+    assert "data" in result
+    url = "local"
+
+    expected_content = f"Q: {question}\nA: {answer}"
+    assert result["data"][0]["content"] == expected_content
+
+    assert result["data"][0]["meta_data"]["url"] == url
+
+    assert result["data"][0]["meta_data"]["question"] == question
+
+    expected_doc_id = hashlib.sha256((expected_content + url).encode()).hexdigest()
+    assert result["doc_id"] == expected_doc_id

tests/loaders/test_local_text.py

@@ -0,0 +1,25 @@
+import hashlib
+import pytest
+from embedchain.loaders.local_text import LocalTextLoader
+
+
+@pytest.fixture
+def text_loader():
+    return LocalTextLoader()
+
+
+def test_load_data(text_loader):
+    mock_content = "This is a sample text content."
+
+    result = text_loader.load_data(mock_content)
+
+    assert "doc_id" in result
+    assert "data" in result
+
+    url = "local"
+    assert result["data"][0]["content"] == mock_content
+
+    assert result["data"][0]["meta_data"]["url"] == url
+
+    expected_doc_id = hashlib.sha256((mock_content + url).encode()).hexdigest()
+    assert result["doc_id"] == expected_doc_id

tests/loaders/test_mdx.py

@@ -0,0 +1,28 @@
+import hashlib
+import pytest
+from unittest.mock import patch, mock_open
+from embedchain.loaders.mdx import MdxLoader
+
+
+@pytest.fixture
+def mdx_loader():
+    return MdxLoader()
+
+
+def test_load_data(mdx_loader):
+    mock_content = "Sample MDX Content"
+
+    # Mock open function to simulate file reading
+    with patch("builtins.open", mock_open(read_data=mock_content)):
+        url = "mock_file.mdx"
+        result = mdx_loader.load_data(url)
+
+        assert "doc_id" in result
+        assert "data" in result
+
+        assert result["data"][0]["content"] == mock_content
+
+        assert result["data"][0]["meta_data"]["url"] == url
+
+        expected_doc_id = hashlib.sha256((mock_content + url).encode()).hexdigest()
+        assert result["doc_id"] == expected_doc_id

tests/loaders/test_notion.py

@@ -0,0 +1,34 @@
+import hashlib
+import os
+import pytest
+from unittest.mock import Mock, patch
+from embedchain.loaders.notion import NotionLoader
+
+
+@pytest.fixture
+def notion_loader():
+    with patch.dict(os.environ, {"NOTION_INTEGRATION_TOKEN": "test_notion_token"}):
+        yield NotionLoader()
+
+
+def test_load_data(notion_loader):
+    source = "https://www.notion.so/Test-Page-1234567890abcdef1234567890abcdef"
+    mock_text = "This is a test page."
+    expected_doc_id = hashlib.sha256((mock_text + source).encode()).hexdigest()
+    expected_data = [
+        {
+            "content": mock_text,
+            "meta_data": {"url": "notion-12345678-90ab-cdef-1234-567890abcdef"},  # formatted_id
+        }
+    ]
+
+    mock_page = Mock()
+    mock_page.text = mock_text
+    mock_documents = [mock_page]
+
+    with patch("embedchain.loaders.notion.NotionPageReader") as mock_reader:
+        mock_reader.return_value.load_data.return_value = mock_documents
+        result = notion_loader.load_data(source)
+
+    assert result["doc_id"] == expected_doc_id
+    assert result["data"] == expected_data

tests/loaders/test_pdf_file.py

@@ -0,0 +1,36 @@
+import pytest
+from langchain.schema import Document
+
+
+def test_load_data(loader, mocker):
+    mocked_pypdfloader = mocker.patch("embedchain.loaders.pdf_file.PyPDFLoader")
+    mocked_pypdfloader.return_value.load_and_split.return_value = [
+        Document(page_content="Page 0 Content", metadata={"source": "example.pdf", "page": 0}),
+        Document(page_content="Page 1 Content", metadata={"source": "example.pdf", "page": 1}),
+    ]
+
+    mock_sha256 = mocker.patch("embedchain.loaders.docs_site_loader.hashlib.sha256")
+    doc_id = "mocked_hash"
+    mock_sha256.return_value.hexdigest.return_value = doc_id
+
+    result = loader.load_data("dummy_url")
+    assert result["doc_id"] is doc_id
+    assert result["data"] == [
+        {"content": "Page 0 Content", "meta_data": {"source": "example.pdf", "page": 0, "url": "dummy_url"}},
+        {"content": "Page 1 Content", "meta_data": {"source": "example.pdf", "page": 1, "url": "dummy_url"}},
+    ]
+
+
+def test_load_data_fails_to_find_data(loader, mocker):
+    mocked_pypdfloader = mocker.patch("embedchain.loaders.pdf_file.PyPDFLoader")
+    mocked_pypdfloader.return_value.load_and_split.return_value = []
+
+    with pytest.raises(ValueError):
+        loader.load_data("dummy_url")
+
+
+@pytest.fixture
+def loader():
+    from embedchain.loaders.pdf_file import PdfFileLoader
+
+    return PdfFileLoader()

tests/loaders/test_web_page.py

@@ -0,0 +1,115 @@
+import hashlib
+import pytest
+from unittest.mock import Mock, patch
+from embedchain.loaders.web_page import WebPageLoader
+
+
+@pytest.fixture
+def web_page_loader():
+    return WebPageLoader()
+
+
+def test_load_data(web_page_loader):
+    page_url = "https://example.com/page"
+    mock_response = Mock()
+    mock_response.status_code = 200
+    mock_response.content = """
+        <html>
+            <head>
+                <title>Test Page</title>
+            </head>
+            <body>
+                <div id="content">
+                    <p>This is some test content.</p>
+                </div>
+            </body>
+        </html>
+    """
+    with patch("embedchain.loaders.web_page.requests.get", return_value=mock_response):
+        result = web_page_loader.load_data(page_url)
+
+    content = web_page_loader._get_clean_content(mock_response.content, page_url)
+    expected_doc_id = hashlib.sha256((content + page_url).encode()).hexdigest()
+    assert result["doc_id"] == expected_doc_id
+
+    expected_data = [
+        {
+            "content": content,
+            "meta_data": {
+                "url": page_url,
+            },
+        }
+    ]
+
+    assert result["data"] == expected_data
+
+
+def test_get_clean_content_excludes_unnecessary_info(web_page_loader):
+    mock_html = """
+        <html>
+        <head>
+            <title>Sample HTML</title>
+            <style>
+                /* Stylesheet to be excluded */
+                .elementor-location-header {
+                    background-color: #f0f0f0;
+                }
+            </style>
+        </head>
+        <body>
+            <header id="header">Header Content</header>
+            <nav class="nav">Nav Content</nav>
+            <aside>Aside Content</aside>
+            <form>Form Content</form>
+            <main>Main Content</main>
+            <footer class="footer">Footer Content</footer>
+            <script>Some Script</script>
+            <noscript>NoScript Content</noscript>
+            <svg>SVG Content</svg>
+            <canvas>Canvas Content</canvas>
+            
+            <div id="sidebar">Sidebar Content</div>
+            <div id="main-navigation">Main Navigation Content</div>
+            <div id="menu-main-menu">Menu Main Menu Content</div>
+            
+            <div class="header-sidebar-wrapper">Header Sidebar Wrapper Content</div>
+            <div class="blog-sidebar-wrapper">Blog Sidebar Wrapper Content</div>
+            <div class="related-posts">Related Posts Content</div>
+        </body>
+        </html>
+    """
+
+    tags_to_exclude = [
+        "nav",
+        "aside",
+        "form",
+        "header",
+        "noscript",
+        "svg",
+        "canvas",
+        "footer",
+        "script",
+        "style",
+    ]
+    ids_to_exclude = ["sidebar", "main-navigation", "menu-main-menu"]
+    classes_to_exclude = [
+        "elementor-location-header",
+        "navbar-header",
+        "nav",
+        "header-sidebar-wrapper",
+        "blog-sidebar-wrapper",
+        "related-posts",
+    ]
+
+    content = web_page_loader._get_clean_content(mock_html, "https://example.com/page")
+
+    for tag in tags_to_exclude:
+        assert tag not in content
+
+    for id in ids_to_exclude:
+        assert id not in content
+
+    for class_name in classes_to_exclude:
+        assert class_name not in content
+
+    assert len(content) > 0

tests/loaders/test_xml.py

@@ -0,0 +1,62 @@
+import tempfile
+
+import pytest
+
+from embedchain.loaders.xml import XmlLoader
+
+# Taken from https://github.com/langchain-ai/langchain/blob/master/libs/langchain/tests/integration_tests/examples/factbook.xml
+SAMPLE_XML = """<?xml version="1.0" encoding="UTF-8"?>
+<factbook>
+  <country>
+    <name>United States</name>
+    <capital>Washington, DC</capital>
+    <leader>Joe Biden</leader>
+    <sport>Baseball</sport>
+  </country>
+  <country>
+    <name>Canada</name>
+    <capital>Ottawa</capital>
+    <leader>Justin Trudeau</leader>
+    <sport>Hockey</sport>
+  </country>
+  <country>
+    <name>France</name>
+    <capital>Paris</capital>
+    <leader>Emmanuel Macron</leader>
+    <sport>Soccer</sport>
+  </country>
+  <country>
+    <name>Trinidad &amp; Tobado</name>
+    <capital>Port of Spain</capital>
+    <leader>Keith Rowley</leader>
+    <sport>Track &amp; Field</sport>
+  </country>
+</factbook>"""
+
+
+@pytest.mark.parametrize("xml", [SAMPLE_XML])
+def test_load_data(xml: str):
+    """
+    Test XML loader
+
+    Tests that XML file is loaded, metadata is correct and content is correct
+    """
+    # Creating temporary XML file
+    with tempfile.NamedTemporaryFile(mode="w+") as tmpfile:
+        tmpfile.write(xml)
+
+        tmpfile.seek(0)
+        filename = tmpfile.name
+
+        # Loading CSV using XmlLoader
+        loader = XmlLoader()
+        result = loader.load_data(filename)
+        data = result["data"]
+
+        # Assertions
+        assert len(data) == 1
+        assert "United States Washington, DC Joe Biden" in data[0]["content"]
+        assert "Canada Ottawa Justin Trudeau" in data[0]["content"]
+        assert "France Paris Emmanuel Macron" in data[0]["content"]
+        assert "Trinidad & Tobado Port of Spain Keith Rowley" in data[0]["content"]
+        assert data[0]["meta_data"]["url"] == filename

tests/loaders/test_youtube_video.py

@@ -0,0 +1,47 @@
+import hashlib
+import pytest
+from unittest.mock import MagicMock, Mock, patch
+from embedchain.loaders.youtube_video import YoutubeVideoLoader
+
+
+@pytest.fixture
+def youtube_video_loader():
+    return YoutubeVideoLoader()
+
+
+def test_load_data(youtube_video_loader):
+    video_url = "https://www.youtube.com/watch?v=VIDEO_ID"
+    mock_loader = Mock()
+    mock_page_content = "This is a YouTube video content."
+    mock_loader.load.return_value = [
+        MagicMock(
+            page_content=mock_page_content,
+            metadata={"url": video_url, "title": "Test Video"},
+        )
+    ]
+
+    with patch("embedchain.loaders.youtube_video.YoutubeLoader.from_youtube_url", return_value=mock_loader):
+        result = youtube_video_loader.load_data(video_url)
+
+    expected_doc_id = hashlib.sha256((mock_page_content + video_url).encode()).hexdigest()
+
+    assert result["doc_id"] == expected_doc_id
+
+    expected_data = [
+        {
+            "content": "This is a YouTube video content.",
+            "meta_data": {"url": video_url, "title": "Test Video"},
+        }
+    ]
+
+    assert result["data"] == expected_data
+
+
+def test_load_data_with_empty_doc(youtube_video_loader):
+    video_url = "https://www.youtube.com/watch?v=VIDEO_ID"
+    mock_loader = Mock()
+    mock_loader.load.return_value = []
+
+    with patch("embedchain.loaders.youtube_video.YoutubeLoader.from_youtube_url", return_value=mock_loader):
+        with pytest.raises(ValueError):
+            youtube_video_loader.load_data(video_url)

tests/models/test_clip_processor.py

@@ -1,51 +1,43 @@
-# import os
-# import tempfile
-# import urllib
+import os
+import tempfile
+import urllib
 
-# import pytest
-# from PIL import Image
+from PIL import Image
 
-# TODO: Uncomment after fixing clip dependency issue
-# from embedchain.models.clip_processor import ClipProcessor
+from embedchain.models.clip_processor import ClipProcessor
 
 
-# class TestClipProcessor:
-#     @pytest.mark.xfail(reason="This test is failing because of the missing CLIP dependency.")
-#     def test_load_model(self):
-#         # Test that the `load_model()` method loads the CLIP model and image preprocessing correctly.
-#         model, preprocess = ClipProcessor.load_model()
-#         assert model is not None
-#         assert preprocess is not None
+class TestClipProcessor:
+    def test_load_model(self):
+        # Test that the `load_model()` method loads the CLIP model and image preprocessing correctly.
+        model = ClipProcessor.load_model()
+        assert model is not None
 
-#     @pytest.mark.xfail(reason="This test is failing because of the missing CLIP dependency.")
-#     def test_get_image_features(self):
-#         # Clone the image to a temporary folder.
-#         with tempfile.TemporaryDirectory() as tmp_dir:
-#             urllib.request.urlretrieve("https://upload.wikimedia.org/wikipedia/en/a/a9/Example.jpg", "image.jpg")
+    def test_get_image_features(self):
+        # Clone the image to a temporary folder.
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            urllib.request.urlretrieve("https://upload.wikimedia.org/wikipedia/en/a/a9/Example.jpg", "image.jpg")
 
-#             image = Image.open("image.jpg")
-#             image.save(os.path.join(tmp_dir, "image.jpg"))
+            image = Image.open("image.jpg")
+            image.save(os.path.join(tmp_dir, "image.jpg"))
 
-#             # Get the image features.
-#             model, preprocess = ClipProcessor.load_model()
-#             ClipProcessor.get_image_features(os.path.join(tmp_dir, "image.jpg"), model, preprocess)
+            # Get the image features.
+            model = ClipProcessor.load_model()
+            ClipProcessor.get_image_features(os.path.join(tmp_dir, "image.jpg"), model)
 
-#             # Delete the temporary file.
-#             os.remove(os.path.join(tmp_dir, "image.jpg"))
+            # Delete the temporary file.
+            os.remove(os.path.join(tmp_dir, "image.jpg"))
 
-#     @pytest.mark.xfail(reason="This test is failing because of the missing CLIP dependency.")
-#     def test_get_text_features(self):
-#         # Test that the `get_text_features()` method returns a list containing the text embedding.
-#         query = "This is a text query."
-#         model, preprocess = ClipProcessor.load_model()
+    def test_get_text_features(self):
+        # Test that the `get_text_features()` method returns a list containing the text embedding.
+        query = "This is a text query."
+        text_features = ClipProcessor.get_text_features(query)
 
-#         text_features = ClipProcessor.get_text_features(query)
+        # Assert that the text embedding is not None.
+        assert text_features is not None
 
-#         # Assert that the text embedding is not None.
-#         assert text_features is not None
+        # Assert that the text embedding is a list of floats.
+        assert isinstance(text_features, list)
 
-#         # Assert that the text embedding is a list of floats.
-#         assert isinstance(text_features, list)
-
-#         # Assert that the text embedding has the correct length.
-#         assert len(text_features) == 512
+        # Assert that the text embedding has the correct length.
+        assert len(text_features) == 512