"""
Haystack Pipeline模块 - 实现基于Haystack的RAG管道
"""
import logging
from typing import Any, Dict, List, Optional
from haystack import Pipeline
from haystack.components.retrievers import InMemoryEmbeddingRetriever
from haystack.document_stores.in_memory import InMemoryDocumentStore
from haystack_integrations.components.embedders.sentence_transformers import SentenceTransformersTextEmbedder
from connectors.base_llm_connector import BaseLLMConnector
from connectors.haystack_llm_connector import ConnectorChatGenerator
from rag.prompt_builder import ChatPromptBuilder
# 配置日志
logger = logging.getLogger(__name__)
[docs]
class RAGPipeline:
"""基于Haystack的RAG管道"""
[docs]
def __init__(
self,
llm_connector: BaseLLMConnector,
document_store: Optional[InMemoryDocumentStore] = None,
retriever_top_k: int = 5,
prompt_builder: Optional[ChatPromptBuilder] = None,
embedding_model: str = "sentence-transformers/all-MiniLM-L6-v2",
):
"""
初始化RAG管道
Args:
llm_connector (BaseLLMConnector): LLM连接器
document_store (InMemoryDocumentStore, optional): 当前Haystack检索器使用内存文档存储，如果为None则创建InMemoryDocumentStore
retriever_top_k (int): 检索器返回的最大文档数量
prompt_builder (ChatPromptBuilder, optional): 提示构建器，如果为None则创建默认构建器
embedding_model (str): 嵌入模型名称
"""
self.llm_connector = llm_connector
self.document_store = document_store or InMemoryDocumentStore()
self.retriever_top_k = retriever_top_k
self.prompt_builder = prompt_builder or ChatPromptBuilder()
self.embedding_model = embedding_model
# 创建查询嵌入器
self.query_embedder = SentenceTransformersTextEmbedder(model=embedding_model, progress_bar=False)
# 创建检索器
self.retriever = InMemoryEmbeddingRetriever(document_store=self.document_store, top_k=retriever_top_k)
# 创建Pipeline
self._create_pipeline()
def _create_pipeline(self) -> None:
"""创建Haystack Pipeline"""
self.pipeline = Pipeline()
# 向Pipeline中添加组件
self.pipeline.add_component("query_embedder", self.query_embedder)
self.pipeline.add_component("retriever", self.retriever)
self.pipeline.add_component("prompt_builder", self.prompt_builder)
self.pipeline.add_component("llm", ConnectorChatGenerator(self.llm_connector))
# 定义组件之间的连接
self.pipeline.connect("query_embedder.embedding", "retriever.query_embedding")
self.pipeline.connect("retriever.documents", "prompt_builder.documents")
self.pipeline.connect("prompt_builder.messages", "llm.messages")
logger.info("成功创建Haystack RAG Pipeline (含查询嵌入器)")
[docs]
def run(
self,
query: str,
chat_history: Optional[List[Dict[str, str]]] = None,
filters: Optional[Dict[str, Any]] = None,
generation_kwargs: Optional[Dict[str, Any]] = None,
top_k: Optional[int] = None,
) -> Dict[str, Any]:
"""
运行RAG管道
Args:
query (str): 用户查询
chat_history (list, optional): 聊天历史记录，格式为[{"role": "user", "content": "消息内容"}]
filters (dict, optional): 检索过滤条件
generation_kwargs (dict, optional): 生成参数
top_k (int, optional): 单次检索数量，省略时保留配置默认值
Returns:
dict: 包含回答和检索文档的结果
"""
try:
if top_k is not None and (type(top_k) is not int or top_k <= 0):
raise ValueError("top_k must be a positive integer")
# 准备输入
inputs: Dict[str, Dict[str, Any]] = {
"query_embedder": {"text": query},
"retriever": {"filters": filters},
"llm": {"generation_kwargs": generation_kwargs or {}},
"prompt_builder": {"query": query, "chat_history": chat_history or []},
}
if top_k is not None:
inputs["retriever"]["top_k"] = top_k
# 运行管道
logger.info(f"运行RAG管道: 查询='{query}'")
logger.debug(f"Pipeline inputs: {inputs}")
output = self.pipeline.run(inputs, include_outputs_from={"retriever"})
# 调试：打印输出结构
logger.debug(f"Pipeline output keys: {list(output.keys())}")
if "retriever" in output:
logger.debug(f"Retriever output documents count: {len(output['retriever']['documents'])}")
# 提取结果
answer = output["llm"]["replies"][0].text
documents = output["retriever"]["documents"] if "retriever" in output else []
# 构建结果
result = {
"answer": answer,
"documents": [{"id": doc.id, "content": doc.content, "metadata": doc.meta} for doc in documents],
"query": query,
}
logger.info(f"RAG管道运行成功: 检索到{len(documents)}个文档")
return result
except Exception as e:
logger.error(f"RAG管道运行失败: {str(e)}")
return {"answer": f"查询处理失败: {str(e)}", "documents": [], "query": query, "error": str(e)}
[docs]
def update_retriever(
self, document_store: Optional[InMemoryDocumentStore] = None, top_k: Optional[int] = None
) -> None:
"""
更新检索器
Args:
document_store (DocumentStore, optional): 新的文档存储
top_k (int, optional): 新的最大检索文档数量
"""
# 更新文档存储
if document_store:
self.document_store = document_store
# 更新top_k
if top_k:
self.retriever_top_k = top_k
# 重新创建检索器
self.retriever = InMemoryEmbeddingRetriever(document_store=self.document_store, top_k=self.retriever_top_k)
# 重新创建Pipeline
self._create_pipeline()
logger.info(f"已更新检索器: top_k={self.retriever_top_k}")
[docs]
def update_llm_connector(self, llm_connector: BaseLLMConnector) -> None:
"""
更新LLM连接器
Args:
llm_connector (BaseLLMConnector): 新的LLM连接器
"""
self.llm_connector = llm_connector
self._create_pipeline()
logger.info(f"已更新LLM连接器: model={llm_connector.model}")
[docs]
class RAGPipelineFactory:
"""RAG管道工厂类，创建基于Haystack的RAG管道"""
[docs]
@staticmethod
def create_pipeline(
llm_connector: BaseLLMConnector,
document_store: Optional[InMemoryDocumentStore] = None,
retriever_top_k: int = 5,
prompt_builder: Optional[ChatPromptBuilder] = None,
config: Optional[Dict[str, Any]] = None,
) -> RAGPipeline:
"""
创建RAG管道
Args:
llm_connector (BaseLLMConnector): LLM连接器
document_store (InMemoryDocumentStore, optional): 当前Haystack检索器使用内存文档存储
retriever_top_k (int): 检索器返回的最大文档数量
prompt_builder (ChatPromptBuilder, optional): 提示构建器
config (dict, optional): 配置参数
Returns:
RAGPipeline: RAG管道实例
"""
# 处理配置参数
if config:
retriever_top_k = config.get("retriever_top_k", config.get("top_k", retriever_top_k))
if isinstance(retriever_top_k, bool) or not isinstance(retriever_top_k, int) or retriever_top_k <= 0:
raise ValueError("retriever_top_k/top_k must be a positive integer")
return RAGPipeline(
llm_connector=llm_connector,
document_store=document_store,
retriever_top_k=retriever_top_k,
prompt_builder=prompt_builder,
)