app.py

import logging
import os
import re
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import gradio as gr
import modelscope_studio.components.antd as antd
import modelscope_studio.components.antdx as antdx
import modelscope_studio.components.base as ms
import modelscope_studio.components.pro as pro
from mem0 import Memory
from modelscope_studio.components.pro.chatbot import (ChatbotBotConfig,
                                                      ChatbotPromptsConfig,
                                                      ChatbotUserConfig,
                                                      ChatbotWelcomeConfig)
from openai import OpenAI


# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('xinyuan_chat.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

@dataclass
class AppConfig:
    """应用配置类"""
    users_file: str = "users.txt"
    memory_path: str = "./faiss_memories"
    model_name: str = "xinyuan-32b-v0609"
    max_tokens: int = 32768
    temperature: float = 0.6
    top_p: float = 0.95
    chatbot_height: int = 1000
    min_username_length: int = 3
    max_memory_results: int = 5

class MemoryManager:
    """记忆管理器"""
    
    def __init__(self, config_path: str):
        self.config = {
            "vector_store": {
                "provider": "faiss",
                "config": {
                    "collection_name": "xinyuan_memories",
                    "path": config_path,
                    "distance_strategy": "euclidean"
                }
            }
        }
        try:
            self.memory = Memory.from_config(self.config)
            logger.info(f"Memory manager initialized with path: {config_path}")
        except Exception as e:
            logger.error(f"Failed to initialize memory manager: {e}")
            raise
    
    def search_memories(self, query: str, user_id: str, limit: int = 5) -> List[Dict[str, Any]]:
        """搜索相关记忆"""
        try:
            if not query or not user_id:
                return []
            
            results = self.memory.search(query=query, user_id=user_id, limit=limit)
            if results and 'results' in results:
                return sorted(results['results'], key=lambda x: x.get('score', 0), reverse=True)
            return []
        except Exception as e:
            logger.error(f"Error searching memories for user {user_id}: {e}")
            return []
    
    def add_memory(self, messages: List[Dict[str, str]], user_id: str) -> bool:
        """添加记忆"""
        try:
            if not messages or not user_id:
                return False
            
            self.memory.add(messages, user_id=user_id)
            logger.info(f"Memory added for user {user_id}")
            return True
        except Exception as e:
            logger.error(f"Error adding memory for user {user_id}: {e}")
            return False

class UserManager:
    """用户管理器"""
    
    def __init__(self, users_file: str, min_username_length: int = 3):
        self.users_file = Path(users_file)
        self.min_username_length = min_username_length
        self._ensure_users_file_exists()
    
    def _ensure_users_file_exists(self):
        """确保用户文件存在"""
        if not self.users_file.exists():
            self.users_file.touch()
            logger.info(f"Created users file: {self.users_file}")
    
    def load_users(self) -> set:
        """加载已注册用户列表"""
        try:
            with open(self.users_file, 'r', encoding='utf-8') as f:
                users = {line.strip() for line in f if line.strip()}
            logger.debug(f"Loaded {len(users)} users")
            return users
        except Exception as e:
            logger.error(f"Error loading users: {e}")
            return set()
    
    def save_user(self, username: str) -> bool:
        """保存新用户到文件"""
        try:
            with open(self.users_file, 'a', encoding='utf-8') as f:
                f.write(f"{username}\n")
            logger.info(f"User {username} saved to file")
            return True
        except Exception as e:
            logger.error(f"Error saving user {username}: {e}")
            return False
    
    def is_valid_username(self, username: str) -> bool:
        """验证用户名是否有效"""
        if not username or not isinstance(username, str):
            return False
        
        # 检查长度
        if len(username) < self.min_username_length:
            return False
        
        # 检查格式：以字母开头，只包含字母、数字和下划线
        return bool(re.match(r'^[a-zA-Z][a-zA-Z0-9_]*$', username))
    
    def login_user(self, username: str) -> Tuple[bool, str]:
        """用户登录验证"""
        if not self.is_valid_username(username):
            return False, "用户名无效！用户名必须以英文字母开头，只能包含英文字母、数字和下划线，且长度至少3位。"
        
        users = self.load_users()
        if username in users:
            logger.info(f"User {username} logged in successfully")
            return True, f"欢迎回来，{username}！"
        else:
            logger.warning(f"Login attempt for unregistered user: {username}")
            return False, f"用户 {username} 未注册，请先注册。"
    
    def register_user(self, username: str) -> Tuple[bool, str]:
        """用户注册"""
        if not self.is_valid_username(username):
            return False, "用户名无效！用户名必须以英文字母开头，只能包含英文字母、数字和下划线，且长度至少3位。"
        
        users = self.load_users()
        if username in users:
            logger.warning(f"Registration attempt for existing user: {username}")
            return False, f"用户名 {username} 已存在，请直接登录。"
        
        if self.save_user(username):
            logger.info(f"User {username} registered successfully")
            return True, f"注册成功！欢迎，{username}！"
        else:
            return False, "注册失败，请稍后重试。"

