Kaynağa Gözat

Merge pull request #12947 from gaby/qdrant-grpc

feat: Add gRPC support to Qdrant
Tim Jaeryang Baek 3 ay önce
ebeveyn
işleme
b6caa95fac

+ 2 - 0
backend/open_webui/config.py

@@ -1694,6 +1694,8 @@ MILVUS_TOKEN = os.environ.get("MILVUS_TOKEN", None)
 QDRANT_URI = os.environ.get("QDRANT_URI", None)
 QDRANT_API_KEY = os.environ.get("QDRANT_API_KEY", None)
 QDRANT_ON_DISK = os.environ.get("QDRANT_ON_DISK", "false").lower() == "true"
+QDRANT_PREFER_GRPC = os.environ.get("QDRANT_PREFER_GRPC", "False").lower() == "true"
+QDRANT_GRPC_PORT = int(os.environ.get("QDRANT_GRPC_PORT", "6334"))
 
 # OpenSearch
 OPENSEARCH_URI = os.environ.get("OPENSEARCH_URI", "https://localhost:9200")

+ 30 - 6
backend/open_webui/retrieval/vector/dbs/qdrant.py

@@ -1,12 +1,19 @@
 from typing import Optional
 import logging
+from urllib.parse import urlparse
 
 from qdrant_client import QdrantClient as Qclient
 from qdrant_client.http.models import PointStruct
 from qdrant_client.models import models
 
 from open_webui.retrieval.vector.main import VectorItem, SearchResult, GetResult
-from open_webui.config import QDRANT_URI, QDRANT_API_KEY, QDRANT_ON_DISK
+from open_webui.config import (
+    QDRANT_URI,
+    QDRANT_API_KEY,
+    QDRANT_ON_DISK,
+    QDRANT_GRPC_PORT,
+    QDRANT_PREFER_GRPC,
+)
 from open_webui.env import SRC_LOG_LEVELS
 
 NO_LIMIT = 999999999
@@ -21,11 +28,28 @@ class QdrantClient:
         self.QDRANT_URI = QDRANT_URI
         self.QDRANT_API_KEY = QDRANT_API_KEY
         self.QDRANT_ON_DISK = QDRANT_ON_DISK
-        self.client = (
-            Qclient(url=self.QDRANT_URI, api_key=self.QDRANT_API_KEY)
-            if self.QDRANT_URI
-            else None
-        )
+        self.PREFER_GRPC = QDRANT_PREFER_GRPC
+        self.GRPC_PORT = QDRANT_GRPC_PORT
+
+        if not self.QDRANT_URI:
+            self.client = None
+            return
+
+        # Unified handling for either scheme
+        parsed = urlparse(self.QDRANT_URI)
+        host = parsed.hostname or self.QDRANT_URI
+        http_port = parsed.port or 6333  # default REST port
+
+        if self.PREFER_GRPC:
+            self.client = Qclient(
+                host=host,
+                port=http_port,
+                grpc_port=self.GRPC_PORT,
+                prefer_grpc=self.PREFER_GRPC,
+                api_key=self.QDRANT_API_KEY,
+            )
+        else:
+            self.client = Qclient(url=self.QDRANT_URI, api_key=self.QDRANT_API_KEY)
 
     def _result_to_get_result(self, points) -> GetResult:
         ids = []