feat: support elasticsearch vector database (#3558)
Co-authored-by: miendinh <miendinh@users.noreply.github.com> Co-authored-by: crazywoola <100913391+crazywoola@users.noreply.github.com> Co-authored-by: crazywoola <427733928@qq.com>pull/7234/head
miendinh
2 years ago
committed by
GitHub
parent
4423710a13
commit
f104b930cf
@ -0,0 +1,191 @@
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
import requests
|
||||
from elasticsearch import Elasticsearch
|
||||
from flask import current_app
|
||||
from pydantic import BaseModel, model_validator
|
||||
|
||||
from core.rag.datasource.entity.embedding import Embeddings
|
||||
from core.rag.datasource.vdb.vector_base import BaseVector
|
||||
from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory
|
||||
from core.rag.datasource.vdb.vector_type import VectorType
|
||||
from core.rag.models.document import Document
|
||||
from models.dataset import Dataset
|
||||
|
||||
|
||||
class ElasticSearchConfig(BaseModel):
|
||||
host: str
|
||||
port: str
|
||||
username: str
|
||||
password: str
|
||||
|
||||
@model_validator(mode='before')
|
||||
def validate_config(cls, values: dict) -> dict:
|
||||
if not values['host']:
|
||||
raise ValueError("config HOST is required")
|
||||
if not values['port']:
|
||||
raise ValueError("config PORT is required")
|
||||
if not values['username']:
|
||||
raise ValueError("config USERNAME is required")
|
||||
if not values['password']:
|
||||
raise ValueError("config PASSWORD is required")
|
||||
return values
|
||||
|
||||
|
||||
class ElasticSearchVector(BaseVector):
|
||||
def __init__(self, index_name: str, config: ElasticSearchConfig, attributes: list):
|
||||
super().__init__(index_name.lower())
|
||||
self._client = self._init_client(config)
|
||||
self._attributes = attributes
|
||||
|
||||
def _init_client(self, config: ElasticSearchConfig) -> Elasticsearch:
|
||||
try:
|
||||
client = Elasticsearch(
|
||||
hosts=f'{config.host}:{config.port}',
|
||||
basic_auth=(config.username, config.password),
|
||||
request_timeout=100000,
|
||||
retry_on_timeout=True,
|
||||
max_retries=10000,
|
||||
)
|
||||
except requests.exceptions.ConnectionError:
|
||||
raise ConnectionError("Vector database connection error")
|
||||
|
||||
return client
|
||||
|
||||
def get_type(self) -> str:
|
||||
return 'elasticsearch'
|
||||
|
||||
def add_texts(self, documents: list[Document], embeddings: list[list[float]], **kwargs):
|
||||
uuids = self._get_uuids(documents)
|
||||
texts = [d.page_content for d in documents]
|
||||
metadatas = [d.metadata for d in documents]
|
||||
|
||||
if not self._client.indices.exists(index=self._collection_name):
|
||||
dim = len(embeddings[0])
|
||||
mapping = {
|
||||
"properties": {
|
||||
"text": {
|
||||
"type": "text"
|
||||
},
|
||||
"vector": {
|
||||
"type": "dense_vector",
|
||||
"index": True,
|
||||
"dims": dim,
|
||||
"similarity": "l2_norm"
|
||||
},
|
||||
}
|
||||
}
|
||||
self._client.indices.create(index=self._collection_name, mappings=mapping)
|
||||
|
||||
added_ids = []
|
||||
for i, text in enumerate(texts):
|
||||
self._client.index(index=self._collection_name,
|
||||
id=uuids[i],
|
||||
document={
|
||||
"text": text,
|
||||
"vector": embeddings[i] if embeddings[i] else None,
|
||||
"metadata": metadatas[i] if metadatas[i] else {},
|
||||
})
|
||||
added_ids.append(uuids[i])
|
||||
|
||||
self._client.indices.refresh(index=self._collection_name)
|
||||
return uuids
|
||||
|
||||
def text_exists(self, id: str) -> bool:
|
||||
return self._client.exists(index=self._collection_name, id=id).__bool__()
|
||||
|
||||
def delete_by_ids(self, ids: list[str]) -> None:
|
||||
for id in ids:
|
||||
self._client.delete(index=self._collection_name, id=id)
|
||||
|
||||
def delete_by_metadata_field(self, key: str, value: str) -> None:
|
||||
query_str = {
|
||||
'query': {
|
||||
'match': {
|
||||
f'metadata.{key}': f'{value}'
|
||||
}
|
||||
}
|
||||
}
|
||||
results = self._client.search(index=self._collection_name, body=query_str)
|
||||
ids = [hit['_id'] for hit in results['hits']['hits']]
|
||||
if ids:
|
||||
self.delete_by_ids(ids)
|
||||
|
||||
def delete(self) -> None:
|
||||
self._client.indices.delete(index=self._collection_name)
|
||||
|
||||
def search_by_vector(self, query_vector: list[float], **kwargs: Any) -> list[Document]:
|
||||
query_str = {
|
||||
"query": {
|
||||
"script_score": {
|
||||
"query": {
|
||||
"match_all": {}
|
||||
},
|
||||
"script": {
|
||||
"source": "cosineSimilarity(params.query_vector, 'vector') + 1.0",
|
||||
"params": {
|
||||
"query_vector": query_vector
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
results = self._client.search(index=self._collection_name, body=query_str)
|
||||
|
||||
docs_and_scores = []
|
||||
for hit in results['hits']['hits']:
|
||||
docs_and_scores.append(
|
||||
(Document(page_content=hit['_source']['text'], metadata=hit['_source']['metadata']), hit['_score']))
|
||||
|
||||
docs = []
|
||||
for doc, score in docs_and_scores:
|
||||
score_threshold = kwargs.get("score_threshold", .0) if kwargs.get('score_threshold', .0) else 0.0
|
||||
if score > score_threshold:
|
||||
doc.metadata['score'] = score
|
||||
docs.append(doc)
|
||||
|
||||
# Sort the documents by score in descending order
|
||||
docs = sorted(docs, key=lambda x: x.metadata['score'], reverse=True)
|
||||
|
||||
return docs
|
||||
def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]:
|
||||
query_str = {
|
||||
"match": {
|
||||
"text": query
|
||||
}
|
||||
}
|
||||
results = self._client.search(index=self._collection_name, query=query_str)
|
||||
docs = []
|
||||
for hit in results['hits']['hits']:
|
||||
docs.append(Document(page_content=hit['_source']['text'], metadata=hit['_source']['metadata']))
|
||||
|
||||
return docs
|
||||
|
||||
def create(self, texts: list[Document], embeddings: list[list[float]], **kwargs):
|
||||
return self.add_texts(texts, embeddings, **kwargs)
|
||||
|
||||
|
||||
class ElasticSearchVectorFactory(AbstractVectorFactory):
|
||||
def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> ElasticSearchVector:
|
||||
if dataset.index_struct_dict:
|
||||
class_prefix: str = dataset.index_struct_dict['vector_store']['class_prefix']
|
||||
collection_name = class_prefix
|
||||
else:
|
||||
dataset_id = dataset.id
|
||||
collection_name = Dataset.gen_collection_name_by_id(dataset_id)
|
||||
dataset.index_struct = json.dumps(
|
||||
self.gen_index_struct_dict(VectorType.ELASTICSEARCH, collection_name))
|
||||
|
||||
config = current_app.config
|
||||
return ElasticSearchVector(
|
||||
index_name=collection_name,
|
||||
config=ElasticSearchConfig(
|
||||
host=config.get('ELASTICSEARCH_HOST'),
|
||||
port=config.get('ELASTICSEARCH_PORT'),
|
||||
username=config.get('ELASTICSEARCH_USERNAME'),
|
||||
password=config.get('ELASTICSEARCH_PASSWORD'),
|
||||
),
|
||||
attributes=[]
|
||||
)
|
||||
@ -0,0 +1,25 @@
|
||||
from core.rag.datasource.vdb.elasticsearch.elasticsearch_vector import ElasticSearchConfig, ElasticSearchVector
|
||||
from tests.integration_tests.vdb.test_vector_store import (
|
||||
AbstractVectorTest,
|
||||
setup_mock_redis,
|
||||
)
|
||||
|
||||
|
||||
class ElasticSearchVectorTest(AbstractVectorTest):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.attributes = ['doc_id', 'dataset_id', 'document_id', 'doc_hash']
|
||||
self.vector = ElasticSearchVector(
|
||||
index_name=self.collection_name.lower(),
|
||||
config=ElasticSearchConfig(
|
||||
host='http://localhost',
|
||||
port='9200',
|
||||
username='elastic',
|
||||
password='elastic'
|
||||
),
|
||||
attributes=self.attributes
|
||||
)
|
||||
|
||||
|
||||
def test_elasticsearch_vector(setup_mock_redis):
|
||||
ElasticSearchVectorTest().run_all_tests()
|
||||
File diff suppressed because it is too large
Load Diff
Loading…
Reference in New Issue