Skip to content

milvusdb

langroid/vector_store/milvusdb.py

MilvusDB(config=None)

Bases: VectorStore

Source code in langroid/vector_store/milvusdb.py
def __init__(self, config: Optional[MilvusDBConfig] = None):
    if config is None:
        config = MilvusDBConfig()
    super().__init__(config)
    self.config: MilvusDBConfig = config
    load_dotenv()
    self.config.uri = self.config.uri or os.getenv("MILVUS_URI", "./milvus.db")
    self.config.token = self.config.token or os.getenv("MILVUS_TOKEN")
    self.config.db_name = self.config.db_name or os.getenv("MILVUS_DB_NAME")
    self.config.metric_type = self.config.metric_type.upper()

    try:
        from pymilvus import MilvusClient
    except ImportError as exc:
        raise LangroidImportError("pymilvus", "milvus") from exc
    if (
        self._is_local_lite_uri(self.config.uri)
        and self._milvus_lite_version() is None
    ):
        raise ValueError(
            "Milvus Lite is unavailable on this platform; set MILVUS_URI "
            "to a Milvus server or Zilliz Cloud endpoint"
        )

    self._milvus_lite_3_0 = self._uses_milvus_lite_3_0(self.config.uri)

    client_kwargs: Dict[str, Any] = {
        "uri": self.config.uri,
        "timeout": self.config.timeout,
    }
    if self.config.token:
        client_kwargs["token"] = self.config.token
    if self.config.db_name:
        client_kwargs["db_name"] = self.config.db_name
    self.client = MilvusClient(**client_kwargs)

    if self.config.collection_name is not None:
        self.create_collection(
            self.config.collection_name,
            replace=self.config.replace_collection,
        )

set_collection(collection_name, replace=False)

Set the active collection and load it when it already exists.

Parameters:

Name Type Description Default
collection_name str

Name of the collection.

required
replace bool

Whether to replace an existing collection.

False
Source code in langroid/vector_store/milvusdb.py
def set_collection(self, collection_name: str, replace: bool = False) -> None:
    """Set the active collection and load it when it already exists.

    Args:
        collection_name: Name of the collection.
        replace: Whether to replace an existing collection.
    """
    super().set_collection(collection_name, replace)
    if not replace and self.client.has_collection(collection_name=collection_name):
        self.create_collection(collection_name, replace=False)