From cf9d972e00a09d543f54a240fd3efb13877aa22f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9D=8E=E5=A6=82=E5=A8=81?= Date: Fri, 19 Dec 2025 14:00:44 +0800 Subject: [PATCH] feat: search --- search.es | 29 +++++++++++++++++++++++++ src/pipeline/core/nodes.py | 36 ++++++++++++++++++++++++++++--- src/pipeline/core/utils.py | 2 -- src/tests/test_nodes.py | 44 ++++++++++++++++++++++++++++---------- 4 files changed, 95 insertions(+), 16 deletions(-) diff --git a/search.es b/search.es index e69de29..0f14b6d 100644 --- a/search.es +++ b/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 \ No newline at end of file diff --git a/src/pipeline/core/nodes.py b/src/pipeline/core/nodes.py index 5d5d871..3cf7d6b 100644 --- a/src/pipeline/core/nodes.py +++ b/src/pipeline/core/nodes.py @@ -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" diff --git a/src/pipeline/core/utils.py b/src/pipeline/core/utils.py index 0cfce81..cb8c4ca 100644 --- a/src/pipeline/core/utils.py +++ b/src/pipeline/core/utils.py @@ -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"): diff --git a/src/tests/test_nodes.py b/src/tests/test_nodes.py index 43a4df8..cdae9f6 100644 --- a/src/tests/test_nodes.py +++ b/src/tests/test_nodes.py @@ -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()