diff --git a/pytest.ini b/pytest.ini index 03f586d..ea76444 100644 --- a/pytest.ini +++ b/pytest.ini @@ -1,2 +1,5 @@ [pytest] -pythonpath = . \ No newline at end of file +pythonpath = . +asyncio_mode = auto +asyncio_default_fixture_loop_scope = function +filterwarnings = ignore:.*Swig.*:DeprecationWarning \ No newline at end of file diff --git a/scripts/run_test.sh b/scripts/run_test.sh index a9b4483..6d78988 100644 --- a/scripts/run_test.sh +++ b/scripts/run_test.sh @@ -1,4 +1,4 @@ #!/usr/bin/env bash export $(cat .env | xargs) -pytest -s -W ignore::DeprecationWarning tests/pipeline/api/test_rag.py::test_chat_message_list +pytest -s -W ignore::DeprecationWarning tests/pipeline/api/test_rag.py::test_chat_message_list_async diff --git a/src/pipeline/config.py b/src/pipeline/config.py index 9ff94fe..646ef3a 100644 --- a/src/pipeline/config.py +++ b/src/pipeline/config.py @@ -1,8 +1,11 @@ from pydantic_settings import BaseSettings +from pydantic import ConfigDict from pydantic import MySQLDsn class Settings(BaseSettings): + model_config = ConfigDict(env_file=".env") + logger_level: str env: str version: str @@ -23,8 +26,7 @@ class Settings(BaseSettings): es_user: str es_password: str mysql_dsn: MySQLDsn - class Config: - env_file = ".env" + # def _read_config() -> Settings: diff --git a/src/pipeline/main.py b/src/pipeline/main.py index d679b99..9c3dc7e 100644 --- a/src/pipeline/main.py +++ b/src/pipeline/main.py @@ -19,7 +19,12 @@ async def lifespan(app: FastAPI): await close_tortoise() -app = FastAPI(title="AI Pipeline", description="轻量级 AI Pipeline", version=config.version, lifespan=lifespan) +app = FastAPI( + title="AI Pipeline", + description="轻量级 AI Pipeline", + version=config.version, + lifespan=lifespan, +) include_router(app) diff --git a/src/pipeline/models/__init__.py b/src/pipeline/models/__init__.py index be236bd..c5c5309 100644 --- a/src/pipeline/models/__init__.py +++ b/src/pipeline/models/__init__.py @@ -5,9 +5,10 @@ from uuid import uuid4 class BaseModel(Model): - id = fields.IntField(pk=True) + id = fields.IntField(primary_key=True) create_at = fields.DatetimeField(auto_now_add=True) update_at = fields.DatetimeField(auto_now=True) + deleted_at = fields.DatetimeField(null=True) is_deleted = fields.IntField(default=0) class Meta: @@ -74,18 +75,18 @@ class User(BaseModel): class Chat(BaseModel): - user_id = fields.IntField(null=True) + user_id = fields.IntField(default=0, db_index=True) title = fields.CharField(max_length=50) - uuid = fields.UUIDField(index=True) + uuid = fields.UUIDField(db_index=True) class Meta: table = "tb_chats" class Message(BaseModel): - user_id = fields.IntField(null=True) - chat_id = fields.IntField(null=True, index=True) - uuid = fields.UUIDField(index=True) + user_id = fields.IntField(default=0, db_index=True) + chat_id = fields.IntField(null=True, db_index=True) + uuid = fields.UUIDField(db_index=True) content = fields.JSONField(null=True) class Meta: @@ -98,3 +99,22 @@ class Message(BaseModel): if "answer" in self.content and self.content["answer"]: llm_messages.append({"role": "assistant", "content": self.content["answer"]}) return llm_messages + + +class File(BaseModel): + user_id = fields.IntField(default=0, db_index=True) + uuid = fields.UUIDField(db_index=True) + # 文件基础属性 + origin_name = fields.CharField(max_length=244, description="原始名字") + stored_name = fields.CharField(max_length=244, description="存储名字") + stored_path = fields.CharField(max_length=500, description="存储相对路径(不包含文件名字)") + ext = fields.CharField(max_length=20, description="文件扩展名") + mine_type = fields.CharField(max_length=100, description="MIME类型") + size = fields.IntField(default=0, description="文件大小(字节)") + hash_md5 = fields.CharField(max_length=32, description="MD5", db_index=True) + # 业务属性 + status = fields.IntField(default=1, description="状态:1正常 2禁用 3删除") + biz_type = fields.CharField(max_length=50, description="业务类型", null=True) + + class Meta: + table = "tb_file" diff --git a/src/pipeline/services/file.py b/src/pipeline/services/file.py new file mode 100644 index 0000000..344c200 --- /dev/null +++ b/src/pipeline/services/file.py @@ -0,0 +1,150 @@ +import os +import hashlib +from uuid import uuid4 +from datetime import datetime, timezone +from typing import Optional + +import aiofiles +import aiofiles.os as aioos + +from src.pipeline.models import File as TFile + + +class FileService: + def __init__(self, base_path: str = "files/uploads"): + """ + base_path: 存储根目录,默认当前工作目录下的 storage/uploads + """ + self.base_path = base_path + os.makedirs(self.base_path, exist_ok=True) + + def _chunk_dir(self, upload_id: str) -> str: + return os.path.join(self.base_path, "chunks", upload_id) + + def _final_dir(self, when: datetime) -> str: + return os.path.join(self.base_path, when.strftime("%Y"), when.strftime("%m"), when.strftime("%d")) + + async def save_chunk(self, upload_id: str, chunk_index: int, chunk_bytes: bytes) -> str: + """ + 保存一个分片到临时目录(异步)。 + 返回保存的分片路径(绝对)。 + """ + chunk_dir = self._chunk_dir(upload_id) + # 使用同步创建目录是可以接受的,但也支持异步创建 + await aioos.makedirs(chunk_dir, exist_ok=True) + part_path = os.path.join(chunk_dir, f"part_{chunk_index}") + async with aiofiles.open(part_path, "wb") as f: + await f.write(chunk_bytes) + return part_path + + async def _async_rmtree(self, path: str) -> None: + """异步递归删除目录及其内容。""" + try: + entries = await aioos.listdir(path) + except FileNotFoundError: + return + for entry in entries: + full = os.path.join(path, entry) + try: + await aioos.remove(full) + except IsADirectoryError: + await self._async_rmtree(full) + try: + await aioos.rmdir(full) + except Exception: + pass + except Exception: + # 可能是目录或其他问题,尝试递归删除 + await self._async_rmtree(full) + try: + await aioos.rmdir(full) + except Exception: + pass + try: + await aioos.rmdir(path) + except Exception: + pass + + async def finalize_upload( + self, + user_id: int, + upload_id: str, + total_chunks: int, + origin_name: str, + mime_type: Optional[str] = None, + biz_type: Optional[str] = None, + ) -> TFile: + """ + 将 upload_id 对应的所有分片合并为最终文件(异步),写入存储并在数据库中创建 File 记录。 + 返回创建的 File ORM 对象。 + """ + chunk_dir = self._chunk_dir(upload_id) + # 检查分片目录是否存在 + try: + await aioos.listdir(chunk_dir) + except FileNotFoundError: + raise FileNotFoundError("chunk directory not found") + + # 确定扩展名和目标路径/文件名 + _, ext = os.path.splitext(origin_name) + ext = ext.lstrip(".") if ext else "" + now = datetime.now(timezone.utc) + final_dir = self._final_dir(now) + await aioos.makedirs(final_dir, exist_ok=True) + stored_name = f"{uuid4().hex}{('.' + ext) if ext else ''}" + final_path = os.path.join(final_dir, stored_name) + + # 合并并计算 md5 与大小(异步读写) + md5 = hashlib.md5() + total_size = 0 + async with aiofiles.open(final_path, "wb") as fw: + for i in range(total_chunks): + part_path = os.path.join(chunk_dir, f"part_{i}") + try: + async with aiofiles.open(part_path, "rb") as fr: + while True: + buf = await fr.read(1024 * 1024) + if not buf: + break + await fw.write(buf) + md5.update(buf) + total_size += len(buf) + except FileNotFoundError: + # 清理已写入的文件 + try: + await aioos.remove(final_path) + except Exception: + pass + raise FileNotFoundError(f"missing chunk: {i}") + + md5_hex = md5.hexdigest() + + # 清理分片目录(异步) + await self._async_rmtree(chunk_dir) + + # stored_path 应为相对路径(不包含文件名),相对于 base_path + stored_path_rel = os.path.relpath(final_dir, self.base_path) + + # 在数据库中创建 File 记录并返回 + file_uuid = uuid4() + file_record = await TFile.create( + user_id=user_id, + uuid=file_uuid, + origin_name=origin_name, + stored_name=stored_name, + stored_path=stored_path_rel, + ext=ext, + mine_type=mime_type or "", + size=total_size, + hash_md5=md5_hex, + biz_type=biz_type or None, + ) + return file_record + + async def abort_upload(self, upload_id: str) -> None: + """删除临时分片目录(中止上传用,异步)。""" + chunk_dir = self._chunk_dir(upload_id) + await self._async_rmtree(chunk_dir) + + +file_service = FileService() diff --git a/tests/pipeline/api/test_rag.py b/tests/pipeline/api/test_rag.py index 55de439..4cf3bf9 100644 --- a/tests/pipeline/api/test_rag.py +++ b/tests/pipeline/api/test_rag.py @@ -2,7 +2,6 @@ from asgi_lifespan import LifespanManager from httpx import AsyncClient, ASGITransport from src.pipeline.main import app from src.pipeline.utils import logger -from fastapi.testclient import TestClient import pytest import pytest_asyncio import json @@ -12,7 +11,7 @@ import json # ---------------------- -@pytest_asyncio.fixture(scope="session") +@pytest_asyncio.fixture async def app_client_async(): async with LifespanManager(app): async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: @@ -21,7 +20,7 @@ async def app_client_async(): @pytest.mark.asyncio async def test_stream_chat(app_client_async): - res = await app_client.post("/api/v1/rag/stream-chat", json={"query": "hello"}) + res = await app_client_async.post("/api/v1/rag/stream-chat", json={"query": "hello"}) assert res.status_code == 200 sse_messages = [] async for line in res.aiter_lines(): @@ -43,20 +42,3 @@ async def test_chat_message_list_async(app_client_async): res = await app_client_async.get("/api/v1/rag/message-list/a1c9108c-9201-4e60-a436-505edec3f47e") assert res.status_code == 200 logger.debug(res.json()) - - -# ---------------------- -# normal -# ---------------------- - - -@pytest.fixture(scope="session") -def app_client(): - client = TestClient(app=app) - return client - - -def test_chat_message_list(app_client): - res = app_client.get("/api/v1/rag/message-list/a1c9108c-9201-4e60-a436-505edec3f47e") - assert res.status_code == 200 - logger.debug(res.json()) diff --git a/tests/pipeline/service/test_file.py b/tests/pipeline/service/test_file.py new file mode 100644 index 0000000..cee1156 --- /dev/null +++ b/tests/pipeline/service/test_file.py @@ -0,0 +1,58 @@ +import hashlib +import os +import pytest +import pytest_asyncio +from src.pipeline.db import Path, init_tortoise, close_tortoise +from src.pipeline.services.file import file_service +from src.pipeline.config import config +from src.pipeline.models import File as TFile +from uuid import uuid4 + + +@pytest_asyncio.fixture +async def init_db(): + db_url = str(config.mysql_dsn) + await init_tortoise(db_url=db_url) + yield + await close_tortoise() + + +@pytest.mark.asyncio +async def test_file_upload(init_db): + upload_id = str(uuid4()) + chunks = [b"part1-", b"part2-", b"part3"] + + # save chunks + for idx, chunk in enumerate(chunks): + part_path = await file_service.save_chunk(upload_id, idx, chunk) + assert os.path.exists(part_path) + + # finalize + origin_name = "hello.txt" + file_record = await file_service.finalize_upload( + user_id=42, + upload_id=upload_id, + total_chunks=len(chunks), + origin_name=origin_name, + mime_type="text/plain", + biz_type="unittest", + ) + + # DB record + assert isinstance(file_record, TFile) + assert file_record.user_id == 42 + assert file_record.origin_name == origin_name + + # file exists on disk + stored_path = Path(file_service.base_path) / Path(file_record.stored_path) + final_path = stored_path / file_record.stored_name + assert final_path.exists() + + # md5 and size correct + combined = b"".join(chunks) + assert file_record.hash_md5 == hashlib.md5(combined).hexdigest() + assert file_record.size == len(combined) + + # chunks directory removed + chunk_dir = Path(file_service.base_path) / "chunks" / upload_id + assert not chunk_dir.exists()