【AI RAG知识库】09.【检索】【节点2】
掌柜智库项目(RAG)实战9. 检索数据节点实现与测试9.2 查询节点搜索向量库 (node_search_embedding)文件:app/query_process/agent/nodes/node_search_embedding.py步骤分解1获取查询上下文从状态state中提取经过改写的问题rewritten_query以及在上一节点确认的商品名称列表item_names。2文本向量化调用嵌入模型接口将改写后的问题转换为向量表示。这里同时生成两种向量稠密向量 (Dense Vector)用于捕捉语义相似度。稀疏向量 (Sparse Vector)用于捕捉关键词匹配。3构建混合检索请求准备 Milvus 混合搜索参数过滤条件如果有确定的商品名构造item_name in [...]过滤器限定搜索范围。检索策略结合稠密向量COSINE 距离和稀疏向量IP 内积进行多路召回。4执行向量检索连接 Milvus 数据库的CHUNKS_COLLECTION集合执行混合检索。加权融合默认使用 0.8/0.2 的权重对稠密和稀疏得分进行重排序Rerank。结果截取返回 Top 5 最相关的文档切片Chunks。5更新状态将检索到的文档切片列表embedding_chunks存入状态供后续节点使用。节点代码实现importsysimportosfromapp.utils.task_utilsimportadd_running_task,add_done_taskfromapp.lm.embedding_utilsimportgenerate_embeddingsfromapp.clients.milvus_utilsimportcreate_hybrid_search_requests,hybrid_search,get_milvus_clientfromapp.core.loggerimportloggerfromdotenvimportload_dotenv,find_dotenv load_dotenv(find_dotenv())defnode_search_embedding(state): 核心节点函数基于已确认商品名改写后的用户问题执行Milvus向量数据库混合检索 流程用户问题向量化 → 构造带商品名过滤的混合搜索请求 → 执行稠密稀疏混合检索 → 返回检索结果 :param state: Dict - 会话状态字典包含上游传递的核心信息关键字段 { session_id: str, # 会话唯一标识 rewritten_query: str, # step3改写后的完整用户问题含商品名 item_names: list[str], # step6已确认的标准化商品名列表 is_stream: bool/None # 是否为流式响应可选 } :return: Dict - 检索结果字典仅包含embedding_chunks字段供下游节点使用 { embedding_chunks: List[Dict] # Milvus检索结果列表无结果则为空列表 # 每个元素为一条匹配的向量数据含业务字段 } logger.info(---search_milvus 开始处理---)add_running_task(state[session_id],sys._getframe().f_code.co_name,state[is_stream])# 1. 从会话状态中提取核心入参为后续检索做准备querystate.get(rewritten_query)# 提取改写后的用户问题含商品名独立完整item_namesstate.get(item_names)# 提取已确认的标准化商品名列表精准过滤用logger.info(f核心入参提取: query{query}, item_names{item_names})# 2. 对改写后的用户问题执行向量化生成BGEM3稠密稀疏向量logger.info(f开始为文本获取嵌入值:{query[:50]}...iflen(query)50elsef开始为“{query}”文本获取嵌入值...)# 调用向量化函数入参为列表支持批量此处仅单条查询# 生成与商品名匹配的语义向量用于后续相似性检索embeddingsgenerate_embeddings([query])dense_vecembeddings.get(dense)[0]sparse_vecembeddings.get(sparse)[0]# 打印稠密/稀疏向量日志便于调试向量生成结果logger.debug(f向量生成成功: dense_dim{len(dense_vec)}, sparse_len{len(sparse_vec)})# 3. 准备Milvus向量数据库连接相关配置指定检索的集合# 从环境变量中获取Milvus中存储「文本片段向量」的集合名表名避免硬编码collection_nameos.environ.get(CHUNKS_COLLECTION)logger.info(f正在连接到 Milvus 并准备集合 {collection_name}...)# 4. 构造Milvus混合搜索请求对象核心步骤# 先通过辅助函数生成商品名过滤表达式精准过滤检索范围# item_name in [苹果15, 华为P60]# 若无商品名直接返回None不做过滤ifnotitem_names:logger.warning(item_names 为空跳过检索返回空结果)return{embedding_chunks:[]}# 对每个商品名添加双引号拼接为Milvus支持的in语法格式quoted, .join(f{v}forvinitem_names)# 构造最终过滤表达式exprfitem_name in [{quoted}]logger.info(f创建搜索请求过滤表达式:{expr})# 构造稠密稀疏混合搜索请求整合向量、过滤条件、搜索参数reqscreate_hybrid_search_requests(dense_vectordense_vec,# 取用户问题的稠密向量单条故取索引0sparse_vectorsparse_vec,# 取用户问题的稀疏向量单条故取索引0exprexpr,# 商品名过滤表达式缩小检索范围仅检索指定商品名的向量limit10# 底层检索返回数量后续会再过滤为5预留更多结果做重排序)# 5. 执行Milvus稠密稀疏混合向量检索核心调用logger.info(开始执行 Milvus 混合检索...)clientget_milvus_client()reshybrid_search(clientclient,collection_namecollection_name,# 检索的目标集合名文本片段向量集合reqsreqs,# 构造好的混合搜索请求对象稠密稀疏ranker_weights(0.8,0.2),# 稠/稀疏向量评分权重配比各占50%可按业务调优norm_scoreTrue,# 开启评分归一化将距离值转为0-1区间的相似度评分limit5,# 最终返回的TOP5相似度最高结果output_fields[chunk_id,content,item_name]# 指定返回的业务字段)# 打印节点处理成功日志输出原始检索结果便于调试hit_countlen(res[0])ifresandlen(res)0else0logger.info(f节点 search_embedding 处理成功检索到{hit_count}条相关片段)ifhit_count0:logger.debug(fTop1 检索结果示例:{res[0][0]})# 标记当前任务完成更新任务状态add_done_task(state[session_id],sys._getframe().f_code.co_name,state.get(is_stream))# 6. 构造并返回结果若检索结果非空取res[0]适配Milvus批量搜索格式否则返回空列表# res[0]为当前单条查询的检索结果包含TOP5匹配的向量数据及业务字段return{embedding_chunks:res[0]ifreselse[]}主流程测试if__name____main__:# 模拟测试数据test_state{session_id:test_search_embedding_001,rewritten_query:HAK 180 烫金机使用说明,# 模拟改写后的查询item_names:[HAK 180 烫金机],# 模拟已确认的商品名is_stream:False}print(\n 开始测试 node_search_embedding 节点...)try:# 执行节点函数resultnode_search_embedding(test_state)logger.info(f检索结果汇总{result})# 验证结果chunksresult.get(embedding_chunks,[])print(f\n 测试完成检索到{len(chunks)}条结果)ifchunks:print(\n Top 1 结果详情:)top1chunks[0]# 打印关键字段注意entity字段可能包含具体业务数据print(fID:{top1.get(id)})print(fDistance:{top1.get(distance)})entitytop1.get(entity,{})print(fItem Name:{entity.get(item_name)})print(fContent Preview:{entity.get(content,)[:100]}...)else:print(\n 警告未检索到任何结果请检查 Milvus 数据或 item_names 是否匹配)exceptExceptionase:logger.error(f测试运行失败:{e},exc_infoTrue)