310 lines
10 KiB
Python
310 lines
10 KiB
Python
import pickle
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
|
|
import gradio as gr
|
|
from theflow.storage import storage
|
|
|
|
from kotaemon.chatbot import ChatConversation
|
|
from kotaemon.contribs.promptui.base import get_component
|
|
from kotaemon.contribs.promptui.export import export
|
|
from kotaemon.contribs.promptui.ui.blocks import ChatBlock
|
|
|
|
from ..logs import ResultLog
|
|
|
|
USAGE_INSTRUCTION = """## How to use:
|
|
|
|
1. Set the desired parameters.
|
|
2. Click "New chat" to start a chat session with the supplied parameters. This
|
|
set of parameters will persist until the end of the chat session. During an
|
|
ongoing chat session, changing the parameters will not take any effect.
|
|
3. Chat and interact with the chat bot on the right panel. You can add any
|
|
additional input (if any), and they will be supplied to the chatbot.
|
|
4. During chat, the log of the chat will show up in the "Output" tabs. This is
|
|
empty by default, so if you want to show the log here, tell the AI developers
|
|
to configure the UI settings.
|
|
5. When finishing chat, select your preference in the radio box. Click "End chat".
|
|
This will save the chat log and the preference to disk.
|
|
6. To compare the result of different run, click "Export" to get an Excel
|
|
spreadsheet summary of different run.
|
|
|
|
## Support:
|
|
|
|
In case of errors, you can:
|
|
|
|
- PromptUI instruction:
|
|
https://github.com/Cinnamon/kotaemon/wiki/Utilities#prompt-engineering-ui
|
|
- Create bug fix and make PR at: https://github.com/Cinnamon/kotaemon
|
|
- Ping any of @john @tadashi @ian @jacky in Slack channel #llm-productization
|
|
|
|
## Contribute:
|
|
|
|
- Follow installation at: https://github.com/Cinnamon/kotaemon/
|
|
"""
|
|
|
|
|
|
def construct_chat_ui(
|
|
config, func_new_chat, func_chat, func_end_chat, func_export_to_excel
|
|
) -> gr.Blocks:
|
|
"""Construct the prompt engineering UI for chat
|
|
|
|
Args:
|
|
config: the UI config
|
|
func_new_chat: the function for starting a new chat session
|
|
func_chat: the function for chatting interaction
|
|
func_end_chat: the function for ending and saving the chat
|
|
func_export_to_excel: the function to export the logs to excel
|
|
|
|
Returns:
|
|
the UI object
|
|
"""
|
|
inputs, outputs, params = [], [], []
|
|
for name, component_def in config.get("inputs", {}).items():
|
|
if "params" not in component_def:
|
|
component_def["params"] = {}
|
|
component_def["params"]["interactive"] = True
|
|
component = get_component(component_def)
|
|
if hasattr(component, "label") and not component.label: # type: ignore
|
|
component.label = name # type: ignore
|
|
|
|
inputs.append(component)
|
|
|
|
for name, component_def in config.get("params", {}).items():
|
|
if "params" not in component_def:
|
|
component_def["params"] = {}
|
|
component_def["params"]["interactive"] = True
|
|
component = get_component(component_def)
|
|
if hasattr(component, "label") and not component.label: # type: ignore
|
|
component.label = name # type: ignore
|
|
|
|
params.append(component)
|
|
|
|
for idx, component_def in enumerate(config.get("outputs", [])):
|
|
if "params" not in component_def:
|
|
component_def["params"] = {}
|
|
component_def["params"]["interactive"] = False
|
|
component = get_component(component_def)
|
|
if hasattr(component, "label") and not component.label: # type: ignore
|
|
component.label = f"Output {idx}" # type: ignore
|
|
|
|
outputs.append(component)
|
|
|
|
sess = gr.State(value=None)
|
|
chatbot = gr.Chatbot(label="Chatbot", show_copy_button=True)
|
|
chat = ChatBlock(
|
|
func_chat, chatbot=chatbot, additional_inputs=[sess], additional_outputs=outputs
|
|
)
|
|
param_state = gr.Textbox(interactive=False)
|
|
|
|
with gr.Blocks(analytics_enabled=False, title="Welcome to PromptUI") as demo:
|
|
sess.render()
|
|
with gr.Accordion(label="HOW TO", open=False):
|
|
gr.Markdown(USAGE_INSTRUCTION)
|
|
with gr.Row():
|
|
run_btn = gr.Button("New chat")
|
|
run_btn.click(
|
|
func_new_chat,
|
|
inputs=params,
|
|
outputs=[
|
|
chat.chatbot,
|
|
chat.chatbot_state,
|
|
chat.saved_input,
|
|
param_state,
|
|
sess,
|
|
*outputs,
|
|
],
|
|
)
|
|
with gr.Accordion(label="End chat", open=False):
|
|
likes = gr.Radio(["like", "dislike", "neutral"], value="neutral")
|
|
save_log = gr.Checkbox(
|
|
value=True,
|
|
label="Save log",
|
|
info="If saved, log can be exported later",
|
|
show_label=True,
|
|
)
|
|
end_btn = gr.Button("End chat")
|
|
end_btn.click(
|
|
func_end_chat,
|
|
inputs=[likes, save_log, sess],
|
|
outputs=[param_state, sess],
|
|
)
|
|
with gr.Accordion(label="Export", open=False):
|
|
exported_file = gr.File(
|
|
label="Output file", show_label=True, height=100
|
|
)
|
|
export_btn = gr.Button("Export")
|
|
export_btn.click(
|
|
func_export_to_excel, inputs=None, outputs=exported_file
|
|
)
|
|
|
|
with gr.Row():
|
|
with gr.Column():
|
|
with gr.Tab("Params"):
|
|
for component in params:
|
|
component.render()
|
|
with gr.Accordion(label="Session state", open=False):
|
|
param_state.render()
|
|
|
|
with gr.Tab("Outputs"):
|
|
for component in outputs:
|
|
component.render()
|
|
with gr.Column():
|
|
chat.render()
|
|
|
|
return demo.queue()
|
|
|
|
|
|
def build_chat_ui(config, pipeline_def):
|
|
"""Build the chat UI
|
|
|
|
Args:
|
|
config: the UI config
|
|
pipeline_def: the pipeline definition
|
|
|
|
Returns:
|
|
the UI object
|
|
"""
|
|
output_dir: Path = Path(storage.url(pipeline_def().config.store_result))
|
|
exported_dir = output_dir.parent / "exported"
|
|
exported_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
resultlog = getattr(pipeline_def, "_promptui_resultlog", ResultLog)
|
|
allowed_resultlog_callbacks = {i for i in dir(resultlog) if not i.startswith("__")}
|
|
|
|
def new_chat(*args):
|
|
"""Start a new chat function
|
|
|
|
Args:
|
|
*args: the pipeline init params
|
|
|
|
Returns:
|
|
new empty states
|
|
"""
|
|
gr.Info("Starting new session...")
|
|
param_dicts = {
|
|
name: value for name, value in zip(config["params"].keys(), args)
|
|
}
|
|
for key in param_dicts.keys():
|
|
if config["params"][key].get("component").lower() == "file":
|
|
param_dicts[key] = param_dicts[key].name
|
|
|
|
# TODO: currently hard-code as ChatConversation
|
|
pipeline = pipeline_def()
|
|
session = ChatConversation(bot=pipeline)
|
|
session.set(param_dicts)
|
|
session.start_session()
|
|
|
|
param_state_str = "\n".join(
|
|
f"- {name}: {value}" for name, value in param_dicts.items()
|
|
)
|
|
|
|
gr.Info("New chat session started.")
|
|
return (
|
|
[],
|
|
[],
|
|
None,
|
|
param_state_str,
|
|
session,
|
|
*[None] * len(config.get("outputs", [])),
|
|
)
|
|
|
|
def chat(message, history, session, *args):
|
|
"""The chat interface
|
|
|
|
# TODO: wrap the input and output of this chat function so that it
|
|
work with more types of chat conversation than simple text
|
|
|
|
Args:
|
|
message: the message from the user
|
|
history: the gradio history of the chat
|
|
session: the chat object session
|
|
*args: the additional inputs
|
|
|
|
Returns:
|
|
the response from the chatbot
|
|
"""
|
|
if session is None:
|
|
raise gr.Error(
|
|
"No active chat session. Please set the params and click New chat"
|
|
)
|
|
|
|
pred = session(message)
|
|
text_response = pred.content
|
|
|
|
additional_outputs = []
|
|
for output_def in config.get("outputs", []):
|
|
value = session.last_run.logs(output_def["step"])
|
|
getter = output_def.get("getter", None)
|
|
if getter and getter in allowed_resultlog_callbacks:
|
|
value = getattr(resultlog, getter)(value)
|
|
additional_outputs.append(value)
|
|
|
|
return text_response, *additional_outputs
|
|
|
|
def end_chat(preference: str, save_log: bool, session):
|
|
"""End the chat session
|
|
|
|
Args:
|
|
preference: the preference of the user
|
|
save_log: whether to save the result
|
|
session: the chat object session
|
|
|
|
Returns:
|
|
the new empty state
|
|
"""
|
|
gr.Info("Ending session...")
|
|
session.end_session()
|
|
output_dir: Path = (
|
|
Path(storage.url(session.config.store_result)) / session.last_run.id()
|
|
)
|
|
|
|
if not save_log:
|
|
if output_dir.exists():
|
|
import shutil
|
|
|
|
shutil.rmtree(output_dir)
|
|
|
|
session = None
|
|
param_state = ""
|
|
gr.Info("End session without saving log.")
|
|
return param_state, session
|
|
|
|
# add preference result to progress
|
|
with (output_dir / "progress.pkl").open("rb") as fi:
|
|
progress = pickle.load(fi)
|
|
progress["preference"] = preference
|
|
with (output_dir / "progress.pkl").open("wb") as fo:
|
|
pickle.dump(progress, fo)
|
|
|
|
# get the original params
|
|
param_dicts = {name: session.getx(name) for name in config["params"].keys()}
|
|
with (output_dir / "params.pkl").open("wb") as fo:
|
|
pickle.dump(param_dicts, fo)
|
|
|
|
session = None
|
|
param_state = ""
|
|
gr.Info("End session and save log.")
|
|
return param_state, session
|
|
|
|
def export_func():
|
|
name = (
|
|
f"{pipeline_def.__module__}.{pipeline_def.__name__}_{datetime.now()}.xlsx"
|
|
)
|
|
path = str(exported_dir / name)
|
|
gr.Info(f"Begin exporting {name}...")
|
|
try:
|
|
export(config=config, pipeline_def=pipeline_def, output_path=path)
|
|
except Exception as e:
|
|
raise gr.Error(f"Failed to export. Please contact project's AIR: {e}")
|
|
gr.Info(f"Exported {name}. Please go to the `Exported file` tab to download")
|
|
return path
|
|
|
|
demo = construct_chat_ui(
|
|
config=config,
|
|
func_new_chat=new_chat,
|
|
func_chat=chat,
|
|
func_end_chat=end_chat,
|
|
func_export_to_excel=export_func,
|
|
)
|
|
return demo
|