""" NOTE: This vector database integration is community-supported and maintained on a best-effort basis. """ from typing import Optional from rexpro_ai.config import ( OPENSEARCH_CERT_VERIFY, OPENSEARCH_PASSWORD, OPENSEARCH_SSL, OPENSEARCH_URI, OPENSEARCH_USERNAME, ) from rexpro_ai.retrieval.vector.main import ( GetResult, SearchResult, VectorDBBase, VectorItem, ) from rexpro_ai.retrieval.vector.utils import process_metadata from opensearchpy import OpenSearch from opensearchpy.helpers import bulk class OpenSearchClient(VectorDBBase): def __init__(self): self.index_prefix = 'rexpro_ai' self.client = OpenSearch( hosts=[OPENSEARCH_URI], use_ssl=OPENSEARCH_SSL, verify_certs=OPENSEARCH_CERT_VERIFY, http_auth=(OPENSEARCH_USERNAME, OPENSEARCH_PASSWORD), ) def _get_index_name(self, collection_name: str) -> str: return f'{self.index_prefix}_{collection_name}' def _result_to_get_result(self, result) -> GetResult: if not result['hits']['hits']: return None ids = [] documents = [] metadatas = [] for hit in result['hits']['hits']: ids.append(hit['_id']) documents.append(hit['_source'].get('text')) metadatas.append(hit['_source'].get('metadata')) return GetResult(ids=[ids], documents=[documents], metadatas=[metadatas]) def _result_to_search_result(self, result) -> SearchResult: if not result['hits']['hits']: return None ids = [] distances = [] documents = [] metadatas = [] for hit in result['hits']['hits']: ids.append(hit['_id']) distances.append(hit['_score']) documents.append(hit['_source'].get('text')) metadatas.append(hit['_source'].get('metadata')) return SearchResult( ids=[ids], distances=[distances], documents=[documents], metadatas=[metadatas], ) def _create_index(self, collection_name: str, dimension: int): body = { 'settings': {'index': {'knn': True}}, 'mappings': { 'properties': { 'id': {'type': 'keyword'}, 'vector': { 'type': 'knn_vector', 'dimension': dimension, # Adjust based on your vector dimensions 'index': True, 'similarity': 'faiss', 'method': { 'name': 'hnsw', 'space_type': 'innerproduct', # Use inner product to approximate cosine similarity 'engine': 'faiss', 'parameters': { 'ef_construction': 128, 'm': 16, }, }, }, 'text': {'type': 'text'}, 'metadata': {'type': 'object'}, } }, } self.client.indices.create(index=self._get_index_name(collection_name), body=body) def _create_batches(self, items: list[VectorItem], batch_size=100): for i in range(0, len(items), batch_size): yield items[i : i + batch_size] def has_collection(self, collection_name: str) -> bool: # has_collection here means has index. # We are simply adapting to the norms of the other DBs. return self.client.indices.exists(index=self._get_index_name(collection_name)) def delete_collection(self, collection_name: str): # delete_collection here means delete index. # We are simply adapting to the norms of the other DBs. self.client.indices.delete(index=self._get_index_name(collection_name)) def search( self, collection_name: str, vectors: list[list[float | int]], filter: Optional[dict] = None, limit: int = 10, ) -> Optional[SearchResult]: try: if not self.has_collection(collection_name): return None query = { 'size': limit, '_source': ['text', 'metadata'], 'query': { 'script_score': { 'query': {'match_all': {}}, 'script': { 'source': '(cosineSimilarity(params.query_value, doc[params.field]) + 1.0) / 2.0', 'params': { 'field': 'vector', 'query_value': vectors[0], }, # Assuming single query vector }, } }, } result = self.client.search(index=self._get_index_name(collection_name), body=query) return self._result_to_search_result(result) except Exception as e: return None def query(self, collection_name: str, filter: dict, limit: Optional[int] = None) -> Optional[GetResult]: if not self.has_collection(collection_name): return None query_body = { 'query': {'bool': {'filter': []}}, '_source': ['text', 'metadata'], } for field, value in filter.items(): query_body['query']['bool']['filter'].append({'term': {'metadata.' + str(field) + '.keyword': value}}) size = limit if limit else 10000 try: result = self.client.search( index=self._get_index_name(collection_name), body=query_body, size=size, ) return self._result_to_get_result(result) except Exception as e: return None def _create_index_if_not_exists(self, collection_name: str, dimension: int): if not self.has_collection(collection_name): self._create_index(collection_name, dimension) def get(self, collection_name: str) -> Optional[GetResult]: query = {'query': {'match_all': {}}, '_source': ['text', 'metadata']} result = self.client.search(index=self._get_index_name(collection_name), body=query) return self._result_to_get_result(result) def insert(self, collection_name: str, items: list[VectorItem]): self._create_index_if_not_exists(collection_name=collection_name, dimension=len(items[0]['vector'])) for batch in self._create_batches(items): actions = [ { '_op_type': 'index', '_index': self._get_index_name(collection_name), '_id': item['id'], '_source': { 'vector': item['vector'], 'text': item['text'], 'metadata': process_metadata(item['metadata']), }, } for item in batch ] bulk(self.client, actions) self.client.indices.refresh(index=self._get_index_name(collection_name)) def upsert(self, collection_name: str, items: list[VectorItem]): self._create_index_if_not_exists(collection_name=collection_name, dimension=len(items[0]['vector'])) for batch in self._create_batches(items): actions = [ { '_op_type': 'update', '_index': self._get_index_name(collection_name), '_id': item['id'], 'doc': { 'vector': item['vector'], 'text': item['text'], 'metadata': process_metadata(item['metadata']), }, 'doc_as_upsert': True, } for item in batch ] bulk(self.client, actions) self.client.indices.refresh(index=self._get_index_name(collection_name)) def delete( self, collection_name: str, ids: Optional[list[str]] = None, filter: Optional[dict] = None, ): if ids: actions = [ { '_op_type': 'delete', '_index': self._get_index_name(collection_name), '_id': id, } for id in ids ] bulk(self.client, actions) elif filter: query_body = { 'query': {'bool': {'filter': []}}, } for field, value in filter.items(): query_body['query']['bool']['filter'].append({'term': {'metadata.' + str(field) + '.keyword': value}}) self.client.delete_by_query(index=self._get_index_name(collection_name), body=query_body) self.client.indices.refresh(index=self._get_index_name(collection_name)) def reset(self): indices = self.client.indices.get(index=f'{self.index_prefix}_*') for index in indices: self.client.indices.delete(index=index)