import uuid from src.pipeline.core.pocket_flow import AsyncBatchNode, AsyncNode from src.pipeline.core.utils import fixed_size_chunk, load_document, logger from src.pipeline.core import llm from src.pipeline.core import es import re class ChunkDocumentsNode(AsyncBatchNode): async def prep_async(self, shared): return shared["documents"] async def exec_async(self, document): """ :param document: {text, file_name} """ # print(f"document: {document}") text = document["text"] # 先将所有制表符等替换为一个空格 text = re.sub(r"[ \t]+", " ", text) # 再将多个空格替换为一个空格 text = re.sub(r" +", " ", text) # 去除首尾空格 text = text.strip() return [ { "text": x, "file_name": document["file_name"], "file_type": document["file_type"], "uuid": uuid.uuid4().hex, } for x in fixed_size_chunk(text, chunk_size=500) ] async def post_async(self, shared, prep_res, exec_res_list): all_chunks = [] for chunks in exec_res_list: all_chunks.extend(chunks) shared["documents"] = all_chunks return "default" class EmbeddingDocumentsNode(AsyncBatchNode): async def prep_async(self, shared): return shared["documents"] async def exec_async(self, document): """ :param document: {text, file_name} """ logger.debug(f"开始 embedding: {document["text"].strip()[:10]}...") res = {**document, "embedding": await llm.client.embedding(document["text"])} logger.debug(f"结束 embedding: {document["text"].strip()[:10]}...") return res async def post_async(self, shared, prep_res, exec_res_list): shared["documents"] = exec_res_list return "default" class ReadDocumentsNode(AsyncBatchNode): async def prep_async(self, shared): return shared["files"] async def exec_async(self, file_path): try: document_text = await load_document(file_path) return { "file_path": file_path, "file_name": file_path.split("/")[-1], "text": document_text, "text_length": len(document_text), "file_type": str(file_path.split(".")[-1]).strip().lower(), "status": "done", "message": "", } except Exception as e: return { "file_path": file_path, "file_name": file_path.split("/")[-1], "text": "", "text_length": 0, "file_type": str(file_path.split(".")[-1]).strip().lower(), "status": "error", "message": str(e), } async def post_async(self, shared, prep_res, exec_res): shared["documents"] = exec_res return "default" class WriteDocumentsToESNode(AsyncBatchNode): async def prep_async(self, shared): index = shared["index"] await es.client.create_index(index) return [ { "index": index, "es_id": x["uuid"], "doc_id": x["file_name"], "embedding": x["embedding"], "title": x["file_name"], "doc_type": x["file_type"], "content": x["text"], } for x in shared["documents"] ] async def exec_async(self, prep_res): await es.client.add_doc(**prep_res) return True async def post_async(self, shared, prep_res, exec_res): return "default" class EmbeddingNode(AsyncNode): async def prep_async(self, shared): return shared["text"] async def exec_async(self, prep_res): return await llm.client.embedding(prep_res) async def post_async(self, shared, prep_res, exec_res): shared["embedding"] = exec_res return "default" class SearchFromESNode(AsyncNode): async def prep_async(self, shared): return { "index": shared["index"], "query_text": shared["text"], "query_vector": shared["embedding"], "top_k": shared.get("top_k", 5), } async def exec_async(self, prep_res): return await es.client.hybrid_search_es(**prep_res) async def post_async(self, shared, prep_res, exec_res): logger.debug(f"SearchFromESNode: length:{len(exec_res)}") shared["results"] = exec_res return "default" class RerankNode(AsyncNode): """ 使用 LLM 对搜索结果进行重排 """ async def prep_async(self, shared): # 准备要重排的数据 return { "query": shared["text"], "results": shared.get("results", [])[: shared["rerank_top_k"]], } async def exec_async(self, prep_res): query = prep_res["query"] results = prep_res["results"] logger.debug(f"query: {query}") logger.debug(f"results: {results}") if not results: return [] # 调用 LLM 进行 rerank,这里假设 llm.client.rerank 接口存在 # 返回格式:[{"es_id":..., "score":..., "rank_score":...}] reranked = await llm.client.rerank(query=query, documents=results) return reranked async def post_async(self, shared, prep_res, exec_res): # 更新 shared 中的结果为重排后的结果 shared["results"] = exec_res return "default"