feat: web search

This commit is contained in:
李如威 2025-12-25 16:48:26 +08:00
parent 1a83148928
commit 5ae6ed73a5
6 changed files with 105 additions and 31 deletions

View File

@ -9,3 +9,4 @@ aiofiles
pillow pillow
loguru loguru
httpx httpx
baidusearch

View File

@ -1,9 +1,11 @@
import json
import uuid import uuid
import re
from src.pipeline.core.pocket_flow import AsyncBatchNode, AsyncNode 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 llm
from src.pipeline.core import es from src.pipeline.core import es
import re from itertools import chain
class ChunkDocumentsNode(AsyncBatchNode): class ChunkDocumentsNode(AsyncBatchNode):
@ -129,36 +131,70 @@ class EmbeddingNode(AsyncNode):
return "default" return "default"
class SearchFromESNode(AsyncNode): class SearchNode(AsyncBatchNode):
async def prep_async(self, shared): async def prep_async(self, shared):
return { tasks = [
"index": shared["index"], {
"query_text": shared["text"], "from": "es",
"query_vector": shared["embedding"], "data": {
"top_k": shared.get("top_k", 5), "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): 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): async def post_async(self, shared, prep_res, exec_res):
logger.debug(f"SearchFromESNode: length:{len(exec_res)}") for _from, _data in exec_res:
shared["results"] = exec_res if _from == 'es':
shared["results"] = _data
if _from == "web":
shared["web_results"] = _data
return "default" return "default"
class RerankNode(AsyncNode): class RerankNode(AsyncBatchNode):
""" """
使用 LLM 对搜索结果进行重排 使用 LLM 对搜索结果进行重排
""" """
async def prep_async(self, shared): async def prep_async(self, shared):
# 准备要重排的数据 # 准备要重排的数据
return { tasks = [
"query": shared["text"], {
"results": shared.get("results", [])[: shared["rerank_top_k"]], "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): async def exec_async(self, prep_res):
query = prep_res["query"] query = prep_res["query"]
@ -167,13 +203,15 @@ class RerankNode(AsyncNode):
logger.debug(f"results: {results}") logger.debug(f"results: {results}")
if not results: if not results:
return [] return []
# 调用 LLM 进行 rerank这里假设 llm.client.rerank 接口存在 # 调用 LLM 进行 rerank这里假设 llm.client.rerank 接口存在
# 返回格式:[{"es_id":..., "score":..., "rank_score":...}] # 返回格式:[{"es_id":..., "score":..., "rerank_score":...}]
reranked = await llm.client.rerank(query=query, documents=results) reranked = await llm.client.rerank(query=query, documents=results)
return reranked return reranked
async def post_async(self, shared, prep_res, exec_res): async def post_async(self, shared, prep_res, exec_res):
# 更新 shared 中的结果为重排后的结果 # 更新 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" return "default"

View File

@ -3,12 +3,13 @@ import docx
import fitz # PyMuPDF import fitz # PyMuPDF
import aiofiles import aiofiles
import io import io
import os import re
import sys import sys
from pathlib import Path from pathlib import Path
from PIL import Image from PIL import Image
from loguru import logger from loguru import logger
from src.pipeline.config import config 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) logger.debug(prompt)
return 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

View File

@ -1,5 +0,0 @@
import pytest
@pytest.mark.asyncio
async def test_embedding():
print("\n1")

View File

@ -48,7 +48,7 @@ async def test_search():
} }
embeddingNode = nodes.EmbeddingNode() embeddingNode = nodes.EmbeddingNode()
searchNode = nodes.SearchFromESNode() searchNode = nodes.SearchNode()
embeddingNode >> searchNode embeddingNode >> searchNode
flow = AsyncFlow(embeddingNode) flow = AsyncFlow(embeddingNode)
@ -77,16 +77,17 @@ async def test_rerank():
logger.debug("search from es") logger.debug("search from es")
shared = { shared = {
"text": "哪里盛产矿石", "text": "山海经中描述了哪里盛产矿石",
"index": "test_kb", "index": "test_kb",
"top_k": 50, "top_k": 10,
"search_web": False,
"rerank_top_k": 5, "rerank_top_k": 5,
"top_n": 3, "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.SearchNode()
rerankNode = nodes.RerankNode() rerankNode = nodes.RerankNode()
embeddingNode >> searchNode >> rerankNode embeddingNode >> searchNode >> rerankNode
# embeddingNode >> searchNode # embeddingNode >> searchNode

10
src/tests/test_utils.py Normal file
View File

@ -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))