From 9686f24b607ba6654ea31580af816bf56df39ce3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9D=8E=E5=A6=82=E5=A8=81?= Date: Wed, 10 Dec 2025 18:34:11 +0800 Subject: [PATCH] feat: embedding --- src/pipeline/config.py | 6 ++++ src/pipeline/llm/__init__.py | 29 +++++++++++++++---- src/pipeline/nodes/chunk_document_node.py | 12 +++----- src/pipeline/nodes/embedding_document_node.py | 19 ++++++++++++ src/tests/test_nodes.py | 8 ++++- 5 files changed, 59 insertions(+), 15 deletions(-) create mode 100644 src/pipeline/nodes/embedding_document_node.py diff --git a/src/pipeline/config.py b/src/pipeline/config.py index 5a0c3e1..5853591 100644 --- a/src/pipeline/config.py +++ b/src/pipeline/config.py @@ -12,6 +12,9 @@ class Config(TypedDict): llm_api_key: str llm_api_host: str llm_model: str + embedding_api_key: str + embedding_api_host: str + embedding_model: str def _read_config() -> Config: @@ -22,6 +25,9 @@ def _read_config() -> Config: "llm_api_host": os.getenv("LLM_API_HOST"), "llm_api_key": os.getenv("LLM_API_KEY"), "llm_model": os.getenv("LLM_MODEL"), + "embedding_api_host": os.getenv("EMBEDDING_API_HOST"), + "embedding_api_key": os.getenv("EMBEDDING_API_KEY"), + "embedding_model": os.getenv("EMBEDDING_MODEL"), } config = _read_config() diff --git a/src/pipeline/llm/__init__.py b/src/pipeline/llm/__init__.py index 0b69fbb..1c1965e 100644 --- a/src/pipeline/llm/__init__.py +++ b/src/pipeline/llm/__init__.py @@ -1,8 +1,25 @@ from src.pipeline.config import config +import httpx +from urllib.parse import urljoin -async def chat_completion(messages, model=None): - model = model or settings.LLM_MODEL - async with httpx.AsyncClient(timeout=60) as client: - r = await client.post(settings.VLLM_CHAT_URL, json={"model": model, "messages": messages}, headers=HEADERS) - r.raise_for_status() - return r.json() +async def get_embedding(text, timeout: int = 30): + try: + async with httpx.AsyncClient(timeout=timeout) as client: + url = config["embedding_api_host"] + body = { + "model": config["embedding_model"], + "input": text, + } + headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {config['embedding_api_key']}", + } + print(url) + print(headers) + res = await client.post(url, headers=headers, json=body) + res.raise_for_status() + data = res.json() + return data["data"][0]["embedding"] + except Exception as e: + print(f"get_embedding[ERROR]: {e}") + return [] diff --git a/src/pipeline/nodes/chunk_document_node.py b/src/pipeline/nodes/chunk_document_node.py index 5b6c528..774feb5 100644 --- a/src/pipeline/nodes/chunk_document_node.py +++ b/src/pipeline/nodes/chunk_document_node.py @@ -1,6 +1,6 @@ +import uuid from src.pipeline.core.pocket_flow import AsyncBatchNode from src.pipeline.core.utils import fixed_size_chunk -import json import re @@ -9,8 +9,7 @@ class ChunkDocumentsNode(AsyncBatchNode): return shared["documents"] async def exec_async(self, document): - """ 简单切片 - :param self + """ :param document: {text, file_name} """ # print(f"document: {document}") @@ -21,16 +20,13 @@ class ChunkDocumentsNode(AsyncBatchNode): text = re.sub(r" +", " ", text) # 去除首尾空格 text = text.strip() - return [{"text": x, "file_name": document["file_name"]} for x in fixed_size_chunk(text)] + return [{"text": x, "file_name": document["file_name"], "uuid": uuid.uuid4().hex} for x in fixed_size_chunk(text)] async def post_async(self, shared, prep_res, exec_res_list): all_chunks = [] for chunks in exec_res_list: all_chunks.extend(chunks) - print(f"all_chunks: {json.dumps(all_chunks, indent=2, ensure_ascii=False)}") - shared["texts"] = all_chunks + shared["documents"] = all_chunks return "default" - - diff --git a/src/pipeline/nodes/embedding_document_node.py b/src/pipeline/nodes/embedding_document_node.py new file mode 100644 index 0000000..472c2cf --- /dev/null +++ b/src/pipeline/nodes/embedding_document_node.py @@ -0,0 +1,19 @@ +from src.pipeline.core.pocket_flow import AsyncBatchNode +from src.pipeline.llm import get_embedding + + +class EmbeddingDocumentsNode(AsyncBatchNode): + async def prep_async(self, shared): + return shared["documents"] + + async def exec_async(self, document): + """ + :param document: {text, file_name} + """ + return {**document, "embedding": await get_embedding(document["text"])} + + async def post_async(self, shared, prep_res, exec_res_list): + + shared["documents"] = exec_res_list + + return "default" diff --git a/src/tests/test_nodes.py b/src/tests/test_nodes.py index dc5ebd4..59707d9 100644 --- a/src/tests/test_nodes.py +++ b/src/tests/test_nodes.py @@ -1,9 +1,12 @@ import pytest +import json from src.pipeline.nodes.read_document_node import ReadDocumentNode from src.pipeline.nodes.chunk_document_node import ChunkDocumentsNode +from src.pipeline.nodes.embedding_document_node import EmbeddingDocumentsNode # from src.pipeline.nodes import ReadDocumentNode, ChunkDocumentsNode from src.pipeline.core.pocket_flow import AsyncFlow + @pytest.mark.asyncio async def test_embedding(): @@ -20,7 +23,10 @@ async def test_embedding(): readNode = ReadDocumentNode() chunkNode = ChunkDocumentsNode() - readNode >> chunkNode + embeddingNode = EmbeddingDocumentsNode() + readNode >> chunkNode >> embeddingNode flow = AsyncFlow(readNode) await flow.run_async(shared) + + print(json.dumps(shared["documents"], indent=4, ensure_ascii=False))