feat: pytest-asyncio
This commit is contained in:
parent
a5f33c88d1
commit
2c8e93dcde
|
|
@ -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
|
||||
|
|
@ -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"""
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
Loading…
Reference in New Issue