|
|
|
@ -1,13 +1,17 @@
|
|
|
|
|
import re
|
|
|
|
|
from typing import io
|
|
|
|
|
|
|
|
|
|
import pandas as pd
|
|
|
|
|
from fastapi import FastAPI, HTTPException, Depends
|
|
|
|
|
from pydantic import BaseModel
|
|
|
|
|
from sqlalchemy import create_engine, Column, Integer, String, SmallInteger, Date
|
|
|
|
|
from sqlalchemy.orm import declarative_base # 更新导入路径
|
|
|
|
|
from sqlalchemy.orm import sessionmaker, Session
|
|
|
|
|
import uvicorn # 导入 uvicorn
|
|
|
|
|
from starlette.staticfiles import StaticFiles
|
|
|
|
|
from starlette.responses import StreamingResponse
|
|
|
|
|
|
|
|
|
|
from Text2Sql.Util.PostgreSQLUtil import PostgreSQLUtil
|
|
|
|
|
from Text2Sql.Util.SaveToExcel import save_to_excel
|
|
|
|
|
from Text2Sql.Util.VannaUtil import VannaUtil
|
|
|
|
|
|
|
|
|
|
# 数据库连接配置
|
|
|
|
@ -131,8 +135,9 @@ def delete_question(question_id: int, db: Session = Depends(get_db)):
|
|
|
|
|
db.commit()
|
|
|
|
|
return {"message": "Question deleted successfully"}
|
|
|
|
|
|
|
|
|
|
@app.post("/questions/generate_sql")
|
|
|
|
|
def generate_sql(question: str):
|
|
|
|
|
# 通过语义生成Excel
|
|
|
|
|
@app.post("/questions/getExcel")
|
|
|
|
|
def getExcel(question: str):
|
|
|
|
|
# 指定学段
|
|
|
|
|
question = '''
|
|
|
|
|
查询:
|
|
|
|
@ -146,12 +151,28 @@ def generate_sql(question: str):
|
|
|
|
|
2、目标数据库是Postgresql 16
|
|
|
|
|
'''
|
|
|
|
|
question = question + common_prompt
|
|
|
|
|
# 开始查询
|
|
|
|
|
print("开始查询...")
|
|
|
|
|
# 获取完整 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"}
|
|
|
|
|
)
|
|
|
|
|
def save_to_excel_stream(data):
|
|
|
|
|
# 将数据保存为Excel文件流
|
|
|
|
|
df = pd.DataFrame(data)
|
|
|
|
|
output = io.BytesIO()
|
|
|
|
|
with pd.ExcelWriter(output, engine='openpyxl') as writer:
|
|
|
|
|
df.to_excel(writer, index=False)
|
|
|
|
|
output.seek(0) # 将指针移动到流的开头
|
|
|
|
|
return output
|
|
|
|
|
# 添加 main 函数
|
|
|
|
|
def main():
|
|
|
|
|
# 开始训练
|
|
|
|
|