langchain/libs/community/langchain_community/retrievers/zilliz.py

import warnings
from typing import Any, Dict, List, Optional

from langchain_core.callbacks import CallbackManagerForRetrieverRun
from langchain_core.documents import Document
from langchain_core.embeddings import Embeddings
from langchain_core.pydantic_v1 import root_validator
from langchain_core.retrievers import BaseRetriever

from langchain_community.vectorstores.zilliz import Zilliz

# TODO: Update to ZillizClient + Hybrid Search when available


class ZillizRetriever(BaseRetriever):
    """`Zilliz API` retriever."""

    embedding_function: Embeddings
    """The underlying embedding function from which documents will be retrieved."""
    collection_name: str = "LangChainCollection"
    """The name of the collection in Zilliz."""
    connection_args: Optional[Dict[str, Any]] = None
    """The connection arguments for the Zilliz client."""
    consistency_level: str = "Session"
    """The consistency level for the Zilliz client."""
    search_params: Optional[dict] = None
    """The search parameters for the Zilliz client."""
    store: Zilliz
    """The underlying Zilliz store."""
    retriever: BaseRetriever
    """The underlying retriever."""

    @root_validator(pre=True)
    def create_client(cls, values: dict) -> dict:
        values["store"] = Zilliz(
            values["embedding_function"],
            values["collection_name"],
            values["connection_args"],
            values["consistency_level"],
        )
        values["retriever"] = values["store"].as_retriever(
            search_kwargs={"param": values["search_params"]}
        )
        return values

    def add_texts(
        self, texts: List[str], metadatas: Optional[List[dict]] = None
    ) -> None:
        """Add text to the Zilliz store

        Args:
            texts (List[str]): The text
            metadatas (List[dict]): Metadata dicts, must line up with existing store
        """
        self.store.add_texts(texts, metadatas)

    def _get_relevant_documents(
        self,
        query: str,
        *,
        run_manager: CallbackManagerForRetrieverRun,
        **kwargs: Any,
    ) -> List[Document]:
        return self.retriever.invoke(
            query, run_manager=run_manager.get_child(), **kwargs
        )


def ZillizRetreiver(*args: Any, **kwargs: Any) -> ZillizRetriever:
    """Deprecated ZillizRetreiver.

    Please use ZillizRetriever ('i' before 'e') instead.

    Args:
        *args:
        **kwargs:

    Returns:
        ZillizRetriever
    """
    warnings.warn(
        "ZillizRetreiver will be deprecated in the future. "
        "Please use ZillizRetriever ('i' before 'e') instead.",
        DeprecationWarning,
    )
    return ZillizRetriever(*args, **kwargs)