Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,40 +1,376 @@
-import altair as alt
-import numpy as np
-import pandas as pd
 import streamlit as st
+import json
+import pandas as pd
+from datetime import datetime
+import time
+from pathlib import Path
+
+# 导入你的MsPatient类和辅助函数
+try:
+    from ms_patient import MsPatient  # 请根据实际路径调整
+    ANNA_AGENT_AVAILABLE = True
+except ImportError:
+    ANNA_AGENT_AVAILABLE = False
+    
+# 导入辅助函数
+from integration_example import (
+    load_dataset, validate_patient_data, initialize_patient_agent,
+    simulate_response, export_chat_history, get_patient_summary
+)
+
+# 页面配置
+st.set_page_config(
+    page_title="AnnaAgent - 心理咨询智能体",
+    page_icon="🧠",
+    layout="wide",
+    initial_sidebar_state="expanded"
+)
+
+# 自定义CSS样式
+st.markdown("""
+<style>
+    .main-header {
+        font-size: 2.5rem;
+        font-weight: 700;
+        color: #2E86AB;
+        text-align: center;
+        margin-bottom: 2rem;
+        text-shadow: 2px 2px 4px rgba(0,0,0,0.1);
+    }
+    
+    .chat-container {
+        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+        border-radius: 15px;
+        padding: 20px;
+        margin: 10px 0;
+        box-shadow: 0 8px 16px rgba(0,0,0,0.2);
+    }
+    
+    .user-message {
+        background: linear-gradient(135deg, #4facfe 0%, #00f2fe 100%);
+        color: white;
+        padding: 15px;
+        border-radius: 15px;
+        margin: 10px 0;
+        box-shadow: 0 4px 8px rgba(0,0,0,0.1);
+        float: right;
+        clear: both;
+        max-width: 70%;
+        margin-left: 30%;
+    }
+    
+    .assistant-message {
+        background: linear-gradient(135deg, #fa709a 0%, #fee140 100%);
+        color: #333;
+        padding: 15px;
+        border-radius: 15px;
+        margin: 10px 0;
+        box-shadow: 0 4px 8px rgba(0,0,0,0.1);
+        float: left;
+        clear: both;
+        max-width: 70%;
+        margin-right: 30%;
+    }
+    
+    .sidebar-info {
+        background: #f0f2f6;
+        padding: 15px;
+        border-radius: 10px;
+        margin: 10px 0;
+        border-left: 4px solid #2E86AB;
+    }
+    
+    .metrics-container {
+        display: flex;
+        justify-content: space-around;
+        margin: 20px 0;
+    }
+    
+    .metric-card {
+        background: white;
+        padding: 20px;
+        border-radius: 10px;
+        text-align: center;
+        box-shadow: 0 4px 8px rgba(0,0,0,0.1);
+        min-width: 120px;
+    }
+    
+    .stButton > button {
+        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+        color: white;
+        border: none;
+        border-radius: 25px;
+        padding: 0.5rem 2rem;
+        font-weight: 600;
+        transition: all 0.3s ease;
+    }
+    
+    .stButton > button:hover {
+        transform: translateY(-2px);
+        box-shadow: 0 6px 12px rgba(0,0,0,0.2);
+    }
+</style>
+""", unsafe_allow_html=True)
+
+# 初始化session state
+if 'messages' not in st.session_state:
+    st.session_state.messages = []
+if 'patient_agent' not in st.session_state:
+    st.session_state.patient_agent = None
+if 'current_patient_data' not in st.session_state:
+    st.session_state.current_patient_data = None
+
+# 标题
+st.markdown('<h1 class="main-header">🧠 AnnaAgent 心理咨询智能体</h1>', unsafe_allow_html=True)
+
+# 侧边栏 - 患者信息和控制
+with st.sidebar:
+    st.markdown("## 📋 患者信息配置")
+    
+    # 数据集加载
+    st.markdown("### 📁 数据集加载")
+    uploaded_file = st.file_uploader(
+        "上传Anna-CPsyCounD数据集文件",
+        type=['json', 'jsonl'],
+        help="请上传包含患者数据的JSON文件"
+    )
+    
+    # 示例数据选择器（如果没有上传文件）
+    if uploaded_file is None:
+        st.markdown("### 🎭 示例患者选择")
+        # 这里提供一些示例数据
+        example_patients = {
+            "示例患者1 - 抑郁症状": {
+                "id": "example_001",
+                "portrait": {
+                    "age": 28,
+                    "gender": "女",
+                    "occupation": "软件工程师",
+                    "marital_status": "未婚",
+                    "symptom": "抑郁情绪、失眠、食欲不振"
+                },
+                "report": {
+                    "chief_complaint": "最近几个月感到情绪低落",
+                    "history": "工作压力大，经常加班",
+                    "mental_status": "情绪低落，思维迟缓"
+                },
+                "conversation": []
+            },
+            "示例患者2 - 焦虑症状": {
+                "id": "example_002", 
+                "portrait": {
+                    "age": 35,
+                    "gender": "男",
+                    "occupation": "销售经理",
+                    "marital_status": "已婚",
+                    "symptom": "焦虑、心慌、担心"
+                },
+                "report": {
+                    "chief_complaint": "经常��到紧张焦虑",
+                    "history": "业绩压力，家庭责任重",
+                    "mental_status": "焦虑不安，注意力不集中"
+                },
+                "conversation": []
+            }
+        }
+        
+        selected_example = st.selectbox(
+            "选择一个示例患者",
+            options=list(example_patients.keys())
+        )
+        
+        if st.button("🚀 启动会话", type="primary"):
+            patient_data = example_patients[selected_example]
+            st.session_state.current_patient_data = patient_data
+            
+            # 初始化患者智能体
+            if ANNA_AGENT_AVAILABLE:
+                agent, message = initialize_patient_agent(patient_data)
+                if agent:
+                    st.session_state.patient_agent = agent
+                    st.success(f"✅ 患者智能体已启动！{message}")
+                else:
+                    st.error(f"❌ 启动失败: {message}")
+                    st.session_state.patient_agent = None
+            else:
+                st.warning("⚠️ 使用模拟模式")
+                st.session_state.patient_agent = None
+            
+            st.rerun()
+    
+    else:
+        # 处理上传的文件
+        try:
+            data = load_dataset(uploaded_file)
+            if data and isinstance(data, list):
+                patients_df = pd.DataFrame(data)
+                st.markdown("### 👥 患者列表")
+                
+                # 显示患者信息表格
+                display_df = patients_df[['id']].copy()
+                if 'portrait' in patients_df.columns:
+                    display_df['年龄'] = patients_df['portrait'].apply(lambda x: x.get('age', 'N/A') if isinstance(x, dict) else 'N/A')
+                    display_df['性别'] = patients_df['portrait'].apply(lambda x: x.get('gender', 'N/A') if isinstance(x, dict) else 'N/A')
+                    display_df['症状'] = patients_df['portrait'].apply(lambda x: x.get('symptom', 'N/A') if isinstance(x, dict) else 'N/A')
+                
+                selected_id = st.selectbox("选择患者ID", options=patients_df['id'].tolist())
+                
+                if st.button("🚀 启动会话", type="primary"):
+                    selected_patient = patients_df[patients_df['id'] == selected_id].iloc[0].to_dict()
+                    st.session_state.current_patient_data = selected_patient
+                    
+                    # 初始化患者智能体
+                    if ANNA_AGENT_AVAILABLE:
+                        agent, message = initialize_patient_agent(selected_patient)
+                        if agent:
+                            st.session_state.patient_agent = agent
+                            st.success(f"✅ 患者智能体已启动！{message}")
+                        else:
+                            st.error(f"❌ 启动失败: {message}")
+                            st.session_state.patient_agent = None
+                    else:
+                        st.warning("⚠️ 使用模拟模式")
+                        st.session_state.patient_agent = None
+                    
+                    st.rerun()
+                    
+        except Exception as e:
+            st.error(f"❌ 文件加载失败: {str(e)})")
+    
+    # 分隔线
+    st.markdown("---")
+    
+    # 当前患者信息显示
+    if st.session_state.current_patient_data:
+        st.markdown("### 👤 当前患者信息")
+        patient_info = st.session_state.current_patient_data
+        
+        if 'portrait' in patient_info:
+            portrait = patient_info['portrait']
+            st.markdown(f"""
+            <div class="sidebar-info">
+                <strong>基本信息:</strong><br>
+                • 年龄: {portrait.get('age', 'N/A')}<br>
+                • 性别: {portrait.get('gender', 'N/A')}<br>
+                • 职业: {portrait.get('occupation', 'N/A')}<br>
+                • 婚姻状态: {portrait.get('marital_status', 'N/A')}<br>
+                • 主要症状: {portrait.get('symptom', 'N/A')}
+            </div>
+            """, unsafe_allow_html=True)
+    
+    # 控制按钮
+    st.markdown("### 🎮 会话控制")
+    
+    col1, col2 = st.columns(2)
+    with col1:
+        if st.button("🔄 重置会话"):
+            st.session_state.messages = []
+            st.rerun()
+    
+    with col2:
+        if st.button("💾 导出记录"):
+            if st.session_state.messages:
+                chat_history = {
+                    "patient_id": st.session_state.current_patient_data.get('id', 'unknown') if st.session_state.current_patient_data else 'unknown',
+                    "timestamp": datetime.now().isoformat(),
+                    "messages": st.session_state.messages
+                }
+                st.download_button(
+                    label="📥 下载聊天记录",
+                    data=json.dumps(chat_history, ensure_ascii=False, indent=2),
+                    file_name=f"chat_history_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json",
+                    mime="application/json"
+                )
+
+# 主聊天界面
+if st.session_state.current_patient_data is None:
+    st.markdown("""
+    <div style='text-align: center; padding: 50px;'>
+        <h3>🎯 欢迎使用AnnaAgent心理咨询智能体</h3>
+        <p>请在左侧面板选择或上传患者数据以开始对话</p>
+    </div>
+    """, unsafe_allow_html=True)
+else:
+    # 显示聊天历史
+    chat_container = st.container()
+    
+    with chat_container:
+        for message in st.session_state.messages:
+            if message["role"] == "user":
+                st.markdown(f"""
+                <div class="user-message">
+                    <strong>咨询师:</strong> {message["content"]}
+                </div>
+                """, unsafe_allow_html=True)
+            else:
+                st.markdown(f"""
+                <div class="assistant-message">
+                    <strong>来访者:</strong> {message["content"]}
+                </div>
+                """, unsafe_allow_html=True)
+    
+    # 输入框
+    with st.form(key="chat_form", clear_on_submit=True):
+        col1, col2 = st.columns([4, 1])
+        
+        with col1:
+            user_input = st.text_area(
+                "输入您的话语:",
+                placeholder="请输入您想对来访者说的话...",
+                height=100,
+                key="user_input"
+            )
+        
+        with col2:
+            st.markdown("<br>", unsafe_allow_html=True)  # 增加一些间距
+            submit_button = st.form_submit_button("💬 发送", type="primary")
+    
+    if submit_button and user_input:
+        # 添加用户消息到聊天历史
+        st.session_state.messages.append({"role": "user", "content": user_input})
+        
+        # 这里调用AnnaAgent生成回复
+        with st.spinner("🤔 来访者正在思考回复..."):
+            try:
+                # 模拟调用（你需要替换为实际的调用）
+                # if st.session_state.patient_agent:
+                #     response = st.session_state.patient_agent.chat(user_input)
+                # else:
+                #     response = "抱歉，智能体尚未初始化。"
+                
+                # 临时模拟回复
+                time.sleep(1)  # 模拟思考时间
+                response = f"谢谢您的关心。您说'{user_input}'，这让我想到..."
+                
+                # 添加助手回复到聊天历史
+                st.session_state.messages.append({"role": "assistant", "content": response})
+                
+            except Exception as e:
+                st.error(f"❌ 生成回复时出错: {str(e)}")
+                st.session_state.messages.append({
+                    "role": "assistant", 
+                    "content": "抱歉，我现在无法回复。请稍后再试。"
+                })
+        
+        st.rerun()
+
+# 底部信息
+st.markdown("---")
+col1, col2, col3 = st.columns(3)
+
+with col1:
+    if st.session_state.messages:
+        st.metric("💬 对话轮数", len(st.session_state.messages) // 2)
+
+with col2:
+    if st.session_state.current_patient_data:
+        st.metric("👤 当前患者", st.session_state.current_patient_data.get('id', 'Unknown'))
+
+with col3:
+    st.metric("🕐 会话时间", f"{datetime.now().strftime('%H:%M:%S')}")
 
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+st.markdown("""
+<div style='text-align: center; color: #666; margin-top: 20px;'>
+    <small>AnnaAgent - 基于大语言模型的心理咨询智能体 | 仅供学术研究使用</small>
+</div>
+""", unsafe_allow_html=True)