feat: 初始化数据库
This commit is contained in:
parent
c3c1d3fbcf
commit
ef68e3cbd8
|
|
@ -10,4 +10,6 @@ aiofiles
|
||||||
pillow
|
pillow
|
||||||
loguru
|
loguru
|
||||||
httpx
|
httpx
|
||||||
baidusearch
|
baidusearch
|
||||||
|
celery[redis]
|
||||||
|
tortoise-orm[asyncmy]
|
||||||
|
|
@ -3,12 +3,16 @@ from fastapi.responses import StreamingResponse
|
||||||
from src.pipeline.services.rag import RagService
|
from src.pipeline.services.rag import RagService
|
||||||
from src.pipeline.schemas.schemas import ChatRequest
|
from src.pipeline.schemas.schemas import ChatRequest
|
||||||
from src.pipeline.utils import logger
|
from src.pipeline.utils import logger
|
||||||
|
from src.pipeline.models import User as TUser
|
||||||
import json
|
import json
|
||||||
|
|
||||||
router = APIRouter()
|
router = APIRouter()
|
||||||
|
|
||||||
@router.post("/stream-chat")
|
@router.post("/stream-chat")
|
||||||
async def stream_chat(body: ChatRequest):
|
async def stream_chat(body: ChatRequest):
|
||||||
|
|
||||||
|
logger.debug(await TUser.all())
|
||||||
|
|
||||||
service = RagService()
|
service = RagService()
|
||||||
|
|
||||||
async def event_generator():
|
async def event_generator():
|
||||||
|
|
|
||||||
|
|
@ -1,10 +1,5 @@
|
||||||
from mimetypes import init
|
|
||||||
from typing import TypedDict
|
|
||||||
from dotenv import load_dotenv
|
|
||||||
import os
|
|
||||||
from pydantic_settings import BaseSettings
|
from pydantic_settings import BaseSettings
|
||||||
|
from pydantic import MySQLDsn
|
||||||
load_dotenv()
|
|
||||||
|
|
||||||
|
|
||||||
class Settings(BaseSettings):
|
class Settings(BaseSettings):
|
||||||
|
|
@ -26,7 +21,7 @@ class Settings(BaseSettings):
|
||||||
es_port: int
|
es_port: int
|
||||||
es_user: str
|
es_user: str
|
||||||
es_password: str
|
es_password: str
|
||||||
|
mysql_dsn: MySQLDsn
|
||||||
class Config:
|
class Config:
|
||||||
env_file = ".env"
|
env_file = ".env"
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -186,5 +186,6 @@ async def init_client():
|
||||||
global client
|
global client
|
||||||
client = AsyncES()
|
client = AsyncES()
|
||||||
|
|
||||||
|
|
||||||
async def close_client():
|
async def close_client():
|
||||||
await client.close()
|
await client.close()
|
||||||
|
|
|
||||||
|
|
@ -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)
|
||||||
|
|
@ -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
|
||||||
|
);
|
||||||
|
|
@ -1,14 +1,16 @@
|
||||||
from fastapi import FastAPI
|
from fastapi import FastAPI
|
||||||
from src.pipeline.api import include_router
|
from src.pipeline.api import include_router
|
||||||
from src.pipeline.config import config
|
from src.pipeline.config import config
|
||||||
from contextlib import asynccontextmanager
|
|
||||||
from src.pipeline.core import llm, es
|
from src.pipeline.core import llm, es
|
||||||
|
from src.pipeline.db import init_tortoise
|
||||||
|
from contextlib import asynccontextmanager
|
||||||
|
|
||||||
|
|
||||||
@asynccontextmanager
|
@asynccontextmanager
|
||||||
async def lifespan(app: FastAPI):
|
async def lifespan(app: FastAPI):
|
||||||
await llm.init_client()
|
await llm.init_client()
|
||||||
await es.init_client()
|
await es.init_client()
|
||||||
|
await init_tortoise(app, db_url=str(config.mysql_dsn))
|
||||||
|
|
||||||
yield
|
yield
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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"
|
||||||
Loading…
Reference in New Issue