initial commit

venv/Lib/site-packages/langchain_community/retrievers/zilliz.py

@@ -0,0 +1,87 @@
+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.retrievers import BaseRetriever
+from pydantic import model_validator
+
+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."""
+
+    @model_validator(mode="before")
+    @classmethod
+    def create_client(cls, values: dict) -> Any:
+        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)