feat: pytest-asyncio

This commit is contained in:
李如威 2025-11-24 15:57:33 +08:00
parent a5f33c88d1
commit 2c8e93dcde
4 changed files with 37 additions and 5 deletions

View File

@ -1,3 +1,5 @@
httpx==0.28.1
chromadb==1.3.4
python-dotenv==1.2.1
python-dotenv==1.2.1
pytest==9.0.1
pytest-asyncio==1.3.0

View File

@ -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"""

View File

@ -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

14
src/tests/test_utils.py Normal file
View File

@ -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)