proper session data and refractor

app.py

@@ -1,10 +1,7 @@
 import json
 import gradio as gr
-import json
-from langchain_legal import run_pipeline
 
-global data
-data = {}
+from langchain_legal import run_pipeline
 
 
 def load_json(path: str) -> dict:
@@ -12,10 +9,9 @@ def load_json(path: str) -> dict:
         return json.load(f)
 
 
-def get_ai_final_summary(text_input, req: gr.Request):
-    global data
-    data[req.client.host] = run_pipeline(text_input)
-    return data[req.client.host]['ai_final_summary']
+def get_ai_final_summary(text_input):
+    data = run_pipeline(text_input)
+    return data, data.get('ai_final_summary', 'No answer yet')
 
 
 STEP_NAMES = [
@@ -26,73 +22,73 @@ STEP_NAMES = [
 ]
 
 
-def get_inital_assesment(req: gr.Request) -> str:
-    client_data = data.get(req.client.host)
-    if client_data is None:
-        return 'No llm answer yet'
-    initial = client_data.get('ai_first_5', None)
-    print(data)
+def get_final_summary(data):
+    summ = data.get('ai_final_summary')
+    if summ is None:
+        return gr.Textbox('No answer yet', label='Final summary')
+    return gr.Textbox(summ, label='Final summary', show_copy_button=True)
+
+
+def get_inital_assesment(data):
+    initial = data.get('ai_first_5', None)
     if initial is None:
-        return 'No llm answer yet'
-    return initial
+        return gr.Textbox('No llm answer yet', show_label=False)
+    return gr.Textbox(initial, show_label=False, lines=25)
 
 
-def get_applicable_eu_regs(summary_tb, json_output, req: gr.Request):
-    client_data = data.get(req.client.host)
-    if client_data is None:
-        return summary_tb, json_output
-    eval = client_data.get('ai_eur_lex_eval', None)
+def get_eu_regs_eval(data):
+    eval = data.get('ai_eur_lex_eval', None)
     if eval is None:
-        return summary_tb, json_output
-    sources = client_data.get('eur_lex_lookup')
-    summary_tb = gr.Textbox(value=eval, show_label=False)
-    if sources is not None:
-        json_output = gr.Json(value=[json.loads(j) for j in sources])
-    return summary_tb, json_output
-
-
-def get_applicable_national_regs(summary_tb, json_output, req: gr.Request):
-    client_data = data.get(req.client.host)
-    if client_data is None:
-        return summary_tb, json_output
-    eval = client_data.get('ai_aus_eval', None)
-    nat = 'aus'
+        return gr.Textbox('No llm answer yet', show_label=False)
+    return gr.Textbox(eval, show_label=False, lines=50)
+
+
+def get_eu_regs_sources(data):
+    lookup = data.get('eur_lex_lookup', None)
+    if lookup is None:
+        return gr.Json({}, show_label=False)
+    return gr.Json(lookup)
+
+
+def get_national_regs_eval(data):
+    eval = data.get('ai_aus_eval', None)
     if eval is None:
-        eval = client_data.get('ai_ger_eval', None)
-        nat = 'ger'
+        eval = data.get('ai_ger_eval')
         if eval is None:
-            return summary_tb, json_output
-    if nat == 'aus':
-        sources = client_data.get('aus_lookup')
-    else:
-        sources = client_data.get('ger_lookup')
-    summary_tb = gr.Textbox(eval, show_label=False)
-    if sources is not None:
-        json_output = gr.Json([json.loads(j) for j in sources])
-    return summary_tb, sources
-
-
-def get_ai_challange(req: gr.Request) -> str:
-    client_data = data.get(req.client.host)
-    if client_data is None:
-        return 'No llm answer yet'
-    challange = client_data.get('ai_challange', None)
-    print(data)
+            return gr.Textbox('No llm answer yet', show_label=False)
+    return gr.Textbox(eval, show_label=False, lines=50)
+
+
+def get_national_regs_sources(data):
+    lookup = data.get('aus_lookup', None)
+    if lookup is None:
+        lookup = data.get('ger_lookup')
+        if lookup is None:
+            return gr.Json({}, show_label=False)
+    return gr.Json(lookup)
+
+
+def get_ai_challange(data):
+    challange = data.get('ai_challange', None)
     if challange is None:
-        return 'No llm answer yet'
-    return challange
+        return gr.Textbox('No llm answer yet', show_label=False)
+    return gr.Textbox(challange, show_label=False, lines=50)
 
 
-CALLBACKS = [
-    get_inital_assesment,
-    get_applicable_eu_regs,
-    get_applicable_national_regs,
-    get_ai_challange,
-]
+def update_ui(data):
+    return (
+        get_inital_assesment(data),
+        get_eu_regs_eval(data),
+        get_eu_regs_sources(data),
+        get_national_regs_eval(data),
+        get_national_regs_sources(data),
+        get_ai_challange(data),
+    )
 
 
 def main():
     with gr.Blocks() as demo:
+        data = gr.State({})
         gr.Markdown("# Lawgarithm's LLM legal research tool")
 
         with gr.Column():
@@ -101,28 +97,46 @@ def main():
                 placeholder='Write your legal case here',
             )
             submit_btn = gr.Button('Submit')
-
-        text_output = gr.Textbox(label='Final Summary')
-
-        submit_btn.click(
-            get_ai_final_summary, inputs=text_input, outputs=text_output
-        )
+            text_output = gr.Textbox(label='Final Summary')
+            submit_btn.click(
+                get_ai_final_summary,
+                inputs=[text_input],
+                outputs=[data, text_output],
+                show_progress='full',
+                queue=True,
+            )
 
         with gr.Tabs():
-            for name, cb in zip(STEP_NAMES, CALLBACKS):
-                with gr.Tab(f'{name}') as t:
-                    tab_text = gr.Textbox(
-                        placeholder='No answer yet', show_label=False
-                    )
-                    if 'regs' not in name:
-                        t.select(cb, outputs=tab_text)
-                    else:
-                        tab_sources = gr.Json({})
-                        t.select(
-                            cb,
-                            inputs=[tab_text, tab_sources],
-                            outputs=[tab_text, tab_sources],
-                        )
+            with gr.Tab(STEP_NAMES[0]):
+                inital_tab_text = gr.Textbox(
+                    placeholder='No answer yet', show_label=False
+                )
+            with gr.Tab(STEP_NAMES[1]):
+                eu_regs_tab_text = gr.Textbox(
+                    placeholder='No answer yet', show_label=False
+                )
+                eu_regs_tab_sources = gr.Json({})
+            with gr.Tab(STEP_NAMES[2]):
+                nat_regs_tab_text = gr.Textbox(
+                    placeholder='No answer yet', show_label=False
+                )
+                nat_regs_tab_sources = gr.Json({})
+            with gr.Tab(STEP_NAMES[2]):
+                comp_assesment_tab_text = gr.Textbox(
+                    placeholder='No answer yet', show_label=False
+                )
+        data.change(
+            update_ui,
+            inputs=data,
+            outputs=[
+                inital_tab_text,
+                eu_regs_tab_text,
+                eu_regs_tab_sources,
+                nat_regs_tab_text,
+                nat_regs_tab_sources,
+                comp_assesment_tab_text,
+            ],
+        )
 
     demo.launch()