feat: add rerank

This commit is contained in:
李如威 2025-12-23 17:51:49 +08:00
parent d815416dda
commit a00847ff82
4 changed files with 70 additions and 5 deletions

View File

@ -1,4 +1,4 @@
#!/usr/bin/env bash #!/usr/bin/env bash
export $(cat .env | xargs) export $(cat .env | xargs)
pytest -s -W ignore::DeprecationWarning src/tests/test_nodes.py pytest -s -W ignore::DeprecationWarning -k test_search src/tests/test_nodes.py

View File

@ -1,3 +1,4 @@
from ast import List
import httpx import httpx
import json import json
from src.pipeline.config import config from src.pipeline.config import config
@ -65,6 +66,9 @@ class AsyncLLm:
return [] return []
async def rerank(self, query: str, documents: List[dict]) -> List[dict]:
return []
async def chat( async def chat(
self, self,
messages: list[dict], messages: list[dict],

View File

@ -145,3 +145,30 @@ class SearchFromESNode(AsyncNode):
async def post_async(self, shared, prep_res, exec_res): async def post_async(self, shared, prep_res, exec_res):
shared["results"] = exec_res shared["results"] = exec_res
return "default" return "default"
class RerankNode(AsyncNode):
"""
使用 LLM 对搜索结果进行重排
"""
async def prep_async(self, shared):
# 准备要重排的数据
return {"query": shared["text"], "results": shared.get("results", [])}
async def exec_async(self, prep_res):
query = prep_res["query"]
results = prep_res["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"

View File

@ -7,7 +7,6 @@ from src.pipeline.core import llm, es, nodes, utils
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_embedding(): async def test_embedding():
return
await llm.init_client() await llm.init_client()
await es.init_client() await es.init_client()
@ -43,7 +42,43 @@ async def test_search():
shared = { shared = {
"text": "哪里盛产矿石", "text": "哪里盛产矿石",
"index": "test_kb", "index": "test_kb",
"top_k": 5, "top_k": 10,
"results": [], # [{es_id, doc_id, title, type, created_at, score, content}]
}
embeddingNode = nodes.EmbeddingNode()
searchNode = nodes.SearchFromESNode()
embeddingNode >> searchNode
flow = AsyncFlow(embeddingNode)
await flow.run_async(shared)
logger.debug(json.dumps({**shared, "embedding": shared["embedding"][:4]}, indent=4, ensure_ascii=False))
request = llm.client.stream_chat(
messages=[
{"role": "system", "content": utils.rag_system_prompt()},
{"role": "user", "content": utils.rag_user_prompt(shared["text"], shared["results"])},
]
)
async for chunk in request:
logger.debug(chunk)
await llm.close_client()
await es.close_client()
@pytest.mark.asyncio
async def test_rerank():
await llm.init_client()
await es.init_client()
logger.debug("search from es")
shared = {
"text": "哪里盛产矿石",
"index": "test_kb",
"top_k": 10,
"results": [], # [{es_id, doc_id, title, type, created_at, score, content}] "results": [], # [{es_id, doc_id, title, type, created_at, score, content}]
} }
@ -59,8 +94,7 @@ async def test_search():
res = await llm.client.chat( res = await llm.client.chat(
messages=[ messages=[
{"role": "system", "content": utils.rag_system_prompt()}, {"role": "system", "content": utils.rag_system_prompt()},
{"role": "system", "content": utils.rag_user_prompt(shared["text"], shared["results"])}, {"role": "user", "content": utils.rag_user_prompt(shared["text"], shared["results"])},
# {"role": "system", "content": "你好"},
] ]
) )