CCCasEEE/internlm_lagent

1import copy2import os3from typing import List4import streamlit as st5from lagent.actions import ArxivSearch, WeatherQuery6from lagent.prompts.parsers import PluginParser7from lagent.agents.stream import INTERPRETER_CN, META_CN, PLUGIN_CN, AgentForInternLM, get_plugin_prompt8from lagent.llms import GPTAPI9 10class SessionState:11    """管理会话状态的类。"""12 13    def init_state(self):14        """初始化会话状态变量。"""15        st.session_state['assistant'] = []  # 助手消息历史16        st.session_state['user'] = []  # 用户消息历史17        # 初始化插件列表18        action_list = [19            ArxivSearch(),20            WeatherQuery(),21        ]22        st.session_state['plugin_map'] = {action.name: action for action in action_list}23        st.session_state['model_map'] = {}  # 存储模型实例24        st.session_state['model_selected'] = None  # 当前选定模型25        st.session_state['plugin_actions'] = set()  # 当前激活插件26        st.session_state['history'] = []  # 聊天历史27        st.session_state['api_base'] = None  # 初始化API base地址28 29    def clear_state(self):30        """清除当前会话状态。"""31        st.session_state['assistant'] = []32        st.session_state['user'] = []33        st.session_state['model_selected'] = None34 35 36class StreamlitUI:37    """管理 Streamlit 界面的类。"""38 39    def __init__(self, session_state: SessionState):40        self.session_state = session_state41        self.plugin_action = []  # 当前选定的插件42        # 初始化提示词43        self.meta_prompt = META_CN44        self.plugin_prompt = PLUGIN_CN45        self.init_streamlit()46 47    def init_streamlit(self):48        """初始化 Streamlit 的 UI 设置。"""49        # st.set_page_config(50        #     layout='wide',51        #     page_title='lagent-web',52        #     page_icon='./docs/imgs/lagent_icon.png'53        # )54        st.header(':robot_face: :blue[Lagent] Web Demo ', divider='rainbow')55 56    def setup_sidebar(self):57        """设置侧边栏，选择模型和插件。"""58        # 模型名称和 API Base 输入框59        model_name = st.sidebar.text_input('模型名称：', value='internlm2.5-latest')60        61        # ================================== 硅基流动的API ==================================62        # 注意，如果采用硅基流动API，模型名称需要更改为：internlm/internlm2_5-7b-chat 或者 internlm/internlm2_5-20b-chat63        # api_base = st.sidebar.text_input(64        #     'API Base 地址：', value='https://api.siliconflow.cn/v1/chat/completions'65        # )66        # ================================== 浦语官方的API ==================================67        api_base = st.sidebar.text_input(68            'API Base 地址：', value='https://internlm-chat.intern-ai.org.cn/puyu/api/v1/chat/completions'69        )70        # ==================================================================================71        # 插件选择72        plugin_name = st.sidebar.multiselect(73            '插件选择',74            options=list(st.session_state['plugin_map'].keys()),75            default=[],76        )77 78        # 根据选择的插件生成插件操作列表79        self.plugin_action = [st.session_state['plugin_map'][name] for name in plugin_name]80 81        # 动态生成插件提示82        if self.plugin_action:83            self.plugin_prompt = get_plugin_prompt(self.plugin_action)84 85        # 清空对话按钮86        if st.sidebar.button('清空对话', key='clear'):87            self.session_state.clear_state()88 89        return model_name, api_base, self.plugin_action90 91    def initialize_chatbot(self, model_name, api_base, plugin_action):92        """初始化 GPTAPI 实例作为 chatbot。"""93        token = os.getenv("token")94        if not token:95            st.error("未检测到环境变量 `token`，请设置环境变量，例如 `export token='your_token_here'` 后重新运行 X﹏X")96            st.stop()  # 停止运行应用97            98        # 创建完整的 meta_prompt，保留原始结构并动态插入侧边栏配置99        meta_prompt = [100            {"role": "system", "content": self.meta_prompt, "api_role": "system"},101            {"role": "user", "content": "", "api_role": "user"},102            {"role": "assistant", "content": self.plugin_prompt, "api_role": "assistant"},103            {"role": "environment", "content": "", "api_role": "environment"}104        ]105 106        api_model = GPTAPI(107            model_type=model_name,108            api_base=api_base,109            key=token,  # 从环境变量中获取授权令牌110            meta_template=meta_prompt,111            max_new_tokens=512,112            temperature=0.8,113            top_p=0.9114        )115        return api_model116 117    def render_user(self, prompt: str):118        """渲染用户输入内容。"""119        with st.chat_message('user'):120            st.markdown(prompt)121 122    def render_assistant(self, agent_return):123        """渲染助手响应内容。"""124        with st.chat_message('assistant'):125            content = getattr(agent_return, "content", str(agent_return))126            st.markdown(content if isinstance(content, str) else str(content))127 128 129def main():130    """主函数，运行 Streamlit 应用。"""131    if 'ui' not in st.session_state:132        session_state = SessionState()133        session_state.init_state()134        st.session_state['ui'] = StreamlitUI(session_state)135    else:136        # st.set_page_config(137        #     layout='wide',138        #     page_title='lagent-web',139        #     page_icon='./docs/imgs/lagent_icon.png'140        # )141        st.header(':robot_face: :blue[Lagent] Web Demo ', divider='rainbow')142 143    # 设置侧边栏并获取模型和插件信息144    model_name, api_base, plugin_action = st.session_state['ui'].setup_sidebar()145    plugins = [dict(type=f"lagent.actions.{plugin.__class__.__name__}") for plugin in plugin_action]146 147    if (148        'chatbot' not in st.session_state or149        model_name != st.session_state['chatbot'].model_type or150        'last_plugin_action' not in st.session_state or151        plugin_action != st.session_state['last_plugin_action'] or152        api_base != st.session_state['api_base']    153    ):154        # 更新 Chatbot155        st.session_state['chatbot'] = st.session_state['ui'].initialize_chatbot(model_name, api_base, plugin_action)156        st.session_state['last_plugin_action'] = plugin_action  # 更新插件状态157        st.session_state['api_base'] = api_base  # 更新 API Base 地址158 159        # 初始化 AgentForInternLM160        st.session_state['agent'] = AgentForInternLM(161            llm=st.session_state['chatbot'],162            plugins=plugins,163            output_format=dict(164                type=PluginParser,165                template=PLUGIN_CN,166                prompt=get_plugin_prompt(plugin_action)167            )168        )169        # 清空对话历史170        st.session_state['session_history'] = []171 172    if 'agent' not in st.session_state:173        st.session_state['agent'] = None174 175    agent = st.session_state['agent']176    for prompt, agent_return in zip(st.session_state['user'], st.session_state['assistant']):177        st.session_state['ui'].render_user(prompt)178        st.session_state['ui'].render_assistant(agent_return)179 180    # 处理用户输入181    if user_input := st.chat_input(''):182        st.session_state['ui'].render_user(user_input)183 184        # 调用模型时确保侧边栏的系统提示词和插件提示词生效185        res = agent(user_input, session_id=0)186        st.session_state['ui'].render_assistant(res)187 188        # 更新会话状态189        st.session_state['user'].append(user_input)190        st.session_state['assistant'].append(copy.deepcopy(res))191 192    st.session_state['last_status'] = None193 194 195if __name__ == '__main__':196    main()197