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
|
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.utils import fixed_size_chunk, load_document, logger
|
||||||
from src.pipeline.core import llm
|
from src.pipeline.core import llm
|
||||||
from src.pipeline.core import es
|
from src.pipeline.core import es
|
||||||
|
|
@ -60,7 +60,7 @@ class EmbeddingDocumentsNode(AsyncBatchNode):
|
||||||
return "default"
|
return "default"
|
||||||
|
|
||||||
|
|
||||||
class ReadDocumentNode(AsyncBatchNode):
|
class ReadDocumentsNode(AsyncBatchNode):
|
||||||
async def prep_async(self, shared):
|
async def prep_async(self, shared):
|
||||||
return shared["files"]
|
return shared["files"]
|
||||||
|
|
||||||
|
|
@ -92,7 +92,7 @@ class ReadDocumentNode(AsyncBatchNode):
|
||||||
return "default"
|
return "default"
|
||||||
|
|
||||||
|
|
||||||
class WriteToElasticsearchNode(AsyncBatchNode):
|
class WriteDocumentsToESNode(AsyncBatchNode):
|
||||||
async def prep_async(self, shared):
|
async def prep_async(self, shared):
|
||||||
index = shared["index"]
|
index = shared["index"]
|
||||||
await es.client.create_index(index)
|
await es.client.create_index(index)
|
||||||
|
|
@ -115,3 +115,33 @@ class WriteToElasticsearchNode(AsyncBatchNode):
|
||||||
|
|
||||||
async def post_async(self, shared, prep_res, exec_res):
|
async def post_async(self, shared, prep_res, exec_res):
|
||||||
return "default"
|
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:
|
async def load_document(path: str) -> str:
|
||||||
suffix = Path(path).suffix.lower()
|
suffix = Path(path).suffix.lower()
|
||||||
|
|
||||||
print(f"读取文件: {path}")
|
|
||||||
|
|
||||||
if suffix == ".txt":
|
if suffix == ".txt":
|
||||||
return await _load_txt(path)
|
return await _load_txt(path)
|
||||||
if suffix in (".md", ".markdown"):
|
if suffix in (".md", ".markdown"):
|
||||||
|
|
|
||||||
|
|
@ -1,31 +1,29 @@
|
||||||
import pytest
|
import pytest
|
||||||
import json
|
import json
|
||||||
from src.pipeline.core.nodes import ReadDocumentNode, ChunkDocumentsNode, EmbeddingDocumentsNode, WriteToElasticsearchNode
|
|
||||||
from src.pipeline.core.pocket_flow import AsyncFlow
|
from src.pipeline.core.pocket_flow import AsyncFlow
|
||||||
from src.pipeline.core.utils import logger
|
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
|
@pytest.mark.asyncio
|
||||||
async def test_embedding():
|
async def test_embedding():
|
||||||
|
return
|
||||||
|
|
||||||
await llm.init_client()
|
await llm.init_client()
|
||||||
await es.init_client()
|
await es.init_client()
|
||||||
|
|
||||||
print("\n\ntest_embedding:\n")
|
logger.debug("file to es")
|
||||||
|
|
||||||
shared = {
|
shared = {
|
||||||
"files": [
|
"files": [ "./files/山海经01.txt"],
|
||||||
"./files/山海经01.txt",
|
"documents": [], # [{text, file_name, file_type, uuid, embedding}]
|
||||||
],
|
|
||||||
"documents": [],
|
|
||||||
"index": "test_kb",
|
"index": "test_kb",
|
||||||
}
|
}
|
||||||
|
|
||||||
readNode = ReadDocumentNode()
|
readNode = nodes.ReadDocumentsNode()
|
||||||
chunkNode = ChunkDocumentsNode()
|
chunkNode = nodes.ChunkDocumentsNode()
|
||||||
embeddingNode = EmbeddingDocumentsNode()
|
embeddingNode = nodes.EmbeddingDocumentsNode()
|
||||||
writeToESNode = WriteToElasticsearchNode()
|
writeToESNode = nodes.WriteDocumentsToESNode()
|
||||||
readNode >> chunkNode >> embeddingNode >> writeToESNode
|
readNode >> chunkNode >> embeddingNode >> writeToESNode
|
||||||
flow = AsyncFlow(readNode)
|
flow = AsyncFlow(readNode)
|
||||||
|
|
||||||
|
|
@ -35,3 +33,27 @@ async def test_embedding():
|
||||||
|
|
||||||
await llm.close_client()
|
await llm.close_client()
|
||||||
await es.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