import sys
sys.path.append("../../Gradio_Config")
import os
from gradio_base import WebUI, UIHelper, PORT, HOST
from gradio_config import GradioConfig as gc
from gradio_config import StateConfig as sc
from typing import List
import gradio as gr
class NovelUI(WebUI):
node2show = {
"Node 1": "Write Character Settings and Script Outlines🖊️",
"Node 2": "Expand the first chapter<br>✍️",
"Node 3": "Expand the second chapter<br>✍️",
"Node 4": "Expand the third chapter<br>✍️",
"state1": "Perform the first plot<br>🎭",
"state2": "Perform the second plot<br>🎭",
"state3": "Perform the third plot<br>🎭",
"state4": "Perform the forth plot<br>🎭"
}
show2node = {}
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 handle_message(self, history:list, record:list, state, agent_name, token, node_name):
RECORDER = True if state//10 ==2 else False
render_data:list = record if RECORDER else history
data:list = self.data_recorder if RECORDER else self.data_history
if state % 10 == 0:
data.append({agent_name: token})
elif state % 10 == 1:
# Same state. Need to add new bubble in same bubble.
data[-1][agent_name] += token
elif state % 10 == 2:
# New state. Need to add new bubble.
render_data.append([None, ""])
data.clear()
data.append({agent_name: token})
else:
assert False, "Invalid state."
render_data = self.render_bubble(render_data, data, node_name, render_node_name=True)
if RECORDER:
record = render_data
else:
history = render_data
return history, record
def update_progress(self, node_name, node_schedule):
DONE = True
node_name = self.node2show[node_name]
for idx, name in enumerate(self.cache['nodes_name']):
name = self.node2show[name]
self.progress_manage['show_type'][idx] = "active-show-up"
self.progress_manage['show_content'][idx] = ("" if name != node_name else "💬",)
if name == node_name:
DONE = False
self.progress_manage['schedule'][idx] = node_schedule
elif DONE:
self.progress_manage['schedule'][idx] = 100
elif DONE == False:
self.progress_manage['schedule'][idx] = 0
if self.cache['nodes_name'].index(self.show2node[node_name]) == len(self.cache['nodes_name']) - 2 and node_schedule == 100:
self.progress_manage['schedule'][-1] = 100
return sc.FORMAT.format(
sc.CSS,
sc.update_states(
current_states=self.progress_manage["schedule"],
current_templates=self.progress_manage["show_type"],
show_content=self.progress_manage["show_content"]
)
)
def __init__(
self,
client_cmd: list,
socket_host: str = HOST,
socket_port: int = PORT,
bufsize: int = 1024,
ui_name: str = "NovelUI"
):
super(NovelUI, self).__init__(client_cmd, socket_host, socket_port, bufsize, ui_name)
self.first_recieve_from_client()
for item in ['agents_name', 'nodes_name', 'output_file_path', 'requirement']:
assert item in self.cache
self.progress_manage = {
"schedule": [None for _ in range(len(self.cache['nodes_name']))],
"show_type": [None for _ in range(len(self.cache['nodes_name']))],
"show_content": [None for _ in range(len(self.cache['nodes_name']))]
}
NovelUI.show2node = {NovelUI.node2show[_]:_ for _ in NovelUI.node2show.keys()}
def construct_ui(self):
with gr.Blocks(css=gc.CSS) as demo:
with gr.Column():
self.progress = gr.HTML(
value=sc.FORMAT.format(
sc.CSS,
sc.create_states([NovelUI.node2show[name] for name in self.cache['nodes_name']], False)
)
)
with gr.Row():
with gr.Column(scale=6):
self.chatbot = gr.Chatbot(
elem_id="chatbot1",
label="Dialog",
height=500
)
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."
)
with gr.Row():
self.text_requirement = gr.Textbox(
placeholder="Requirement of the novel",
value=self.cache['requirement'],
scale=9
)
self.btn_start = gr.Button(
value="Start",
scale=1
)
self.btn_reset = gr.Button(
value="Restart",
visible=False
)
with gr.Column(scale=5):
self.chat_record = gr.Chatbot(
elem_id="chatbot1",
label="Record",
visible=False
)
self.file_show = gr.File(
value=[],
label="FileList",
visible=False
)
self.chat_show = gr.Chatbot(
elem_id="chatbot1",
label="FileRead",
visible=False
)
# ===============Event Listener===============
self.btn_start.click(
fn=self.btn_start_when_click,
inputs=[self.text_requirement, self.text_api],
outputs=[self.chatbot, self.chat_record, self.btn_start, self.text_requirement]
).then(
fn=self.btn_start_after_click,
inputs=[self.chatbot, self.chat_record],
outputs=[self.progress, self.chatbot, self.chat_record, self.chat_show, self.btn_start, self.btn_reset, self.text_requirement, self.file_show]
)
self.btn_reset.click(
fn=self.btn_reset_when_click,
inputs=[],
outputs=[self.progress, self.chatbot, self.chat_record, self.chat_show, self.btn_start, self.btn_reset, self.text_requirement, self.file_show]
).then(
fn=self.btn_reset_after_click,
inputs=[],
outputs=[self.progress, self.chatbot, self.chat_record, self.chat_show, self.btn_start, self.btn_reset, self.text_requirement, self.file_show]
)
self.file_show.select(
fn=self.file_when_select,
inputs=[self.file_show],
outputs=[self.chat_show]
)
# ===========================================
self.demo = demo
def btn_start_when_click(self, text_requirement:str, api_key:str):
"""
inputs=[self.text_requirement],
outputs=[self.chatbot, self.chat_record, self.btn_start, self.text_requirement]
"""
history = [[UIHelper.wrap_css(content=text_requirement, name="User"), None]]
yield history,\
gr.Chatbot.update(visible=True),\
gr.Button.update(interactive=False, value="Running"),\
gr.Textbox.update(value="", interactive=False)
self.send_start_cmd({'requirement': text_requirement, "api_key": api_key})
return
def btn_start_after_click(self, history:List, record):
def walk_file():
print("file:", self.cache['output_file_path'])
files = []
for _ in os.listdir(self.cache['output_file_path']):
if os.path.isfile(self.cache['output_file_path']+'/'+_):
files.append(self.cache['output_file_path']+'/'+_)
return files
"""
inputs=[self.chatbot, self.chat_record],
outputs=[self.progress, self.chatbot, self.chat_record, self.chat_show, self.btn_start, self.btn_reset, self.text_requirement, self.file_show]
"""
self.data_recorder = list()
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, node_schedule = data
assert isinstance(state, int)
fs:List = walk_file()
# 10/11/12 -> history
# 20/21/22 -> recorder
# 99 -> finish
# 30 -> register new agent
assert state in [10, 11, 12, 20, 21, 22, 99, 30]
if state == 30:
# register new agent.
gc.add_agent(eval(agent_name))
continue
if state == 99:
# finish
yield gr.HTML.update(value=self.update_progress(node_name, node_schedule)),\
history,\
gr.Chatbot.update(visible=True, value=record),\
gr.Chatbot.update(visible=True),\
gr.Button.update(visible=True, interactive=False, value="Done"),\
gr.Button.update(visible=True, interactive=True),\
gr.Textbox.update(visible=True, interactive=False),\
gr.File.update(value=fs, visible=True, interactive=True)
return
history, record = self.handle_message(history, record, state, agent_name, token, node_name)
# [self.progress, self.chatbot, self.chat_record, self.chat_show, self.btn_start, self.btn_reset, self.text_requirement, self.file_show]
yield gr.HTML.update(value=self.update_progress(node_name, node_schedule)),\
history,\
gr.Chatbot.update(visible=True, value=record),\
gr.Chatbot.update(visible=False),\
gr.Button.update(visible=True, interactive=False),\
gr.Button.update(visible=False, interactive=True),\
gr.Textbox.update(visible=True, interactive=False),\
gr.File.update(value=fs, visible=True, interactive=True)
def btn_reset_when_click(self):
"""
inputs=[],
outputs=[self.progress, self.chatbot, self.chat_record, self.chat_show, self.btn_start, self.btn_reset, self.text_requirement, self.file_show]
"""
return gr.HTML.update(value=sc.create_states(states_name=self.cache['nodes_name'])),\
gr.Chatbot.update(value=None),\
gr.Chatbot.update(value=None, visible=False),\
gr.Chatbot.update(value=None, visible=False),\
gr.Button.update(value="Restarting...", visible=True, interactive=False),\
gr.Button.update(value="Restarting...", visible=True, interactive=False),\
gr.Textbox.update(value="Restarting...", interactive=False, visible=True),\
gr.File.update(visible=False)
def btn_reset_after_click(self):
"""
inputs=[],
outputs=[self.progress, self.chatbot, self.chat_record, self.chat_show, self.btn_start, self.btn_reset, self.text_requirement, self.file_show]
"""
self.reset()
self.first_recieve_from_client(reset_mode=True)
return gr.HTML.update(value=sc.create_states(states_name=self.cache['nodes_name'])),\
gr.Chatbot.update(value=None),\
gr.Chatbot.update(value=None, visible=False),\
gr.Chatbot.update(value=None, visible=False),\
gr.Button.update(value="Start", visible=True, interactive=True),\
gr.Button.update(value="Restart", visible=False, interactive=False),\
gr.Textbox.update(value="", interactive=True, visible=True),\
gr.File.update(visible=False)
def file_when_select(self, file_obj):
"""
inputs=[self.file_show],
outputs=[self.chat_show]
"""
CODE_PREFIX = "```json\n{}\n```"
with open(file_obj.name, "r", encoding='utf-8') as f:
contents = f.readlines()
codes = "".join(contents)
return [[CODE_PREFIX.format(codes),None]]
if __name__ == '__main__':
ui = NovelUI(client_cmd=["python","gradio_backend.py"])
ui.construct_ui()
ui.run()