ailine/frontend/frontend.py

"""
Streamlit 前端 - 支持模型选择
"""

# 标准库
import os
import uuid

# 第三方库
from dotenv import load_dotenv
import requests
import streamlit as st

# 加载 .env 文件
load_dotenv()

# 后端 API 地址配置
# 优先级：环境变量 API_URL > Docker 内部服务名 > 本地开发地址
API_URL = os.getenv("API_URL", "http://localhost:8001/chat")

st.set_page_config(page_title="AI 个人助手", page_icon="🤖")
st.title("🤖 个人生活与数据分析助手")

# 模型选项（与后端支持的模型名称一致）
MODEL_OPTIONS = {
    "zhipu": "智谱 GLM-4.7-Flash（在线）",
    "deepseek": "DeepSeek V3.2（在线）",
    "local": "本地 vLLM（Gemma-4）"
}

# 初始化会话状态
if "messages" not in st.session_state:
    st.session_state.messages = []
if "thread_id" not in st.session_state:
    st.session_state.thread_id = str(uuid.uuid4())
if "selected_model" not in st.session_state:
    st.session_state.selected_model = "zhipu"

# 侧边栏：模型选择和会话管理
with st.sidebar:
    st.header("⚙️ 设置")
    
    # 模型选择
    selected_model_key = st.selectbox(
        "选择大模型",
        options=list(MODEL_OPTIONS.keys()),
        format_func=lambda x: MODEL_OPTIONS[x],
        index=0
    )
    st.session_state.selected_model = selected_model_key
    
    # 会话信息显示
    st.write(f"当前会话 ID: `{st.session_state.thread_id[:8]}...`")
    
    # 新会话按钮
    if st.button("🔄 新会话"):
        st.session_state.thread_id = str(uuid.uuid4())
        st.session_state.messages = []
        st.rerun()

# 显示历史消息
for msg in st.session_state.messages:
    with st.chat_message(msg["role"]):
        st.markdown(msg["content"])

# 用户输入
if prompt := st.chat_input("请输入您的问题..."):
    # 显示用户消息
    with st.chat_message("user"):
        st.markdown(prompt)
    st.session_state.messages.append({"role": "user", "content": prompt})

    # 调用后端 API（携带模型参数）
    with st.chat_message("assistant"):
        with st.spinner("思考中..."):
            try:
                response = requests.post(
                    API_URL,
                    json={
                        "message": prompt,
                        "thread_id": st.session_state.thread_id,
                        "model": st.session_state.selected_model
                    },
                    timeout=60
                )
                response.raise_for_status()
                data = response.json()
                reply = data["reply"]
                model_used = data["model_used"]
                input_tokens = data.get("input_tokens", 0)
                output_tokens = data.get("output_tokens", 0)
                total_tokens = data.get("total_tokens", 0)
                elapsed_time = data.get("elapsed_time", 0.0)
                
                # 显示回复
                st.markdown(reply)
                
                # 显示使用的模型和性能指标
                stats_text = f"🤖 模型: {MODEL_OPTIONS.get(model_used, model_used)}"
                stats_text += f" | ⏱️ 耗时: {elapsed_time:.2f}s"
                if total_tokens > 0:
                    stats_text += f" | 📊 Tokens: {input_tokens}(输入) + {output_tokens}(输出) = {total_tokens}(总计)"
                st.caption(stats_text)
                
                st.session_state.messages.append({"role": "assistant", "content": reply})
            except Exception as e:
                error_msg = f"请求失败: {e}"
                st.error(error_msg)
                st.session_state.messages.append({"role": "assistant", "content": error_msg})