197 lines
6.9 KiB
Python
197 lines
6.9 KiB
Python
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
|