zencorn commited on
Commit
3634d3e
·
1 Parent(s): ac819bc

Updates Files

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