LiSiyi13146413708/Baseline_Agent

1import os
2import uuid
3import shutil
4import subprocess
5import sys
6from pathlib import Path
7import zipfile
8import io
9import streamlit as st
10from streamlit.components.v1 import html as st_html
11
12BASE_DIR = Path(__file__).parent.resolve()
13UPLOAD_DIR = BASE_DIR / "uploads"
14OUTPUTS_DIR = BASE_DIR / "frontend_outputs"
15UPLOAD_DIR.mkdir(exist_ok=True)
16OUTPUTS_DIR.mkdir(exist_ok=True)
17
18
19st.set_page_config(page_title="技术路线图生成器", layout="centered")
20st.title("技术路线图生成器")
21st.write("请放入您需要生成技术路线图的文件")
22
23uploaded = st.file_uploader("拖拽或点击上传文件", type=None)
24
25if uploaded is not None:
26  job_id = uuid.uuid4().hex
27  filename = uploaded.name
28  saved_path = UPLOAD_DIR / f"{job_id}_{filename}"
29  with open(saved_path, "wb") as fh:
30    fh.write(uploaded.getbuffer())
31
32  out_dir = OUTPUTS_DIR / f"out_{job_id}"
33  if out_dir.exists():
34    shutil.rmtree(out_dir)
35  out_dir.mkdir(parents=True, exist_ok=True)
36
37  with st.spinner('技术路线图正在生成，请耐心等待...'):
38    # main.py 期望 --output 是文件前缀（不含扩展名），不是目录路径。
39    # 这里构造一个输出前缀：out_dir/roadmap
40    output_prefix = out_dir / 'roadmap'
41    cmd = [
42      sys.executable,
43      str(Path(__file__).parent / 'main.py'),
44      '--input', str(saved_path),
45      '--output', str(output_prefix),
46    ]
47    try:
48      proc = subprocess.run(cmd, capture_output=True, text=True, timeout=60*60)
49    except subprocess.TimeoutExpired:
50      st.error('处理超时')
51      st.stop()
52
53  if proc.returncode != 0:
54    st.error('子进程返回错误')
55    st.text(proc.stdout)
56    st.text(proc.stderr)
57  else:
58    st.success('生成完成')
59
60    # 查找可预览文件（优先 HTML）
61    preview_path = None
62    for root, dirs, files in os.walk(out_dir):
63      for name in files:
64        lower = name.lower()
65        if lower.endswith(('.html', '.htm')):
66          preview_path = Path(root) / name
67          break
68      if preview_path:
69        break
70
71    if preview_path and preview_path.exists():
72      content = preview_path.read_text(encoding='utf-8', errors='ignore')
73      st_html(content, height=600, scrolling=True)
74    else:
75      # 查找文本类型
76      text_path = None
77      for root, dirs, files in os.walk(out_dir):
78        for name in files:
79          if name.lower().endswith(('.txt', '.md')):
80            text_path = Path(root) / name
81            break
82        if text_path:
83          break
84
85      if text_path and text_path.exists():
86        text = text_path.read_text(encoding='utf-8', errors='ignore')
87        st.code(text)
88      else:
89        # 打包为 zip 并提供下载
90        buf = io.BytesIO()
91        with zipfile.ZipFile(buf, 'w', zipfile.ZIP_DEFLATED) as zf:
92          for root, dirs, files in os.walk(out_dir):
93            for name in files:
94              file_path = Path(root) / name
95              arcname = file_path.relative_to(out_dir)
96              zf.write(file_path, arcname)
97        buf.seek(0)
98        st.download_button('下载生成结果（ZIP）', data=buf, file_name=f'result_{job_id}.zip')
99
100    # 显示子进程输出以便调试
101    if proc.stdout:
102      st.subheader('子进程输出')
103      st.text(proc.stdout)
104    if proc.stderr:
105      st.subheader('子进程错误输出')
106      st.text(proc.stderr)
107
108