Milvus Search

Inputs

Outputs

Source

python

# Input preparation
inp = __ivy_ctx__["nodes"][__ivy_node_id__]["input"]
milvus_client = inp["milvus_client"]
collection_name = inp["collection_name"]
embedding = inp["embedding"]

# Compute
from pymilvus import MilvusClient, MilvusException
from typing import List, Tuple, Optional

class MilvusSearch:
    """
    A class to perform search operations on a Milvus collection.

    Attributes:
        milvus_client (MilvusClient): The Milvus client instance used for communication.
        collection_name (str): The name of the Milvus collection to search in.

    Methods:
        search(embedding: List[float], limit: int = 3, search_params: Optional[dict] = None) -> str:
            Searches for similar vectors in the Milvus collection and returns the results as a formatted string.
    """

    def __init__(self, milvus_client: MilvusClient, collection_name: str):
        """
        Initializes the MilvusSearch with a Milvus client and a collection name.

        Args:
            milvus_client (MilvusClient): The Milvus client instance.
            collection_name (str): The name of the Milvus collection.

        Raises:
            ValueError: If `collection_name` is empty.
        """
        if not isinstance(collection_name, str) or not collection_name.strip():
            raise ValueError("collection_name must be a non-empty string.")
        
        self.milvus_client = milvus_client
        self.collection_name = collection_name

    def search(self, embedding: List[float], limit: int = 3, search_params: Optional[dict] = None) -> str:
        """
        Searches for similar vectors in the Milvus collection.

        Args:
            embedding (List[float]): The embedding vector to search for.
            limit (int, optional): The number of top results to return. Default is 3.
            search_params (Optional[dict], optional): Optional search parameters. Default is `{"metric_type": "IP", "params": {}}`.

        Returns:
            str: A formatted string containing the search results.

        Raises:
            ValueError: If the embedding is empty or contains invalid data.
            RuntimeError: If the search operation fails due to a MilvusException or unexpected error.
        """
        if not isinstance(embedding, list) or not all(isinstance(x, (int, float)) for x in embedding):
            raise ValueError("Embedding must be a list of numeric values.")
        if not embedding:
            raise ValueError("Embedding cannot be empty.")

        if search_params is None:
            search_params = {"metric_type": "IP", "params": {}}

        try:
            # Perform the search operation
            search_res = self.milvus_client.search(
                collection_name=self.collection_name,
                data=[embedding],
                limit=limit,
                search_params=search_params,
                output_fields=["text"],
            )

            # Extract and return the results
            search_results = [(res["entity"]["text"], res["distance"]) for res in search_res[0]]
            context = "\n".join([res[0] for res in search_results])
            
            return context
        except MilvusException as e:
            raise RuntimeError(f"Search operation failed in collection '{self.collection_name}': {e}")
        except Exception as e:
            raise RuntimeError(f"Unexpected error during search: {e}")

milvus_search = MilvusSearch(milvus_client=milvus_client, collection_name=collection_name)

context = milvus_search.search(embedding)

# Output collection (runner reads __ivy_ctx__)
out = __ivy_ctx__["nodes"][__ivy_node_id__]["output"]
out["context"] = context

Tests

Requires: python:3.11

• basic-search

  Search with valid embedding vector (should succeed).

  searchbasic
• error-empty-embedding

  Empty embedding should raise an error.

  error-handlingvalidation