218 lines
6.7 KiB
Python
218 lines
6.7 KiB
Python
import json
|
||
import uuid
|
||
import re
|
||
from src.pipeline.core.pocket_flow import AsyncBatchNode, AsyncNode
|
||
from src.pipeline.core.utils import fixed_size_chunk, load_document, logger, baidu_search_async
|
||
from src.pipeline.core import llm
|
||
from src.pipeline.core import es
|
||
from itertools import chain
|
||
|
||
|
||
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 SearchNode(AsyncBatchNode):
|
||
|
||
async def prep_async(self, shared):
|
||
tasks = [
|
||
{
|
||
"from": "es",
|
||
"data": {
|
||
"index": shared["index"],
|
||
"query_text": shared["text"],
|
||
"query_vector": shared["embedding"],
|
||
"top_k": shared.get("top_k", 5),
|
||
},
|
||
}
|
||
]
|
||
if shared.get("search_web", False):
|
||
tasks.append(
|
||
{
|
||
"from": "web",
|
||
"data": {
|
||
"query": shared["text"],
|
||
"max_results": shared.get("top_k", 5),
|
||
},
|
||
}
|
||
)
|
||
return tasks
|
||
|
||
async def exec_async(self, prep_res):
|
||
_data = prep_res["data"]
|
||
_from = prep_res["from"]
|
||
if _from == "es":
|
||
return ("es", await es.client.hybrid_search_es(**_data))
|
||
if _from == "web":
|
||
return ("web", await baidu_search_async(**_data))
|
||
|
||
async def post_async(self, shared, prep_res, exec_res):
|
||
for _from, _data in exec_res:
|
||
if _from == 'es':
|
||
shared["results"] = _data
|
||
if _from == "web":
|
||
shared["web_results"] = _data
|
||
return "default"
|
||
|
||
|
||
class RerankNode(AsyncBatchNode):
|
||
"""
|
||
使用 LLM 对搜索结果进行重排
|
||
"""
|
||
|
||
async def prep_async(self, shared):
|
||
# 准备要重排的数据
|
||
tasks = [
|
||
{
|
||
"query": shared["text"],
|
||
"results": shared.get("results", [])[: shared["rerank_top_k"]],
|
||
}
|
||
]
|
||
if shared.get("search_web", False):
|
||
tasks.append(
|
||
{
|
||
"query": shared["text"],
|
||
"results": shared.get("web_results", [])[: shared["rerank_top_k"]],
|
||
}
|
||
)
|
||
return tasks
|
||
|
||
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":..., "rerank_score":...}]
|
||
reranked = await llm.client.rerank(query=query, documents=results)
|
||
return reranked
|
||
|
||
async def post_async(self, shared, prep_res, exec_res):
|
||
# 更新 shared 中的结果为重排后的结果
|
||
results = list(chain.from_iterable(exec_res))
|
||
logger.debug(results)
|
||
results.sort(key=lambda x: x["rerank_score"], reverse=True)
|
||
shared["results"] = results
|
||
return "default"
|