import os
import shutil
from pathlib import Path
from flask import Flask, render_template_string, send_file
from update_predictions import Config, main_task, get_business_info
app = Flask(__name__)
# 统一配置（与update_predictions.py保持一致）
Config = {
"REPO_PATH": Path(__file__).parent.resolve(),
"MODEL_PATH": os.path.join("/tmp", "Kronos_model"),
# 关键修改：HTML和图表都存放在/tmp目录（可写）
"HTML_PATH": os.path.join("/tmp", "index.html"),
"CHART_PATH": os.path.join("/tmp", "prediction_chart.png")
}
# 确保/tmp目录下有模板文件
def initialize_html_template():
# 源模板路径（项目中的原始index.html）
src_html = Config["REPO_PATH"] / "templates" / "index.html"
# 目标路径（/tmp下的可写副本）
dest_html = Path(Config["HTML_PATH"])
# 如果目标不存在，从源模板复制
if not dest_html.exists():
# 确保/tmp目录存在
dest_html.parent.mkdir(parents=True, exist_ok=True)
if src_html.exists():
shutil.copy2(src_html, dest_html)
print(f"已复制模板文件到可写目录：{dest_html}")
else:
# 如果源模板也不存在，创建基础HTML
base_html = """
清华大学Kronos股票K型大模型上证指数预测
清华大学Kronos股票K型大模型上证指数概率预测
最后更新时间：未更新（UTC）
同 步 网 站：火狼工具站
"""
with open(dest_html, 'w', encoding='utf-8') as f:
f.write(base_html)
print(f"已创建基础模板文件：{dest_html}")
# 初始化模型（全局一次）
from update_predictions import load_local_model as load_model
predictor = load_model()
# 初始化HTML模板
initialize_html_template()
@app.route("/")
def index():
# 获取当前业务日
current_business_date, _ = get_business_info()
# 容错处理：检查配置键是否存在，不存在则执行主任务
# 同时处理键存在但值为None的情况（首次运行）
if ("LAST_INFERENCED_BUSINESS_DATE" not in Config
or Config["LAST_INFERENCED_BUSINESS_DATE"] != current_business_date):
main_task(predictor)
# 读取并返回 HTML
try:
with open(Config["HTML_PATH"], 'r', encoding='utf-8') as f:
html_content = f.read()
return render_template_string(html_content)
except FileNotFoundError:
# 处理HTML文件不存在的情况
return "预测页面生成中，请稍后刷新", 202
except Exception as e:
return f"页面加载错误：{str(e)}", 500
@app.route("/prediction_chart.png")
def get_chart():
# 提供图表访问
chart_path = Path(Config["CHART_PATH"])
if chart_path.exists():
return send_file(chart_path, mimetype='image/png')
else:
return "预测图表生成中，请稍后刷新", 202
if __name__ == "__main__":
app.run(host="0.0.0.0", port=7860)
