feat: add rerank
This commit is contained in:
parent
d815416dda
commit
a00847ff82
|
|
@ -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
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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],
|
||||||
|
|
|
||||||
|
|
@ -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"
|
||||||
|
|
|
||||||
|
|
@ -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": "你好"},
|
|
||||||
]
|
]
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue