feat: web search
This commit is contained in:
parent
1a83148928
commit
5ae6ed73a5
|
|
@ -9,3 +9,4 @@ aiofiles
|
|||
pillow
|
||||
loguru
|
||||
httpx
|
||||
baidusearch
|
||||
|
|
@ -1,9 +1,11 @@
|
|||
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
|
||||
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
|
||||
import re
|
||||
from itertools import chain
|
||||
|
||||
|
||||
class ChunkDocumentsNode(AsyncBatchNode):
|
||||
|
|
@ -129,36 +131,70 @@ class EmbeddingNode(AsyncNode):
|
|||
return "default"
|
||||
|
||||
|
||||
class SearchFromESNode(AsyncNode):
|
||||
class SearchNode(AsyncBatchNode):
|
||||
|
||||
async def prep_async(self, shared):
|
||||
return {
|
||||
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):
|
||||
return await es.client.hybrid_search_es(**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):
|
||||
logger.debug(f"SearchFromESNode: length:{len(exec_res)}")
|
||||
shared["results"] = 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(AsyncNode):
|
||||
class RerankNode(AsyncBatchNode):
|
||||
"""
|
||||
使用 LLM 对搜索结果进行重排
|
||||
"""
|
||||
|
||||
async def prep_async(self, shared):
|
||||
# 准备要重排的数据
|
||||
return {
|
||||
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"]
|
||||
|
|
@ -167,13 +203,15 @@ class RerankNode(AsyncNode):
|
|||
logger.debug(f"results: {results}")
|
||||
if not results:
|
||||
return []
|
||||
|
||||
# 调用 LLM 进行 rerank,这里假设 llm.client.rerank 接口存在
|
||||
# 返回格式:[{"es_id":..., "score":..., "rank_score":...}]
|
||||
# 返回格式:[{"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 中的结果为重排后的结果
|
||||
shared["results"] = exec_res
|
||||
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"
|
||||
|
|
|
|||
|
|
@ -3,12 +3,13 @@ import docx
|
|||
import fitz # PyMuPDF
|
||||
import aiofiles
|
||||
import io
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from PIL import Image
|
||||
from loguru import logger
|
||||
from src.pipeline.config import config
|
||||
from baidusearch.baidusearch import search
|
||||
|
||||
# -----------------------------
|
||||
# 日志
|
||||
|
|
@ -187,3 +188,31 @@ def rag_user_prompt(query: str, documents: list[dict]) -> str:
|
|||
"""
|
||||
logger.debug(prompt)
|
||||
return prompt
|
||||
|
||||
# -----------------------------
|
||||
# 其他工具
|
||||
# -----------------------------
|
||||
|
||||
|
||||
async def baidu_search_async(query: str, max_results: int = 5):
|
||||
"""
|
||||
异步调用 baidusearch(内部用 asyncio.to_thread 封装同步函数)
|
||||
返回结构化搜索结果列表
|
||||
"""
|
||||
|
||||
def sync_search():
|
||||
return list(search(query, num_results=max_results))
|
||||
|
||||
results = await asyncio.to_thread(sync_search)
|
||||
docs = []
|
||||
for r in results:
|
||||
docs.append(
|
||||
{
|
||||
"title": r.get("title"),
|
||||
"content": re.sub(r"\s+", " ", r.get("abstract", "")),
|
||||
"url": r.get("url"),
|
||||
"type": "web",
|
||||
"score": 0.5,
|
||||
}
|
||||
)
|
||||
return docs
|
||||
|
|
|
|||
|
|
@ -1,5 +0,0 @@
|
|||
import pytest
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_embedding():
|
||||
print("\n1")
|
||||
|
|
@ -48,7 +48,7 @@ async def test_search():
|
|||
}
|
||||
|
||||
embeddingNode = nodes.EmbeddingNode()
|
||||
searchNode = nodes.SearchFromESNode()
|
||||
searchNode = nodes.SearchNode()
|
||||
embeddingNode >> searchNode
|
||||
flow = AsyncFlow(embeddingNode)
|
||||
|
||||
|
|
@ -77,16 +77,17 @@ async def test_rerank():
|
|||
logger.debug("search from es")
|
||||
|
||||
shared = {
|
||||
"text": "哪里盛产矿石",
|
||||
"text": "山海经中描述了哪里盛产矿石",
|
||||
"index": "test_kb",
|
||||
"top_k": 50,
|
||||
"top_k": 10,
|
||||
"search_web": False,
|
||||
"rerank_top_k": 5,
|
||||
"top_n": 3,
|
||||
"results": [], # [{es_id, doc_id, title, type, created_at, score, content}]
|
||||
}
|
||||
|
||||
embeddingNode = nodes.EmbeddingNode()
|
||||
searchNode = nodes.SearchFromESNode()
|
||||
searchNode = nodes.SearchNode()
|
||||
rerankNode = nodes.RerankNode()
|
||||
embeddingNode >> searchNode >> rerankNode
|
||||
# embeddingNode >> searchNode
|
||||
|
|
|
|||
|
|
@ -0,0 +1,10 @@
|
|||
import pytest
|
||||
import json
|
||||
from src.pipeline.core.utils import logger, baidu_search_async
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_search_web():
|
||||
query = "孙悟空哪里出生"
|
||||
logger.debug(f"query: {query} ...")
|
||||
results = await baidu_search_async(query, max_results=5)
|
||||
logger.debug(json.dumps(results, indent=4, ensure_ascii=False))
|
||||
Loading…
Reference in New Issue