diff --git a/requirements.txt b/requirements.txt index a5a3c35..ac3b269 100644 --- a/requirements.txt +++ b/requirements.txt @@ -10,4 +10,6 @@ aiofiles pillow loguru httpx -baidusearch \ No newline at end of file +baidusearch +celery[redis] +tortoise-orm[asyncmy] \ No newline at end of file diff --git a/src/pipeline/api/rag.py b/src/pipeline/api/rag.py index 1f316a3..1c83817 100644 --- a/src/pipeline/api/rag.py +++ b/src/pipeline/api/rag.py @@ -3,12 +3,16 @@ from fastapi.responses import StreamingResponse from src.pipeline.services.rag import RagService from src.pipeline.schemas.schemas import ChatRequest from src.pipeline.utils import logger +from src.pipeline.models import User as TUser import json router = APIRouter() @router.post("/stream-chat") async def stream_chat(body: ChatRequest): + + logger.debug(await TUser.all()) + service = RagService() async def event_generator(): diff --git a/src/pipeline/config.py b/src/pipeline/config.py index cf6c7e2..3985649 100644 --- a/src/pipeline/config.py +++ b/src/pipeline/config.py @@ -1,10 +1,5 @@ -from mimetypes import init -from typing import TypedDict -from dotenv import load_dotenv -import os from pydantic_settings import BaseSettings - -load_dotenv() +from pydantic import MySQLDsn class Settings(BaseSettings): @@ -26,7 +21,7 @@ class Settings(BaseSettings): es_port: int es_user: str es_password: str - + mysql_dsn: MySQLDsn class Config: env_file = ".env" diff --git a/src/pipeline/core/es.py b/src/pipeline/core/es.py index 5cf57dc..bcd8a2b 100644 --- a/src/pipeline/core/es.py +++ b/src/pipeline/core/es.py @@ -186,5 +186,6 @@ async def init_client(): global client client = AsyncES() + async def close_client(): await client.close() diff --git a/src/pipeline/db/__init__.py b/src/pipeline/db/__init__.py new file mode 100644 index 0000000..9d19ffa --- /dev/null +++ b/src/pipeline/db/__init__.py @@ -0,0 +1,26 @@ +from pathlib import Path +from tortoise import Tortoise, generate_config +from tortoise.contrib.fastapi import RegisterTortoise +from fastapi import FastAPI +from src.pipeline.utils import logger + + +DEFAULT_MODELS = {"models": ["src.pipeline.models"]} +DEFAULT_INIT_SQL = "src/pipeline/db/init.sql" + + +async def init_tortoise(app: FastAPI, db_url: str, modules: dict = DEFAULT_MODELS, schema_sql: str = DEFAULT_INIT_SQL): + config = generate_config(db_url, app_modules=modules) + await RegisterTortoise(app, config=config) + + if schema_sql: + conn = Tortoise.get_connection("default") + existing = await conn.execute_query_dict("SHOW TABLES") + if not existing: + logger.debug(f"数据库初始化文件: {schema_sql}") + sql_file = Path(schema_sql) + if sql_file.exists(): + logger.debug("初始化数据库") + sql = sql_file.read_text() + logger.debug(f'sql:\n{sql}') + await conn.execute_script(sql) diff --git a/src/pipeline/db/init.sql b/src/pipeline/db/init.sql new file mode 100644 index 0000000..92d5dcd --- /dev/null +++ b/src/pipeline/db/init.sql @@ -0,0 +1,9 @@ +-- tb_users table +CREATE TABLE IF NOT EXISTS tb_users ( + id INT AUTO_INCREMENT PRIMARY KEY, + username VARCHAR(50) NOT NULL UNIQUE, + password VARCHAR(128) NOT NULL, + user_status TINYINT(1) NOT NULL DEFAULT 1, -- 0=禁用, 1=启用 + create_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + update_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP +); \ No newline at end of file diff --git a/src/pipeline/main.py b/src/pipeline/main.py index 3304dd4..b05ac31 100644 --- a/src/pipeline/main.py +++ b/src/pipeline/main.py @@ -1,14 +1,16 @@ from fastapi import FastAPI from src.pipeline.api import include_router from src.pipeline.config import config -from contextlib import asynccontextmanager from src.pipeline.core import llm, es +from src.pipeline.db import init_tortoise +from contextlib import asynccontextmanager @asynccontextmanager async def lifespan(app: FastAPI): await llm.init_client() await es.init_client() + await init_tortoise(app, db_url=str(config.mysql_dsn)) yield diff --git a/src/pipeline/models/__init__.py b/src/pipeline/models/__init__.py new file mode 100644 index 0000000..3827fba --- /dev/null +++ b/src/pipeline/models/__init__.py @@ -0,0 +1,21 @@ +from email.policy import default +from tortoise import fields +from tortoise.models import Model + + +class BaseModel(Model): + id = fields.IntField(pk=True) + create_at = fields.DatetimeField(auto_now_add=True) + update_at = fields.DatetimeField(auto_now=True) + + class Meta: + abstract = True + + +class User(BaseModel): + username = fields.CharField(max_length=50, unique=True) + password = fields.CharField(max_length=128) + user_status = fields.IntField(default=1) # 0=禁用, 1=启用 + + class Meta: + table = "tb_users" diff --git a/src/pipeline/tasks/rag.py b/src/pipeline/tasks/rag.py new file mode 100644 index 0000000..e69de29