Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,376 +1,181 @@
 import streamlit as st
-import json
-import pandas as pd
-from datetime import datetime
+from openai import OpenAI
 import time
-from pathlib import Path
+from datasets import load_dataset # 导入Hugging Face datasets库
+
+# =======================================================================
+# 1. 导入您的类和数据
+# =======================================================================
+
+# 请将 'your_agent_file' 替换为包含 MsPatient 类的实际文件名
+from ms_patient import MsPatient 
+
+# 使用Streamlit的缓存功能来加载和缓存数据集
+@st.cache_data
+def load_hf_dataset():
+    """
+    从Hugging Face Hub加载并缓存数据集。
+    这将返回一个字典列表，每个字典代表一个病人数据。
+    """
+    try:
+        # 加载'train'分割部分的数据
+        dataset = load_dataset("sci-m-wang/Anna-CPsyCounD", split='train')
+        # 转换为pandas DataFrame再转为字典列表，方便处理
+        return dataset.to_pandas().to_dict('records')
+    except Exception as e:
+        st.error(f"从Hugging Face加载数据集失败: {e}")
+        return []
+
+# 加载数据
+ALL_PATIENTS = load_hf_dataset()
 
-# 导入你的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
-)
 
-# 页面配置
+# =======================================================================
+# 2. Streamlit 应用界面
+# =======================================================================
+
+# --- 页面配置 ---
 st.set_page_config(
-    page_title="AnnaAgent - 心理咨询智能体",
-    page_icon="🧠",
-    layout="wide",
-    initial_sidebar_state="expanded"
+    page_title="与Anna对话",
+    page_icon="�",
+    layout="wide"
 )
 
-# 自定义CSS样式
+# --- 自定义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%;
+    /* 主聊天容器 */
+    .st-emotion-cache-1y4p8pa { 
+        padding-top: 2rem;
     }
-    
-    .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%;
+    /* 聊天消息 */
+    .st-chat-message {
+        border-radius: 0.8rem;
+        padding: 0.9rem 1.2rem;
+        box-shadow: 0 2px 5px rgba(0,0,0,0.05);
+        background-color: #ffffff;
     }
