Browse Source

Merge pull request #8562 from NYU-ITS/milvus_custom_db

[feat] Milvus: add new config var, MILVUS_DB
Timothy Jaeryang Baek 3 months ago
parent
commit
d1bde9f348

+ 1 - 0
backend/open_webui/config.py

@@ -1194,6 +1194,7 @@ CHROMA_HTTP_SSL = os.environ.get("CHROMA_HTTP_SSL", "false").lower() == "true"
 # Milvus
 # Milvus
 
 
 MILVUS_URI = os.environ.get("MILVUS_URI", f"{DATA_DIR}/vector_db/milvus.db")
 MILVUS_URI = os.environ.get("MILVUS_URI", f"{DATA_DIR}/vector_db/milvus.db")
+MILVUS_DB = os.environ.get("MILVUS_DB", "default")
 
 
 # Qdrant
 # Qdrant
 QDRANT_URI = os.environ.get("QDRANT_URI", None)
 QDRANT_URI = os.environ.get("QDRANT_URI", None)

+ 2 - 2
backend/open_webui/retrieval/vector/dbs/milvus.py

@@ -6,14 +6,14 @@ from typing import Optional
 
 
 from open_webui.retrieval.vector.main import VectorItem, SearchResult, GetResult
 from open_webui.retrieval.vector.main import VectorItem, SearchResult, GetResult
 from open_webui.config import (
 from open_webui.config import (
-    MILVUS_URI,
+    MILVUS_URI, MILVUS_DB,
 )
 )
 
 
 
 
 class MilvusClient:
 class MilvusClient:
     def __init__(self):
     def __init__(self):
         self.collection_prefix = "open_webui"
         self.collection_prefix = "open_webui"
-        self.client = Client(uri=MILVUS_URI)
+        self.client = Client(uri=MILVUS_URI, database=MILVUS_DB)
 
 
     def _result_to_get_result(self, result) -> GetResult:
     def _result_to_get_result(self, result) -> GetResult:
         ids = []
         ids = []