Add support for OpenSearch as vector database (#725)
This commit is contained in:
Deshraj Yadav
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
@@ -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):
|
||||
|
||||
196
embedchain/vectordb/opensearch.py
Normal file
196
embedchain/vectordb/opensearch.py
Normal file
@@ -0,0 +1,196 @@
|
||||
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,
|
||||
)
|
||||
docs = docsearch.similarity_search(
|
||||
input_query,
|
||||
search_type="script_scoring",
|
||||
space_type="cosinesimil",
|
||||
vector_field="embeddings",
|
||||
text_field="text",
|
||||
metadata_field="metadata",
|
||||
)
|
||||
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
|
||||
Reference in New Issue
Block a user