class ChatManager:
    """聊天管理器"""
    
    def __init__(self, config: AppConfig, memory_manager: MemoryManager):
        self.config = config
        self.memory_manager = memory_manager
        self.client = self._initialize_openai_client()
    
    def _initialize_openai_client(self) -> OpenAI:
        """初始化OpenAI客户端"""
        try:
            # 可以根据需要配置API密钥和基础URL
            gw_api_key = os.getenv("GW_API_KEY")
            client = OpenAI(
                base_url='https://api.geniuworks.com/v2',
                api_key=gw_api_key,
            )
            logger.info("OpenAI client initialized successfully")
            return client
        except Exception as e:
            logger.error(f"Failed to initialize OpenAI client: {e}")
            raise
    
    def format_history(self, sender_value: str, history: List[Dict], username: Optional[str] = None) -> List[Dict[str, str]]:
        """格式化聊天历史"""
        messages = []
        
        # 添加系统提示
        if username:
            system_prompt = f"""You are Xinyuan, a large language model trained by Cylingo Group. You are a helpful assistant. 目前和你聊天的用户是{username}."""
            
            # 搜索相关记忆
            if sender_value:
                related_memories = self.memory_manager.search_memories(
                    query=sender_value, 
                    user_id=username, 
                    limit=self.config.max_memory_results
                )
                
                if related_memories:
                    memory_content = "\n相关记忆：\n"
                    for idx, memory in enumerate(related_memories):
                        memory_content += f"记忆{idx + 1}：{memory.get('memory', '')} (相关度: {memory.get('score', 0):.3f})\n"
                    system_prompt += memory_content
            
            messages.append({"role": "system", "content": system_prompt})
        
        # 添加历史对话
        for item in history:
            if item.get("role") == "user":
                messages.append({"role": "user", "content": item.get("content", "")})
            elif item.get("role") == "assistant" and item.get("content"):
                # 提取助手回复的文本内容
                content_list = item.get("content", [])
                if content_list and len(content_list) > 1:
                    assistant_content = content_list[-1].get("content", "")
                    if assistant_content:
                        messages.append({"role": "assistant", "content": assistant_content})
        
        return messages
    
    def create_chat_completion(self, messages: List[Dict[str, str]]) -> Any:
        """创建聊天完成请求"""
        try:
            return self.client.chat.completions.create(
                model=self.config.model_name,
                messages=messages,
                stream=True,
                max_tokens=self.config.max_tokens,
                temperature=self.config.temperature,
                top_p=self.config.top_p,
            )
        except Exception as e:
            logger.error(f"Error creating chat completion: {e}")
            raise

# 全局配置和管理器实例
config = AppConfig()
memory_manager = MemoryManager(config.memory_path)
user_manager = UserManager(config.users_file, config.min_username_length)
chat_manager = ChatManager(config, memory_manager)

# Gradio界面函数
def handle_auth(username: str, is_register: bool) -> Tuple:
    """处理认证逻辑"""
    try:
        if is_register:
            success, message = user_manager.register_user(username)
        else:
            success, message = user_manager.login_user(username)
        
        if success:
            return (
                gr.update(visible=False),  # 隐藏登录界面
                gr.update(visible=True),   # 显示聊天界面
                gr.update(message=message, type="success", visible=True),
                username
            )
        else:
            return (
                gr.update(visible=True),   # 保持登录界面可见
                gr.update(visible=False),  # 隐藏聊天界面
                gr.update(message=message, type="error", visible=True),
                ""
            )
    except Exception as e:
        logger.error(f"Error in handle_auth: {e}")
        return (
            gr.update(visible=True),
            gr.update(visible=False),
            gr.update(message="系统错误，请稍后重试。", type="error", visible=True),
            ""
        )

def prompt_select(e: gr.EventData) -> gr.update:
    """处理提示选择"""
    try:
        return gr.update(value=e._data["payload"][0]["value"]["description"])
    except Exception as e:
        logger.error(f"Error in prompt_select: {e}")
        return gr.update(value="")

def clear() -> gr.update:
    """清空聊天记录"""
    return gr.update(value=None)

def retry(chatbot_value: List, e: gr.EventData, username: Optional[str] = None):
    """重试功能"""
    try:
        index = e._data["payload"][0]["index"]
        chatbot_value = chatbot_value[:index]
        
        yield gr.update(value=None, loading=True), gr.update(value=chatbot_value), gr.update(disabled=True)
        
        for chunk in submit(None, chatbot_value, username):
            yield chunk
    except Exception as e:
        logger.error(f"Error in retry: {e}")
        yield gr.update(value=None, loading=False), gr.update(value=chatbot_value), gr.update(disabled=False)

