feat: add rerank

This commit is contained in:
李如威 2025-12-24 16:08:11 +08:00
parent a00847ff82
commit 1a83148928
5 changed files with 81 additions and 8 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 -k test_search src/tests/test_nodes.py pytest -s -W ignore::DeprecationWarning src/tests/test_nodes.py::test_rerank

View File

@ -14,6 +14,9 @@ class Config(TypedDict):
llm_api_key: str llm_api_key: str
llm_api_host: str llm_api_host: str
llm_model: str llm_model: str
rerank_api_key: str
rerank_api_host: str
rerank_model: str
embedding_api_key: str embedding_api_key: str
embedding_api_host: str embedding_api_host: str
embedding_model: str embedding_model: str
@ -33,6 +36,9 @@ def _read_config() -> Config:
"llm_api_host": os.getenv("LLM_API_HOST"), "llm_api_host": os.getenv("LLM_API_HOST"),
"llm_api_key": os.getenv("LLM_API_KEY"), "llm_api_key": os.getenv("LLM_API_KEY"),
"llm_model": os.getenv("LLM_MODEL"), "llm_model": os.getenv("LLM_MODEL"),
"rerank_api_host": os.getenv("RERANK_API_HOST"),
"rerank_api_key": os.getenv("RERANK_API_KEY"),
"rerank_model": os.getenv("RERANK_MODEL"),
"embedding_api_host": os.getenv("EMBEDDING_API_HOST"), "embedding_api_host": os.getenv("EMBEDDING_API_HOST"),
"embedding_api_key": os.getenv("EMBEDDING_API_KEY"), "embedding_api_key": os.getenv("EMBEDDING_API_KEY"),
"embedding_model": os.getenv("EMBEDDING_MODEL"), "embedding_model": os.getenv("EMBEDDING_MODEL"),

View File

@ -1,4 +1,3 @@
from ast import List
import httpx import httpx
import json import json
from src.pipeline.config import config from src.pipeline.config import config
@ -46,6 +45,23 @@ class AsyncLLm:
"Authorization": f"Bearer {self.embedding_api_key}", "Authorization": f"Bearer {self.embedding_api_key}",
}, },
) )
# rerank
self.rerank_api = config["rerank_api_host"].rstrip("/") + "/score"
self.rerank_model = config["rerank_model"]
self.rerank_api_key = config["rerank_api_key"]
self.rerank_client = httpx.AsyncClient(
http2=False,
trust_env=False,
timeout=httpx.Timeout(timeout),
limits=httpx.Limits(
max_connections=max_connections,
max_keepalive_connections=max_keepalive,
),
headers={
"Content-Type": "application/json",
"Authorization": f"Bearer {self.rerank_api_key}",
},
)
async def embedding(self, text: str) -> list[float]: async def embedding(self, text: str) -> list[float]:
try: try:
@ -66,7 +82,48 @@ class AsyncLLm:
return [] return []
async def rerank(self, query: str, documents: List[dict]) -> List[dict]: async def rerank(self, query: str, documents: list[dict]) -> list[dict]:
logger.debug("\n\n default:\n\n" + "\n".join([str({"content": x["content"][:10], "score": x["score"]}) for x in documents]))
if not documents:
return []
text_1 = []
text_2 = []
for doc in documents:
text_1.append(query)
text_2.append(doc.get("content", "")[:1024])
# Qwen3 专用
payload = {
"model": self.rerank_model,
"text_1": text_1,
"text_2": text_2,
}
try:
resp = await self.rerank_client.post(self.rerank_api, json=payload)
resp.raise_for_status()
result = resp.json()
scores = [item["score"] for item in result["data"]]
reranked = []
for doc, score in zip(documents, scores):
d = dict(doc)
d["rerank_score"] = float(score)
reranked.append(d)
reranked.sort(key=lambda x: x["rerank_score"], reverse=True)
logger.debug("\n\n rerank: \n\n" + "\n".join([str({"content": x["content"][:10], "score": x["score"], "rerank_score": x["rerank_score"]}) for x in reranked]))
return reranked
except httpx.HTTPStatusError as e:
logger.error(e)
logger.error(f"Rerank HTTP error: {e.response.text}")
except Exception as e:
logger.exception("Rerank request failed")
return [] return []
async def chat( async def chat(

View File

@ -143,6 +143,7 @@ class SearchFromESNode(AsyncNode):
return await es.client.hybrid_search_es(**prep_res) return await es.client.hybrid_search_es(**prep_res)
async def post_async(self, shared, prep_res, exec_res): async def post_async(self, shared, prep_res, exec_res):
logger.debug(f"SearchFromESNode: length:{len(exec_res)}")
shared["results"] = exec_res shared["results"] = exec_res
return "default" return "default"
@ -154,12 +155,16 @@ class RerankNode(AsyncNode):
async def prep_async(self, shared): async def prep_async(self, shared):
# 准备要重排的数据 # 准备要重排的数据
return {"query": shared["text"], "results": shared.get("results", [])} return {
"query": shared["text"],
"results": shared.get("results", [])[: shared["rerank_top_k"]],
}
async def exec_async(self, prep_res): async def exec_async(self, prep_res):
query = prep_res["query"] query = prep_res["query"]
results = prep_res["results"] results = prep_res["results"]
logger.debug(f"query: {query}")
logger.debug(f"results: {results}")
if not results: if not results:
return [] return []

View File

@ -1,3 +1,4 @@
import asyncio
import pytest import pytest
import json import json
from src.pipeline.core.pocket_flow import AsyncFlow from src.pipeline.core.pocket_flow import AsyncFlow
@ -78,13 +79,17 @@ async def test_rerank():
shared = { shared = {
"text": "哪里盛产矿石", "text": "哪里盛产矿石",
"index": "test_kb", "index": "test_kb",
"top_k": 10, "top_k": 50,
"rerank_top_k": 5,
"top_n": 3,
"results": [], # [{es_id, doc_id, title, type, created_at, score, content}] "results": [], # [{es_id, doc_id, title, type, created_at, score, content}]
} }
embeddingNode = nodes.EmbeddingNode() embeddingNode = nodes.EmbeddingNode()
searchNode = nodes.SearchFromESNode() searchNode = nodes.SearchFromESNode()
embeddingNode >> searchNode rerankNode = nodes.RerankNode()
embeddingNode >> searchNode >> rerankNode
# embeddingNode >> searchNode
flow = AsyncFlow(embeddingNode) flow = AsyncFlow(embeddingNode)
await flow.run_async(shared) await flow.run_async(shared)
@ -94,7 +99,7 @@ async def test_rerank():
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": "user", "content": utils.rag_user_prompt(shared["text"], shared["results"])}, {"role": "user", "content": utils.rag_user_prompt(shared["text"], shared["results"][:shared["top_n"]])},
] ]
) )