feat: test
This commit is contained in:
parent
2050622dea
commit
bc29326c2c
|
|
@ -0,0 +1,3 @@
|
|||
VERSION=1.0.0
|
||||
HOST=0.0.0.0
|
||||
PORT=8011
|
||||
|
|
@ -188,3 +188,4 @@ local_config.py
|
|||
# docker_image_build_tmp/
|
||||
# logs/
|
||||
# results/
|
||||
files/
|
||||
|
|
|
|||
|
|
@ -0,0 +1,2 @@
|
|||
[pytest]
|
||||
pythonpath = .
|
||||
|
|
@ -6,4 +6,5 @@ python-dotenv
|
|||
pytest-asyncio
|
||||
pymupdf
|
||||
python-docx
|
||||
scikit-learn
|
||||
scikit-learn
|
||||
aiofiles
|
||||
|
|
@ -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
|
||||
|
|
@ -0,0 +1,3 @@
|
|||
#!/usr/bin/env bash
|
||||
export $(cat .env | xargs)
|
||||
pytest -s -W ignore::DeprecationWarning src/tests/test_nodes.py
|
||||
|
|
@ -0,0 +1,5 @@
|
|||
|
||||
from fastapi import FastAPI
|
||||
|
||||
def include_router(app:FastAPI):
|
||||
pass
|
||||
|
|
@ -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()
|
||||
|
|
@ -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 {}
|
||||
|
|
@ -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}")
|
||||
|
|
@ -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()
|
||||
|
|
@ -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
|
||||
|
|
@ -0,0 +1,5 @@
|
|||
import pytest
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_embedding():
|
||||
print("\n1")
|
||||
|
|
@ -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)
|
||||
Loading…
Reference in New Issue