|
|
import re
|
|
|
|
|
|
import uvicorn # 导入 uvicorn
|
|
|
from fastapi import FastAPI, HTTPException, Depends
|
|
|
from sqlalchemy.orm import Session
|
|
|
from starlette.responses import StreamingResponse
|
|
|
from Text2Sql.Util.PostgreSQLUtil import PostgreSQLUtil
|
|
|
from Text2Sql.Util.SaveToExcel import save_to_excel_stream
|
|
|
from Text2Sql.Util.VannaUtil import VannaUtil
|
|
|
from Model.biModel import *
|
|
|
|
|
|
# 初始化 FastAPI
|
|
|
app = FastAPI()
|
|
|
|
|
|
@app.get("/")
|
|
|
def read_root():
|
|
|
return {"message": "Hello, World!"}
|
|
|
|
|
|
|
|
|
# 获取数据库会话
|
|
|
def get_db():
|
|
|
db = SessionLocal()
|
|
|
try:
|
|
|
yield db
|
|
|
finally:
|
|
|
db.close()
|
|
|
|
|
|
|
|
|
# 创建记录
|
|
|
@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
|
|
|
@app.post("/questions/get_excel")
|
|
|
def get_excel(question: str):
|
|
|
# 指定学段
|
|
|
# question = '''
|
|
|
# 查询:
|
|
|
# 1、发布时间是2024年度
|
|
|
# 2、每个学段，每个科目，上传课程数量，按由多到少排序
|
|
|
# 3、字段名: 学段,科目,排名,课程数量
|
|
|
# '''
|
|
|
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)
|
|
|
# 将数据保存为Excel文件流
|
|
|
excel_stream = save_to_excel_stream(_data)
|
|
|
# 返回Excel文件流
|
|
|
return StreamingResponse(
|
|
|
excel_stream,
|
|
|
media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
|
|
|
headers={"Content-Disposition": "attachment; filename=导出信息.xlsx"}
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
# 确保直接运行脚本时启动 FastAPI 应用
|
|
|
if __name__ == "__main__":
|
|
|
vn = VannaUtil()
|
|
|
uvicorn.run("app:app", host="0.0.0.0", port=8000, reload=True)
|
