from rag.db import MilvusOperate from rag.load_model import * from rag.llm import VllmApi from config import redis_config_dict import json import re import gc import redis from utils.get_logger import setup_logger import random import concurrent.futures logger = setup_logger(__name__) rerank_model_mapping = { "bce_rerank_model": (bce_rerank_tokenizer, bce_rerank_base_model), "rerank": (bce_rerank_tokenizer, bce_rerank_base_model) } redis_host = redis_config_dict.get("host") redis_port = redis_config_dict.get("port") redis_db = redis_config_dict.get("db") class ChatRetrieverRag: def __init__(self, chat_json: dict = None, chat_id: str=None): self.chat_id = chat_id self.redis_client = redis.StrictRedis(host=redis_host, port=redis_port, db=redis_db) if not chat_id: self.vllm_client = VllmApi(chat_json) def rerank_result(self,query, top_k, hybrid_search_result, rerank_embedding_name): rerank_list = [] for result in hybrid_search_result: rerank_list.append([query, result["content"]]) tokenizer, rerank_model = rerank_model_mapping.get(rerank_embedding_name) # 重排序 rerank_model.eval() inputs = tokenizer(rerank_list, padding=True, truncation=True, max_length=512, return_tensors="pt") inputs = {k: v.to(device) for k, v in inputs.items()} with torch.no_grad(): logits = rerank_model(**inputs, return_dict=True).logits.view(-1,).float() logits = logits.detach().cpu() scores = torch.sigmoid(logits).tolist() # scores_list = scores.tolist() logger.info(f"重排序的得分：{scores}") sorted_pairs = sorted(zip(scores, hybrid_search_result), key=lambda x: x[0], reverse=True) sorted_scores, sorted_search = zip(*sorted_pairs) sorted_scores_list = list(sorted_scores) sorted_search_list = list(sorted_search) for score, search in zip(sorted_scores_list, sorted_search_list): search["rerank_score"] = score del inputs, logits gc.collect() torch.cuda.empty_cache() search_result = random.sample(sorted_search_list, top_k) return search_result def request_milvus(self, collection_name_list, rag_embedding_name, query, top, mode): search_multi_collection_list = [] search_doc_id_to_knowledge_id_dict = {} def query_collection(collection_name): return collection_name, MilvusOperate( collection_name=collection_name, embedding_name=rag_embedding_name )._search(query, k=top, mode=mode) max_workers = min(10, len(collection_name_list)) with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: future_to_name = { executor.submit(query_collection, name): name for name in collection_name_list } for future in concurrent.futures.as_completed(future_to_name): collection_name = future_to_name[future] try: _, hybrid_search_result = future.result() for result in hybrid_search_result: doc_id = result.get("doc_id") if doc_id not in search_doc_id_to_knowledge_id_dict: search_doc_id_to_knowledge_id_dict[doc_id] = collection_name search_multi_collection_list.extend(hybrid_search_result) except Exception as e: logger.error(f"查询集合 {collection_name} 时出错: {str(e)}") continue return search_multi_collection_list, search_doc_id_to_knowledge_id_dict def retriever_result(self, chat_json): mode_search_mode = { "embedding": "dense", "keyword": "sparse", "mixed": "hybrid", } collection_name_list = chat_json.get("knowledgeIds") # 如果指定多个knowledge传list top = chat_json.get("sliceCount", 5) retriever_info = json.loads(chat_json.get("knowledgeInfo", "{}")) rag_embedding_name = chat_json.get("embeddingId") mode_rag = retriever_info.get("recall_method") mode = mode_search_mode.get(mode_rag, "hybrid") rerank_embedding_name = retriever_info.get("rerank_model_name") rerank_status = retriever_info.get("rerank_status") query = chat_json.get("query") # search_multi_collection_list = [] # search_doc_id_to_knowledge_id_dict = {} # for collection_name in collection_name_list: # hybrid_search_result = MilvusOperate( # collection_name=collection_name,embedding_name=rag_embedding_name)._search(query, k=top, mode=mode) # for result in hybrid_search_result: # doc_id = result.get("doc_id") # if doc_id not in search_doc_id_to_knowledge_id_dict: # search_doc_id_to_knowledge_id_dict[doc_id] = collection_name # search_multi_collection_list.extend(hybrid_search_result) search_multi_collection_list, search_doc_id_to_knowledge_id_dict = self.request_milvus(collection_name_list, rag_embedding_name, query, top, mode) # logger.info(f"根据{collection_name_list}检索到结果：{search_multi_collection_list}") if len(search_multi_collection_list) <= 0: rerank_result_list = [] elif rerank_status: rerank_result_list = self.rerank_result(query, top, search_multi_collection_list, rerank_embedding_name) else: for result in search_multi_collection_list: result["rerank_score"] = 1 # rerank_result_list = search_multi_collection_list rerank_result_list = random.sample(search_multi_collection_list, top) return rerank_result_list, search_doc_id_to_knowledge_id_dict async def generate_sync_response(self, chat_json): """非流式生成响应""" try: logger.info(f"非流式rag聊天的请求参数：{chat_json}") # 1. 检索相关文档 retriever_result_list, search_doc_id_to_knowledge_id_dict = self.retriever_result(chat_json) logger.info(f"向量库中获取的最终结果：{retriever_result_list}") # 2. 解析检索结果 chunk_content, knowledge_info_dict = self.parse_retriever_list( retriever_result_list, search_doc_id_to_knowledge_id_dict ) # 3. 构建提示词 prompt = chat_json.get("prompt") query = chat_json.get("query") prompt = prompt.replace("{知识}", chunk_content).replace("{用户}", query) logger.info(f"请求的提示词：{prompt}") # 4. 调用LLM temperature = float(chat_json.get("temperature", 0.6)) top_p = float(chat_json.get("topP", 0.7)) max_token = chat_json.get("maxToken", 4096) enable_think = chat_json.get("enable_think", False) model = chat_json.get("model", "DeepSeek-R1-Distill-Qwen-14B") # 非流式调用 chat_resp = "" async for chunk in self.vllm_client.chat( prompt, model, temperature=temperature, top_p=top_p, max_tokens=max_token, stream=False, # 设置为False非流式 history=[], enable_think=enable_think ): chat_resp = chunk.choices[0].message.content # 非流式只返回一次完整内容 # 记录检索数据 chat_id = chunk.get("id") self.redis_client.set(chat_id, json.dumps(knowledge_info_dict)) logger.info(f"LLM生成的回答：{chat_resp[:100]}...") # 5. 返回结果 return { "code": 200, "data": { "answer": chat_resp, "chat_id": chat_id } } except Exception as e: logger.error(f"非流式RAG聊天出错：{e}") return { "code": 500, "message": str(e) } async def generate_rag_response(self, chat_json, chunk_content): # logger.info(f"rag聊天的请求参数：{chat_json}") # retriever_result_list = self.retriever_result(chat_json) # logger.info(f"向量库中获取的最终结果：{retriever_result_list}") prompt = chat_json.get("prompt") query = chat_json.get("query") # chunk_content = "" # for retriever in retriever_result_list: # chunk_content += retriever["content"] prompt = prompt.replace("{知识}", chunk_content).replace("{用户}", query) # logger.info(f"请求的提示词：{prompt}") temperature = float(chat_json.get("temperature", 0.6)) top_p = float(chat_json.get("topP", 0.7)) max_token = chat_json.get("maxToken", 4096) enable_think = chat_json.get("enable_think", False) # 调用模型获取返回的结果 model = chat_json.get("model", "DeepSeek-R1-Distill-Qwen-14B") # model = "Qwen3-Coder-30B-loft" chat_resp = "" async for chunk in self.vllm_client.chat(prompt, model, temperature=temperature, top_p=top_p, max_tokens=max_token, stream=True, history=[], enable_think=enable_think): if chunk.get("event") == "add": chat_resp += chunk.get("data") chunk["data"] = chat_resp yield chunk else: yield chunk def parse_retriever_list(self, retriever_result_list, search_doc_id_to_knowledge_id_dict): doc_info = {} chunk_content = "" # 组织成每个doc对应的json格式 for retriever in retriever_result_list: chunk_text = retriever["content"] chunk_content += chunk_text doc_id = retriever["doc_id"] chunk_id = retriever["chunk_id"] rerank_score = retriever["rerank_score"] doc_name = retriever["metadata"]["source"] if doc_id in doc_info: c_d = { "chunk_id": chunk_id, "rerank_score": rerank_score, "chunk_len": len(chunk_text) } doc_info[doc_id]["chunk_list"].append(c_d) # doc_info[doc_id]["rerank_score"].append(rerank_score) else: c_d = { "chunk_id": chunk_id, "rerank_score": rerank_score, "chunk_len": len(chunk_text) } doc_info[doc_id] = { "doc_name": doc_name, "chunk_list": [c_d], } # doc_list = [] knowledge_info_dict = {} """ { "knowledge_id": [{}, {}], "knowledge_id" } """ for k, v in doc_info.items(): d = {} knowledge_id = search_doc_id_to_knowledge_id_dict.get(k) d["doc_id"] = k d["doc_name"] = v.get("doc_name") d["chunk_nums"] = len(v.get("chunk_list")) d["chunk_info_list"] = v.get("chunk_list") # d["chunk_len"] = len(v.get("chunk_id_list")) if knowledge_id in knowledge_info_dict: knowledge_info_dict[knowledge_id].append(d) else: knowledge_info_dict[knowledge_id] = [d] # doc_list.append(d) return chunk_content, knowledge_info_dict async def generate_event(self, chat_json, request): chat_id = "" try: logger.info(f"rag聊天的请求参数：{chat_json}") # knowledge_id = chat_json.get("knowledgeIds") retriever_result_list, search_doc_id_to_knowledge_id_dict = self.retriever_result(chat_json) logger.info(f"向量库中获取的最终结果：{retriever_result_list}") chunk_content, knowledge_info_dict = self.parse_retriever_list(retriever_result_list, search_doc_id_to_knowledge_id_dict) first = True async for event in self.generate_rag_response(chat_json, chunk_content): chat_id = event.get("id") if first: #json_dict = {"knowledge_id": knowledge_id, "doc": doc_list} self.redis_client.set(chat_id, json.dumps(knowledge_info_dict)) # logger.info(f"返回检索出的切片信息：{json_dict}") # yield {"id": chat_id, "event": "json", "data": json_dict} first = False yield event # yield json.dumps(event, ensure_ascii=False) if await request.is_disconnected(): logger.info(f"chat id:{chat_id}连接中断") yield {"id": chat_id, "event": "interrupted", "data": ""} return # yield json.dumps({"id": chat_id, "event": "interrupted", "data": ""}, ensure_ascii=False) except Exception as e: logger.info(f"执行出错：{e}") yield {"id": chat_id, "event": "finish", "data": ""} return # yield json.dumps({"id": chat_id, "event": "finish", "data": ""}, ensure_ascii=False) async def query_redis_chunk_by_chat_id(self, chat_id): chunk_redis_str = self.redis_client.get(chat_id) chunk_json = json.loads(chunk_redis_str) logger.info(f"redis中存储的chunk 信息：{chunk_json}") #knowledge_id = chunk_json.get("knowledge_id") knowledge_id = list(chunk_json.keys())[0] if chunk_json else None chunk_list = [] #for chunk in chunk_json["doc"]: # doc_chunk_list = chunk.get("chunk_info_list") # chunk_list.extend(doc_chunk_list) for knowledge_id_key, doc_list in chunk_json.items(): for doc in doc_list: doc_chunk_list = doc.get("chunk_info_list") if doc_chunk_list: chunk_list.extend(doc_chunk_list) if len(chunk_list) > 3: random_chunk_list = random.sample(chunk_list, 3) else: random_chunk_list = chunk_list chunk_id_list = [chunk["chunk_id"] for chunk in random_chunk_list] return knowledge_id, chunk_id_list async def generate_relevant_query(self, query_json): chat_id = query_json.get("chat_id") if chat_id: knowledge_id, chunk_id_list = await self.query_redis_chunk_by_chat_id(chat_id) chunk_content_list = MilvusOperate(collection_name=knowledge_id)._search_by_chunk_id_list(chunk_id_list) chunk_content = "\n".join(chunk_content_list) # generate_query_prompt = "请根据下面提供的信息，帮我生成" query_messages = query_json.get("messages") user_dict = query_messages.pop() messages = [{"role": "assistant", "content": chunk_content}, user_dict] else: messages = query_json.get("messages") model = query_json.get("model") # model = "Qwen3-Coder-30B-loft" query_result = self.vllm_client.chat(model=model, stream=False, history=messages) async for result in query_result: # result_json = json.loads(result) # logger.info(f"生成的问题：{result_json}") # result_str = result_json.get("choices", [{}])[0].get("message", {}).get("content", "").strip() result = result.strip() logger.info(f"模型生成的问题：{result}") try: if "```json" in result: json_pattern = r'```json\s(.*?)```' matches = re.findall(json_pattern, result, re.DOTALL) result = matches[0] query_json = json.loads(result) except Exception as e: query_json = eval(result) query_list = query_json.get("问题") return {"code": 200, "data": query_list} async def search_slice(self): try: chunk_redis_str = self.redis_client.get(self.chat_id) chunk_json = json.loads(chunk_redis_str) chunk_json["code"] = 200 except Exception as e: logger.error(f"查询redis报错：{e}") chunk_json = { "code": 500, "message": str(e) } return chunk_json