backend/rexpro_ai/retrieval/vector/dbs/milvus.py

"""
NOTE: This vector database integration is community-supported and maintained on a best-effort basis.
"""

import json
import logging
from typing import Optional

from rexpro_ai.config import (
    MILVUS_DB,
    MILVUS_DISKANN_MAX_DEGREE,
    MILVUS_DISKANN_SEARCH_LIST_SIZE,
    MILVUS_HNSW_EFCONSTRUCTION,
    MILVUS_HNSW_M,
    MILVUS_INDEX_TYPE,
    MILVUS_IVF_FLAT_NLIST,
    MILVUS_METRIC_TYPE,
    MILVUS_TOKEN,
    MILVUS_URI,
)
from rexpro_ai.retrieval.vector.main import (
    GetResult,
    SearchResult,
    VectorDBBase,
    VectorItem,
)
from rexpro_ai.retrieval.vector.utils import process_metadata
from pymilvus import Collection, DataType, FieldSchema, connections
from pymilvus import MilvusClient as Client

log = logging.getLogger(__name__)


class MilvusClient(VectorDBBase):
    def __init__(self):
        self.collection_prefix = 'rexpro_ai'
        if MILVUS_TOKEN is None:
            self.client = Client(uri=MILVUS_URI, db_name=MILVUS_DB)
        else:
            self.client = Client(uri=MILVUS_URI, db_name=MILVUS_DB, token=MILVUS_TOKEN)

    def _result_to_get_result(self, result) -> GetResult:
        ids = []
        documents = []
        metadatas = []
        for match in result:
            _ids = []
            _documents = []
            _metadatas = []
            for item in match:
                _ids.append(item.get('id'))
                _documents.append(item.get('data', {}).get('text'))
                _metadatas.append(item.get('metadata'))
            ids.append(_ids)
            documents.append(_documents)
            metadatas.append(_metadatas)
        return GetResult(
            **{
                'ids': ids,
                'documents': documents,
                'metadatas': metadatas,
            }
        )

    def _result_to_search_result(self, result) -> SearchResult:
        ids = []
        distances = []
        documents = []
        metadatas = []
        for match in result:
            _ids = []
            _distances = []
            _documents = []
            _metadatas = []
            for item in match:
                _ids.append(item.get('id'))
                # normalize milvus score from [-1, 1] to [0, 1] range
                # https://milvus.io/docs/de/metric.md
                _dist = (item.get('distance') + 1.0) / 2.0
                _distances.append(_dist)
                _documents.append(item.get('entity', {}).get('data', {}).get('text'))
                _metadatas.append(item.get('entity', {}).get('metadata'))
            ids.append(_ids)
            distances.append(_distances)
            documents.append(_documents)
            metadatas.append(_metadatas)
        return SearchResult(
            **{
                'ids': ids,
                'distances': distances,
                'documents': documents,
                'metadatas': metadatas,
            }
        )

    def _create_collection(self, collection_name: str, dimension: int):
        schema = self.client.create_schema(
            auto_id=False,
            enable_dynamic_field=True,
        )
        schema.add_field(
            field_name='id',
            datatype=DataType.VARCHAR,
            is_primary=True,
            max_length=65535,
        )
        schema.add_field(
            field_name='vector',
            datatype=DataType.FLOAT_VECTOR,
            dim=dimension,
            description='vector',
        )
        schema.add_field(field_name='data', datatype=DataType.JSON, description='data')
        schema.add_field(field_name='metadata', datatype=DataType.JSON, description='metadata')

        index_params = self.client.prepare_index_params()

        # Use configurations from config.py
        index_type = MILVUS_INDEX_TYPE.upper()
        metric_type = MILVUS_METRIC_TYPE.upper()

        log.info(f'Using Milvus index type: {index_type}, metric type: {metric_type}')

        index_creation_params = {}
        if index_type == 'HNSW':
            index_creation_params = {
                'M': MILVUS_HNSW_M,
                'efConstruction': MILVUS_HNSW_EFCONSTRUCTION,
            }
            log.info(f'HNSW params: {index_creation_params}')
        elif index_type == 'IVF_FLAT':
            index_creation_params = {'nlist': MILVUS_IVF_FLAT_NLIST}
            log.info(f'IVF_FLAT params: {index_creation_params}')
        elif index_type == 'DISKANN':
            index_creation_params = {
                'max_degree': MILVUS_DISKANN_MAX_DEGREE,
                'search_list_size': MILVUS_DISKANN_SEARCH_LIST_SIZE,
            }
            log.info(f'DISKANN params: {index_creation_params}')
        elif index_type in ['FLAT', 'AUTOINDEX']:
            log.info(f'Using {index_type} index with no specific build-time params.')
        else:
            log.warning(
                f"Unsupported MILVUS_INDEX_TYPE: '{index_type}'. "
                f'Supported types: HNSW, IVF_FLAT, DISKANN, FLAT, AUTOINDEX. '
                f'Milvus will use its default for the collection if this type is not directly supported for index creation.'
            )
            # For unsupported types, pass the type directly to Milvus; it might handle it or use a default.
            # If Milvus errors out, the user needs to correct the MILVUS_INDEX_TYPE env var.

        index_params.add_index(
            field_name='vector',
            index_type=index_type,
            metric_type=metric_type,
            params=index_creation_params,
        )

        self.client.create_collection(
            collection_name=f'{self.collection_prefix}_{collection_name}',
            schema=schema,
            index_params=index_params,
        )
        log.info(
            f"Successfully created collection '{self.collection_prefix}_{collection_name}' with index type '{index_type}' and metric '{metric_type}'."
        )

    def has_collection(self, collection_name: str) -> bool:
        # Check if the collection exists based on the collection name.
        collection_name = collection_name.replace('-', '_')
        return self.client.has_collection(collection_name=f'{self.collection_prefix}_{collection_name}')

    def delete_collection(self, collection_name: str):
        # Delete the collection based on the collection name.
        collection_name = collection_name.replace('-', '_')
        return self.client.drop_collection(collection_name=f'{self.collection_prefix}_{collection_name}')

    def search(
        self,
        collection_name: str,
        vectors: list[list[float | int]],
        filter: Optional[dict] = None,
        limit: int = 10,
    ) -> Optional[SearchResult]:
        # Search for the nearest neighbor items based on the vectors and return 'limit' number of results.
        collection_name = collection_name.replace('-', '_')
        # For some index types like IVF_FLAT, search params like nprobe can be set.
        # Example: search_params = {"nprobe": 10} if using IVF_FLAT
        # For simplicity, not adding configurable search_params here, but could be extended.
        result = self.client.search(
            collection_name=f'{self.collection_prefix}_{collection_name}',
            data=vectors,
            limit=limit,
            output_fields=['data', 'metadata'],
            # search_params=search_params # Potentially add later if needed
        )
        return self._result_to_search_result(result)

    def query(self, collection_name: str, filter: dict, limit: int = -1):
        connections.connect(uri=MILVUS_URI, token=MILVUS_TOKEN, db_name=MILVUS_DB)

        collection_name = collection_name.replace('-', '_')
        if not self.has_collection(collection_name):
            log.warning(f'Query attempted on non-existent collection: {self.collection_prefix}_{collection_name}')
            return None

        filter_expressions = []
        for key, value in filter.items():
            if isinstance(value, str):
                filter_expressions.append(f'metadata["{key}"] == "{value}"')
            else:
                filter_expressions.append(f'metadata["{key}"] == {value}')

        filter_string = ' && '.join(filter_expressions)

        collection = Collection(f'{self.collection_prefix}_{collection_name}')
        collection.load()

        try:
            log.info(
                f"Querying collection {self.collection_prefix}_{collection_name} with filter: '{filter_string}', limit: {limit}"
            )

            iterator = collection.query_iterator(
                expr=filter_string,
                output_fields=[
                    'id',
                    'data',
                    'metadata',
                ],
                limit=limit if limit > 0 else -1,
            )

            all_results = []
            while True:
                batch = iterator.next()
                if not batch:
                    iterator.close()
                    break
                all_results.extend(batch)

            log.debug(f'Total results from query: {len(all_results)}')
            return self._result_to_get_result([all_results] if all_results else [[]])

        except Exception as e:
            log.exception(
                f"Error querying collection {self.collection_prefix}_{collection_name} with filter '{filter_string}' and limit {limit}: {e}"
            )
            return None

    def get(self, collection_name: str) -> Optional[GetResult]:
        # Get all the items in the collection. This can be very resource-intensive for large collections.
        collection_name = collection_name.replace('-', '_')
        log.warning(
            f"Fetching ALL items from collection '{self.collection_prefix}_{collection_name}'. This might be slow for large collections."
        )
        # Using query with a trivial filter to get all items.
        # This will use the paginated query logic.
        return self.query(collection_name=collection_name, filter={}, limit=-1)

    def insert(self, collection_name: str, items: list[VectorItem]):
        # Insert the items into the collection, if the collection does not exist, it will be created.
        collection_name = collection_name.replace('-', '_')
        if not self.client.has_collection(collection_name=f'{self.collection_prefix}_{collection_name}'):
            log.info(f'Collection {self.collection_prefix}_{collection_name} does not exist. Creating now.')
            if not items:
                log.error(
                    f'Cannot create collection {self.collection_prefix}_{collection_name} without items to determine dimension.'
                )
                raise ValueError('Cannot create Milvus collection without items to determine vector dimension.')
            self._create_collection(collection_name=collection_name, dimension=len(items[0]['vector']))

        log.info(f'Inserting {len(items)} items into collection {self.collection_prefix}_{collection_name}.')
        return self.client.insert(
            collection_name=f'{self.collection_prefix}_{collection_name}',
            data=[
                {
                    'id': item['id'],
                    'vector': item['vector'],
                    'data': {'text': item['text']},
                    'metadata': process_metadata(item['metadata']),
                }
                for item in items
            ],
        )

    def upsert(self, collection_name: str, items: list[VectorItem]):
        # Update the items in the collection, if the items are not present, insert them. If the collection does not exist, it will be created.
        collection_name = collection_name.replace('-', '_')
        if not self.client.has_collection(collection_name=f'{self.collection_prefix}_{collection_name}'):
            log.info(f'Collection {self.collection_prefix}_{collection_name} does not exist for upsert. Creating now.')
            if not items:
                log.error(
                    f'Cannot create collection {self.collection_prefix}_{collection_name} for upsert without items to determine dimension.'
                )
                raise ValueError(
                    'Cannot create Milvus collection for upsert without items to determine vector dimension.'
                )
            self._create_collection(collection_name=collection_name, dimension=len(items[0]['vector']))

        log.info(f'Upserting {len(items)} items into collection {self.collection_prefix}_{collection_name}.')
        return self.client.upsert(
            collection_name=f'{self.collection_prefix}_{collection_name}',
            data=[
                {
                    'id': item['id'],
                    'vector': item['vector'],
                    'data': {'text': item['text']},
                    'metadata': process_metadata(item['metadata']),
                }
                for item in items
            ],
        )

    def delete(
        self,
        collection_name: str,
        ids: Optional[list[str]] = None,
        filter: Optional[dict] = None,
    ):
        # Delete the items from the collection based on the ids or filter.
        collection_name = collection_name.replace('-', '_')
        if not self.has_collection(collection_name):
            log.warning(f'Delete attempted on non-existent collection: {self.collection_prefix}_{collection_name}')
            return None

        if ids:
            log.info(f'Deleting items by IDs from {self.collection_prefix}_{collection_name}. IDs: {ids}')
            return self.client.delete(
                collection_name=f'{self.collection_prefix}_{collection_name}',
                ids=ids,
            )
        elif filter:
            filter_string = ' && '.join([f'metadata["{key}"] == {json.dumps(value)}' for key, value in filter.items()])
            log.info(
                f'Deleting items by filter from {self.collection_prefix}_{collection_name}. Filter: {filter_string}'
            )
            return self.client.delete(
                collection_name=f'{self.collection_prefix}_{collection_name}',
                filter=filter_string,
            )
        else:
            log.warning(
                f'Delete operation on {self.collection_prefix}_{collection_name} called without IDs or filter. No action taken.'
            )
            return None

    def reset(self):
        # Resets the database. This will delete all collections and item entries that match the prefix.
        log.warning(f"Resetting Milvus: Deleting all collections with prefix '{self.collection_prefix}'.")
        collection_names = self.client.list_collections()
        deleted_collections = []
        for collection_name_full in collection_names:
            if collection_name_full.startswith(self.collection_prefix):
                try:
                    self.client.drop_collection(collection_name=collection_name_full)
                    deleted_collections.append(collection_name_full)
                    log.info(f'Deleted collection: {collection_name_full}')
                except Exception as e:
                    log.error(f'Error deleting collection {collection_name_full}: {e}')
        log.info(f'Milvus reset complete. Deleted collections: {deleted_collections}')