feat: test

This commit is contained in:
李如威 2025-12-09 17:25:04 +08:00
parent 2050622dea
commit bc29326c2c
17 changed files with 189 additions and 1 deletions

3
.env.example Normal file
View File

@ -0,0 +1,3 @@
VERSION=1.0.0
HOST=0.0.0.0
PORT=8011

1
.gitignore vendored
View File

@ -188,3 +188,4 @@ local_config.py
# docker_image_build_tmp/
# logs/
# results/
files/

View File

@ -0,0 +1,5 @@
## 测试
```python
pytest -s src/test/test.py
```

2
pytest.ini Normal file
View File

@ -0,0 +1,2 @@
[pytest]
pythonpath = .

View File

@ -7,3 +7,4 @@ pytest-asyncio
pymupdf
python-docx
scikit-learn
aiofiles

3
scripts/run_debug.sh Normal file
View File

@ -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

3
scripts/run_test.sh Normal file
View File

@ -0,0 +1,3 @@
#!/usr/bin/env bash
export $(cat .env | xargs)
pytest -s -W ignore::DeprecationWarning src/tests/test_nodes.py

0
src/__init__.py Normal file
View File

View File

@ -0,0 +1,5 @@
from fastapi import FastAPI
def include_router(app:FastAPI):
pass

View File

@ -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()

View File

@ -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 {}

View File

@ -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}")

View File

@ -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()

View File

@ -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

View File

@ -0,0 +1,5 @@
import pytest
@pytest.mark.asyncio
async def test_embedding():
print("\n1")

22
src/tests/test_nodes.py Normal file
View File

@ -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)