-    
-    .sidebar-info {
-        background: #f0f2f6;
-        padding: 15px;
-        border-radius: 10px;
-        margin: 10px 0;
-        border-left: 4px solid #2E86AB;
+    .st-chat-message[data-testid="chat-message-container-user"] {
+        background-color: #dcf8c6;
     }
-    
-    .metrics-container {
-        display: flex;
-        justify-content: space-around;
-        margin: 20px 0;
+    /* 侧边栏 */
+    .st-sidebar {
+        background-color: #f8f9fa;
+        border-right: 1px solid #e9ecef;
     }
-    
-    .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;
+    .st-sidebar h2 {
+        color: #343a40;
     }
-    
-    .stButton > button {
-        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
-        color: white;
-        border: none;
-        border-radius: 25px;
-        padding: 0.5rem 2rem;
+    .st-expanderHeader {
+        font-size: 1.1rem;
         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)
+# --- 初始化 Session State ---
+if "patient_agent" not in st.session_state:
+    st.session_state.patient_agent = None
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+if "selected_patient_id" not in st.session_state:
+    st.session_state.selected_patient_id = None
+if "openai_client" not in st.session_state:
+    st.session_state.openai_client = None
+if "model_name" not in st.session_state:
+    st.session_state.model_name = "gpt-4o-mini" # 默认模型
 
-# 侧边栏 - 患者信息和控制
+# --- 侧边栏 ---
 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
+    st.title("👩 AnnaAgent 设置")
+    st.markdown("---")
+
+    # API Key 输入
+    with st.expander("🔑 API 设置", expanded=True):
+        api_key = st.text_input("输入您的 OpenAI API Key", type="password", help="您的API Key将仅用于本次会话，不会被储存。")
+        base_url = st.text_input("API Base URL (可选)", value="https://api.openai.com/v1")
+        model_name = st.text_input("模型名称", value=st.session_state.model_name)
+
+        if st.button("连接模型"):
+            if api_key:
+                try:
+                    st.session_state.openai_client = OpenAI(api_key=api_key, base_url=base_url)
+                    st.session_state.model_name = model_name
+                    st.success("连接成功！")
+                    if st.session_state.patient_agent:
+                        st.session_state.patient_agent.client = st.session_state.openai_client
+                except Exception as e:
+                    st.error(f"连接失败: {e}")
             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.warning("请输入API Key。")
+
     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 not ALL_PATIENTS:
+        st.error("无法加载病人数据，请检查网络连接或数据集名称。")
+    else:
+        patient_options = {p["id"]: f"{p['portrait']['gender']}，{p['portrait']['age']}岁 - {p['portrait']['symptom']}" for p in ALL_PATIENTS}
+        selected_id = st.selectbox(
+            "选择一位病人进行对话",
+            options=list(patient_options.keys()),
+            format_func=lambda x: patient_options[x]
+        )
+
+        # 当选择的病人变化时，重置状态
+        if st.session_state.selected_patient_id != selected_id:
+            st.session_state.selected_patient_id = selected_id
+            selected_patient_data = next((p for p in ALL_PATIENTS if p["id"] == selected_id), None)
+            
+            with st.spinner("正在生成病人角色..."):
+                st.session_state.patient_agent = MsPatient(
+                    portrait=selected_patient_data["portrait"],
+                    report=selected_patient_data["report"],
+                    previous_conversations=selected_patient_data["conversation"],
+                    language="Chinese",
+                    client=st.session_state.openai_client
                 )
+            st.session_state.messages = [{"role": "assistant", "content": "你好，医生..."}]
+            st.rerun()
 
-# 主聊天界面
-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})
+    # 显示病人信息
+    if st.session_state.patient_agent:
+        st.markdown("---")
+        st.subheader("病人信息")
+        agent = st.session_state.patient_agent
+        st.info(f"""
+        **基本情况**: {agent.portrait['gender']}, {agent.portrait['age']}岁, {agent.portrait['occupation']}, {agent.portrait['marital_status']}
         
-        # 这里调用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()
+        **近期状态**: {agent.status}
+        """)
 
-# 底部信息
-st.markdown("---")
-col1, col2, col3 = st.columns(3)
+        with st.expander("查看完整系统提示 (System Prompt)"):
+            st.code(agent.get_system_prompt(), language='markdown')
 
-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'))
+# --- 主聊天界面 ---
+st.title("💬 与 Anna 对话")
+st.caption("这是一个模拟心理咨询的AI Agent。由 `MsPatient` 类驱动。")
 
-with col3:
-    st.metric("🕐 会话时间", f"{datetime.now().strftime('%H:%M:%S')}")
+# 显示聊天记录
+for message in st.session_state.messages:
+    avatar = "👩" if message["role"] == "assistant" else "🧑‍⚕️"
+    with st.chat_message(message["role"], avatar=avatar):
+        st.markdown(message["content"])
 
-st.markdown("""
-<div style='text-align: center; color: #666; margin-top: 20px;'>
-    <small>AnnaAgent - 基于大语言模型的心理咨询智能体 | 仅供学术研究使用</small>
-</div>
-""", unsafe_allow_html=True)
+# 获取用户输入
+if prompt := st.chat_input("请输入您想说的话..."):
+    st.session_state.messages.append({"role": "user", "content": prompt})
+    with st.chat_message("user", avatar="🧑‍⚕️"):
+        st.markdown(prompt)
+
+    if st.session_state.patient_agent:
+        with st.chat_message("assistant", avatar="👩"):
+            with st.spinner("Anna正在思考..."):
+                response = st.session_state.patient_agent.chat(prompt)
+                st.markdown(response)
+        
+        st.session_state.messages.append({"role": "assistant", "content": response})
+    else:
+        st.warning("请先在左侧选择一位病人并配置API Key。")