ailine/frontend/components/chat_area.py

"""
中间聊天区组件
包含模型选择、消息显示和输入框
"""

import re
import streamlit as st

# 使用绝对导入
from frontend.state import AppState
from frontend.api_client import api_client
from frontend.config import config


def render_chat_area():
    """渲染中间聊天区域"""
    # 顶部：极简模型选择器（可选放在顶部中间）
    _render_model_selector()
    
    # 使用空白占位符或者不需要 divider 让界面更干净
    st.write("")
    
    # 渲染历史消息
    _render_chat_history()
    
    # 输入框和流式响应处理
    _render_input_and_response()


def _render_model_selector():
    """渲染模型选择器，极简风格"""
    col_empty1, col_model, col_empty2 = st.columns([1, 2, 1])
    
    with col_model:
        selected_model = st.selectbox(
            "选择模型",
            options=list(config.model_options.keys()),
            format_func=lambda x: config.model_options[x],
            index=_get_model_index(),
            label_visibility="collapsed" # 隐藏标签，只显示下拉框，更加现代
        )
        AppState.set_selected_model(selected_model)


def _get_model_index() -> int:
    """
    获取当前选中模型的索引
    
    Returns:
        模型索引
    """
    current_model = AppState.get_selected_model()
    model_keys = list(config.model_options.keys())
    return model_keys.index(current_model) if current_model in model_keys else 0


def _render_chat_history():
    """渲染历史聊天消息"""
    messages = AppState.get_messages()
    for msg in messages:
        with st.chat_message(msg["role"]):
            content = msg["content"]
            
            # 1. 尝试解析我们在前端流式结束后存入的 ```thought 格式
            if "```thought\n" in content:
                parts = content.split("```thought\n")
                if parts[0].strip():
                    st.markdown(parts[0])
                
                for part in parts[1:]:
                    if "\n```\n" in part:
                        thought, rest = part.split("\n```\n", 1)
                        with st.expander("🤔 思考过程", expanded=False):
                            st.markdown(thought)
                        if rest.strip():
                            st.markdown(rest)
                    else:
                        st.markdown("```thought\n" + part)
                        
            # 2. 尝试解析从后端原始加载的历史记录中包含的 <think> 标签
            elif "<think>" in content and "</think>" in content:
                # 提取思考内容和剩余正文
                thought_match = re.search(r'<think>(.*?)</think>', content, re.DOTALL)
                if thought_match:
                    thought = thought_match.group(1).strip()
                    rest = re.sub(r'<think>.*?</think>', '', content, flags=re.DOTALL).strip()
                    
                    with st.expander("🤔 思考过程", expanded=False):
                        st.markdown(thought)
                    if rest:
                        st.markdown(rest)
                else:
                    st.markdown(content)
            else:
                st.markdown(content)


def _render_input_and_response():
    """渲染输入框并处理用户输入与AI响应"""
    if prompt := st.chat_input("请输入您的问题...", key="chat_input"):
        # 显示用户消息
        with st.chat_message("user"):
            st.markdown(prompt)
        AppState.add_message("user", prompt)
        
        # 流式调用 AI 回复
        _handle_ai_response()


