[OpenSearch] Add support for filtering docs based on app_id in opensearch db (#729)

Signed-off-by: Osama Nabih <nabih_the_4th@yahoo.com>

Makefile

@@ -7,8 +7,19 @@ PROJECT_NAME := embedchain
 .PHONY: install format lint clean test ci_lint ci_test
 
 install:
-	$(PIP) install --upgrade pip
-	$(PIP) install -e .[dev]
+	poetry install
+
+install_es:
+	poetry install --extras elasticsearch
+
+install_opensearch:
+	poetry install --extras opensearch
+
+shell:
+	poetry shell
+
+py_shell:
+	poetry run python
 
 format:
 	$(PYTHON) -m black .

docs/advanced/interface_types.mdx

@@ -70,6 +70,6 @@ app.reset()
 Counts the number of embeddings (chunks) in the database.
 
 ```python
-print(app.count())
+print(app.db.count())
 # returns: 481
 ```

docs/advanced/vector_database.mdx

@@ -2,7 +2,7 @@
 title: '💾 Vector Database'
 ---
 
-We support `Chroma` and `Elasticsearch` as two vector database. 
+We support `Chroma`, `Elasticsearch` and `OpenSearch` as vector databases.
 `Chroma` is used as a default database.
 
 ## Elasticsearch
@@ -22,13 +22,13 @@ Please note that the key needs certain privileges. For testing you can just togg
 2. Load the app
 ```python
 from embedchain import CustomApp
-from embedchain.embedder.openai import OpenAiEmbedder
+from embedchain.embedder.openai import OpenAIEmbedder
 from embedchain.llm.openai import OpenAILlm
 from embedchain.vectordb.elasticsearch import ElasticsearchDB
 
 es_app = CustomApp(
     llm=OpenAILlm(),
-    embedder=OpenAiEmbedder(),
+    embedder=OpenAIEmbedder(),
     db=ElasticsearchDB(),
 )
 ```
@@ -45,7 +45,7 @@ import os
 
 from embedchain import CustomApp
 from embedchain.config import CustomAppConfig, ElasticsearchDBConfig
-from embedchain.embedder.openai import OpenAiEmbedder
+from embedchain.embedder.openai import OpenAIEmbedder
 from embedchain.llm.openai import OpenAILlm
 from embedchain.vectordb.elasticsearch import ElasticsearchDB
 
@@ -61,10 +61,58 @@ es_config = ElasticsearchDBConfig(
 es_app = CustomApp(
     config=CustomAppConfig(log_level="INFO"),
     llm=OpenAILlm(),
-    embedder=OpenAiEmbedder(),
+    embedder=OpenAIEmbedder(),
     db=ElasticsearchDB(config=es_config),
 )
 ```
 3. This should log your connection details to the console.
 4. Alternatively to a URL, you `ElasticsearchDBConfig` accepts `es_url` as a list of nodes url with different hosts and ports.
 5. Additionally we can pass named parameters supported by Python Elasticsearch client.
+
+
+## OpenSearch 🔍
+
+To use OpenSearch as a vector database with a CustomApp, follow these simple steps:
+
+1. Set the `OPENAI_API_KEY` environment variable:
+
+```
+OPENAI_API_KEY=sk-xxxx
+```
+
+2. Define the OpenSearch configuration in your Python code:
+
+```python
+from embedchain import CustomApp
+from embedchain.config import OpenSearchDBConfig
+from embedchain.embedder.openai import OpenAIEmbedder
+from embedchain.llm.openai import OpenAILlm
+from embedchain.vectordb.opensearch import OpenSearchDB
+
+opensearch_url = "https://localhost:9200"
+http_auth = ("username", "password")
+
+db_config = OpenSearchDBConfig(
+    opensearch_url=opensearch_url,
+    http_auth=http_auth,
+    collection_name="embedchain-app",
+    use_ssl=True,
+    timeout=30,
+)
+db = OpenSearchDB(config=db_config)
+```
+
+2. Instantiate the app and add data:
+
+```python
+app = CustomApp(llm=OpenAILlm(), embedder=OpenAIEmbedder(), db=db)
+app.add("https://en.wikipedia.org/wiki/Elon_Musk")
+app.add("https://www.forbes.com/profile/elon-musk")
+app.add("https://www.britannica.com/biography/Elon-Musk")
+```
+
+3. You're all set! Start querying using the following command:
+
+```python
+app.query("What is the net worth of Elon Musk?")
+```

docs/integration/langsmith.mdx

@@ -0,0 +1,51 @@
+---
+title: '🛠️ LangSmith'
+description: 'Integrate with Langsmith to debug and monitor your LLM app'
+---
+
+Embedchain now supports integration with [LangSmith](https://www.langchain.com/langsmith).
+
+To use langsmith, you need to do the following steps
+
+1. Have an account on langsmith and keep the environment variables in handy
+2. Set the environments variables in your app so that embedchain has context about it.
+3. Just use embedchain and everything will be logged to LangSmith, so that you can better test and monitor your application.
+
+Lets cover each step in detail.
+
+* First make sure that you a LangSmith account created and have all the necessary variables handy. LangSmith has a [good documentation](https://docs.smith.langchain.com/) on how to get started with their service.
+
+* Once you have the account setup, we will need the following environment variables
+
+```bash
+export LANGCHAIN_TRACING_V2=true
+export LANGCHAIN_ENDPOINT=https://api.smith.langchain.com
+export LANGCHAIN_API_KEY=<your-api-key>
+export LANGCHAIN_PROJECT=<your-project>  # if not specified, defaults to "default"
+```
+
+If you are using Python, you can use the following code to set environment variables
+
+```python
+import os
+
+os.environ['LANGCHAIN_TRACING_V2'] = 'true'
+os.environ['LANGCHAIN_ENDPOINT'] = 'https://api.smith.langchain.com'
+os.environ['LANGCHAIN_API_KEY'] = <your-api-key>
+os.environ['LANGCHAIN_PROJECT] = <your-project>
+```
+
+* Now create an app using embedchain and everything will be automatically visible in the LangSmith
+
+
+```python
+from embedchain import App
+
+app = App()
+app.add("https://en.wikipedia.org/wiki/Elon_Musk")
+app.query("How many companies did Elon found?")
+```
+
+* Now the entire log for this will be visible in langsmith.
+
+<img src="/images/langsmith.png"/>

docs/mint.json

@@ -19,6 +19,10 @@
     {
       "name": "Discord",
       "url": "https://discord.gg/6PzXDgEjG5"
+    },
+    {
+      "name":"Slack",
+      "url":"https://join.slack.com/t/embedchain/shared_invite/zt-22uwz3c46-Zg7cIh5rOBteT_xe1jwLDw"
     }
   ],
   "topbarCtaButton": {
@@ -38,6 +42,10 @@
       "group": "Examples",
       "pages": ["examples/full_stack", "examples/api_server", "examples/discord_bot", "examples/slack_bot", "examples/telegram_bot", "examples/whatsapp_bot", "examples/poe_bot"]
     },
+    {
+      "group": "Integration",
+      "pages": ["integration/langsmith"]
+    },
     {
       "group": "Contribution Guidelines",
       "pages": ["contribution/dev", "contribution/docs"]
@@ -48,7 +56,8 @@
     "twitter": "https://twitter.com/embedchain",
     "github": "https://github.com/embedchain/embedchain",
     "linkedin": "https://www.linkedin.com/company/embedchain",
-    "website": "https://embedchain.ai"
+    "website": "https://embedchain.ai",
+    "slack":"https://join.slack.com/t/embedchain/shared_invite/zt-22uwz3c46-Zg7cIh5rOBteT_xe1jwLDw"
   },
   "backgroundImage": "/background.png",
   "isWhiteLabeled": true

embedchain/__init__.py

@@ -2,7 +2,7 @@ import importlib.metadata
 
 __version__ = importlib.metadata.version(__package__ or __name__)
 
-from embedchain.apps.App import App  # noqa: F401
+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

embedchain/apps/Llama2App.py

@@ -2,7 +2,7 @@ from typing import Optional
 
 from embedchain.apps.custom_app import CustomApp
 from embedchain.config import CustomAppConfig
-from embedchain.embedder.openai import OpenAiEmbedder
+from embedchain.embedder.openai import OpenAIEmbedder
 from embedchain.helper.json_serializable import register_deserializable
 from embedchain.llm.llama2 import Llama2Llm
 from embedchain.vectordb.chroma import ChromaDB
@@ -29,5 +29,5 @@ class Llama2App(CustomApp):
             config = CustomAppConfig()
 
         super().__init__(
-            config=config, llm=Llama2Llm(), db=ChromaDB(), embedder=OpenAiEmbedder(), system_prompt=system_prompt
+            config=config, llm=Llama2Llm(), db=ChromaDB(), embedder=OpenAIEmbedder(), system_prompt=system_prompt
         )

embedchain/apps/PersonApp.py

@@ -1,6 +1,6 @@
 from string import Template
 
-from embedchain.apps.App import App
+from embedchain.apps.app import App
 from embedchain.apps.open_source_app import OpenSourceApp
 from embedchain.config import BaseLlmConfig
 from embedchain.config.apps.base_app_config import BaseAppConfig

embedchain/apps/app.py

@@ -3,7 +3,7 @@ from typing import Optional
 from embedchain.config import (AppConfig, BaseEmbedderConfig, BaseLlmConfig,
                                ChromaDbConfig)
 from embedchain.embedchain import EmbedChain
-from embedchain.embedder.openai import OpenAiEmbedder
+from embedchain.embedder.openai import OpenAIEmbedder
 from embedchain.helper.json_serializable import register_deserializable
 from embedchain.llm.openai import OpenAILlm
 from embedchain.vectordb.chroma import ChromaDB
@@ -48,7 +48,7 @@ class App(EmbedChain):
             config = AppConfig()
 
         llm = OpenAILlm(config=llm_config)
-        embedder = OpenAiEmbedder(config=BaseEmbedderConfig(model="text-embedding-ada-002"))
+        embedder = OpenAIEmbedder(config=BaseEmbedderConfig(model="text-embedding-ada-002"))
         database = ChromaDB(config=chromadb_config)
 
         super().__init__(config, llm, db=database, embedder=embedder, system_prompt=system_prompt)

embedchain/bots/base.py

@@ -2,7 +2,7 @@ from typing import Any
 
 from embedchain import CustomApp
 from embedchain.config import AddConfig, CustomAppConfig, LlmConfig
-from embedchain.embedder.openai import OpenAiEmbedder
+from embedchain.embedder.openai import OpenAIEmbedder
 from embedchain.helper.json_serializable import (JSONSerializable,
                                                  register_deserializable)
 from embedchain.llm.openai import OpenAILlm
@@ -12,7 +12,7 @@ from embedchain.vectordb.chroma import ChromaDB
 @register_deserializable
 class BaseBot(JSONSerializable):
     def __init__(self):
-        self.app = CustomApp(config=CustomAppConfig(), llm=OpenAILlm(), db=ChromaDB(), embedder=OpenAiEmbedder())
+        self.app = CustomApp(config=CustomAppConfig(), llm=OpenAILlm(), db=ChromaDB(), embedder=OpenAIEmbedder())
 
     def add(self, data: Any, config: AddConfig = None):
         """

embedchain/chunkers/docs_site.py

@@ -3,7 +3,7 @@ from typing import Optional
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 
 from embedchain.chunkers.base_chunker import BaseChunker
-from embedchain.config.AddConfig import ChunkerConfig
+from embedchain.config.add_config import ChunkerConfig
 from embedchain.helper.json_serializable import register_deserializable
 
 

embedchain/chunkers/docx_file.py

@@ -3,7 +3,7 @@ from typing import Optional
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 
 from embedchain.chunkers.base_chunker import BaseChunker
-from embedchain.config.AddConfig import ChunkerConfig
+from embedchain.config.add_config import ChunkerConfig
 from embedchain.helper.json_serializable import register_deserializable
 
 

embedchain/chunkers/mdx.py

@@ -3,7 +3,7 @@ from typing import Optional
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 
 from embedchain.chunkers.base_chunker import BaseChunker
-from embedchain.config.AddConfig import ChunkerConfig
+from embedchain.config.add_config import ChunkerConfig
 from embedchain.helper.json_serializable import register_deserializable
 
 

embedchain/chunkers/notion.py

@@ -3,7 +3,7 @@ from typing import Optional
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 
 from embedchain.chunkers.base_chunker import BaseChunker
-from embedchain.config.AddConfig import ChunkerConfig
+from embedchain.config.add_config import ChunkerConfig
 from embedchain.helper.json_serializable import register_deserializable
 
 

embedchain/chunkers/pdf_file.py

@@ -3,7 +3,7 @@ from typing import Optional
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 
 from embedchain.chunkers.base_chunker import BaseChunker
-from embedchain.config.AddConfig import ChunkerConfig
+from embedchain.config.add_config import ChunkerConfig
 from embedchain.helper.json_serializable import register_deserializable
 
 

embedchain/chunkers/qna_pair.py

@@ -3,7 +3,7 @@ from typing import Optional
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 
 from embedchain.chunkers.base_chunker import BaseChunker
-from embedchain.config.AddConfig import ChunkerConfig
+from embedchain.config.add_config import ChunkerConfig
 from embedchain.helper.json_serializable import register_deserializable
 
 

embedchain/chunkers/table.py

@@ -3,7 +3,7 @@ from typing import Optional
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 
 from embedchain.chunkers.base_chunker import BaseChunker
-from embedchain.config.AddConfig import ChunkerConfig
+from embedchain.config.add_config import ChunkerConfig
 
 
 class TableChunker(BaseChunker):

embedchain/chunkers/text.py

@@ -3,7 +3,7 @@ from typing import Optional
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 
 from embedchain.chunkers.base_chunker import BaseChunker
-from embedchain.config.AddConfig import ChunkerConfig
+from embedchain.config.add_config import ChunkerConfig
 from embedchain.helper.json_serializable import register_deserializable
 
 

embedchain/chunkers/web_page.py

@@ -3,7 +3,7 @@ from typing import Optional
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 
 from embedchain.chunkers.base_chunker import BaseChunker
-from embedchain.config.AddConfig import ChunkerConfig
+from embedchain.config.add_config import ChunkerConfig
 from embedchain.helper.json_serializable import register_deserializable
 
 

embedchain/chunkers/youtube_video.py

@@ -3,7 +3,7 @@ from typing import Optional
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 
 from embedchain.chunkers.base_chunker import BaseChunker
-from embedchain.config.AddConfig import ChunkerConfig
+from embedchain.config.add_config import ChunkerConfig
 from embedchain.helper.json_serializable import register_deserializable
 
 

embedchain/config/__init__.py

@@ -1,13 +1,14 @@
 # flake8: noqa: F401
 
-from .AddConfig import AddConfig, ChunkerConfig
+from .add_config import AddConfig, ChunkerConfig
 from .apps.app_config import AppConfig
 from .apps.custom_app_config import CustomAppConfig
 from .apps.open_source_app_config import OpenSourceAppConfig
-from .BaseConfig import BaseConfig
-from .embedder.BaseEmbedderConfig import BaseEmbedderConfig
-from .embedder.BaseEmbedderConfig import BaseEmbedderConfig as EmbedderConfig
+from .base_config import BaseConfig
+from .embedder.base import BaseEmbedderConfig
+from .embedder.base import BaseEmbedderConfig as EmbedderConfig
 from .llm.base_llm_config import BaseLlmConfig
 from .llm.base_llm_config import BaseLlmConfig as LlmConfig
-from .vectordbs.ChromaDbConfig import ChromaDbConfig
-from .vectordbs.ElasticsearchDBConfig import ElasticsearchDBConfig
+from .vectordb.chroma import ChromaDbConfig
+from .vectordb.elasticsearch import ElasticsearchDBConfig
+from .vectordb.opensearch import OpenSearchDBConfig

embedchain/config/add_config.py

@@ -1,6 +1,6 @@
 from typing import Callable, Optional
 
-from embedchain.config.BaseConfig import BaseConfig
+from embedchain.config.base_config import BaseConfig
 from embedchain.helper.json_serializable import register_deserializable
 
 

embedchain/config/apps/base_app_config.py

@@ -1,7 +1,7 @@
 import logging
 from typing import Optional
 
-from embedchain.config.BaseConfig import BaseConfig
+from embedchain.config.base_config import BaseConfig
 from embedchain.helper.json_serializable import JSONSerializable
 from embedchain.vectordb.base import BaseVectorDB
 

embedchain/config/llm/base_llm_config.py

@@ -2,7 +2,7 @@ import re
 from string import Template
 from typing import Any, Dict, Optional
 
-from embedchain.config.BaseConfig import BaseConfig
+from embedchain.config.base_config import BaseConfig
 from embedchain.helper.json_serializable import register_deserializable
 
 DEFAULT_PROMPT = """

embedchain/config/vectordb/base.py

@@ -1,6 +1,6 @@
 from typing import Optional
 
-from embedchain.config.BaseConfig import BaseConfig
+from embedchain.config.base_config import BaseConfig
 
 
 class BaseVectorDbConfig(BaseConfig):

embedchain/config/vectordb/chroma.py

@@ -1,6 +1,6 @@
 from typing import Optional
 
-from embedchain.config.vectordbs.BaseVectorDbConfig import BaseVectorDbConfig
+from embedchain.config.vectordb.base import BaseVectorDbConfig
 from embedchain.helper.json_serializable import register_deserializable
 
 

embedchain/config/vectordb/elasticsearch.py

@@ -1,7 +1,7 @@
 import os
 from typing import Dict, List, Optional, Union
 
-from embedchain.config.vectordbs.BaseVectorDbConfig import BaseVectorDbConfig
+from embedchain.config.vectordb.base import BaseVectorDbConfig
 from embedchain.helper.json_serializable import register_deserializable
 
 

embedchain/config/vectordb/opensearch.py

@@ -0,0 +1,37 @@
+from typing import Dict, Optional, Tuple
+
+from embedchain.config.vectordb.base import BaseVectorDbConfig
+from embedchain.helper.json_serializable import register_deserializable
+
+
+@register_deserializable
+class OpenSearchDBConfig(BaseVectorDbConfig):
+    def __init__(
+        self,
+        opensearch_url: str,
+        http_auth: Tuple[str, str],
+        vector_dimension: int = 1536,
+        collection_name: Optional[str] = None,
+        dir: Optional[str] = None,
+        **extra_params: Dict[str, any],
+    ):
+        """
+        Initializes a configuration class instance for an OpenSearch client.
+
+        :param collection_name: Default name for the collection, defaults to None
+        :type collection_name: Optional[str], optional
+        :param opensearch_url: URL of the OpenSearch domain
+        :type opensearch_url: str, Eg, "http://localhost:9200"
+        :param http_auth: Tuple of username and password
+        :type http_auth: Tuple[str, str], Eg, ("username", "password")
+        :param vector_dimension: Dimension of  the vector, defaults to 1536 (openai embedding model)
+        :type vector_dimension: int, optional
+        :param dir: Path to the database directory, where the database is stored, defaults to None
+        :type dir: Optional[str], optional
+        """
+        self.opensearch_url = opensearch_url
+        self.http_auth = http_auth
+        self.vector_dimension = vector_dimension
+        self.extra_params = extra_params
+
+        super().__init__(collection_name=collection_name, dir=dir)

embedchain/data_formatter/data_formatter.py

@@ -10,7 +10,7 @@ from embedchain.chunkers.text import TextChunker
 from embedchain.chunkers.web_page import WebPageChunker
 from embedchain.chunkers.youtube_video import YoutubeVideoChunker
 from embedchain.config import AddConfig
-from embedchain.config.AddConfig import ChunkerConfig, LoaderConfig
+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

embedchain/embedchain.py

@@ -21,7 +21,8 @@ from embedchain.embedder.base import BaseEmbedder
 from embedchain.helper.json_serializable import JSONSerializable
 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.models.data_type import (DataType, DirectDataType,
+                                         IndirectDataType, SpecialDataType)
 from embedchain.utils import detect_datatype
 from embedchain.vectordb.base import BaseVectorDB
 
@@ -60,16 +61,13 @@ class EmbedChain(JSONSerializable):
         """
 
         self.config = config
-
-        # Add subclasses
-        ## Llm
+        # Llm
         self.llm = llm
-        ## Database
         # Database has support for config assignment for backwards compatibility
         if db is None and (not hasattr(self.config, "db") or self.config.db is None):
             raise ValueError("App requires Database.")
         self.db = db or self.config.db
-        ## Embedder
+        # Embedder
         if embedder is None:
             raise ValueError("App requires Embedder.")
         self.embedder = embedder
@@ -255,7 +253,6 @@ class EmbedChain(JSONSerializable):
         )
         return self.add(source=source, data_type=data_type, metadata=metadata, config=config)
 
-
     def _get_existing_doc_id(self, chunker: BaseChunker, src: Any):
         """
         Get id of existing document for a given source, based on the data type
@@ -394,10 +391,10 @@ class EmbedChain(JSONSerializable):
             return list(documents), metadatas, ids, 0
 
         # Count before, to calculate a delta in the end.
-        chunks_before_addition = self.count()
+        chunks_before_addition = self.db.count()
 
         self.db.add(documents=documents, metadatas=metadatas, ids=ids)
-        count_new_chunks = self.count() - chunks_before_addition
+        count_new_chunks = self.db.count() - chunks_before_addition
         print((f"Successfully saved {src} ({chunker.data_type}). New chunks count: {count_new_chunks}"))
         return list(documents), metadatas, ids, count_new_chunks
 

embedchain/embedder/base.py

@@ -1,6 +1,6 @@
 from typing import Any, Callable, Optional
 
-from embedchain.config.embedder.BaseEmbedderConfig import BaseEmbedderConfig
+from embedchain.config.embedder.base import BaseEmbedderConfig
 
 try:
     from chromadb.api.types import Documents, Embeddings

embedchain/embedder/openai.py

@@ -16,7 +16,7 @@ except RuntimeError:
     from chromadb.utils import embedding_functions
 
 
-class OpenAiEmbedder(BaseEmbedder):
+class OpenAIEmbedder(BaseEmbedder):
     def __init__(self, config: Optional[BaseEmbedderConfig] = None):
         super().__init__(config=config)
         if self.config.model is None:

embedchain/llm/azure_openai.py

@@ -12,10 +12,10 @@ class AzureOpenAILlm(BaseLlm):
         super().__init__(config=config)
 
     def get_llm_model_answer(self, prompt):
-        return AzureOpenAILlm._get_azure_openai_answer(prompt=prompt, config=self.config)
+        return AzureOpenAILlm._get_answer(prompt=prompt, config=self.config)
 
     @staticmethod
-    def _get_azure_openai_answer(prompt: str, config: BaseLlmConfig) -> str:
+    def _get_answer(prompt: str, config: BaseLlmConfig) -> str:
         from langchain.chat_models import AzureChatOpenAI
 
         if not config.deployment_name:

embedchain/llm/openai.py

@@ -1,6 +1,7 @@
 from typing import Optional
 
-import openai
+from langchain.chat_models import ChatOpenAI
+from langchain.schema import HumanMessage, SystemMessage
 
 from embedchain.config import BaseLlmConfig
 from embedchain.helper.json_serializable import register_deserializable
@@ -12,31 +13,32 @@ class OpenAILlm(BaseLlm):
     def __init__(self, config: Optional[BaseLlmConfig] = None):
         super().__init__(config=config)
 
-    # NOTE: This class does not use langchain. One reason is that `top_p` is not supported.
-
     def get_llm_model_answer(self, prompt):
-        messages = []
-        if self.config.system_prompt:
-            messages.append({"role": "system", "content": self.config.system_prompt})
-        messages.append({"role": "user", "content": prompt})
-        response = openai.ChatCompletion.create(
-            model=self.config.model or "gpt-3.5-turbo-0613",
-            messages=messages,
-            temperature=self.config.temperature,
-            max_tokens=self.config.max_tokens,
-            top_p=self.config.top_p,
-            stream=self.config.stream,
-        )
+        response = OpenAILlm._get_answer(prompt, self.config)
 
         if self.config.stream:
-            return self._stream_llm_model_response(response)
+            return response
         else:
-            return response["choices"][0]["message"]["content"]
+            return response.content
 
-    def _stream_llm_model_response(self, response):
-        """
-        This is a generator for streaming response from the OpenAI completions API
-        """
-        for line in response:
-            chunk = line["choices"][0].get("delta", {}).get("content", "")
-            yield chunk
+    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 = {
+            "model": config.model or "gpt-3.5-turbo-0613",
+            "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 = ChatOpenAI(**kwargs, streaming=config.stream, callbacks=[StreamingStdOutCallbackHandler()])
+        else:
+            chat = ChatOpenAI(**kwargs)
+        return chat(messages)

embedchain/models/__init__.py

@@ -1,4 +1,4 @@
-from .EmbeddingFunctions import EmbeddingFunctions  # noqa: F401
-from .Providers import Providers  # noqa: F401
+from .embedding_functions import EmbeddingFunctions  # noqa: F401
+from .providers import Providers  # noqa: F401
 from .vector_databases import VectorDatabases  # noqa: F401
-from .VectorDimensions import VectorDimensions  # noqa: F401
+from .vector_dimensions import VectorDimensions  # noqa: F401

embedchain/models/vector_databases.py

@@ -4,3 +4,4 @@ from enum import Enum
 class VectorDatabases(Enum):
     CHROMADB = "CHROMADB"
     ELASTICSEARCH = "ELASTICSEARCH"
+    OPENSEARCH = "OPENSEARCH"

embedchain/vectordb/base.py

@@ -1,4 +1,4 @@
-from embedchain.config.vectordbs.BaseVectorDbConfig import BaseVectorDbConfig
+from embedchain.config.vectordb.base import BaseVectorDbConfig
 from embedchain.embedder.base import BaseEmbedder
 from embedchain.helper.json_serializable import JSONSerializable
 

embedchain/vectordb/base_vector_db.py

@@ -1,50 +0,0 @@
-from embedchain.config.vectordbs.BaseVectorDbConfig import BaseVectorDbConfig
-from embedchain.embedder.base_embedder import BaseEmbedder
-from embedchain.helper_classes.json_serializable import JSONSerializable
-
-
-class BaseVectorDB(JSONSerializable):
-    """Base class for vector database."""
-
-    def __init__(self, config: BaseVectorDbConfig):
-        self.client = self._get_or_create_db()
-        self.config: BaseVectorDbConfig = config
-
-    def _initialize(self):
-        """
-        This method is needed because `embedder` attribute needs to be set externally before it can be initialized.
-
-        So it's can't be done in __init__ in one step.
-        """
-        raise NotImplementedError
-
-    def _get_or_create_db(self):
-        """Get or create the database."""
-        raise NotImplementedError
-
-    def _get_or_create_collection(self):
-        raise NotImplementedError
-
-    def _set_embedder(self, embedder: BaseEmbedder):
-        self.embedder = embedder
-
-    def get(self):
-        raise NotImplementedError
-
-    def add(self):
-        raise NotImplementedError
-
-    def query(self):
-        raise NotImplementedError
-
-    def count(self):
-        raise NotImplementedError
-
-    def delete(self):
-        raise NotImplementedError
-
-    def reset(self):
-        raise NotImplementedError
-
-    def set_collection_name(self, name: str):
-        raise NotImplementedError

embedchain/vectordb/chroma.py

@@ -63,7 +63,9 @@ class ChromaDB(BaseVectorDB):
         This method is needed because `embedder` attribute needs to be set externally before it can be initialized.
         """
         if not self.embedder:
-            raise ValueError("Embedder not set. Please set an embedder with `set_embedder` before initialization.")
+            raise ValueError(
+                "Embedder not set. Please set an embedder with `_set_embedder()` function before initialization."
+            )
         self._get_or_create_collection(self.config.collection_name)
 
     def _get_or_create_db(self):

embedchain/vectordb/opensearch.py

@@ -0,0 +1,203 @@
+import logging
+from typing import Dict, List, Optional, Set
+
+try:
+    from opensearchpy import OpenSearch
+    from opensearchpy.helpers import bulk
+except ImportError:
+    raise ImportError(
+        "OpenSearch requires extra dependencies. Install with `pip install --upgrade embedchain[opensearch]`"
+    ) from None
+
+from langchain.embeddings.openai import OpenAIEmbeddings
+from langchain.vectorstores import OpenSearchVectorSearch
+
+from embedchain.config import OpenSearchDBConfig
+from embedchain.helper.json_serializable import register_deserializable
+from embedchain.vectordb.base import BaseVectorDB
+
+
+@register_deserializable
+class OpenSearchDB(BaseVectorDB):
+    """
+    OpenSearch as vector database
+    """
+
+    def __init__(self, config: OpenSearchDBConfig):
+        """OpenSearch as vector database.
+
+        :param config: OpenSearch domain config
+        :type config: OpenSearchDBConfig
+        """
+        if config is None:
+            raise ValueError("OpenSearchDBConfig is required")
+        self.config = config
+        self.client = OpenSearch(
+            hosts=[self.config.opensearch_url],
+            http_auth=self.config.http_auth,
+            **self.config.extra_params,
+        )
+        info = self.client.info()
+        logging.info(f"Connected to {info['version']['distribution']}. Version: {info['version']['number']}")
+        # Remove auth credentials from config after successful connection
+        super().__init__(config=self.config)
+
+    def _initialize(self):
+        logging.info(self.client.info())
+        index_name = self._get_index()
+        if self.client.indices.exists(index=index_name):
+            print(f"Index '{index_name}' already exists.")
+            return
+
+        index_body = {
+            "settings": {"knn": True},
+            "mappings": {
+                "properties": {
+                    "text": {"type": "text"},
+                    "embeddings": {
+                        "type": "knn_vector",
+                        "index": False,
+                        "dimension": self.config.vector_dimension,
+                    },
+                }
+            },
+        }
+        self.client.indices.create(index_name, body=index_body)
+        print(self.client.indices.get(index_name))
+
+    def _get_or_create_db(self):
+        """Called during initialization"""
+        return self.client
+
+    def _get_or_create_collection(self, name):
+        """Note: nothing to return here. Discuss later"""
+
+    def get(
+        self, ids: Optional[List[str]] = None, where: Optional[Dict[str, any]] = None, limit: Optional[int] = None
+    ) -> Set[str]:
+        """
+        Get existing doc ids present in vector database
+
+        :param ids: _list of doc ids to check for existence
+        :type ids: List[str]
+        :param where: to filter data
+        :type where: Dict[str, any]
+        :return: ids
+        :type: Set[str]
+        """
+        if ids:
+            query = {"query": {"bool": {"must": [{"ids": {"values": ids}}]}}}
+        else:
+            query = {"query": {"bool": {"must": []}}}
+        if "app_id" in where:
+            app_id = where["app_id"]
+            query["query"]["bool"]["must"].append({"term": {"metadata.app_id": app_id}})
+
+        # OpenSearch syntax is different from Elasticsearch
+        response = self.client.search(index=self._get_index(), body=query, _source=False, size=limit)
+        docs = response["hits"]["hits"]
+        ids = [doc["_id"] for doc in docs]
+        return {"ids": set(ids)}
+
+    def add(self, documents: List[str], metadatas: List[object], ids: List[str]):
+        """add data in vector database
+
+        :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]
+        """
+
+        docs = []
+        embeddings = self.embedder.embedding_fn(documents)
+        for id, text, metadata, embeddings in zip(ids, documents, metadatas, embeddings):
+            docs.append(
+                {
+                    "_index": self._get_index(),
+                    "_id": id,
+                    "_source": {"text": text, "metadata": metadata, "embeddings": embeddings},
+                }
+            )
+        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]) -> List[str]:
+        """
+        query contents from vector data base based on vector similarity
+
+        :param input_query: list of query string
+        :type input_query: List[str]
+        :param n_results: no of similar documents to fetch from database
+        :type n_results: int
+        :param where: Optional. to filter data
+        :type where: Dict[str, any]
+        :return: Database contents that are the result of the query
+        :rtype: List[str]
+        """
+        embeddings = OpenAIEmbeddings()
+        docsearch = OpenSearchVectorSearch(
+            index_name=self._get_index(),
+            embedding_function=embeddings,
+            opensearch_url=f"{self.config.opensearch_url}",
+            http_auth=self.config.http_auth,
+            use_ssl=True,
+        )
+
+        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}}]}}
+        docs = docsearch.similarity_search(
+            input_query,
+            search_type="script_scoring",
+            space_type="cosinesimil",
+            vector_field="embeddings",
+            text_field="text",
+            metadata_field="metadata",
+            pre_filter=pre_filter,
+            k=n_results,
+        )
+        contents = [doc.page_content for doc in docs]
+        return contents
+
+    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
+
+    def count(self) -> int:
+        """
+        Count number of documents/chunks embedded in the database.
+
+        :return: number of documents
+        :rtype: int
+        """
+        query = {"query": {"match_all": {}}}
+        response = self.client.count(index=self._get_index(), body=query)
+        doc_count = response["count"]
+        return doc_count
+
+    def reset(self):
+        """
+        Resets the database. Deletes all embeddings irreversibly.
+        """
+        # Delete all data from the database
+        if self.client.indices.exists(index=self._get_index()):
+            # delete index in Es
+            self.client.indices.delete(index=self._get_index())
+
+    def _get_index(self) -> str:
+        """Get the OpenSearch index for a collection
+
+        :return: OpenSearch index
+        :rtype: str
+        """
+        return self.config.collection_name

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "embedchain"
-version = "0.0.60"
+version = "0.0.63"
 description = "embedchain is a framework to easily create LLM powered bots over any dataset"
 authors = ["Taranjeet Singh"]
 license = "Apache License"
@@ -98,6 +98,7 @@ torch = { version = ">=2.0.0, !=2.0.1", optional = true }
 # Torch 2.0.1 is not compatible with poetry (https://github.com/pytorch/pytorch/issues/100974)
 gpt4all = { version = "1.0.8", optional = true }
 # 1.0.9 is not working for some users (https://github.com/nomic-ai/gpt4all/issues/1394)
+opensearch-py = { version = "2.3.1", optional = true }
 elasticsearch = { version = "^8.9.0", optional = true }
 flask = { version = "^2.3.3", optional = true }
 twilio = { version = "^8.5.0", optional = true }
@@ -123,6 +124,7 @@ streamlit = ["streamlit"]
 community = ["llama-hub"]
 opensource = ["sentence-transformers", "torch", "gpt4all"]
 elasticsearch = ["elasticsearch"]
+opensearch = ["opensearch-py"]
 poe = ["fastapi-poe"]
 discord = ["discord"]
 slack = ["slack-sdk", "flask"]

tests/llm/test_query.py

@@ -46,41 +46,29 @@ class TestApp(unittest.TestCase):
         self.assertEqual(input_query_arg, "Test query")
         mock_answer.assert_called_once()
 
-    @patch("openai.ChatCompletion.create")
-    def test_query_config_app_passing(self, mock_create):
-        mock_create.return_value = {"choices": [{"message": {"content": "response"}}]}  # Mock response
+    @patch("embedchain.llm.openai.OpenAILlm._get_answer")
+    def test_query_config_app_passing(self, mock_get_answer):
+        mock_get_answer.return_value = MagicMock()
+        mock_get_answer.return_value.content = "Test answer"
 
         config = AppConfig(collect_metrics=False)
         chat_config = BaseLlmConfig(system_prompt="Test system prompt")
         app = App(config=config, llm_config=chat_config)
+        answer = app.llm.get_llm_model_answer("Test query")
 
-        app.llm.get_llm_model_answer("Test query")
-
-        # Test system_prompt: Check that the 'create' method was called with the correct 'messages' argument
-        messages_arg = mock_create.call_args.kwargs["messages"]
-        self.assertTrue(messages_arg[0].get("role"), "system")
-        self.assertEqual(messages_arg[0].get("content"), "Test system prompt")
-        self.assertTrue(messages_arg[1].get("role"), "user")
-        self.assertEqual(messages_arg[1].get("content"), "Test query")
-
-        # TODO: Add tests for other config variables
-
-    @patch("openai.ChatCompletion.create")
-    def test_app_passing(self, mock_create):
-        mock_create.return_value = {"choices": [{"message": {"content": "response"}}]}  # Mock response
+        self.assertEqual(app.llm.config.system_prompt, "Test system prompt")
+        self.assertEqual(answer, "Test answer")
 
+    @patch("embedchain.llm.openai.OpenAILlm._get_answer")
+    def test_app_passing(self, mock_get_answer):
+        mock_get_answer.return_value = MagicMock()
+        mock_get_answer.return_value.content = "Test answer"
         config = AppConfig(collect_metrics=False)
         chat_config = BaseLlmConfig()
         app = App(config=config, llm_config=chat_config, system_prompt="Test system prompt")
-
+        answer = app.llm.get_llm_model_answer("Test query")
         self.assertEqual(app.llm.config.system_prompt, "Test system prompt")
-
-        app.llm.get_llm_model_answer("Test query")
-
-        # Test system_prompt: Check that the 'create' method was called with the correct 'messages' argument
-        messages_arg = mock_create.call_args.kwargs["messages"]
-        self.assertTrue(messages_arg[0].get("role"), "system")
-        self.assertEqual(messages_arg[0].get("content"), "Test system prompt")
+        self.assertEqual(answer, "Test answer")
 
     @patch("chromadb.api.models.Collection.Collection.add", MagicMock)
     def test_query_with_where_in_params(self):