159 lines
4.3 KiB
Python
159 lines
4.3 KiB
Python
import asyncio
|
|
from platform import node
|
|
import pytest
|
|
import json
|
|
from src.pipeline.core.pocket_flow import AsyncFlow
|
|
from src.pipeline.utils import logger
|
|
from src.pipeline.core import llm, es, nodes
|
|
from src.pipeline.core.skills import load_skills
|
|
from src.pipeline import utils
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
async def init_llm():
|
|
await llm.init_client()
|
|
yield
|
|
await llm.close_client()
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
async def init_es():
|
|
await es.init_client()
|
|
yield
|
|
await es.close_client()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_embedding(init_llm, init_es):
|
|
logger.debug("file to es")
|
|
|
|
shared = {
|
|
"files": [ "./files/西游记.txt"],
|
|
"documents": [], # [{text, file_name, file_type, uuid, embedding}]
|
|
"index": "test_kb_1",
|
|
}
|
|
|
|
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))
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_search(init_llm, init_es):
|
|
await llm.init_client()
|
|
await es.init_client()
|
|
logger.debug("search from es")
|
|
|
|
shared = {
|
|
"text": "哪里盛产矿石",
|
|
"index": "test_kb",
|
|
"top_k": 10,
|
|
"results": [], # [{es_id, doc_id, title, type, created_at, score, content}]
|
|
}
|
|
|
|
embeddingNode = nodes.EmbeddingNode()
|
|
searchNode = nodes.SearchNode()
|
|
embeddingNode >> searchNode
|
|
flow = AsyncFlow(embeddingNode)
|
|
|
|
await flow.run_async(shared)
|
|
|
|
logger.debug(json.dumps({**shared, "embedding": shared["embedding"][:4]}, indent=4, ensure_ascii=False))
|
|
|
|
request = llm.client.stream_chat(
|
|
messages=[
|
|
{"role": "system", "content": utils.rag_system_prompt()},
|
|
{"role": "user", "content": utils.rag_user_prompt(shared["text"], shared["results"])},
|
|
]
|
|
)
|
|
|
|
async for chunk in request:
|
|
logger.debug(chunk)
|
|
|
|
await llm.close_client()
|
|
await es.close_client()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_rerank(init_llm, init_es):
|
|
await llm.init_client()
|
|
await es.init_client()
|
|
logger.debug("search from es")
|
|
|
|
shared = {
|
|
"text": "山海经中描述了哪里盛产矿石",
|
|
"index": "test_kb",
|
|
"top_k": 10,
|
|
"search_web": False,
|
|
"rerank_top_k": 5,
|
|
"top_n": 3,
|
|
"results": [], # [{es_id, doc_id, title, type, created_at, score, content}]
|
|
}
|
|
|
|
embeddingNode = nodes.EmbeddingNode()
|
|
searchNode = nodes.SearchNode()
|
|
rerankNode = nodes.RerankNode()
|
|
embeddingNode >> searchNode >> rerankNode
|
|
# 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": "user", "content": utils.rag_user_prompt(shared["text"], shared["results"][:shared["top_n"]])},
|
|
]
|
|
)
|
|
|
|
logger.debug(res)
|
|
|
|
await llm.close_client()
|
|
await es.close_client()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_agent(init_llm, init_es):
|
|
await llm.init_client()
|
|
await es.init_client()
|
|
|
|
shared = {
|
|
"query": "山海经中描述了哪里盛产矿石",
|
|
"skills": await load_skills("./src/skills"),
|
|
}
|
|
|
|
decideNode = nodes.DecideSkillNode()
|
|
refineNode = nodes.RefineSkillNode()
|
|
disambiguateNode = nodes.DisambiguateSkillNode()
|
|
runNode = nodes.ExecuteSkillNode()
|
|
|
|
decideNode >> refineNode
|
|
|
|
refineNode >> runNode
|
|
|
|
disambiguateNode >> runNode
|
|
|
|
refineNode - "disambiguate" >> disambiguateNode
|
|
|
|
flow = AsyncFlow(decideNode)
|
|
await flow.run_async(shared)
|
|
|
|
res = await llm.client.chat(
|
|
messages=[
|
|
{"role": "system", "content": utils.rag_system_prompt()},
|
|
{"role": "user", "content": utils.rag_user_prompt(shared["query"], shared["results"])},
|
|
]
|
|
)
|
|
|
|
logger.debug(res)
|
|
await llm.close_client()
|
|
await es.close_client()
|