From bc29326c2c5b149b4c694af162da5d61d625de42 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9D=8E=E5=A6=82=E5=A8=81?= Date: Tue, 9 Dec 2025 17:25:04 +0800 Subject: [PATCH] feat: test --- .env.example | 3 + .gitignore | 1 + README.md | 5 ++ .../rag/nodes.py => docker/Dockerfile | 0 pytest.ini | 2 + requirements.txt | 3 +- scripts/run_debug.sh | 3 + scripts/run_test.sh | 3 + src/__init__.py | 0 src/pipeline/api/__init__.py | 5 ++ src/pipeline/config.py | 27 +++++++++ src/pipeline/core/nodes.py | 30 ++++++++++ src/pipeline/core/utils.py | 56 +++++++++++++++++++ src/pipeline/llm/__init__.py | 8 +++ src/pipeline/main.py | 17 ++++++ src/tests/test.py | 5 ++ src/tests/test_nodes.py | 22 ++++++++ 17 files changed, 189 insertions(+), 1 deletion(-) create mode 100644 .env.example rename src/pipeline/rag/nodes.py => docker/Dockerfile (100%) create mode 100644 pytest.ini create mode 100644 scripts/run_debug.sh create mode 100644 scripts/run_test.sh create mode 100644 src/__init__.py create mode 100644 src/pipeline/core/nodes.py create mode 100644 src/pipeline/core/utils.py create mode 100644 src/tests/test_nodes.py diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..da4a7e6 --- /dev/null +++ b/.env.example @@ -0,0 +1,3 @@ +VERSION=1.0.0 +HOST=0.0.0.0 +PORT=8011 \ No newline at end of file diff --git a/.gitignore b/.gitignore index 6493694..5fb2783 100644 --- a/.gitignore +++ b/.gitignore @@ -188,3 +188,4 @@ local_config.py # docker_image_build_tmp/ # logs/ # results/ +files/ diff --git a/README.md b/README.md index e69de29..9f75ed4 100644 --- a/README.md +++ b/README.md @@ -0,0 +1,5 @@ +## 测试 + +```python +pytest -s src/test/test.py +``` \ No newline at end of file diff --git a/src/pipeline/rag/nodes.py b/docker/Dockerfile similarity index 100% rename from src/pipeline/rag/nodes.py rename to docker/Dockerfile diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..03f586d --- /dev/null +++ b/pytest.ini @@ -0,0 +1,2 @@ +[pytest] +pythonpath = . \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index bd9b848..ea487e2 100644 --- a/requirements.txt +++ b/requirements.txt @@ -6,4 +6,5 @@ python-dotenv pytest-asyncio pymupdf python-docx -scikit-learn \ No newline at end of file +scikit-learn +aiofiles \ No newline at end of file diff --git a/scripts/run_debug.sh b/scripts/run_debug.sh new file mode 100644 index 0000000..6ebdbdf --- /dev/null +++ b/scripts/run_debug.sh @@ -0,0 +1,3 @@ +#!/usr/bin/env bash +export $(cat .env | xargs) +uvicorn src.pipeline.main:app --host ${HOST:-0.0.0.0} --port ${PORT:-8000} --reload diff --git a/scripts/run_test.sh b/scripts/run_test.sh new file mode 100644 index 0000000..904a5f4 --- /dev/null +++ b/scripts/run_test.sh @@ -0,0 +1,3 @@ +#!/usr/bin/env bash +export $(cat .env | xargs) +pytest -s -W ignore::DeprecationWarning src/tests/test_nodes.py diff --git a/src/__init__.py b/src/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/pipeline/api/__init__.py b/src/pipeline/api/__init__.py index e69de29..668b471 100644 --- a/src/pipeline/api/__init__.py +++ b/src/pipeline/api/__init__.py @@ -0,0 +1,5 @@ + +from fastapi import FastAPI + +def include_router(app:FastAPI): + pass \ No newline at end of file diff --git a/src/pipeline/config.py b/src/pipeline/config.py index e69de29..5a0c3e1 100644 --- a/src/pipeline/config.py +++ b/src/pipeline/config.py @@ -0,0 +1,27 @@ +from mimetypes import init +from typing import TypedDict +from dotenv import load_dotenv +import os + +load_dotenv() + +class Config(TypedDict): + version: str + port: int + host: str + llm_api_key: str + llm_api_host: str + llm_model: str + + +def _read_config() -> Config: + return { + "host": os.getenv("HOST"), + "port": int(os.getenv("PORT")), + "version": os.getenv("VERSION"), + "llm_api_host": os.getenv("LLM_API_HOST"), + "llm_api_key": os.getenv("LLM_API_KEY"), + "llm_model": os.getenv("LLM_MODEL"), + } + +config = _read_config() diff --git a/src/pipeline/core/nodes.py b/src/pipeline/core/nodes.py new file mode 100644 index 0000000..bd32c06 --- /dev/null +++ b/src/pipeline/core/nodes.py @@ -0,0 +1,30 @@ +from src.pipeline.core.pocket_flow import AsyncBatchNode +from src.pipeline.core.utils import load_document + +class ReadDocumentNode(AsyncBatchNode): + async def prep_async(self, shared): + return shared["files"] + + async def exec_async(self, file_path): + try: + document_text = await load_document(file_path) + return { + "file_path": file_path, + "text": document_text, + "text_length": len(document_text), + "status": "done", + "message": "", + } + except Exception as e: + return { + "file_path": file_path, + "text": "", + "text_length": 0, + "status": "error", + "message": str(e), + } + + async def post_async(self, shared, prep_res, exec_res): + + print([{**x, "text": x["text"][:5] + "..."} for x in exec_res]) + return {} diff --git a/src/pipeline/core/utils.py b/src/pipeline/core/utils.py new file mode 100644 index 0000000..0930e96 --- /dev/null +++ b/src/pipeline/core/utils.py @@ -0,0 +1,56 @@ +import asyncio +import docx +import fitz # PyMuPDF +from pathlib import Path +import aiofiles + +# ----------------------------- +# 文件类型加载器 +# ----------------------------- + + +async def load_txt(path: str) -> str: + async with aiofiles.open(path, "r", encoding="utf-8") as f: + return await f.read() + + +async def load_md(path: str) -> str: + async with aiofiles.open(path, "r", encoding="utf-8") as f: + return await f.read() + + +async def load_docx(path: str) -> str: + # docx 读取是阻塞 I/O → 放入线程池 + def _read(): + doc = docx.Document(path) + return "\n".join(p.text for p in doc.paragraphs) + + return await asyncio.to_thread(_read) + + +async def load_pdf(path: str) -> str: + # PyMuPDF 也是阻塞 → 放入线程池 + def _read(): + pdf = fitz.open(path) + return "\n".join(page.get_text() for page in pdf) + + return await asyncio.to_thread(_read) + + +# ----------------------------- +# 统一调度器 +# ----------------------------- + + +async def load_document(path: str) -> str: + suffix = Path(path).suffix.lower() + if suffix == ".txt": + return await load_txt(path) + if suffix in (".md", ".markdown"): + return await load_md(path) + if suffix == ".docx": + return await load_docx(path) + if suffix == ".pdf": + return await load_pdf(path) + + raise ValueError(f"Unsupported file type: {suffix}") diff --git a/src/pipeline/llm/__init__.py b/src/pipeline/llm/__init__.py index e69de29..0b69fbb 100644 --- a/src/pipeline/llm/__init__.py +++ b/src/pipeline/llm/__init__.py @@ -0,0 +1,8 @@ +from src.pipeline.config import config + +async def chat_completion(messages, model=None): + model = model or settings.LLM_MODEL + async with httpx.AsyncClient(timeout=60) as client: + r = await client.post(settings.VLLM_CHAT_URL, json={"model": model, "messages": messages}, headers=HEADERS) + r.raise_for_status() + return r.json() diff --git a/src/pipeline/main.py b/src/pipeline/main.py index e69de29..cada61f 100644 --- a/src/pipeline/main.py +++ b/src/pipeline/main.py @@ -0,0 +1,17 @@ +from fastapi import FastAPI +from src.pipeline.api import include_router +from src.pipeline.config import config +app = FastAPI( + title="AI Pipeline", + description="轻量级 AI Pipeline", + version=config["version"], +) + +include_router(app) + +@app.get("/") +async def healthz(): + return {"status": "running", "version": config["version"]} + +if __name__ == '__main__': + pass diff --git a/src/tests/test.py b/src/tests/test.py index e69de29..b8b7c46 100644 --- a/src/tests/test.py +++ b/src/tests/test.py @@ -0,0 +1,5 @@ +import pytest + +@pytest.mark.asyncio +async def test_embedding(): + print("\n1") \ No newline at end of file diff --git a/src/tests/test_nodes.py b/src/tests/test_nodes.py new file mode 100644 index 0000000..4e97cbb --- /dev/null +++ b/src/tests/test_nodes.py @@ -0,0 +1,22 @@ +import pytest +from src.pipeline.core.nodes import ReadDocumentNode +from src.pipeline.core.pocket_flow import AsyncFlow + + +@pytest.mark.asyncio +async def test_embedding(): + + print("\n\ntest_embedding:\n") + + shared = { + "files": [ + "./files/大白智问-API接入文档-V1.2.2.pdf", + "./files/我来帮您创建一个美观简洁的微信小程序订单详情页面。首先让我了解一下当前的项目结构.md", + "./files/deepsearch状态.txt", + ] + } + + readNode = ReadDocumentNode() + flow = AsyncFlow(readNode) + + await flow.run_async(shared)