feat: search
This commit is contained in:
parent
fcaac9c0f1
commit
cf9d972e00
29
search.es
29
search.es
|
|
@ -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
|
||||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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"):
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Reference in New Issue