Spaces:
Running
Running
zencorn
commited on
Commit
·
3634d3e
1
Parent(s):
ac819bc
Updates Files
Browse files
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()
|