ai_pipeline/src/tests/test_nodes.py

72 lines
2.0 KiB
Python

import pytest
import json
from src.pipeline.core.pocket_flow import AsyncFlow
from src.pipeline.core.utils import logger
from src.pipeline.core import llm, es, nodes, utils
@pytest.mark.asyncio
async def test_embedding():
return
await llm.init_client()
await es.init_client()
logger.debug("file to es")
shared = {
"files": [ "./files/山海经01.txt"],
"documents": [], # [{text, file_name, file_type, uuid, embedding}]
"index": "test_kb",
}
readNode = nodes.ReadDocumentsNode()
chunkNode = nodes.ChunkDocumentsNode()
embeddingNode = nodes.EmbeddingDocumentsNode()
writeToESNode = nodes.WriteDocumentsToESNode()
readNode >> chunkNode >> embeddingNode
# >> writeToESNode
flow = AsyncFlow(readNode)
await flow.run_async(shared)
logger.debug(json.dumps([{**x, "embedding":x["embedding"][:4]} for x in shared["documents"]], indent=4, ensure_ascii=False))
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",
"top_k": 1,
"results": [], # [{es_id, doc_id, title, type, created_at, score, content}]
}
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))
res = await llm.client.chat(
messages=[
{"role": "system", "content": utils.rag_system_prompt()},
{"role": "system", "content": utils.rag_user_prompt(shared["text"], shared["results"])},
# {"role": "system", "content": "你好"},
]
)
logger.debug(res)
await llm.close_client()
await es.close_client()