mem0ai--mem0/mem0/vector_stores/milvus.py

import logging
import re
from typing import Dict, Optional

from pydantic import BaseModel

from mem0.configs.vector_stores.milvus import MetricType
from mem0.vector_stores.base import VectorStoreBase

try:
    import pymilvus  # noqa: F401
except ImportError:
    raise ImportError("The 'pymilvus' library is required. Please install it using 'pip install pymilvus'.")

from pymilvus import (
    CollectionSchema,
    DataType,
    FieldSchema,
    Function,
    FunctionType,
    MilvusClient,
)

logger = logging.getLogger(__name__)


class OutputData(BaseModel):
    id: Optional[str]  # memory id
    score: Optional[float]  # distance
    payload: Optional[Dict]  # metadata


class MilvusDB(VectorStoreBase):
    def __init__(
        self,
        url: str,
        token: str,
        collection_name: str,
        embedding_model_dims: int,
        metric_type: MetricType,
        db_name: str,
    ) -> None:
        """Initialize the MilvusDB database.

        Args:
            url (str): Full URL for Milvus/Zilliz server.
            token (str): Token/api_key for Zilliz server / for local setup defaults to None.
            collection_name (str): Name of the collection (defaults to mem0).
            embedding_model_dims (int): Dimensions of the embedding model (defaults to 1536).
            metric_type (MetricType): Metric type for similarity search (defaults to L2).
            db_name (str): Name of the database (defaults to "").
        """
        self.collection_name = collection_name
        self.embedding_model_dims = embedding_model_dims
        self.metric_type = metric_type
        self.client = MilvusClient(uri=url, token=token, db_name=db_name)
        # Whether this collection has the `text` + `sparse` fields for v3 BM25.
        # Pre-v3 collections lack them; writing a top-level `text` field is rejected.
        self._has_bm25_schema = False
        self.create_col(
            collection_name=self.collection_name,
            vector_size=self.embedding_model_dims,
            metric_type=self.metric_type,
        )

    def create_col(
        self,
        collection_name: str,
        vector_size: int,
        metric_type: MetricType = MetricType.COSINE,
    ) -> None:
        """Create a new collection with index_type AUTOINDEX.

        Args:
            collection_name (str): Name of the collection (defaults to mem0).
            vector_size (int): Dimensions of the embedding model (defaults to 1536).
            metric_type (MetricType, optional): etric type for similarity search. Defaults to MetricType.COSINE.
        """

        if self.client.has_collection(collection_name):
            logger.info(f"Collection {collection_name} already exists. Skipping creation.")
            desc = self.client.describe_collection(collection_name=collection_name)
            field_names = {f.get("name") for f in desc.get("fields", [])}
            self._has_bm25_schema = "text" in field_names and "sparse" in field_names
            if not self._has_bm25_schema:
                logger.warning(
                    f"Collection '{collection_name}' predates v3 hybrid search (no 'text'/'sparse' fields). "
                    "BM25 keyword scoring will be disabled for this collection; semantic search works normally. "
                    "To enable hybrid search, use a fresh collection."
                )
        else:
            fields = [
                FieldSchema(name="id", dtype=DataType.VARCHAR, is_primary=True, max_length=512),
                FieldSchema(name="vectors", dtype=DataType.FLOAT_VECTOR, dim=vector_size),
                FieldSchema(name="metadata", dtype=DataType.JSON),
                # Text field for BM25 full-text search (auto-tokenized by Milvus analyzer)
                FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=65535, enable_analyzer=True),
                # Sparse vector field populated automatically by the BM25 function below
                FieldSchema(name="sparse", dtype=DataType.SPARSE_FLOAT_VECTOR),
            ]

            schema = CollectionSchema(fields, enable_dynamic_field=True)

            # Add BM25 function so Milvus auto-generates sparse vectors from the text field
            bm25_function = Function(
                name="bm25",
                input_field_names=["text"],
                output_field_names=["sparse"],
                function_type=FunctionType.BM25,
            )
            schema.add_function(bm25_function)

            index_params = self.client.prepare_index_params()
            index_params.add_index(
                field_name="vectors", metric_type=metric_type, index_type="AUTOINDEX", index_name="vector_index"
            )
            index_params.add_index(
                field_name="sparse",
                index_type="SPARSE_INVERTED_INDEX",
                metric_type="BM25",
                index_name="sparse_index",
            )
            self.client.create_collection(collection_name=collection_name, schema=schema, index_params=index_params)
            self._has_bm25_schema = True

    def insert(self, ids, vectors, payloads, **kwargs: Optional[dict[str, any]]):
        """Insert vectors into a collection.

        Args:
            vectors (List[List[float]]): List of vectors to insert.
            payloads (List[Dict], optional): List of payloads corresponding to vectors.
            ids (List[str], optional): List of IDs corresponding to vectors.
        """
        # Batch insert all records at once for better performance and consistency.
        # Only include the `text` field when the collection's schema has it — legacy
        # collections created pre-v3 reject unknown top-level fields.
        def _build_record(idx, embedding, metadata):
            record = {"id": idx, "vectors": embedding, "metadata": metadata}
            if self._has_bm25_schema:
                # Populate the text field for BM25 sparse search; prefer lemmatized text, fall back to raw data
                record["text"] = (metadata.get("text_lemmatized") or metadata.get("data", ""))[:65535] if metadata else ""
            return record

        data = [_build_record(idx, embedding, metadata) for idx, embedding, metadata in zip(ids, vectors, payloads)]
        self.client.insert(collection_name=self.collection_name, data=data, **kwargs)

    _SAFE_FILTER_KEY = re.compile(r"^[a-zA-Z_][a-zA-Z0-9_]*$")

    def _create_filter(self, filters: dict):
        """Prepare filters for efficient query.

        Args:
            filters (dict): filters [user_id, agent_id, run_id]

        Returns:
            str: formated filter.
        """
        operands = []
        for key, value in filters.items():
            if not self._SAFE_FILTER_KEY.match(key):
                raise ValueError(f"Invalid filter key: {key!r}")
            if value == "*":
                # Wildcard - match any value (MilvusDB doesn't have direct wildcard, so we skip this filter)
                continue
            elif isinstance(value, str):
                escaped = value.replace("\\", "\\\\").replace('"', '\\"')
                operands.append(f'(metadata["{key}"] == "{escaped}")')
            elif isinstance(value, (int, float, bool)):
                operands.append(f'(metadata["{key}"] == {value})')
            else:
                raise ValueError(
                    f"Filter value for {key!r} must be str, int, float, or bool, "
                    f"got {type(value).__name__}"
                )

        return " and ".join(operands)

    def _parse_output(self, data: list):
        """
        Parse the output data.

        Args:
            data (Dict): Output data.

        Returns:
            List[OutputData]: Parsed output data.
        """
        memory = []

        for value in data:
            uid = value.get("id")
            raw_distance = value.get("distance")
            metadata = value.get("entity", {}).get("metadata")

            if raw_distance is not None and self.metric_type in (MetricType.L2, "L2"):
                score = 1.0 / (1.0 + raw_distance)
            else:
                score = raw_distance

            memory_obj = OutputData(id=uid, score=score, payload=metadata)
            memory.append(memory_obj)

        return memory

    def search(self, query: str, vectors: list, top_k: int = 5, filters: dict = None) -> list:
        """
        Search for similar vectors.

        Args:
            query (str): Query.
            vectors (List[float]): Query vector.
            top_k (int, optional): Number of results to return. Defaults to 5.
            filters (Dict, optional): Filters to apply to the search. Defaults to None.

        Returns:
            list: Search results.
        """
        query_filter = self._create_filter(filters) if filters else None
        # v3 collections carry both a dense `vectors` field and a sparse `sparse`
        # field (for BM25), which makes anns_field ambiguous — Milvus rejects the
        # query otherwise with "multiple anns_fields exist". Legacy single-vector
        # collections don't need the hint, so only pass it when the hybrid schema
        # is present.
        search_kwargs = {
            "collection_name": self.collection_name,
            "data": [vectors],
            "limit": top_k,
            "filter": query_filter,
            "output_fields": ["*"],
        }
        if self._has_bm25_schema:
            search_kwargs["anns_field"] = "vectors"
        hits = self.client.search(**search_kwargs)
        result = self._parse_output(data=hits[0])
        return result

    def keyword_search(self, query, top_k=5, filters=None):
        """
        Search for memories using BM25-based full-text search via Milvus sparse vector support.

        Milvus 2.5+ supports native BM25 via full-text search with a SPARSE_FLOAT_VECTOR field.
        This method attempts to use that capability. If the collection does not have a sparse
        field configured, it returns None gracefully.

        Args:
            query (str): The text query for keyword-based search.
            top_k (int, optional): Number of results to return. Defaults to 5.
            filters (dict, optional): Filters to apply to the search. Defaults to None.

        Returns:
            list: Search results in the same format as search(), or None if sparse search
                  is not supported on this collection.
        """
        if not self._has_bm25_schema:
            return None
        try:
            query_filter = self._create_filter(filters) if filters else None
            hits = self.client.search(
                collection_name=self.collection_name,
                data=[query],
                anns_field="sparse",
                limit=top_k,
                filter=query_filter,
                output_fields=["*"],
            )
            result = self._parse_output(data=hits[0])
            return result
        except Exception as e:
            logger.debug(f"Keyword search not available for collection {self.collection_name}: {e}")
            return None

    def delete(self, vector_id):
        """
        Delete a vector by ID.

        Args:
            vector_id (str): ID of the vector to delete.
        """
        self.client.delete(collection_name=self.collection_name, ids=[vector_id])

    def update(self, vector_id=None, vector=None, payload=None):
        """
        Update a vector and its payload.

        Args:
            vector_id (str): ID of the vector to update.
            vector (List[float], optional): Updated vector.
            payload (Dict, optional): Updated payload.
        """
        if vector is None or payload is None:
            existing = self.client.get(collection_name=self.collection_name, ids=vector_id)
            if not existing:
                raise ValueError(f"Vector with id {vector_id} not found in collection {self.collection_name}")
            if vector is None:
                vector = existing[0].get("vectors")
                if vector is None:
                    raise ValueError(f"Existing record {vector_id} has no vector data")
            if payload is None:
                payload = existing[0].get("metadata")

        text = ""
        if payload:
            text = (payload.get("text_lemmatized") or payload.get("data", ""))[:65535]
        schema = {"id": vector_id, "vectors": vector, "metadata": payload, "text": text}
        self.client.upsert(collection_name=self.collection_name, data=schema)

    def get(self, vector_id) -> Optional[OutputData]:
        """
        Retrieve a vector by ID.

        Args:
            vector_id (str): ID of the vector to retrieve.

        Returns:
            Optional[OutputData]: Retrieved vector, or None if the ID is not found.
        """
        result = self.client.get(collection_name=self.collection_name, ids=vector_id)
        if not result:
            return None
        output = OutputData(
            id=result[0].get("id", None),
            score=None,
            payload=result[0].get("metadata", None),
        )
        return output

    def list_cols(self):
        """
        List all collections.

        Returns:
            List[str]: List of collection names.
        """
        return self.client.list_collections()

    def delete_col(self):
        """Delete a collection."""
        return self.client.drop_collection(collection_name=self.collection_name)

    def col_info(self):
        """
        Get information about a collection.

        Returns:
            Dict[str, Any]: Collection information.
        """
        return self.client.get_collection_stats(collection_name=self.collection_name)

    def list(self, filters: dict = None, top_k: int = 100) -> list:
        """
        List all vectors in a collection.

        Args:
            filters (Dict, optional): Filters to apply to the list.
            top_k (int, optional): Number of vectors to return. Defaults to 100.

        Returns:
            List[OutputData]: List of vectors.
        """
        query_filter = self._create_filter(filters) if filters else None
        result = self.client.query(collection_name=self.collection_name, filter=query_filter, limit=top_k)
        memories = []
        for data in result:
            obj = OutputData(id=data.get("id"), score=None, payload=data.get("metadata"))
            memories.append(obj)
        return [memories]

    def reset(self):
        """Reset the index by deleting and recreating it."""
        logger.warning(f"Resetting index {self.collection_name}...")
        self.delete_col()
        self.create_col(self.collection_name, self.embedding_model_dims, self.metric_type)