|
|
import uuid
|
|
|
|
|
|
import uvicorn # 导入 uvicorn
|
|
|
from fastapi import FastAPI, HTTPException, Depends, Form
|
|
|
from sqlalchemy.orm import Session
|
|
|
from starlette.responses import JSONResponse
|
|
|
from starlette.staticfiles import StaticFiles
|
|
|
|
|
|
from Model.biModel import *
|
|
|
from Text2Sql.Util.PostgreSQLUtil import PostgreSQLUtil
|
|
|
from Text2Sql.Util.SaveToExcel import save_to_excel
|
|
|
from Text2Sql.Util.VannaUtil import VannaUtil
|
|
|
|
|
|
# 初始化 FastAPI
|
|
|
app = FastAPI()
|
|
|
# 配置静态文件目录
|
|
|
app.mount("/static", StaticFiles(directory="static"), name="static")
|
|
|
vn = VannaUtil()
|
|
|
@app.get("/")
|
|
|
def read_root():
|
|
|
return {"message": "Welcome to Vanna AI SQL !"}
|
|
|
|
|
|
|
|
|
# 创建记录
|
|
|
@app.post("/questions/", response_model=TBIQuestionCreate)
|
|
|
def create_question(question: TBIQuestionCreate, db: Session = Depends(get_db)):
|
|
|
db_question = TBIQuestion(**question.dict())
|
|
|
db.add(db_question)
|
|
|
db.commit()
|
|
|
db.refresh(db_question)
|
|
|
return db_question
|
|
|
|
|
|
|
|
|
# 读取记录
|
|
|
@app.get("/questions/{question_id}", response_model=TBIQuestionCreate)
|
|
|
def read_question(question_id: int, db: Session = Depends(get_db)):
|
|
|
db_question = db.query(TBIQuestion).filter(TBIQuestion.id == question_id).first()
|
|
|
if db_question is None:
|
|
|
raise HTTPException(status_code=404, detail="Question not found")
|
|
|
return db_question
|
|
|
|
|
|
|
|
|
# 更新记录
|
|
|
@app.put("/questions/{question_id}", response_model=TBIQuestionCreate)
|
|
|
def update_question(question_id: int, question: TBIQuestionUpdate, db: Session = Depends(get_db)):
|
|
|
db_question = db.query(TBIQuestion).filter(TBIQuestion.id == question_id).first()
|
|
|
if db_question is None:
|
|
|
raise HTTPException(status_code=404, detail="Question not found")
|
|
|
for key, value in question.dict().items():
|
|
|
if value is not None:
|
|
|
setattr(db_question, key, value)
|
|
|
db.commit()
|
|
|
db.refresh(db_question)
|
|
|
return db_question
|
|
|
|
|
|
|
|
|
# 删除记录
|
|
|
@app.delete("/questions/{question_id}")
|
|
|
def delete_question(question_id: int, db: Session = Depends(get_db)):
|
|
|
db_question = db.query(TBIQuestion).filter(TBIQuestion.id == question_id).first()
|
|
|
if db_question is None:
|
|
|
raise HTTPException(status_code=404, detail="Question not found")
|
|
|
db.delete(db_question)
|
|
|
db.commit()
|
|
|
return {"message": "Question deleted successfully"}
|
|
|
|
|
|
|
|
|
# 通过语义生成Excel
|
|
|
# http://10.10.21.20:8000/questions/get_excel
|
|
|
# 参数：question
|
|
|
@app.post("/questions/get_excel")
|
|
|
def get_excel(question: str = Form(...)):
|
|
|
|
|
|
common_prompt = '''
|
|
|
返回的信息要求：
|
|
|
1、行政区划为NULL 或者是空字符的不参加统计
|
|
|
2、目标数据库是Postgresql 16
|
|
|
'''
|
|
|
question = question + common_prompt
|
|
|
# 获取完整 SQL
|
|
|
sql = vn.generate_sql(question)
|
|
|
print("生成的查询 SQL:\n", sql)
|
|
|
# 执行SQL查询
|
|
|
with PostgreSQLUtil() as db:
|
|
|
_data = db.execute_query(sql)
|
|
|
# 在static目录下，生成一个guid号的临时文件
|
|
|
uuidStr = str(uuid.uuid4())
|
|
|
filename = f"static/{uuidStr}.xlsx"
|
|
|
save_to_excel(_data, filename)
|
|
|
# 返回静态文件URL
|
|
|
return {"success": True, "message": "Excel文件生成成功", "download_url": f"/static/{uuidStr}.xlsx"}
|
|
|
|
|
|
|
|
|
|
|
|
# 确保直接运行脚本时启动 FastAPI 应用
|
|
|
if __name__ == "__main__":
|
|
|
uvicorn.run("app:app", host="0.0.0.0", port=8000, reload=True)
|
