add Milvus db

config.py

@@ -60,9 +60,21 @@ CHUNK_SIZE = 250
 CHUNK_OVERLAP = 50  # 添加重叠以保持上下文连贯性，提升检索准确率
 
 # 向量数据库配置
+VECTOR_STORE_TYPE = os.environ.get("VECTOR_STORE_TYPE", "chroma")  # 可选: "chroma", "milvus"
 COLLECTION_NAME = "rag-chroma"
 EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"  # HuggingFace嵌入模型
 
+# Milvus 配置 (仅当 VECTOR_STORE_TYPE="milvus" 时生效)
+# 1. Milvus Lite (本地文件模式): 仅需设置 MILVUS_URI，无需 User/Password。适合 Kaggle/本地开发。
+# 2. Zilliz Cloud (云服务): 需要设置 MILVUS_URI (https://...) 和 MILVUS_PASSWORD (API Key/Token)。需官网注册。
+# 3. Milvus Server (Docker/K8s): 需要设置 HOST/PORT，可选 User/Password。
+MILVUS_HOST = os.environ.get("MILVUS_HOST", "localhost")
+MILVUS_PORT = os.environ.get("MILVUS_PORT", "19530")
+MILVUS_USER = os.environ.get("MILVUS_USER", "")      # 仅在自建 Server 开启认证或使用 Zilliz Cloud 时需要
+MILVUS_PASSWORD = os.environ.get("MILVUS_PASSWORD", "") # Zilliz Cloud 的 API Key 也填在这里
+# Milvus Lite 配置: 如果设置了 MILVUS_URI (如 "./milvus_demo.db")，将优先使用本地文件模式
+MILVUS_URI = os.environ.get("MILVUS_URI", "./milvus_rag.db")
+
 # 搜索配置
 WEB_SEARCH_RESULTS_COUNT = 3
 

document_processor.py

@@ -25,6 +25,13 @@ from config import (
     KEYWORD_SEARCH_K,
     BM25_K1,
     BM25_B,
+    # 向量库配置
+    VECTOR_STORE_TYPE,
+    MILVUS_HOST,
+    MILVUS_PORT,
+    MILVUS_USER,
+    MILVUS_PASSWORD,
+    MILVUS_URI,
     # 查询扩展配置
     ENABLE_QUERY_EXPANSION,
     QUERY_EXPANSION_MODEL,
@@ -236,12 +243,66 @@ class DocumentProcessor:
             os.makedirs(persist_directory, exist_ok=True)
             print(f"💾 使用默认持久化目录: {persist_directory}")
         
-        self.vectorstore = Chroma.from_documents(
-            documents=doc_splits,
-            collection_name=COLLECTION_NAME,
-            embedding=self.embeddings,
-            persist_directory=persist_directory  # 添加持久化目录
-        )
+        if VECTOR_STORE_TYPE.lower() == "milvus":
+            try:
+                from langchain_community.vectorstores import Milvus
+                
+                # 准备连接参数
+                connection_args = {}
+                
+                # 优先使用 URI (支持 Milvus Lite 本地文件 或 Zilliz Cloud)
+                # 只要 MILVUS_URI 被设置（config中默认是 ./milvus_rag.db），且不是空字符串
+                if MILVUS_URI and len(MILVUS_URI.strip()) > 0:
+                    # 判断是本地文件还是云服务
+                    is_local_file = not (MILVUS_URI.startswith("http://") or MILVUS_URI.startswith("https://"))
+                    mode_name = "Lite (Local File)" if is_local_file else "Cloud (HTTP)"
+                    
+                    print(f"🔄 正在连接 Milvus {mode_name} ({MILVUS_URI})...")
+                    connection_args["uri"] = MILVUS_URI
+                    
+                    # 如果是云服务，通常需要 token (使用 password 字段作为 token)
+                    if not is_local_file and MILVUS_PASSWORD:
+                         connection_args["token"] = MILVUS_PASSWORD
+                else:
+                    # 传统的 Host/Port 连接
+                    print(f"🔄 正在连接 Milvus Server ({MILVUS_HOST}:{MILVUS_PORT})...")
+                    connection_args = {
+                        "host": MILVUS_HOST,
+                        "port": MILVUS_PORT,
+                        "user": MILVUS_USER,
+                        "password": MILVUS_PASSWORD
+                    }
+
+                self.vectorstore = Milvus.from_documents(
+                    documents=doc_splits,
+                    embedding=self.embeddings,
+                    collection_name=COLLECTION_NAME,
+                    connection_args=connection_args,
+                    drop_old=True  # 重新创建索引
+                )
+                print("✅ Milvus 向量数据库初始化成功")
+            except ImportError:
+                print("❌ 未安装 pymilvus，请运行: pip install pymilvus")
+                raise
+            except Exception as e:
+                print(f"❌ Milvus 连接失败: {e}")
+                print("⚠️ 回退到 Chroma 数据库...")
+                # Fallback to Chroma
+                self.vectorstore = Chroma.from_documents(
+                    documents=doc_splits,
+                    collection_name=COLLECTION_NAME,
+                    embedding=self.embeddings,
+                    persist_directory=persist_directory
+                )
+        else:
+            # Default: Chroma
+            self.vectorstore = Chroma.from_documents(
+                documents=doc_splits,
+                collection_name=COLLECTION_NAME,
+                embedding=self.embeddings,
+                persist_directory=persist_directory  # 添加持久化目录
+            )
+            
         self.retriever = self.vectorstore.as_retriever()
         
         # 如果启用混合检索，创建BM25检索器和集成检索器

requirements.txt

@@ -12,6 +12,7 @@ langchain-ollama>=0.1.0
 
 # 向量数据库和嵌入
 chromadb>=0.4.0
+pymilvus>=2.4.2  # Milvus 客户端及 Lite 本地模式支持
 sentence-transformers>=2.2.0
 torch>=2.0.0
 transformers>=4.30.0