feat: web search
This commit is contained in:
parent
1a83148928
commit
5ae6ed73a5
|
|
@ -9,3 +9,4 @@ aiofiles
|
||||||
pillow
|
pillow
|
||||||
loguru
|
loguru
|
||||||
httpx
|
httpx
|
||||||
|
baidusearch
|
||||||
|
|
@ -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 = [
|
||||||
|
{
|
||||||
|
"from": "es",
|
||||||
|
"data": {
|
||||||
"index": shared["index"],
|
"index": shared["index"],
|
||||||
"query_text": shared["text"],
|
"query_text": shared["text"],
|
||||||
"query_vector": shared["embedding"],
|
"query_vector": shared["embedding"],
|
||||||
"top_k": shared.get("top_k", 5),
|
"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"],
|
"query": shared["text"],
|
||||||
"results": shared.get("results", [])[: shared["rerank_top_k"]],
|
"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"
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
|
||||||
|
|
@ -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()
|
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
|
||||||
|
|
|
||||||
|
|
@ -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