import asyncio from platform import node 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 from src.pipeline.core.skills import load_skills @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()