app.py

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()