callanwu's picture
modify Coder label
f138f52
import sys
sys.path.append("../../agents")
import os
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
from Agent import Agent
import time
import json
from utils import cos_sim
from design_states import get_desgin_states, get_cot_result, gen_coder_task
from gen_utils import *
import openai
import torch
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 CodeUI(WebUI):
def render_and_register_ui(self):
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 __init__(
self,
client_cmd: list,
socket_host: str = HOST,
socket_port: int = PORT,
bufsize: int = 1024,
ui_name: str = "CodeUI"
):
super(CodeUI, self).__init__(client_cmd, socket_host, socket_port, bufsize, ui_name)
self.first_recieve_from_client()
self.data_history = list()
self.caller = 0
def load_sop_fn(self,sop):
return sop.name
def generate_sop_stage_1(self,api_key,proxy,target):
os.environ["API_KEY"] = api_key
# os.environ["PROXY"] = proxy
self.software = "You are a software,aim to write a snake game with python"
self.debate = "Simulate a debate competition"
self.ecological_environment = "Simulate the interactions and competition among different organisms within an ecosystem"
self.software = get_embedding(self.software,api_key)
self.debate = get_embedding(self.debate,api_key)
self.ecological_environment = get_embedding(self.ecological_environment,api_key)
self.embeddings = torch.cat([self.software,self.debate,self.ecological_environment],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:
self.index = top_k_idx
else:
self.index = 0
target_processed = get_cot_result(target)
print("finished!!!!")
return target_processed,self.target_finish_flag.update(visible = True)
def generate_sop_stage_2(self,target):
design_states = get_desgin_states(target,self.index)
root = design_states[0]["state_name"]
self.SOP["root"] = root
return design_states,self.state_finish_flag.update(visible = True)
def generate_sop_stage_3(self,design_states):
agents = get_agents(design_states,self.index)
relations = get_relations(design_states)
self.SOP["relations"] = relations
self.SOP["agents"] = agents
return agents, self.agent_relation_finish_flag.update(visible = True),self.reminder.update(visible = True)
def generate_sop_stage_4(self,agents, need_coder,design_states):
states = gen_states(design_states,self.index)
if "Coder" in need_coder:
agents["coder"] = {"style": "professional", "roles": {}}
for state_name, state in states.items():
if state_name != "end_state":
agents["coder"]["roles"][state_name] = "coder"
state["roles"].append("coder")
task = gen_coder_task(state["environment_prompt"])
now_coder = self.coder.copy()
now_coder["task"]["task"] = task
state["agent_states"]["coder"] = now_coder
state["controller"]["max_chat_nums"] = str(
int(state["controller"]["max_chat_nums"])+2)
for name, agent in state["agent_states"].items():
if name != "coder":
agent["rule"]["rule"] += "\nEvaluate the code of the coder and provide feedback and response as concise as possible.It is best not to exceed 100 words"
agent["task"]["task"] += "\nEvaluate the code of the coder and provide feedback."
self.SOP["states"] = states
# 将字典写入JSON文件
file_name = 'generated_sop.json'
with open(file_name, "w") as json_file:
json.dump(self.SOP, json_file ,indent=4,ensure_ascii=False)
return file_name
def construct_ui(self):
with gr.Blocks(css=gc.CSS) as demo:
with gr.Tab(label="SOP generation") as tab1:
self.coder = {
"task": {
"task": ""
},
"rule": {"rule": "1.write code that conforms to standards like PEP8, is modular, easy to read, and maintainable.\n 2.The output strictly follows the following format:<title>{the file name}</title>\n<python>{the target code}</python>\n3.Please carefully modify the code based on feedback from others.\n4.Output the code only."},
"last": {
"last_prompt": "The output strictly follows the following format:<title>{the file name}</title>\n<python>{the target code}</python>,Output the code only."
}
}
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():
api_key = gr.Textbox(label="api_key")
proxy = gr.Textbox(label="proxy",visible=False)
with gr.Row():
requirement = gr.Textbox(value ="A software company aim to write a mine sweeping game",label="requirement")
with gr.Row():
need_coder = gr.CheckboxGroup(["Coder"],label="Please check this option if your multi-agent system aims to produce code.")
with gr.Row():
self.target_finish_flag = gr.Label(value = "The process of completing requirement handling is finished.",visible=False)
with gr.Row():
self.state_finish_flag = gr.Label(value = "The process of determining the state is completed.",visible=False)
with gr.Row():
self.agent_relation_finish_flag = gr.Label(value = "The process of initializing the agent and relation is completed.",visible=False)
with gr.Row():
self.reminder = gr.Markdown("""Generating SOP...""",visible=False)
generated_sop = gr.File(label="generated_file")
generate_button = gr.Button(label="Generate")
target_processed = gr.State()
design_states = gr.State()
agents = gr.State()
generate_button.click(self.generate_sop_stage_1,[api_key,proxy,requirement],[target_processed,self.target_finish_flag]).then(
self.generate_sop_stage_2, [target_processed], [design_states,self.state_finish_flag]).then(
self.generate_sop_stage_3, [design_states], [agents,self.agent_relation_finish_flag,self.reminder]).then(
self.generate_sop_stage_4, [agents, need_coder,design_states], [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.Row():
with gr.Column():
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.radio_mode = gr.Radio(
[Client.SINGLE_MODE],
value=Client.SINGLE_MODE,
interactive=True,
label = Client.MODE_LABEL,
info = Client.MODE_INFO
)
self.chatbot = gr.Chatbot(
elem_id="chatbot1"
)
self.btn_next = gr.Button(
value="Next Agent",
visible=False, elem_id="btn"
)
with gr.Row():
self.text_requirement = gr.Textbox(
value=self.cache['requirement'],
placeholder="Please enter your content",
scale=9,
)
self.btn_start = gr.Button(
value="Start!",
scale=1
)
self.btn_reset = gr.Button(
value="Restart",
visible=False
)
with gr.Column():
self.file = gr.File(visible=False)
self.chat_code_show = gr.Chatbot(
elem_id="chatbot1",
visible=False
)
self.btn_start.click(
fn=self.btn_send_when_click,
inputs=[self.chatbot, self.text_requirement, self.radio_mode, self.text_api,uploaded_sop],
outputs=[self.chatbot, self.btn_start, self.text_requirement, self.btn_reset]
).then(
fn=self.btn_send_after_click,
inputs=[self.file, self.chatbot, self.chat_code_show, self.btn_start, self.btn_reset, self.text_requirement],
outputs=[self.file, self.chatbot, self.chat_code_show, self.btn_start, self.btn_reset, self.text_requirement, self.btn_next]
)
self.text_requirement.submit(
fn=self.btn_send_when_click,
inputs=[self.chatbot, self.text_requirement, self.text_api,uploaded_sop],
outputs=[self.chatbot, self.btn_start, self.text_requirement, self.btn_reset]
).then(
fn=self.btn_send_after_click,
inputs=[self.file, self.chatbot, self.chat_code_show, self.btn_start, self.btn_reset, self.text_requirement],
outputs=[self.file, self.chatbot, self.chat_code_show, self.btn_start, self.btn_reset, self.text_requirement, self.btn_next]
)
self.btn_reset.click(
fn=self.btn_reset_when_click,
inputs=[],
outputs=[self.file, self.chatbot, self.chat_code_show, self.btn_start, self.btn_reset, self.text_requirement, self.btn_next]
).then(
fn=self.btn_reset_after_click,
inputs=[self.file, self.chatbot, self.chat_code_show, self.btn_start, self.btn_reset, self.text_requirement],
outputs=[self.file, self.chatbot, self.chat_code_show, self.btn_start, self.btn_reset, self.text_requirement, self.btn_next]
)
self.file.select(
fn=self.file_when_select,
inputs=[self.file],
outputs=[self.chat_code_show]
)
self.btn_next.click(
fn = self.btn_next_when_click,
inputs=[],
outputs=[self.btn_next]
).then(
fn=self.btn_send_after_click,
inputs=[self.file, self.chatbot, self.chat_code_show, self.btn_start, self.btn_reset, self.text_requirement],
outputs=[self.file, self.chatbot, self.chat_code_show, self.btn_start, self.btn_reset, self.text_requirement, self.btn_next]
)
self.demo = demo
def handle_message(self, history:list, 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 btn_send_when_click(self, chatbot, text_requirement, mode, api, sop):
"""
inputs=[self.chatbot, self.text_requirement, radio, text_api],
outputs=[self.chatbot, self.btn_start, self.text_requirement, self.btn_reset]
"""
chatbot = [[UIHelper.wrap_css(content=text_requirement, name="User"), None]]
yield chatbot,\
gr.Button.update(visible=True, interactive=False, value="Running"),\
gr.Textbox.update(visible=True, interactive=False, value=""),\
gr.Button.update(visible=False, interactive=False)
self.send_start_cmd({'requirement': text_requirement, "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
def new_render_and_register_ui(self,agent_names):
gc.add_agent(agent_names, 0)
def btn_send_after_click(
self,
file,
history,
show_code,
btn_send,
btn_reset,
text_requirement
):
"""
outputs=[self.file, self.chatbot, self.chat_code_show, self.btn_start, self.btn_reset, self.text_requirement, self.btn_next]
"""
if self.caller == 0:
self.data_history = list()
self.caller = 0
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
assert isinstance(state, int)
assert state in [10, 11, 12, 99, 98]
if state == 99:
# finish
fs = [self.cache['pwd']+'/output_code/'+_ for _ in os.listdir(self.cache['pwd']+'/output_code')]
yield gr.File.update(value=fs, visible=True, interactive=True),\
history, \
gr.Chatbot.update(visible=True),\
gr.Button.update(visible=True, interactive=True, value="Start"),\
gr.Button.update(visible=True, interactive=True),\
gr.Textbox.update(visible=True, interactive=True, placeholder="Please input your requirement", value=""),\
gr.Button.update(visible=False)
return
elif state == 98:
yield gr.File.update(visible=False),\
history, \
gr.Chatbot.update(visible=False),\
gr.Button.update(visible=True, interactive=False),\
gr.Button.update(visible=True, interactive=True),\
gr.Textbox.update(visible=True, interactive=False),\
gr.Button.update(visible=True, value=f"Next Agent: 🤖{agent_name} | Next Node: ⭕{node_name}")
return
history = self.handle_message(history, state, agent_name, token, node_name)
yield gr.File.update(visible=False),\
history, \
gr.Chatbot.update(visible=False),\
gr.Button.update(visible=True, interactive=False),\
gr.Button.update(visible=False, interactive=False),\
gr.Textbox.update(visible=True, interactive=False),\
gr.Button.update(visible=False)
def btn_reset_when_click(self):
"""
inputs = []
outputs = [self.file, self.chatbot, self.chat_code_show, self.btn_start, self.btn_reset, self.text_requirement, self.btn_next]
"""
return gr.File.update(visible=False),\
None, None, gr.Button.update(value="Restarting...", interactive=False),\
gr.Button.update(value="Restarting...", interactive=False),\
gr.Textbox.update(value="Restarting", interactive=False),\
gr.Button.update(visible=False)
def btn_reset_after_click(
self,
file,
chatbot,
show_code,
btn_send,
btn_reset,
text_requirement
):
self.reset()
self.first_recieve_from_client(reset_mode=True)
return gr.File.update(value=None, visible=False),\
gr.Chatbot.update(value=None, visible=True),\
gr.Chatbot.update(value=None, visible=False),\
gr.Button.update(value="Start", visible=True, interactive=True),\
gr.Button.update(value="Restart", interactive=False, visible=False),\
gr.Textbox.update(value=self.cache['requirement'], interactive=True, visible=True),\
gr.Button.update(visible=False)
def file_when_select(self, file):
CODE_PREFIX = "```python\n{}\n```"
with open(file.name, "r", encoding='utf-8') as f:
contents = f.readlines()
codes = "".join(contents)
return [[CODE_PREFIX.format(codes),None]]
def btn_next_when_click(self):
self.caller = 1 # it will remain the value in self.data_history
self.send_message("nothing")
time.sleep(0.5)
yield gr.Button.update(visible=False)
return
if __name__ == '__main__':
ui = CodeUI(client_cmd=["python3","gradio_backend.py"])
ui.construct_ui()
ui.run(share=True)