fix: change milvus init args from (host, port) to (url, token) (#8019)

Signed-off-by: ChengZi <chen.zhang@zilliz.com>

api/.env.example

@@ -107,11 +107,10 @@ QDRANT_GRPC_ENABLED=false
 QDRANT_GRPC_PORT=6334
 
 # Milvus configuration
-MILVUS_HOST=127.0.0.1
-MILVUS_PORT=19530
+MILVUS_URI=http://127.0.0.1:19530
+MILVUS_TOKEN=
 MILVUS_USER=root
 MILVUS_PASSWORD=Milvus
-MILVUS_SECURE=false
 
 # MyScale configuration
 MYSCALE_HOST=127.0.0.1

api/configs/middleware/vdb/milvus_config.py

@@ -1,6 +1,6 @@
 from typing import Optional
 
-from pydantic import Field, PositiveInt
+from pydantic import Field
 from pydantic_settings import BaseSettings
 
 
@@ -9,14 +9,14 @@ class MilvusConfig(BaseSettings):
     Milvus configs
     """
 
-    MILVUS_HOST: Optional[str] = Field(
-        description="Milvus host",
-        default=None,
+    MILVUS_URI: Optional[str] = Field(
+        description="Milvus uri",
+        default="http://127.0.0.1:19530",
     )
 
-    MILVUS_PORT: PositiveInt = Field(
-        description="Milvus RestFul API port",
-        default=9091,
+    MILVUS_TOKEN: Optional[str] = Field(
+        description="Milvus token",
+        default=None,
     )
 
     MILVUS_USER: Optional[str] = Field(
@@ -29,11 +29,6 @@ class MilvusConfig(BaseSettings):
         default=None,
     )
 
-    MILVUS_SECURE: bool = Field(
-        description="whether to use SSL connection for Milvus",
-        default=False,
-    )
-
     MILVUS_DATABASE: str = Field(
         description="Milvus database, default to `default`",
         default="default",

api/core/rag/datasource/vdb/milvus/milvus_vector.py

@@ -1,10 +1,9 @@
 import json
 import logging
 from typing import Any, Optional
-from uuid import uuid4
 
 from pydantic import BaseModel, model_validator
-from pymilvus import MilvusClient, MilvusException, connections
+from pymilvus import MilvusClient, MilvusException
 from pymilvus.milvus_client import IndexParams
 
 from configs import dify_config
@@ -21,20 +20,17 @@ logger = logging.getLogger(__name__)
 
 
 class MilvusConfig(BaseModel):
-    host: str
-    port: int
+    uri: str
+    token: Optional[str] = None
     user: str
     password: str
-    secure: bool = False
     batch_size: int = 100
     database: str = "default"
 
     @model_validator(mode='before')
     def validate_config(cls, values: dict) -> dict:
-        if not values.get('host'):
-            raise ValueError("config MILVUS_HOST is required")
-        if not values.get('port'):
-            raise ValueError("config MILVUS_PORT is required")
+        if not values.get('uri'):
+            raise ValueError("config MILVUS_URI is required")
         if not values.get('user'):
             raise ValueError("config MILVUS_USER is required")
         if not values.get('password'):
@@ -43,11 +39,10 @@ class MilvusConfig(BaseModel):
 
     def to_milvus_params(self):
         return {
-            'host': self.host,
-            'port': self.port,
+            'uri': self.uri,
+            'token': self.token,
             'user': self.user,
             'password': self.password,
-            'secure': self.secure,
             'db_name': self.database,
         }
 
@@ -111,32 +106,14 @@ class MilvusVector(BaseVector):
             return None
 
     def delete_by_metadata_field(self, key: str, value: str):
-        alias = uuid4().hex
-        if self._client_config.secure:
-            uri = "https://" + str(self._client_config.host) + ":" + str(self._client_config.port)
-        else:
-            uri = "http://" + str(self._client_config.host) + ":" + str(self._client_config.port)
-        connections.connect(alias=alias, uri=uri, user=self._client_config.user, password=self._client_config.password,
-                            db_name=self._client_config.database)
-
-        from pymilvus import utility
-        if utility.has_collection(self._collection_name, using=alias):
+        if self._client.has_collection(self._collection_name):
 
             ids = self.get_ids_by_metadata_field(key, value)
             if ids:
                 self._client.delete(collection_name=self._collection_name, pks=ids)
 
     def delete_by_ids(self, ids: list[str]) -> None:
-        alias = uuid4().hex
-        if self._client_config.secure:
-            uri = "https://" + str(self._client_config.host) + ":" + str(self._client_config.port)
-        else:
-            uri = "http://" + str(self._client_config.host) + ":" + str(self._client_config.port)
-        connections.connect(alias=alias, uri=uri, user=self._client_config.user, password=self._client_config.password,
-                            db_name=self._client_config.database)
-
-        from pymilvus import utility
-        if utility.has_collection(self._collection_name, using=alias):
+        if self._client.has_collection(self._collection_name):
 
             result = self._client.query(collection_name=self._collection_name,
                                         filter=f'metadata["doc_id"] in {ids}',
@@ -146,29 +123,11 @@ class MilvusVector(BaseVector):
                 self._client.delete(collection_name=self._collection_name, pks=ids)
 
     def delete(self) -> None:
-        alias = uuid4().hex
-        if self._client_config.secure:
-            uri = "https://" + str(self._client_config.host) + ":" + str(self._client_config.port)
-        else:
-            uri = "http://" + str(self._client_config.host) + ":" + str(self._client_config.port)
-        connections.connect(alias=alias, uri=uri, user=self._client_config.user, password=self._client_config.password,
-                            db_name=self._client_config.database)
-
-        from pymilvus import utility
-        if utility.has_collection(self._collection_name, using=alias):
-            utility.drop_collection(self._collection_name, None, using=alias)
+        if self._client.has_collection(self._collection_name):
+            self._client.drop_collection(self._collection_name, None)
 
     def text_exists(self, id: str) -> bool:
-        alias = uuid4().hex
-        if self._client_config.secure:
-            uri = "https://" + str(self._client_config.host) + ":" + str(self._client_config.port)
-        else:
-            uri = "http://" + str(self._client_config.host) + ":" + str(self._client_config.port)
-        connections.connect(alias=alias, uri=uri, user=self._client_config.user, password=self._client_config.password,
-                            db_name=self._client_config.database)
-
-        from pymilvus import utility
-        if not utility.has_collection(self._collection_name, using=alias):
+        if not self._client.has_collection(self._collection_name):
             return False
 
         result = self._client.query(collection_name=self._collection_name,
@@ -210,15 +169,7 @@ class MilvusVector(BaseVector):
             if redis_client.get(collection_exist_cache_key):
                 return
             # Grab the existing collection if it exists
-            from pymilvus import utility
-            alias = uuid4().hex
-            if self._client_config.secure:
-                uri = "https://" + str(self._client_config.host) + ":" + str(self._client_config.port)
-            else:
-                uri = "http://" + str(self._client_config.host) + ":" + str(self._client_config.port)
-            connections.connect(alias=alias, uri=uri, user=self._client_config.user,
-                                password=self._client_config.password, db_name=self._client_config.database)
-            if not utility.has_collection(self._collection_name, using=alias):
+            if not self._client.has_collection(self._collection_name):
                 from pymilvus import CollectionSchema, DataType, FieldSchema
                 from pymilvus.orm.types import infer_dtype_bydata
 
@@ -263,11 +214,7 @@ class MilvusVector(BaseVector):
             redis_client.set(collection_exist_cache_key, 1, ex=3600)
 
     def _init_client(self, config) -> MilvusClient:
-        if config.secure:
-            uri = "https://" + str(config.host) + ":" + str(config.port)
-        else:
-            uri = "http://" + str(config.host) + ":" + str(config.port)
-        client = MilvusClient(uri=uri, user=config.user, password=config.password, db_name=config.database)
+        client = MilvusClient(uri=config.uri, user=config.user, password=config.password, db_name=config.database)
         return client
 
 
@@ -285,11 +232,10 @@ class MilvusVectorFactory(AbstractVectorFactory):
         return MilvusVector(
             collection_name=collection_name,
             config=MilvusConfig(
-                host=dify_config.MILVUS_HOST,
-                port=dify_config.MILVUS_PORT,
+                uri=dify_config.MILVUS_URI,
+                token=dify_config.MILVUS_TOKEN,
                 user=dify_config.MILVUS_USER,
                 password=dify_config.MILVUS_PASSWORD,
-                secure=dify_config.MILVUS_SECURE,
                 database=dify_config.MILVUS_DATABASE,
             )
         )

api/tests/integration_tests/vdb/milvus/test_milvus.py

@@ -12,8 +12,7 @@ class MilvusVectorTest(AbstractVectorTest):
         self.vector = MilvusVector(
             collection_name=self.collection_name,
             config=MilvusConfig(
-                host="localhost",
-                port=19530,
+                uri="http://localhost:19530",
                 user="root",
                 password="Milvus",
             ),

api/tests/unit_tests/core/rag/datasource/vdb/milvus/test_milvus.py

@@ -5,7 +5,7 @@ from core.rag.datasource.vdb.milvus.milvus_vector import MilvusConfig
 
 
 def test_default_value():
-    valid_config = {"host": "localhost", "port": 19530, "user": "root", "password": "Milvus"}
+    valid_config = {"uri": "http://localhost:19530", "user": "root", "password": "Milvus"}
 
     for key in valid_config:
         config = valid_config.copy()
@@ -15,5 +15,4 @@ def test_default_value():
         assert e.value.errors()[0]["msg"] == f"Value error, config MILVUS_{key.upper()} is required"
 
     config = MilvusConfig(**valid_config)
-    assert config.secure is False
     assert config.database == "default"

docker-legacy/docker-compose.yaml

@@ -128,16 +128,14 @@ services:
       # The Qdrant server gRPC mode PORT.
       QDRANT_GRPC_PORT: 6334
       # Milvus configuration Only available when VECTOR_STORE is `milvus`.
-      # The milvus host.
-      MILVUS_HOST: 127.0.0.1
-      # The milvus host.
-      MILVUS_PORT: 19530
+      # The milvus uri.
+      MILVUS_URI: http://127.0.0.1:19530
+      # The milvus token.
+      MILVUS_TOKEN: ''
       # The milvus username.
       MILVUS_USER: root
       # The milvus password.
       MILVUS_PASSWORD: Milvus
-      # The milvus tls switch.
-      MILVUS_SECURE: 'false'
       # relyt configurations
       RELYT_HOST: db
       RELYT_PORT: 5432
@@ -308,16 +306,14 @@ services:
       # The Qdrant server gRPC mode PORT.
       QDRANT_GRPC_PORT: 6334
       # Milvus configuration Only available when VECTOR_STORE is `milvus`.
-      # The milvus host.
-      MILVUS_HOST: 127.0.0.1
-      # The milvus host.
-      MILVUS_PORT: 19530
+      # The milvus uri.
+      MILVUS_URI: http://127.0.0.1:19530
+      # The milvus token.
+      MILVUS_PORT: ''
       # The milvus username.
       MILVUS_USER: root
       # The milvus password.
       MILVUS_PASSWORD: Milvus
-      # The milvus tls switch.
-      MILVUS_SECURE: 'false'
       # Mail configuration, support: resend
       MAIL_TYPE: ''
       # default send from email address, if not specified

docker/.env.example

@@ -335,16 +335,14 @@ QDRANT_GRPC_ENABLED=false
 QDRANT_GRPC_PORT=6334
 
 # Milvus configuration Only available when VECTOR_STORE is `milvus`.
-# The milvus host.
-MILVUS_HOST=127.0.0.1
-# The milvus host.
-MILVUS_PORT=19530
+# The milvus uri.
+MILVUS_URI=http://127.0.0.1:19530
+# The milvus token.
+MILVUS_TOKEN=
 # The milvus username.
 MILVUS_USER=root
 # The milvus password.
 MILVUS_PASSWORD=Milvus
-# The milvus tls switch.
-MILVUS_SECURE=false
 
 # MyScale configuration, only available when VECTOR_STORE is `myscale`
 # For multi-language support, please set MYSCALE_FTS_PARAMS with referring to:

docker/README.md

@@ -83,7 +83,7 @@ The `.env.example` file provided in the Docker setup is extensive and covers a w
 
 7. **Vector Database Configuration**:
     - `VECTOR_STORE`: Type of vector database (e.g., `weaviate`, `milvus`).
-    - Specific settings for each vector store like `WEAVIATE_ENDPOINT`, `MILVUS_HOST`.
+    - Specific settings for each vector store like `WEAVIATE_ENDPOINT`, `MILVUS_URI`.
 
 8. **CORS Configuration**:
     - `WEB_API_CORS_ALLOW_ORIGINS`, `CONSOLE_CORS_ALLOW_ORIGINS`: Settings for cross-origin resource sharing.

docker/docker-compose.yaml

@@ -89,11 +89,10 @@ x-shared-env: &shared-api-worker-env
   QDRANT_CLIENT_TIMEOUT: ${QDRANT_CLIENT_TIMEOUT:-20}
   QDRANT_GRPC_ENABLED: ${QDRANT_GRPC_ENABLED:-false}
   QDRANT_GRPC_PORT: ${QDRANT_GRPC_PORT:-6334}
-  MILVUS_HOST: ${MILVUS_HOST:-127.0.0.1}
-  MILVUS_PORT: ${MILVUS_PORT:-19530}
+  MILVUS_URI: ${MILVUS_URI:-http://127.0.0.1:19530}
+  MILVUS_TOKEN: ${MILVUS_TOKEN:-}
   MILVUS_USER: ${MILVUS_USER:-root}
   MILVUS_PASSWORD: ${MILVUS_PASSWORD:-Milvus}
-  MILVUS_SECURE: ${MILVUS_SECURE:-false}
   MYSCALE_HOST: ${MYSCALE_HOST:-myscale}
   MYSCALE_PORT: ${MYSCALE_PORT:-8123}
   MYSCALE_USER: ${MYSCALE_USER:-default}