余明辉 commited on
Commit
95cf203
·
1 Parent(s): 452b700
Files changed (1) hide show
  1. agent_api_web_demo.py +0 -199
agent_api_web_demo.py DELETED
@@ -1,199 +0,0 @@
1
-
2
- import copy
3
- import os
4
- from typing import List
5
- import streamlit as st
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
- # model_name = st.sidebar.text_input('模型名称:', value='internlm/internlm2_5-7b-chat')
62
-
63
-
64
- # ================================== 硅基流动的API ==================================
65
- # 注意,如果采用硅基流动API,模型名称需要更改为:internlm/internlm2_5-7b-chat 或者 internlm/internlm2_5-20b-chat
66
- # api_base = st.sidebar.text_input(
67
- # 'API Base 地址:', value='https://api.siliconflow.cn/v1/chat/completions'
68
- # )
69
- # ================================== 浦语官方的API ==================================
70
- api_base = st.sidebar.text_input(
71
- 'API Base 地址:', value='https://internlm-chat.intern-ai.org.cn/puyu/api/v1/chat/completions'
72
- )
73
- # ==================================================================================
74
- # 插件选择
75
- plugin_name = st.sidebar.multiselect(
76
- '插件选择',
77
- options=list(st.session_state['plugin_map'].keys()),
78
- default=[],
79
- )
80
-
81
- # 根据选择的插件生成插件操作列表
82
- self.plugin_action = [st.session_state['plugin_map'][name] for name in plugin_name]
83
-
84
- # 动态生成插件提示
85
- if self.plugin_action:
86
- self.plugin_prompt = get_plugin_prompt(self.plugin_action)
87
-
88
- # 清空对话按钮
89
- if st.sidebar.button('清空对话', key='clear'):
90
- self.session_state.clear_state()
91
-
92
- return model_name, api_base, self.plugin_action
93
-
94
- def initialize_chatbot(self, model_name, api_base, plugin_action):
95
- """初始化 GPTAPI 实例作为 chatbot。"""
96
- token = os.getenv("token")
97
- if not token:
98
- st.error("未检测到环境变量 `token`,请设置环境变量,例如 `export token='your_token_here'` 后重新运行 X﹏X")
99
- st.stop() # 停止运行应用
100
-
101
- # 创建完整的 meta_prompt,保留原始结构并动态插入侧边栏配置
102
- meta_prompt = [
103
- {"role": "system", "content": self.meta_prompt, "api_role": "system"},
104
- {"role": "user", "content": "", "api_role": "user"},
105
- {"role": "assistant", "content": self.plugin_prompt, "api_role": "assistant"},
106
- {"role": "environment", "content": "", "api_role": "environment"}
107
- ]
108
-
109
- api_model = GPTAPI(
110
- model_type=model_name,
111
- api_base=api_base,
112
- key=token, # 从环境变量中获取授权令牌
113
- meta_template=meta_prompt,
114
- max_new_tokens=512,
115
- temperature=0.8,
116
- top_p=0.9
117
- )
118
- return api_model
119
-
120
- def render_user(self, prompt: str):
121
- """渲染用户输入内容。"""
122
- with st.chat_message('user'):
123
- st.markdown(prompt)
124
-
125
- def render_assistant(self, agent_return):
126
- """渲染助手响应内容。"""
127
- with st.chat_message('assistant'):
128
- content = getattr(agent_return, "content", str(agent_return))
129
- st.markdown(content if isinstance(content, str) else str(content))
130
-
131
-
132
- def main():
133
- """主函数,运行 Streamlit 应用。"""
134
- if 'ui' not in st.session_state:
135
- session_state = SessionState()
136
- session_state.init_state()
137
- st.session_state['ui'] = StreamlitUI(session_state)
138
- else:
139
- # st.set_page_config(
140
- # layout='wide',
141
- # page_title='lagent-web',
142
- # page_icon='./docs/imgs/lagent_icon.png'
143
- # )
144
- st.header(':robot_face: :blue[Lagent] Web Demo ', divider='rainbow')
145
-
146
- # 设置侧边栏并获取模型和插件信息
147
- model_name, api_base, plugin_action = st.session_state['ui'].setup_sidebar()
148
- plugins = [dict(type=f"lagent.actions.{plugin.__class__.__name__}") for plugin in plugin_action]
149
-
150
- if (
151
- 'chatbot' not in st.session_state or
152
- model_name != st.session_state['chatbot'].model_type or
153
- 'last_plugin_action' not in st.session_state or
154
- plugin_action != st.session_state['last_plugin_action'] or
155
- api_base != st.session_state['api_base']
156
- ):
157
- # 更新 Chatbot
158
- st.session_state['chatbot'] = st.session_state['ui'].initialize_chatbot(model_name, api_base, plugin_action)
159
- st.session_state['last_plugin_action'] = plugin_action # 更新插件状态
160
- st.session_state['api_base'] = api_base # 更新 API Base 地址
161
-
162
- # 初始化 AgentForInternLM
163
- st.session_state['agent'] = AgentForInternLM(
164
- llm=st.session_state['chatbot'],
165
- plugins=plugins,
166
- output_format=dict(
167
- type=PluginParser,
168
- template=PLUGIN_CN,
169
- prompt=get_plugin_prompt(plugin_action)
170
- )
171
- )
172
- # 清空对话历史
173
- st.session_state['session_history'] = []
174
-
175
- if 'agent' not in st.session_state:
176
- st.session_state['agent'] = None
177
-
178
- agent = st.session_state['agent']
179
- for prompt, agent_return in zip(st.session_state['user'], st.session_state['assistant']):
180
- st.session_state['ui'].render_user(prompt)
181
- st.session_state['ui'].render_assistant(agent_return)
182
-
183
- # 处理用户输入
184
- if user_input := st.chat_input(''):
185
- st.session_state['ui'].render_user(user_input)
186
-
187
- # 调用模型时确保侧边栏的系统提示词和插件提示词生效
188
- res = agent(user_input, session_id=0)
189
- st.session_state['ui'].render_assistant(res)
190
-
191
- # 更新会话状态
192
- st.session_state['user'].append(user_input)
193
- st.session_state['assistant'].append(copy.deepcopy(res))
194
-
195
- st.session_state['last_status'] = None
196
-
197
-
198
- if __name__ == '__main__':
199
- main()