Create app.py

app.py

@@ -0,0 +1,437 @@
+import os
+import re
+import time
+
+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
+
+config = {
+    "vector_store": {
+        "provider": "faiss",
+        "config": {
+            "collection_name": "test",
+            "path": "./faiss_memories",
+            "distance_strategy": "euclidean"
+        }
+    }
+}
+
+m = Memory.from_config(config)
+
+gw_api_key = os.getenv("GW_API_KEY")
+
+client = OpenAI(
+    base_url='https://api.geniuworks.com/v2',
+    api_key=gw_api_key, 
+)
+
+model = "xinyuan-32b-v0609"
+# model = "gpt-4.1-2025-04-14"
+
+# 用户管理相关函数
+USERS_FILE = "users.txt"
+
+def load_users():
+    """加载已注册用户列表"""
+    if not os.path.exists(USERS_FILE):
+        return set()
+    with open(USERS_FILE, 'r', encoding='utf-8') as f:
+        return set(line.strip() for line in f if line.strip())
+
+def save_user(username):
+    """保存新用户到文件"""
+    with open(USERS_FILE, 'a', encoding='utf-8') as f:
+        f.write(username + '\n')
+
+def is_valid_username(username):
+    """验证用户名是否有效（仅英文字母和数字）"""
+    if not username:
+        return False
+    return bool(re.match(r'^[a-zA-Z][a-zA-Z0-9_]*$', username)) and len(username) >= 3
+
+def login_user(username):
+    """用户登录验证"""
+    if not is_valid_username(username):
+        return False, "用户名无效！用户名必须以英文字母开头，只能包含英文字母、数字和下划线，且长度至少3位。"
+    
+    users = load_users()
+    if username in users:
+        return True, f"欢迎回来，{username}！"
+    else:
+        return False, f"用户 {username} 未注册，请先注册。"
+
+def register_user(username):
+    """用户注册"""
+    if not is_valid_username(username):
+        return False, "用户名无效！用户名必须以英文字母开头，只能包含英文字母、数字和下划线，且长度至少3位。"
+    
+    users = load_users()
+    if username in users:
+        return False, f"用户名 {username} 已存在，请直接登录。"
+    
+    save_user(username)
+    return True, f"注册成功！欢迎，{username}！"
+
+def handle_auth(username, is_register):
+    """处理认证逻辑"""
+    if is_register:
+        success, message = register_user(username)
+    else:
+        success, message = 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),  # 显示错误消息
+            ""
+        )
+
+def prompt_select(e: gr.EventData):
+    return gr.update(value=e._data["payload"][0]["value"]["description"])
+
+
+def clear():
+    return gr.update(value=None)
+
+
+def retry(chatbot_value, e: gr.EventData, username=None):
+    index = e._data["payload"][0]["index"]
+    chatbot_value = chatbot_value[:index]
+
+    yield gr.update(loading=True), gr.update(value=chatbot_value), gr.update(
+        disabled=True)
+    for chunk in submit(None, chatbot_value, username):
+        yield chunk
+
+
+def cancel(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)
+
+
+def format_history(sender_value, history, username=None):
+    messages = []
+    
+    # 添加系统提示，包含用户名信息
+    if username:
+        system_prompt = f"""You are Xinyuan, a large language model trained by Cylingo Group. You are a helpful assistant. 目前和你聊天的用户是{username}."""
+        messages.append({"role": "system", "content": system_prompt})
+    
+    related_memories = m.search(query=sender_value, user_id=username)
+    print(related_memories)
+    
+    related_memories_content = ""
+
+# {'results': [{'id': '8de25384-f210-4442-a04f-cd6c7796a5b7', 'memory': 'Loves sci-fi movies', 'hash': '1110b1af77367917ea2022355a16f187', 'metadata': None, 'score': 0.1812809524839618, 'created_at': '2025-08-05T23:54:13.694114-07:00', 'updated_at': None, 'user_id': 'alice'}, {'id': 'a4aa36b6-0595-492c-b6b1-5013511820d1', 'memory': 'Not a big fan of thriller movies', 'hash': '028dfab4483f28980e292f62578d3293', 'metadata': None, 'score': 0.17128575336629281, 'created_at': '2025-08-05T23:54:13.691791-07:00', 'updated_at': None, 'user_id': 'alice'}, {'id': 'a736ea22-3042-4275-ab9b-596324348119', 'memory': 'Planning to watch a movie tonight', 'hash': 'bf55418607cfdca4afa311b5fd8496bd', 'metadata': None, 'score': 0.1213398963070364, 'created_at': '2025-08-05T23:54:13.687585-07:00', 'updated_at': None, 'user_id': 'alice'}]}
+    # 将related_memories按照score排序
+    if related_memories and 'results' in related_memories:
+        related_memories_list = sorted(related_memories['results'], key=lambda x: x['score'], reverse=True)
+        
+        for id, item in enumerate(related_memories_list):
+            # 将score添加到memory中
+            related_memories_content += f"相关记忆{id}：\n内容：{item['memory']}\n相关度：{item['score']}\n\n"
+    if related_memories_content:
+        system_prompt += f"\n相关记忆：\n{related_memories_content}"
+        messages.insert(0, {"role": "system", "content": system_prompt})
+    
+    for item in history:
+        if item["role"] == "user":
+            messages.append({"role": "user", "content": item["content"]})
+        elif item["role"] == "assistant":
+            # ignore thought message
+            messages.append({
+                "role": "assistant",
+                "content": item["content"][-1]["content"]
+            })
+
+        
+    print(related_memories)
+    print(messages)
+    return messages
+
+
+def submit(sender_value, chatbot_value, username=None):
+    if sender_value is not None:
+        chatbot_value.append({
+            "role": "user",
+            "content": sender_value,
+        })
+    history_messages = format_history(sender_value, chatbot_value, username)
+    chatbot_value.append({
+        "role": "assistant",
+        "content": [],
+        "loading": True,
+        "status": "pending"
+    })
+    yield {
+        sender: gr.update(value=None, loading=True),
+        clear_btn: gr.update(disabled=True),
+        chatbot: gr.update(value=chatbot_value)
+    }
+
+    try:
+        response = client.chat.completions.create(model=model,
+                                                    messages=history_messages,
+                                                    stream=True,
+                                                    max_tokens=32768,
+                                                    temperature=0.6,
+                                                    top_p=0.95,
+                                                    )
+        
+        thought_done = False
+        start_time = time.time()
+        message_content = chatbot_value[-1]["content"]
+        # thought content
+        message_content.append({
+            "copyable": False,
+            "editable": False,
+            "type": "tool",
+            "content": "",
+            "options": {
+                "title": "Thinking..."
+            }
+        })
+        # content
+        message_content.append({
+            "type": "text",
+            "content": "",
+        })
+        
+        # 收集完整的助手响应内容用于保存到内存
+        full_assistant_content = ""
+        
+        for chunk in response:
+            try:
+                reasoning_content = chunk.choices[0].delta.reasoning_content
+            except:
+                reasoning_content = ""
+            try:
+                content = chunk.choices[0].delta.content
+            except:
+                content = ""
+            chatbot_value[-1]["loading"] = False
+            message_content[-2]["content"] += reasoning_content or ""
+            message_content[-1]["content"] += content or ""
+            
+            # 收集助手的实际回复内容（不包括思考过程）
+            if content:
+                full_assistant_content += content
+
+            if content and not thought_done:
+                thought_done = True
+                thought_cost_time = "{:.2f}".format(time.time() - start_time)
+                message_content[-2]["options"][
+                    "title"] = f"End of Thought ({thought_cost_time}s)"
+                message_content[-2]["options"]["status"] = "done"
+
+            yield {chatbot: gr.update(value=chatbot_value)}
+
+        # 在流式响应完成后保存到内存
+        if username and sender_value and full_assistant_content:
+            memory_messages = [
+                {'role': 'user', 'content': sender_value},
+                {'role': 'assistant', 'content': full_assistant_content}
+            ]
+            m.add(memory_messages, user_id=username)
+
+        chatbot_value[-1]["footer"] = "{:.2f}".format(time.time() -
+                                                        start_time) + 's'
+        chatbot_value[-1]["status"] = "done"
+        yield {
+            clear_btn: gr.update(disabled=False),
+            sender: gr.update(loading=False),
+            chatbot: gr.update(value=chatbot_value),
+        }
+    except Exception as e:
+        chatbot_value[-1]["loading"] = False
+        chatbot_value[-1]["status"] = "done"
+        chatbot_value[-1]["content"] = "Failed to respond, please try again."
+        yield {
+            clear_btn: gr.update(disabled=False),
+            sender: gr.update(loading=False),
+            chatbot: gr.update(value=chatbot_value),
+        }
+        raise e
+
+
+with gr.Blocks() 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=1000,
+            welcome_config=ChatbotWelcomeConfig(
+                variant="borderless",
+                icon="./xinyuan.png",
+                title=f"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):
+        return handle_auth(username, False)
+    
+    def handle_register(username):
+        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):
+        if username:
+            return f"**当前用户: {username}**"
+        return ""
+    
+    # 登录按钮事件
+    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])
+
+if __name__ == "__main__":
+    demo.queue().launch()