def cancel(chatbot_value: List) -> Tuple:
    """取消当前对话"""
    try:
        if chatbot_value:
            chatbot_value[-1]["loading"] = False
            chatbot_value[-1]["status"] = "done"
            chatbot_value[-1]["footer"] = "Chat completion paused"
        
        return (
            gr.update(value=chatbot_value),
            gr.update(loading=False),
            gr.update(disabled=False)
        )
    except Exception as e:
        logger.error(f"Error in cancel: {e}")
        return (
            gr.update(value=chatbot_value),
            gr.update(loading=False),
            gr.update(disabled=False)
        )

def submit(sender_value: Optional[str], chatbot_value: List, username: Optional[str] = None):
    """提交聊天消息"""
    start_time = time.time()
    
    try:
        # 添加用户消息
        if sender_value is not None:
            chatbot_value.append({
                "role": "user",
                "content": sender_value,
            })
        
        # 格式化历史消息
        history_messages = chat_manager.format_history(sender_value, chatbot_value, username)
        
        # 添加助手消息占位符
        chatbot_value.append({
            "role": "assistant",
            "content": [],
            "loading": True,
            "status": "pending"
        })
        
        # 更新UI状态
        yield (
            gr.update(value=None, loading=True),  # sender
            gr.update(value=chatbot_value),       # chatbot
            gr.update(disabled=True)              # clear_btn
        )
        
        # 创建聊天完成请求
        response = chat_manager.create_chat_completion(history_messages)
        
        # 处理流式响应
        thought_done = False
        message_content = chatbot_value[-1]["content"]
        
        # 初始化消息内容结构
        message_content.append({
            "copyable": False,
            "editable": False,
            "type": "tool",
            "content": "",
            "options": {"title": "Thinking..."}
        })
        message_content.append({
            "type": "text",
            "content": "",
        })
        
        full_assistant_content = ""
        
        # 处理流式响应
        for chunk in response:
            try:
                reasoning_content = getattr(chunk.choices[0].delta, 'reasoning_content', None) or ""
                content = getattr(chunk.choices[0].delta, 'content', None) or ""
                
                chatbot_value[-1]["loading"] = False
                message_content[-2]["content"] += reasoning_content
                message_content[-1]["content"] += content
                
                if content:
                    full_assistant_content += content
                
                if content and not thought_done:
                    thought_done = True
                    thought_cost_time = f"{time.time() - start_time:.2f}"
                    message_content[-2]["options"]["title"] = f"End of Thought ({thought_cost_time}s)"
                    message_content[-2]["options"]["status"] = "done"
                
                yield (
                    gr.update(),                      # sender
                    gr.update(value=chatbot_value),   # chatbot
                    gr.update()                       # clear_btn
                )
                
            except Exception as chunk_error:
                logger.error(f"Error processing chunk: {chunk_error}")
                continue
        
        # 保存到记忆
        if username and sender_value and full_assistant_content:
            memory_messages = [
                {'role': 'user', 'content': sender_value},
                {'role': 'assistant', 'content': full_assistant_content}
            ]
            memory_manager.add_memory(memory_messages, username)
        
        # 完成响应
        total_time = f"{time.time() - start_time:.2f}s"
        chatbot_value[-1]["footer"] = total_time
        chatbot_value[-1]["status"] = "done"
        
        yield (
            gr.update(loading=False),             # sender
            gr.update(value=chatbot_value),       # chatbot
            gr.update(disabled=False)             # clear_btn
        )
        
    except Exception as e:
        logger.error(f"Error in submit: {e}")
        
        # 错误处理
        if chatbot_value:
            chatbot_value[-1]["loading"] = False
            chatbot_value[-1]["status"] = "done"
            chatbot_value[-1]["content"] = "抱歉，处理您的请求时出现错误，请稍后重试。"
        
        yield (
            gr.update(loading=False),             # sender
            gr.update(value=chatbot_value),       # chatbot
            gr.update(disabled=False)             # clear_btn
        )

