本地RAG尝试

app/rag/retriever.py

@@ -0,0 +1,144 @@
+"""
+Qdrant 向量检索器
+
+提供基础向量检索、混合检索（Dense + BM25）功能。
+"""
+
+import os
+from typing import List, Dict, Any, Optional
+from langchain_qdrant import Qdrant
+from langchain.embeddings.base import Embeddings
+from langchain.retrievers import ContextualCompressionRetriever
+from langchain.retrievers.document_compressors import DocumentCompressorPipeline
+from langchain.retrievers import EnsembleRetriever
+from qdrant_client import QdrantClient
+from qdrant_client.http import models
+
+
+def create_qdrant_client(
+    url: Optional[str] = None,
+    api_key: Optional[str] = None,
+) -> QdrantClient:
+    """
+    创建 Qdrant 客户端
+    
+    Args:
+        url: Qdrant 服务地址，默认从环境变量 QDRANT_URL 读取
+        api_key: API 密钥，默认从环境变量 QDRANT_API_KEY 读取
+    
+    Returns:
+        QdrantClient 实例
+    """
+    url = url or os.getenv("QDRANT_URL", "http://localhost:6333")
+    api_key = api_key or os.getenv("QDRANT_API_KEY")
+    
+    client_args = {"url": url}
+    if api_key:
+        client_args["api_key"] = api_key
+    
+    return QdrantClient(**client_args)
+
+
+def create_base_retriever(
+    collection_name: str,
+    embeddings: Embeddings,
+    search_kwargs: Optional[Dict[str, Any]] = None,
+    client: Optional[QdrantClient] = None,
+) -> Qdrant:
+    """
+    创建基础向量检索器
+    
+    Args:
+        collection_name: Qdrant 集合名称
+        embeddings: 嵌入模型
+        search_kwargs: 搜索参数，默认 {"k": 20}
+        client: Qdrant 客户端，如果为 None 则自动创建
+    
+    Returns:
+        Qdrant 检索器实例
+    """
+    if client is None:
+        client = create_qdrant_client()
+    
+    search_kwargs = search_kwargs or {"k": 20}
+    
+    # 创建 Qdrant 检索器
+    retriever = Qdrant.from_existing_collection(
+        embedding=embeddings,
+        collection_name=collection_name,
+        client=client,
+        content_payload_key="content",  # 假设存储的文本字段名为 "content"
+        metadata_payload_key="metadata",  # 元数据字段名
+    )
+    
+    return retriever.as_retriever(search_kwargs=search_kwargs)
+
+
+def create_hybrid_retriever(
+    collection_name: str,
+    embeddings: Embeddings,
+    dense_k: int = 10,
+    sparse_k: int = 10,
+    client: Optional[QdrantClient] = None,
+) -> ContextualCompressionRetriever:
+    """
+    创建混合检索器（Dense Vector + BM25）
+    
+    Args:
+        collection_name: Qdrant 集合名称
+        embeddings: 嵌入模型
+        dense_k: 向量检索返回数量
+        sparse_k: BM25 检索返回数量
+        client: Qdrant 客户端
+    
+    Returns:
+        混合检索器
+    """
+    if client is None:
+        client = create_qdrant_client()
+    
+    # 基础检索器（Qdrant 支持混合检索）
+    base_retriever = Qdrant.from_existing_collection(
+        embedding=embeddings,
+        collection_name=collection_name,
+        client=client,
+        content_payload_key="content",
+        metadata_payload_key="metadata",
+    )
+    
+    # 配置混合检索参数
+    search_kwargs = {
+        "k": dense_k + sparse_k,  # 总返回数量
+        "score_threshold": 0.3,   # 相似度阈值
+    }
+    
+    return base_retriever.as_retriever(search_kwargs=search_kwargs)
+
+
+def create_ensemble_retriever(
+    retrievers: List[Any],
+    weights: Optional[List[float]] = None,
+    c: int = 60,
+) -> EnsembleRetriever:
+    """
+    创建集成检索器，支持倒数排名融合 (RRF)
+    
+    Args:
+        retrievers: 检索器列表
+        weights: 检索器权重
+        c: RRF 常数（通常为60）
+    
+    Returns:
+        集成检索器
+    """
+    if weights is None:
+        weights = [1.0 / len(retrievers)] * len(retrievers)
+    
+    ensemble = EnsembleRetriever(
+        retrievers=retrievers,
+        weights=weights,
+        c=c,
+        search_type="rrf",  # 使用倒数排名融合
+    )
+    
+    return ensemble