open-webui/backend/open_webui/retrieval/vector/dbs/opensearch.py

from opensearchpy import OpenSearch
from opensearchpy.helpers import bulk
from typing import Optional

from open_webui.retrieval.vector.main import VectorItem, SearchResult, GetResult
from open_webui.config import (
    OPENSEARCH_URI,
    OPENSEARCH_SSL,
    OPENSEARCH_CERT_VERIFY,
    OPENSEARCH_USERNAME,
    OPENSEARCH_PASSWORD,
)


class OpenSearchClient:
    def __init__(self):
        self.index_prefix = "open_webui"
        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]], limit: int
    ) -> 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",
                            "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({
                "match": {
                    "metadata." + str(field): value
                }
            })

        size = limit if limit else 10

        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": item["metadata"],
                    },
                }
                for item in batch
            ]
            bulk(self.client, actions)

    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": item["metadata"],
                    },
                    "doc_as_upsert": True,
                }
                for item in batch
            ]
            bulk(self.client, actions)

    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({
                    "match": {
                        "metadata." + str(field): value
                    }
                })
            self.client.delete_by_query(index=self._get_index_name(collection_name), body=query_body)
                
    def reset(self):
        indices = self.client.indices.get(index=f"{self.index_prefix}_*")
        for index in indices:
            self.client.indices.delete(index=index)