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/
|
# docker_image_build_tmp/
|
||||||
# logs/
|
# logs/
|
||||||
# results/
|
# results/
|
||||||
|
files/
|
||||||
|
|
|
||||||
|
|
@ -0,0 +1,2 @@
|
||||||
|
[pytest]
|
||||||
|
pythonpath = .
|
||||||
|
|
@ -7,3 +7,4 @@ pytest-asyncio
|
||||||
pymupdf
|
pymupdf
|
||||||
python-docx
|
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