Update src/webui.py

src/webui.py

@@ -5,7 +5,6 @@ import gradio as gr
 from pipeline import Pipeline
 from models import *
 
-
 examples = [
     {
         "task": "NER",
@@ -16,6 +15,7 @@ examples = [
         "constraint": """["nationality", "country capital", "place of death", "children", "location contains", "place of birth", "place lived", "administrative division of country", "country of administrative divisions", "company", "neighborhood of", "company founders"]""",
         "file_path": None,
         "update_case": False,
+        "truth": "",
     },
     {
         "task": "RE",
@@ -26,6 +26,7 @@ examples = [
         "constraint": """["nationality", "country capital", "place of death", "children", "location contains", "place of birth", "place lived", "administrative division of country", "country of administrative divisions", "company", "neighborhood of", "company founders"]""",
         "file_path": None,
         "update_case": False,
+        "truth": "",
     },
     {
         "task": "EE",
@@ -36,6 +37,7 @@ examples = [
         "constraint": """{"phishing": ["damage amount", "attack pattern", "tool", "victim", "place", "attacker", "purpose", "trusted entity", "time"], "data breach": ["damage amount", "attack pattern", "number of data", "number of victim", "tool", "compromised data", "victim", "place", "attacker", "purpose", "time"], "ransom": ["damage amount", "attack pattern", "payment method", "tool", "victim", "place", "attacker", "price", "time"], "discover vulnerability": ["vulnerable system", "vulnerability", "vulnerable system owner", "vulnerable system version", "supported platform", "common vulnerabilities and exposures", "capabilities", "time", "discoverer"], "patch vulnerability": ["vulnerable system", "vulnerability", "issues addressed", "vulnerable system version", "releaser", "supported platform", "common vulnerabilities and exposures", "patch number", "time", "patch"]}""",
         "file_path": None,
         "update_case": False,
+        "truth": "",
     },
     {
         "task": "Base",
@@ -46,6 +48,7 @@ examples = [
         "constraint": "",
         "text": "",
         "update_case": False,
+        "truth": "",
     },
     {
         "task": "Base",
@@ -56,6 +59,7 @@ examples = [
         "constraint": "",
         "text": "",
         "update_case": False,
+        "truth": "",
     },
 ]
 
@@ -111,15 +115,15 @@ def create_interface():
                 extraction_Agent_gr = gr.Dropdown(choices=["NOT REQUIRED", "extract_information_direct", "extract_information_with_case"], value="NOT REQUIRED", label="🤖 Select your Extraction-Agent", visible=False)
                 reflection_agent_gr = gr.Dropdown(choices=["NOT REQUIRED", "reflect_with_case"], value="NOT REQUIRED", label="🤖 Select your Reflection-Agent", visible=False)
 
-        with gr.Row():
-            use_file_gr = gr.Checkbox(label="📂 Use File", value=True)
-            update_case_gr = gr.Checkbox(label="💰 Update Case", value=False)
-
+        use_file_gr = gr.Checkbox(label="📂 Use File", value=True)
         file_path_gr = gr.File(label="📖 Upload a File", visible=True)
         text_gr = gr.Textbox(label="📖 Text", placeholder="Enter your Text", visible=False)
         instruction_gr = gr.Textbox(label="🕹️ Instruction", visible=True)
         constraint_gr = gr.Textbox(label="🕹️ Constraint", visible=False)
 
+        update_case_gr = gr.Checkbox(label="💰 Update Case", value=False)
+        truth_gr = gr.Textbox(label="⛳ Truth", visible=False)
+
         def customized_mode(mode):
             if mode == "customized":
                 return gr.update(visible=True), gr.update(visible=True), gr.update(visible=True)
@@ -142,6 +146,12 @@ def create_interface():
             else:
                 return gr.update(visible=True), gr.update(visible=False)
 
+        def update_case(update_case):
+            if update_case:
+                return gr.update(visible=True)
+            else:
+                return gr.update(visible=False)
+
         def start_with_example():
             example_index = random.randint(0, len(examples) - 1)
             example = examples[example_index]
@@ -154,12 +164,13 @@ def create_interface():
                 gr.update(value=example["instruction"], visible=example["task"] == "Base"),
                 gr.update(value=example["constraint"], visible=example["task"] in ["NER", "RE", "EE"]),
                 gr.update(value=example["update_case"]),
+                gr.update(value=example["truth"]),
                 gr.update(value="NOT REQUIRED", visible=False),
                 gr.update(value="NOT REQUIRED", visible=False),
                 gr.update(value="NOT REQUIRED", visible=False),
             )
 
-        def submit(model, api_key, base_url, task, mode, instruction, constraint, text, use_file, file_path, update_case, schema_agent, extraction_Agent, reflection_agent):
+        def submit(model, api_key, base_url, task, mode, instruction, constraint, text, use_file, file_path, update_case, truth, schema_agent, extraction_Agent, reflection_agent):
             try:
                 if base_url == "DEFAULT" or base_url == "":
                     if model in ["gpt-3.5-turbo", "gpt-4o", "gpt-4o-mini"]:
@@ -184,6 +195,8 @@ def create_interface():
                 else:
                     text = text
                     file_path = None
+                if not update_case:
+                    truth = ""
 
                 agent3 = {}
                 if mode == "customized":
@@ -206,8 +219,8 @@ def create_interface():
                     three_agents=agent3,
                     isgui=True,
                     update_case=update_case,
+                    truth=truth,
                     output_schema="",
-                    truth="",
                     show_trajectory=False,
                 )
 
@@ -232,6 +245,7 @@ def create_interface():
                 gr.update(value="", visible=False),  # text
                 gr.update(value=None, visible=True),  # file_path
                 gr.update(value=False),  # update_case
+                gr.update(value=""),  # truth
                 gr.update(value=""),
                 gr.update(value=""),
                 gr.update(value="", visible=False),  # error_output
@@ -257,6 +271,7 @@ def create_interface():
         task_gr.change(fn=update_fields, inputs=task_gr, outputs=[instruction_gr, constraint_gr])
         mode_gr.change(fn=customized_mode, inputs=mode_gr, outputs=[schema_agent_gr, extraction_Agent_gr, reflection_agent_gr])
         use_file_gr.change(fn=update_input_fields, inputs=use_file_gr, outputs=[text_gr, file_path_gr])
+        update_case_gr.change(fn=update_case, inputs=update_case_gr, outputs=[truth_gr])
 
         example_button_gr.click(
             fn=start_with_example,
@@ -270,6 +285,7 @@ def create_interface():
                 instruction_gr,
                 constraint_gr,
                 update_case_gr,
+                truth_gr,
                 schema_agent_gr,
                 extraction_Agent_gr,
                 reflection_agent_gr,
@@ -289,6 +305,7 @@ def create_interface():
                 use_file_gr,
                 file_path_gr,
                 update_case_gr,
+                truth_gr,
                 schema_agent_gr,
                 extraction_Agent_gr,
                 reflection_agent_gr,
@@ -310,6 +327,7 @@ def create_interface():
                 text_gr,
                 file_path_gr,
                 update_case_gr,
+                truth_gr,
                 py_output_gr,
                 json_output_gr,
                 error_output_gr,