From 2c8e93dcdef62472f867e0b00529a9a9d57c9001 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9D=8E=E5=A6=82=E5=A8=81?= Date: Mon, 24 Nov 2025 15:57:33 +0800 Subject: [PATCH] feat: pytest-asyncio --- requirements.txt | 4 +++- src/sample_ai/nodes.py | 13 ++++++++++++- src/sample_ai/utils.py | 11 ++++++++--- src/tests/test_utils.py | 14 ++++++++++++++ 4 files changed, 37 insertions(+), 5 deletions(-) create mode 100644 src/tests/test_utils.py diff --git a/requirements.txt b/requirements.txt index fcecb49..1ebc069 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,5 @@ httpx==0.28.1 chromadb==1.3.4 -python-dotenv==1.2.1 \ No newline at end of file +python-dotenv==1.2.1 +pytest==9.0.1 +pytest-asyncio==1.3.0 \ No newline at end of file diff --git a/src/sample_ai/nodes.py b/src/sample_ai/nodes.py index b38c7a2..eaa6c4f 100644 --- a/src/sample_ai/nodes.py +++ b/src/sample_ai/nodes.py @@ -19,7 +19,7 @@ class ChunkDocumentsNode(AsyncBatchNode): text = re.sub(r' +', ' ', text) # 去除首尾空格 text = text.strip() - return fixed_size_chunk(text, chunk_size=100) + return fixed_size_chunk(text, chunk_size=500, overlap=50) async def post_async(self, shared, prep_res, exec_res_list): """Store chunked texts in the shared store""" @@ -36,6 +36,17 @@ class ChunkDocumentsNode(AsyncBatchNode): return "default" +class EmbedDocumentsNode(AsyncBatchNode): + async def prep_async(self, shared): + return shared["texts"] + + async def exec_async(self, text): + return fixed_size_chunk(text, chunk_size=500, overlap=50) + + async def post_async(self, shared, prep_res, exec_res_list): + return "default" + + class CreateIndexNode(Node): def prep(self, shared): """Get embeddings from shared store""" diff --git a/src/sample_ai/utils.py b/src/sample_ai/utils.py index 8d0ac7f..744e52d 100644 --- a/src/sample_ai/utils.py +++ b/src/sample_ai/utils.py @@ -3,6 +3,8 @@ import httpx import asyncio import os import json +from dotenv import load_dotenv +load_dotenv() MAIN_HTTP_URL = os.environ.get("MAIN_HTTP_URL") MAIN_HTTP_KEY = os.environ.get("MAIN_HTTP_KEY") @@ -43,10 +45,13 @@ def _get_content(resp: dict) -> str: return "" -def fixed_size_chunk(text, chunk_size=2000): +def fixed_size_chunk(text, chunk_size=2000, overlap=50): chunks = [] - for i in range(0, len(text), chunk_size): - chunks.append(text[i : i + chunk_size]) + start = 0 + while start < len(text): + end = start + chunk_size + chunks.append(text[max(0, start - overlap) : min(len(text), end + overlap)]) + start += chunk_size return chunks diff --git a/src/tests/test_utils.py b/src/tests/test_utils.py new file mode 100644 index 0000000..a9d16d8 --- /dev/null +++ b/src/tests/test_utils.py @@ -0,0 +1,14 @@ +import pytest + +from sample_ai.utils import get_embedding, call_llm + +@pytest.mark.asyncio +async def test_embedidng(): + res = await get_embedding("你好") + print(res) + + +@pytest.mark.asyncio +async def test_chat(): + res = await call_llm("你好") + print(res)