# 创建Gradio界面
def create_interface():
    """创建Gradio界面"""
    with gr.Blocks(title="Xinyuan 聊天助手") as demo, ms.Application(), antdx.XProvider():
        # 状态变量
        current_user = gr.State("")
        
        # 登录界面
        with antd.Flex(vertical=True, gap="large", elem_id="login_container") as login_container:
            with antd.Card(title="欢迎使用 Xinyuan 聊天助手"):
                with antd.Flex(vertical=True, gap="middle"):
                    antd.Typography.Title("用户登录/注册", level=3)
                    antd.Typography.Text("请输入您的英文用户名（3位以上，仅支持英文字母、数字和下划线）")
                    
                    username_input = antd.Input(
                        placeholder="请输入用户名（如：john_doe）",
                        size="large"
                    )
                    
                    with antd.Flex(gap="small"):
                        login_btn = antd.Button("登录", type="primary", size="large")
                        register_btn = antd.Button("注册", size="large")
                    
                    auth_message = antd.Alert(
                        message="请输入用户名",
                        type="info",
                        visible=False
                    )
        
        # 聊天界面
        with antd.Flex(vertical=True, gap="middle", visible=False) as chat_container:
            # 用户信息栏
            with antd.Flex(justify="space-between", align="center"):
                user_info = gr.Markdown("")
                logout_btn = antd.Button("退出登录", size="small")
            
            # 聊天机器人组件
            chatbot = pro.Chatbot(
                height=config.chatbot_height,
                welcome_config=ChatbotWelcomeConfig(
                    variant="borderless",
                    icon="./xinyuan.png",
                    title="Hello, I'm Xinyuan👋",
                    description="You can input text to get started.",
                    prompts=ChatbotPromptsConfig(
                        title="How can I help you today?",
                        styles={
                            "list": {"width": '100%'},
                            "item": {"flex": 1},
                        },
                        items=[
                            {
                                "label": "💝 心理学与实际应用",
                                "children": [
                                    {"description": "课题分离是什么意思？"},
                                    {"description": "回避型依恋和焦虑型依恋有什么区别？还有其他依恋类型吗？"},
                                    {"description": "为什么我背单词的时候总是只记得开头和结尾，中间全忘了？"}
                                ]
                            },
                            {
                                "label": "👪 儿童教育与发展",
                                "children": [
                                    {"description": "什么是正念养育？"},
                                    {"description": "2岁孩子分离焦虑严重，送托育中心天天哭闹怎么办？"},
                                    {"description": "4岁娃说话不清还爱打人，是心理问题还是欠管教？"}
                                ]
                            }
                        ]
                    )
                ),
                user_config=ChatbotUserConfig(
                    avatar="https://api.dicebear.com/7.x/miniavs/svg?seed=3",
                    variant="shadow"
                ),
                bot_config=ChatbotBotConfig(
                    header='Xinyuan',
                    avatar="./xinyuan.png",
                    actions=["copy", "retry"],
                    variant="shadow"
                ),
            )
            
            # 发送器组件
            with antdx.Sender() as sender:
                with ms.Slot("prefix"):
                    with antd.Button(value=None, color="default", variant="text") as clear_btn:
                        with ms.Slot("icon"):
                            antd.Icon("ClearOutlined")
        
        # 事件处理函数
        def handle_login(username: str):
            return handle_auth(username, False)
        
        def handle_register(username: str):
            return handle_auth(username, True)
        
        def handle_logout():
            return (
                gr.update(visible=True),   # 显示登录界面
                gr.update(visible=False),  # 隐藏聊天界面
                gr.update(message="已退出登录", type="info", visible=True),
                gr.update(value=""),       # 清空用户名输入
                "",                        # 清空用户信息显示
                ""                         # 清空当前用户状态
            )
        
        def update_user_info(username: str) -> str:
            return f"**当前用户: {username}**" if username else ""
        
        # 绑定事件
        login_btn.click(
            fn=handle_login,
            inputs=[username_input],
            outputs=[login_container, chat_container, auth_message, current_user]
        ).then(
            fn=update_user_info,
            inputs=[current_user],
            outputs=[user_info]
        )
        
        register_btn.click(
            fn=handle_register,
            inputs=[username_input],
            outputs=[login_container, chat_container, auth_message, current_user]
        ).then(
            fn=update_user_info,
            inputs=[current_user],
            outputs=[user_info]
        )
        
        logout_btn.click(
            fn=handle_logout,
            outputs=[login_container, chat_container, auth_message, username_input, user_info, current_user]
        )
        
        # 聊天功能事件绑定
        clear_btn.click(fn=clear, outputs=[chatbot])
        
        submit_event = sender.submit(
            fn=submit,
            inputs=[sender, chatbot, current_user],
            outputs=[sender, chatbot, clear_btn]
        )
        
        sender.cancel(
            fn=cancel,
            inputs=[chatbot],
            outputs=[chatbot, sender, clear_btn],
            cancels=[submit_event],
            queue=False
        )
        
        chatbot.retry(
            fn=retry,
            inputs=[chatbot, current_user],
            outputs=[sender, chatbot, clear_btn]
        )
        
        chatbot.welcome_prompt_select(fn=prompt_select, outputs=[sender])
    
    return demo

def main():
    """主函数"""
    try:
        logger.info("Starting Xinyuan Chat Application")
        demo = create_interface()
        demo.queue().launch()
    except Exception as e:
        logger.error(f"Failed to start application: {e}")
        raise

if __name__ == "__main__":
    main()