110 lines
3.6 KiB
Python
110 lines
3.6 KiB
Python
"""
|
||
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})
|