feat: search

This commit is contained in:
李如威 2025-12-19 14:00:44 +08:00
parent fcaac9c0f1
commit cf9d972e00
4 changed files with 95 additions and 16 deletions

View File

@ -0,0 +1,29 @@
// Shift + Command + P > ES: Set Host : http://elastic:12345@localhost:9210
//
POST /_analyze
{
"analyzer": "ik_smart",
"text": "上经常下雨却没有花草树木,蝮虫很多。 \n\n 又东三百七十里,曰瞿 q*)父之山,无草木,多金玉。 \n 【译文】再往东三百七十里是座瞿q*)父山,山上没有花草树木,但有丰富的金属矿物和玉石。 \n\n 又东四百里,曰句余之山,无草木,多金玉。 \n 【译文】再往东四百里,是座句余山,山上没有花草树木,但有丰富的金属矿物和玉石。 \n\n 又东五百里曰浮玉之山北望具区①东望诸■p0。有兽焉其状如虎而牛尾其音如吠犬其名曰彘zh@),是食人。苕水出于其阴,北流注于具区。其中多鮆 z9鱼②。 \n 【注释】①具区:就是现在江苏境内的太湖,②鮆鱼:即■鱼,又叫鲚鱼,鮤鱼,头生长得很长而狭薄,大的有一尺多长。 \n 【译文】再往东五百里,是座浮玉山,在山上向北可以望见具区泽,向东可以望见诸■水,山中有一种野兽,形状像老虎却长着牛的尾巴,发出的叫声如同狗叫,名称是彘,是能吃人的。苕水从这座山的北麓发源,向北流入具区泽。它里面生长着很多鮆鱼。 \n\n 又东五百里,曰成山,四方而三坛,其上多金玉,其下多青雘 hu^。■zhu#水出焉而南流注于虖h&shu^),其中多黄金。 \n 【译文】再往东五百里,是座成山,呈现四方形而像三层土坛,山上盛产金属矿物和玉石,山下多产青雘。■水从这座山发源,然后向南流入虖勺水,水中有丰富的黄金。 \n\n 又东五百里,曰会 gu@j@)之"
}
//
GET /test_kb/_search
{
"size": 5,
"query": {
"match": {
"content": "玉石"
}
},
"_source": [
"title",
"content",
"type",
"created_at"
]
}
//
// DELETE /test_kb

View File

@ -1,5 +1,5 @@
import uuid
from src.pipeline.core.pocket_flow import AsyncBatchNode
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 import llm
from src.pipeline.core import es
@ -60,7 +60,7 @@ class EmbeddingDocumentsNode(AsyncBatchNode):
return "default"
class ReadDocumentNode(AsyncBatchNode):
class ReadDocumentsNode(AsyncBatchNode):
async def prep_async(self, shared):
return shared["files"]
@ -92,7 +92,7 @@ class ReadDocumentNode(AsyncBatchNode):
return "default"
class WriteToElasticsearchNode(AsyncBatchNode):
class WriteDocumentsToESNode(AsyncBatchNode):
async def prep_async(self, shared):
index = shared["index"]
await es.client.create_index(index)
@ -115,3 +115,33 @@ class WriteToElasticsearchNode(AsyncBatchNode):
async def post_async(self, shared, prep_res, exec_res):
return "default"
class EmbeddingNode(AsyncNode):
async def prep_async(self, shared):
return shared["text"]
async def exec_async(self, prep_res):
return await llm.client.embedding(prep_res)
async def post_async(self, shared, prep_res, exec_res):
shared["embedding"] = exec_res
return "default"
class SearchFromESNode(AsyncNode):
async def prep_async(self, shared):
return {
"index": shared["index"],
"query_text": shared["text"],
"query_vector": shared["embedding"],
"top_k": shared.get("top_k", 5),
}
async def exec_async(self, prep_res):
return await es.client.hybrid_search_es(**prep_res)
async def post_async(self, shared, prep_res, exec_res):
shared["results"] = exec_res
return "default"

View File

@ -118,8 +118,6 @@ async def _load_pdf(path: str) -> str:
async def load_document(path: str) -> str:
suffix = Path(path).suffix.lower()
print(f"读取文件: {path}")
if suffix == ".txt":
return await _load_txt(path)
if suffix in (".md", ".markdown"):

View File

@ -1,31 +1,29 @@
import pytest
import json
from src.pipeline.core.nodes import ReadDocumentNode, ChunkDocumentsNode, EmbeddingDocumentsNode, WriteToElasticsearchNode
from src.pipeline.core.pocket_flow import AsyncFlow
from src.pipeline.core.utils import logger
from src.pipeline.core import llm, es
from src.pipeline.core import llm, es, nodes
@pytest.mark.asyncio
async def test_embedding():
return
await llm.init_client()
await es.init_client()
print("\n\ntest_embedding:\n")
logger.debug("file to es")
shared = {
"files": [
"./files/山海经01.txt",
],
"documents": [],
"files": [ "./files/山海经01.txt"],
"documents": [], # [{text, file_name, file_type, uuid, embedding}]
"index": "test_kb",
}
readNode = ReadDocumentNode()
chunkNode = ChunkDocumentsNode()
embeddingNode = EmbeddingDocumentsNode()
writeToESNode = WriteToElasticsearchNode()
readNode = nodes.ReadDocumentsNode()
chunkNode = nodes.ChunkDocumentsNode()
embeddingNode = nodes.EmbeddingDocumentsNode()
writeToESNode = nodes.WriteDocumentsToESNode()
readNode >> chunkNode >> embeddingNode >> writeToESNode
flow = AsyncFlow(readNode)
@ -35,3 +33,27 @@ async def test_embedding():
await llm.close_client()
await es.close_client()
@pytest.mark.asyncio
async def test_search():
await llm.init_client()
await es.init_client()
logger.debug("search from es")
shared = {
"text": "那座山盛产金属矿物",
"index": "test_kb",
}
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))
await llm.close_client()
await es.close_client()