From 0940f01634262468253e1925e48679f45da87527 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=91=86=E8=90=8C=E9=97=B7=E6=B2=B9=E7=93=B6?=
 <253605712@qq.com>
Date: Sun, 28 Apr 2024 15:33:32 +0800
Subject: [PATCH] enhancement:support  Qdrant gRPC mode (#3929)

---
 api/.env.example                                    |  4 +++-
 api/config.py                                       |  3 +++
 api/core/rag/datasource/vdb/qdrant/qdrant_vector.py |  6 +++++-
 api/core/rag/datasource/vdb/vector_factory.py       |  4 +++-
 docker/docker-compose.middleware.yaml               |  1 +
 docker/docker-compose.qdrant.yaml                   |  1 +
 docker/docker-compose.yaml                          | 11 ++++++++++-
 7 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/api/.env.example b/api/.env.example
index 0c30c341a6..1e348d77ba 100644
--- a/api/.env.example
+++ b/api/.env.example
@@ -75,6 +75,8 @@ WEAVIATE_BATCH_SIZE=100
 QDRANT_URL=http://localhost:6333
 QDRANT_API_KEY=difyai123456
 QDRANT_CLIENT_TIMEOUT=20
+QDRANT_GRPC_ENABLED=false
+QDRANT_GRPC_PORT=6334
 
 # Milvus configuration
 MILVUS_HOST=127.0.0.1
@@ -148,4 +150,4 @@ API_TOOL_DEFAULT_CONNECT_TIMEOUT=10
 API_TOOL_DEFAULT_READ_TIMEOUT=60
 
 # Log file path
-LOG_FILE=
\ No newline at end of file
+LOG_FILE=
diff --git a/api/config.py b/api/config.py
index 0de1a01881..12a57e9371 100644
--- a/api/config.py
+++ b/api/config.py
@@ -37,6 +37,7 @@ DEFAULTS = {
     'WEAVIATE_GRPC_ENABLED': 'True',
     'WEAVIATE_BATCH_SIZE': 100,
     'QDRANT_CLIENT_TIMEOUT': 20,
+    'QDRANT_GRPC_ENABLED': 'False',
     'CELERY_BACKEND': 'database',
     'LOG_LEVEL': 'INFO',
     'LOG_FILE': '',
@@ -226,6 +227,8 @@ class Config:
         self.QDRANT_URL = get_env('QDRANT_URL')
         self.QDRANT_API_KEY = get_env('QDRANT_API_KEY')
         self.QDRANT_CLIENT_TIMEOUT = get_env('QDRANT_CLIENT_TIMEOUT')
+        self.QDRANT_GRPC_ENABLED = get_env('QDRANT_GRPC_ENABLED')
+        self.QDRANT_GRPC_PORT = get_env('QDRANT_GRPC_PORT')
 
         # milvus / zilliz setting
         self.MILVUS_HOST = get_env('MILVUS_HOST')
diff --git a/api/core/rag/datasource/vdb/qdrant/qdrant_vector.py b/api/core/rag/datasource/vdb/qdrant/qdrant_vector.py
index 3e39978596..0f6bff6dd7 100644
--- a/api/core/rag/datasource/vdb/qdrant/qdrant_vector.py
+++ b/api/core/rag/datasource/vdb/qdrant/qdrant_vector.py
@@ -36,6 +36,8 @@ class QdrantConfig(BaseModel):
     api_key: Optional[str]
     timeout: float = 20
     root_path: Optional[str]
+    grpc_port: int = 6334
+    prefer_grpc: bool = False
 
     def to_qdrant_params(self):
         if self.endpoint and self.endpoint.startswith('path:'):
@@ -51,7 +53,9 @@ class QdrantConfig(BaseModel):
                 'url': self.endpoint,
                 'api_key': self.api_key,
                 'timeout': self.timeout,
-                'verify': self.endpoint.startswith('https')
+                'verify': self.endpoint.startswith('https'),
+                'grpc_port': self.grpc_port,
+                'prefer_grpc': self.prefer_grpc
             }
 
 
diff --git a/api/core/rag/datasource/vdb/vector_factory.py b/api/core/rag/datasource/vdb/vector_factory.py
index e7d3ca42b2..59f86b5b5f 100644
--- a/api/core/rag/datasource/vdb/vector_factory.py
+++ b/api/core/rag/datasource/vdb/vector_factory.py
@@ -86,7 +86,9 @@ class Vector:
                     endpoint=config.get('QDRANT_URL'),
                     api_key=config.get('QDRANT_API_KEY'),
                     root_path=current_app.root_path,
-                    timeout=config.get('QDRANT_CLIENT_TIMEOUT')
+                    timeout=config.get('QDRANT_CLIENT_TIMEOUT'),
+                    grpc_port=config.get('QDRANT_GRPC_PORT'),
+                    prefer_grpc=config.get('QDRANT_GRPC_ENABLED')
                 )
             )
         elif vector_type == "milvus":
diff --git a/docker/docker-compose.middleware.yaml b/docker/docker-compose.middleware.yaml
index f85f188bb9..dffb0a48f0 100644
--- a/docker/docker-compose.middleware.yaml
+++ b/docker/docker-compose.middleware.yaml
@@ -80,3 +80,4 @@ services:
   #     QDRANT_API_KEY: 'difyai123456'
   #   ports:
   #     - "6333:6333"
+  #     - "6334:6334"
diff --git a/docker/docker-compose.qdrant.yaml b/docker/docker-compose.qdrant.yaml
index 0f72b5f10f..04d34fac64 100644
--- a/docker/docker-compose.qdrant.yaml
+++ b/docker/docker-compose.qdrant.yaml
@@ -10,3 +10,4 @@ services:
       QDRANT_API_KEY: 'difyai123456'
     ports:
       - "6333:6333"
+      - "6334:6334"
diff --git a/docker/docker-compose.yaml b/docker/docker-compose.yaml
index 056dc10630..e03ea153ad 100644
--- a/docker/docker-compose.yaml
+++ b/docker/docker-compose.yaml
@@ -99,8 +99,12 @@ services:
       QDRANT_URL: http://qdrant:6333
       # The Qdrant API key.
       QDRANT_API_KEY: difyai123456
-      # The Qdrant clinet timeout setting.
+      # The Qdrant client timeout setting.
       QDRANT_CLIENT_TIMEOUT: 20
+      # The Qdrant client enable gRPC mode.
+      QDRANT_GRPC_ENABLED: 'false'
+      # 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
@@ -214,6 +218,10 @@ services:
       QDRANT_API_KEY: difyai123456
       # The Qdrant clinet timeout setting.
       QDRANT_CLIENT_TIMEOUT: 20
+      # The Qdrant client enable gRPC mode.
+      QDRANT_GRPC_ENABLED: 'false'
+      # 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
@@ -354,6 +362,7 @@ services:
   #   # uncomment to expose qdrant port to host
   #   # ports:
   #   #  - "6333:6333"
+  #   #  - "6334:6334"
 
   # The nginx reverse proxy.
   # used for reverse proxying the API service and Web service.