def _handle_ai_response():
    """处理 AI 流式响应 (适配 LangGraph v2 事件格式)"""
    with st.chat_message("assistant"):
        # 用于容纳思考过程的占位符（只有在使用 DeepSeek reasoner 时才显示）
        thought_placeholder = st.empty()
        message_placeholder = st.empty()
        tool_status_placeholder = st.empty()
        
        raw_text = ""
        api_thought = ""
        display_text = ""
        display_thought = ""
        rag_sources = None  # 存储 RAG 检索来源信息
        
        # 调用流式 API
        stream = api_client.chat_stream(
            message=AppState.get_messages()[-1]["content"],
            thread_id=AppState.get_current_thread_id(),
            model=AppState.get_selected_model(),
            user_id=AppState.get_user_id()
        )
        
        # 消费流式响应 (v2 格式)
        for event in stream:
            event_type = event.get("type")
            
            # [DEBUG] 可以在前端终端看到接收到的事件
            import logging
            if event_type == "llm_token":
                logging.debug(f"[Frontend Stream] token: {repr(event.get('token'))}, reasoning: {repr(event.get('reasoning_token'))}")
            
            # 1. 处理 LLM Token 流 (打字机效果)
            if event_type == "llm_token":
                # 确保只处理来自 LLM 的 token，避免将工具的输出作为 token 显示
                if event.get("node") == "llm_call":
                    token = str(event.get("token", ""))
                    reasoning_token = str(event.get("reasoning_token", ""))
                    
                    if reasoning_token:
                        api_thought += reasoning_token
                    if token:
                        raw_text += token
                        
                    display_thought = api_thought
                    display_text = raw_text
                    is_thinking = False
                    
                    # 1. 原生 API 推理模式 (如 DeepSeek-Reasoner)
                    if api_thought:
                        is_thinking = not bool(raw_text.strip())
                        
                    # 2. 本地模型 <think> 标签模式 (如 Gemma, 本地 DeepSeek)
                    if "<think>" in raw_text:
                        think_match = re.search(r'<think>(.*?)(</think>|$)', raw_text, re.DOTALL)
                        if think_match:
                            display_thought = think_match.group(1).strip()
                            is_thinking = "</think>" not in raw_text
                            
                            # 正文部分应该是除去了整个 <think>...</think> 块后的剩余内容
                            # 注意：流式输出时可能 </think> 还没出来，此时也要把 <think> 到末尾的部分剔除，只显示正文
                            if is_thinking:
                                display_text = re.sub(r'<think>.*$', '', raw_text, flags=re.DOTALL).strip()
                            else:
                                display_text = re.sub(r'<think>.*?</think>', '', raw_text, flags=re.DOTALL).strip()
                    elif "<" in raw_text and "think" in raw_text and not raw_text.startswith("<think>"):
                        # 处理一种特殊情况：模型正在输出 <think> 标签的过程中（例如刚输出了 "<thin"）
                        # 此时正则表达式匹配不到完整的 "<think>"，会导致残缺的标签显示在正文中
                        # 我们做个简单拦截：如果在开头发现了不完整的标签，暂时不显示它
                        if re.match(r'^<t[hink>]*$', raw_text):
                            display_text = ""
                            is_thinking = True
                    
                    # 渲染思考过程
                    if display_thought:
                        # 使用 st.empty 的特殊方式来避免闪烁和嵌套
                        # Streamlit 无法在流式中动态切换 expander 的 expanded 状态
                        # 最好的方法是直接写一个 markdown 组件，使用 info 的样式来模拟
                        if is_thinking:
                            # 正在思考时，直接显示内容，不要用 expander
                            thought_placeholder.info(f"**🤔 思考过程 (正在思考...)**\n\n{display_thought}▌")
                        else:
                            # 思考完毕后，将 placeholder 替换为空，等待最终替换为折叠面板
                            thought_placeholder.info(f"**🤔 思考过程**\n\n{display_thought}")
                    
                    # 渲染正式回复
                    if display_text or not is_thinking:
                        cursor = "▌" if not is_thinking else ""
                        message_placeholder.markdown(display_text + cursor)
            
            # 2. 处理状态更新 (节点完成、工具结果等)
            elif event_type == "state_update":
                # state_update 可能包含多种数据，常见的是 messages 更新
                data = event.get("data", {})
                messages_update = event.get("messages", [])
                
                if not messages_update and isinstance(data, dict):
                    for node_name, node_data in data.items():
                        if isinstance(node_data, dict) and "messages" in node_data:
                            messages_update = node_data["messages"]
                
                # 如果更新中包含 messages，说明某个节点输出了完整消息
                # 但我们已经在用 token 流构建回复，这里可以用来检测工具调用结果
                if messages_update:
                    # 检查最后一条消息是否来自工具
                    last_msg = messages_update[-1] if messages_update else {}
                    if isinstance(last_msg, dict) and last_msg.get("role") == "tool":
                        tool_name = last_msg.get("name", "unknown")
                        tool_content = last_msg.get("content", "")
                        # 存储 RAG 检索结果
                        if tool_name == "search_knowledge_base":
                            # 尝试解析 tool_content，它可能是 JSON 字符串
                            sources = []
                            try:
                                if isinstance(tool_content, str):
                                    import json
                                    data = json.loads(tool_content)
                                else:
                                    data = tool_content
                                # 提取来源列表
                                if isinstance(data, dict) and "sources" in data:
                                    sources = data["sources"]
                                else:
                                    sources = [str(data)]
                            except Exception:
                                sources = [str(tool_content)]
                            rag_sources = sources
                        tool_status_placeholder.success(f"✅ 工具 {tool_name} 执行完成")
                        # 短暂显示后清除，保持界面清爽
                        import time
                        time.sleep(0.5)
                        tool_status_placeholder.empty()
            
            # 3. 处理自定义事件 (你在后端通过 get_stream_writer 发送的)
            elif event_type == "custom":
                custom_data = event.get("data", {})
                # 检查是否是完成事件
                if custom_data.get("type") == "done":
                    _show_completion_stats(custom_data)
                # 其他自定义事件，比如工具调用状态
                elif "type" in custom_data:
                    custom_type = custom_data["type"]
                    if custom_type == "tool_start":
                        tool_name = custom_data.get("tool", "unknown")
                        tool_status_placeholder.info(f"🔧 调用工具: {tool_name}...")
                    elif custom_type == "tool_end":
                        tool_name = custom_data.get("tool", "unknown")
                        tool_status_placeholder.success(f"✅ 工具 {tool_name} 完成")
                        tool_status_placeholder.empty()
                    elif "status" in custom_data:
                        status_msg = custom_data.get("status", "")
                        tool_status_placeholder.info(f"🔧 {status_msg}")
            
            # 4. 处理错误
            elif event_type == "error":
                st.error(f"❌ 错误: {event.get('message', '未知错误')}")
                break  # 发生错误时停止处理
            
            # 注意：v2 格式中没有固定的 "done" 事件，流结束即代表完成
            # 统计信息 (token_usage, elapsed_time) 通常会在最后的 state_update 中携带
            # 如果后端在最终状态里返回了这些信息，可以在此处理
        
        # 流结束后，移除光标并保存完整回复
        display_text = raw_text
        display_thought = api_thought
        
        # 最后的标签清理，以防未闭合
        if "<think>" in raw_text:
            think_match = re.search(r'<think>(.*?)(</think>|$)', raw_text, re.DOTALL)
            if think_match:
                display_thought = think_match.group(1).strip()
                display_text = re.sub(r'<think>.*?(</think>|$)', '', raw_text, flags=re.DOTALL).strip()
                
        if display_thought:
            # 只有在最终结束时，才把它放进折叠面板
            with thought_placeholder.container():
                with st.expander("🤔 思考过程", expanded=False):
                    st.markdown(display_thought)
        else:
            thought_placeholder.empty()
                    
        # 移除光标
        message_placeholder.markdown(display_text)
        
        # 显示 RAG 检索来源（如果有）
        if rag_sources:
            with st.expander("🔍 检索来源", expanded=False):
                # 格式化来源列表
                if isinstance(rag_sources, list):
                    for i, source in enumerate(rag_sources, 1):
                        if isinstance(source, dict):
                            content = source.get("page_content", source.get("content", str(source)))
                            metadata = source.get("metadata", {})
                            filename = metadata.get("filename", metadata.get("source", "未知文件"))
                            page = metadata.get("page", metadata.get("page_number", ""))
                            if page:
                                source_info = f"**来源 {i}:** {filename} (第{page}页)"
                            else:
                                source_info = f"**来源 {i}:** {filename}"
                            st.markdown(source_info)
                            # 显示内容预览（前200字符）
                            preview = content[:200] + "..." if len(content) > 200 else content
                            st.markdown(f"> {preview}")
                            st.markdown("---")
                        else:
                            st.markdown(f"**来源 {i}:** {str(source)}")
                else:
                    st.markdown(str(rag_sources))
        
        # 拼装包含思考过程的完整内容，以便后续在历史中正确渲染
        final_content = display_text
        if display_thought:
            final_content = f"```thought\n{display_thought}\n```\n\n" + display_text
            
        AppState.add_message("assistant", final_content)
        tool_status_placeholder.empty()
        
        # 消息发送完毕后，静默刷新历史记录列表
        # （因为可能生成了新对话，或者旧对话摘要已更新）
        from frontend.components.sidebar import _refresh_threads
        _refresh_threads()
        
        # 强制重绘页面，使侧边栏立即显示最新记录
        st.rerun()


def _show_completion_stats(event: dict):
    """
    显示对话完成统计信息
    
    Args:
        event: 完成事件数据
    """
    token_usage = event.get("token_usage", {})
    elapsed = event.get("elapsed_time", 0)
    
    if token_usage:
        total_tokens = token_usage.get("total_tokens", 0)
        st.caption(f"📊 消耗 {total_tokens} tokens | ⏱️ {elapsed:.2f}s")