ai_pipeline/src/pipeline/core/nodes.py

180 lines
5.5 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

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"