import sys
import os
import argparse
from gradio_base import WebUI, UIHelper, PORT, HOST, Client
from gradio_config import GradioConfig as gc
from typing import List, Tuple, Any
import gradio as gr
import time
from Agent import Agent
from design_states import get_desgin_states,get_cot_result
from gen_utils import *
from utils import get_embedding,cos_sim
import torch
import json
import openai
def get_embedding(sentence,api_key):
openai.api_key = api_key
embedding_model = openai.Embedding
embed = embedding_model.create(
model="text-embedding-ada-002",
input=sentence
)
embed = embed["data"][0]["embedding"]
embed = torch.tensor(embed,dtype=torch.float32)
if len(embed.shape)==1:
embed = embed.unsqueeze(0)
return embed
class GeneralUI(WebUI):
def render_and_register_ui(self):
# bind the agent with avatar
self.agent_name:list = [self.cache["agents_name"]] if isinstance(self.cache["agents_name"], str) else self.cache['agents_name']
gc.add_agent(self.agent_name)
def handle_message(self, history, state, agent_name, token, node_name):
if state % 10 == 0:
self.data_history.append({agent_name: token})
elif state % 10 == 1:
# Same state. Need to add new bubble in same bubble.
self.data_history[-1][agent_name] += token
elif state % 10 == 2:
# New state. Need to add new bubble.
history.append([None, ""])
self.data_history.clear()
self.data_history.append({agent_name: token})
else:
assert False, "Invalid state."
render_data = self.render_bubble(history, self.data_history, node_name, render_node_name= True)
return render_data
def __init__(
self,
client_cmd: list,
socket_host: str = HOST,
socket_port: int = PORT,
bufsize: int = 1024,
ui_name: str = "GeneralUI"
):
super(GeneralUI, self).__init__(client_cmd, socket_host, socket_port, bufsize, ui_name)
self.first_recieve_from_client()
self.current_node_name = ""
self.data_history = None
for _ in ['agents_name', 'api_key']:
assert _ in self.cache
def generate_sop(self,api_key,proxy,target):
os.environ["API_KEY"] = api_key
# os.environ["PROXY"] = proxy
self.design_assistant = "An assistant that can help users create content such as articles, blogs, advertising copy, etc"
self.tutor = "A tutor who provides personalized learning resources for students to help them understand complex concepts and problems"
self.online_medical_consultant = "An online medical consultant who offers preliminary medical advice to patients and answers common questions about diseases, symptoms, and treatments."
self.online_legal_consultant = "An online legal advisor who can respond to inquiries related to legal matters, providing basic legal information and advice."
self.online_financial_advisor = "An online financial advisor who can analyze financial markets and data, offering investment advice and market forecasts to users."
self.virtual_tour_guide = "A virtual tour guide providing destination information, travel recommendations, and virtual travel experiences for travelers."
self.design_assistant = get_embedding(self.design_assistant,api_key)
self.tutor = get_embedding(self.tutor,api_key)
self.online_medical_consultant = get_embedding(self.online_medical_consultant,api_key)
self.online_legal_consultant = get_embedding(self.online_legal_consultant,api_key)
self.online_financial_advisor = get_embedding(self.online_financial_advisor,api_key)
self.virtual_tour_guide = get_embedding(self.virtual_tour_guide,api_key)
self.embeddings = torch.cat([self.design_assistant,self.tutor,self.online_medical_consultant,self.online_legal_consultant,self.online_financial_advisor,self.virtual_tour_guide],dim = 0)
self.SOP["config"]["API_KEY"] = api_key
# self.SOP["config"]["PROXY"] = proxy
target_tensor = get_embedding(target,api_key)
sim_scores = cos_sim(target_tensor, self.embeddings)[0]
top_k_score, top_k_idx = torch.topk(sim_scores,k = 1)
if top_k_score > 0.7:
index = top_k_idx
else:
index = 0
target = get_cot_result(target)
design_states = get_desgin_states(target,index)
root = design_states[0]["state_name"]
agents = get_agents(design_states)
relations = get_relations(design_states)
states = gen_states(design_states)
for state_name,state_dict in states.items():
state_dict["begin_role"] = list(agents.keys())[0]
state_dict["begin_query"] = "Now that we are in the **{}**, I'm glad to offer you assistance.".format(state_name)
self.SOP["root"] = root
self.SOP["relations"] = relations
self.SOP["agents"] = agents
self.SOP["states"] = states
# 将字典写入JSON文件
print(self.SOP)
file_name = 'generated_sop.json'
with open(file_name, "w",encoding="utf-8") as json_file:
json.dump(self.SOP, json_file ,indent=4,ensure_ascii=False)
return file_name
def load_sop_fn(self,sop):
return sop.name
def construct_ui(self):
with gr.Blocks(css=gc.CSS) as demo:
with gr.Tab(label="SOP generation") as tab1:
self.SOP = {
"config": {
"API_KEY": "sk-********",
"MAX_CHAT_HISTORY": "5",
"User_Names": '["User"]',
},
"root": "state1",
"relations": {
"state1": {"0": "state1", "1": "state2"},
"state2": {"0": "state2", "1": "end_state"},
},
"agents": None,
"states": None,
}
gr.Markdown("""# Generate Agent""")
with gr.Row():
self.api_key_sop_generation = gr.Textbox(label="api_key")
self.proxy_sop_generation = gr.Textbox(label="proxy",visible=False)
with gr.Row():
self.requirement_sop_generation = gr.Textbox(value ="a shopping assistant help customer to buy the commodity",label="requirement")
with gr.Row():
self.generated_sop = gr.File(label="generated_file")
self.generate_button = gr.Button(label="Generate")
self.generate_button.click(fn = self.generate_sop,inputs=[self.api_key_sop_generation,self.proxy_sop_generation,self.requirement_sop_generation],outputs=[self.generated_sop])
with gr.Tab(label="Chat") as tab2:
uploaded_sop = gr.State()
with gr.Row():
sop = gr.File(label="upload your custmized SOP")
load_sop_btn = gr.Button(value="Load SOP")
load_sop_btn.click(self.load_sop_fn, sop,uploaded_sop)
with gr.Column():
self.radio_mode = gr.Radio(
[Client.SINGLE_MODE],
label = Client.MODE_LABEL,
info = Client.MODE_INFO,
value= Client.SINGLE_MODE,
interactive=True
# label="Select the execution mode",
# info="Single mode refers to when the current agent output ends, it will stop running until you click to continue. Auto mode refers to when you complete the input, all agents will continue to output until the task ends."
)
self.text_api = gr.Textbox(
value = self.cache["api_key"],
placeholder="openai key",
label="Please input valid openai key for gpt-3.5-turbo-16k."
)
self.btn_start = gr.Button(
value="Start😁(Click here to start!)",
)
self.chatbot = gr.Chatbot(
elem_id="chatbot1",
label="Dialog",
visible=False,
height=700
)
self.btn_next = gr.Button(
value="Next Agent Start",
visible=False
)
with gr.Row():
self.text_input = gr.Textbox(
placeholder="Please enter your content.",
label="Input",
scale=9,
visible=False
)
self.btn_send = gr.Button(
value="Send",
visible=False
)
self.btn_reset = gr.Button(
value="Restart",
visible=False
)
all_components = [self.btn_start, self.btn_send, self.btn_reset, self.chatbot, self.text_input, self.btn_next]
self.btn_start.click(
fn = self.btn_start_when_click,
inputs=[self.radio_mode, self.text_api,uploaded_sop],
outputs=[self.btn_start, self.btn_send, self.btn_reset, self.chatbot, self.text_input, self.btn_next, self.radio_mode, self.text_api]
).then(
fn = self.btn_start_after_click,
inputs=[self.chatbot],
outputs=all_components
)
self.btn_send.click(
fn=self.btn_send_when_click,
inputs=[self.text_input, self.chatbot],
outputs=all_components
).then(
fn=self.btn_send_after_click,
inputs=[self.text_input, self.chatbot],
outputs=all_components
)
self.text_input.submit(
fn=self.btn_send_when_click,
inputs=[self.text_input, self.chatbot],
outputs=all_components
).then(
fn=self.btn_send_after_click,
inputs=[self.text_input, self.chatbot],
outputs=all_components
)
self.btn_reset.click(
fn=self.btn_reset_when_click,
inputs=[],
outputs=all_components
).then(
fn=self.btn_reset_after_click,
inputs=[],
outputs=[self.btn_start, self.btn_send, self.btn_reset, self.chatbot, self.text_input, self.btn_next, self.radio_mode, self.text_api]
)
self.btn_next.click(
fn=self.btn_next_when_click,
inputs=[self.chatbot],
outputs=all_components
).then(
fn=self.btn_next_after_click,
inputs=[self.chatbot],
outputs=all_components
)
self.demo = demo
def btn_start_when_click(self, mode, api,sop):
"""
inputs=[mode, api]
outputs=[self.btn_start, self.btn_send, self.btn_reset, self.chatbot, self.text_input, self.btn_next, self.radio_mode]
"""
print("server: send ", mode, api)
self.send_start_cmd({"mode": mode, "api_key":api,"uploaded_sop": sop})
agents,roles_to_names,names_to_roles = Agent.from_config(str(sop))
agents_name = []
for i in names_to_roles :
for j in names_to_roles[i]:
agents_name.append(j+"("+names_to_roles[i][j]+")")
self.new_render_and_register_ui(agents_name)
return gr.Button.update(visible=False), \
gr.Button.update(visible=False),\
gr.Button.update(visible=False),\
gr.Chatbot.update(visible=True),\
gr.Textbox.update(visible=False),\
gr.Button.update(visible=False),\
gr.Radio.update(visible=False),\
gr.Textbox.update(visible=False)
def new_render_and_register_ui(self,agent_names):
gc.add_agent(agent_names, 0)
def btn_start_after_click(self, history):
"""
inputs=[self.chatbot]
outputs=[self.btn_start, self.btn_send, self.btn_reset, self.chatbot, self.text_input, self.btn_next]
"""
if self.data_history is None:
self.data_history = list()
receive_server = self.receive_server
while True:
data_list: List = receive_server.send(None)
for item in data_list:
data = eval(item)
assert isinstance(data, list)
state, agent_name, token, node_name = data
self.current_node_name = node_name
assert isinstance(state, int)
assert state in [10, 11, 12, 30, 99, 98]
if state == 99:
# finish
yield gr.Button.update(visible=False),\
gr.Button.update(visible=True, interactive=False),\
gr.Button.update(visible=True, interactive=True),\
history,\
gr.Textbox.update(visible=True, interactive=False),\
gr.Button.update(visible=False)
return
elif state == 98:
# single mode
yield gr.Button.update(visible=False), \
gr.Button.update(visible=False),\
gr.Button.update(visible=True),\
history,\
gr.Textbox.update(visible=False),\
gr.Button.update(visible=True, value=f"Next Agent: 🤖{agent_name} | Next Node: ⭕{node_name}")
return
elif state == 30:
# user input
yield gr.Button.update(visible=False), \
gr.Button.update(visible=True),\
gr.Button.update(visible=True),\
history,\
gr.Textbox.update(visible=True, value=""),\
gr.Button.update(visible=False)
return
history = self.handle_message(history, state, agent_name, token, node_name)
yield gr.Button.update(visible=False), \
gr.Button.update(visible=False),\
gr.Button.update(visible=False),\
history,\
gr.Textbox.update(visible=False),\
gr.Button.update(visible=False)
def btn_send_when_click(self, text_input, history):
'''
inputs=[self.text_input, self.chatbot]
outputs=[self.btn_start, self.btn_send, self.btn_reset, self.chatbot, self.text_input, self.btn_next]
'''
history = self.handle_message(history, 10, 'User', text_input, self.current_node_name)
self.send_message("<USER>"+text_input+self.SIGN["SPLIT"])
yield gr.Button.update(visible=False), \
gr.Button.update(visible=False),\
gr.Button.update(visible=False),\
history,\
gr.Textbox.update(visible=False),\
gr.Button.update(visible=False)
return
def btn_send_after_click(self, text_input, history):
'''
inputs=[self.text_input, self.chatbot]
outputs=[self.btn_start, self.btn_send, self.btn_reset, self.chatbot, self.text_input, self.btn_next]
'''
yield from self.btn_start_after_click(history=history)
return
def btn_reset_when_click(self):
"""
outputs=[self.btn_start, self.btn_send, self.btn_reset, self.chatbot, self.text_input, self.btn_next]
"""
return gr.Button.update(interactive=False), gr.Button.update(interactive=False), gr.Button.update(interactive=False, value="Restarting....."), gr.Chatbot.update(label="Dialog"), \
gr.Textbox.update(interactive=False), gr.Button.update(visible=False)
def btn_reset_after_click(self):
"""
outputs=[self.btn_start, self.btn_send, self.btn_reset, self.chatbot, self.text_input, self.btn_next, self.radio_mode]
"""
self.reset()
self.first_recieve_from_client(reset_mode=True)
self.current_node_name = ""
self.data_history = None
return gr.Button.update(interactive=True, visible=True), \
gr.Button.update(interactive=True, visible=False), \
gr.Button.update(interactive=True, value="Restart", visible=False), \
gr.Chatbot.update(label="Dialog", visible=False, value=None), \
gr.Textbox.update(interactive=True, visible=False),\
gr.Button.update(visible=False),\
gr.Radio.update(visible=True), \
gr.Textbox.update(visible=True)
def btn_next_when_click(self, history):
"""
outputs=[self.btn_start, self.btn_send, self.btn_reset, self.chatbot, self.text_input, self.btn_next]
"""
yield gr.Button.update(visible=False), \
gr.Button.update(visible=False),\
gr.Button.update(visible=False),\
history,\
gr.Textbox.update(visible=False),\
gr.Button.update(visible=False)
self.send_message("nothing")
return
def btn_next_after_click(self, history):
time.sleep(1)
yield from self.btn_start_after_click(history=history)
return
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='A demo of chatbot')
parser.add_argument('--agent', type=str, help='path to SOP json')
args = parser.parse_args()
ui = GeneralUI(client_cmd=["python","gradio_backend.py"])
ui.construct_ui()
ui.run()