Complete refactor of app.py state to dictionary

app.py

@@ -40,35 +40,20 @@ theme = gr.themes.Soft(
     secondary_hue="slate",
     neutral_hue="slate",
     font=[gr.themes.GoogleFont('Inter'), 'sans-serif'],
-).set(
-    button_primary_background_fill='linear-gradient(90deg, #4f46e5 0%, #3b82f6 100%)',
-    button_primary_background_fill_hover='linear-gradient(90deg, #4338ca 0%, #2563eb 100%)',
 )
 
-class ResearchState:
-    def __init__(self):
-        self.initial_topic = ""
-        self.suggestions = []
-        self.final_topic = ""
-        self.research_plan = ""
-        self.subtasks = []
-        self.final_report = ""
-        self.pdf_data = None
-        self.pdf_path = ""
-
 def start_clarification(topic, hf_key, state):
     if not topic:
         return gr.update(), "Please enter a topic.", state, gr.update()
     if not hf_key:
         return gr.update(), "Please provide an HF Token.", state, gr.update()
     
-    state.initial_topic = topic
+    state["initial_topic"] = topic
     clarifier = Clarifier(model_name=CLARIFIER_MODEL, hf_key=hf_key)
     try:
         suggestions = clarifier.get_suggestions(topic)
-        state.suggestions = suggestions
+        state["suggestions"] = suggestions
         
-        # Format suggestions for the UI
         suggestion_md = "### Refine Your Topic\n\nChoose one of the suggested directions or enter a custom one below:\n\n"
         for i, s in enumerate(suggestions):
             suggestion_md += f"**Option {i+1}: {s['title']}**\n{s['description']}\n\n"
@@ -78,21 +63,21 @@ def start_clarification(topic, hf_key, state):
         return gr.update(), f"Error: {e}", state, gr.update()
 
 def select_suggestion(index, custom_topic, state):
-    if custom_topic.strip():
-        state.final_topic = custom_topic
-    elif index is not None and 0 <= int(index)-1 < len(state.suggestions):
-        sug = state.suggestions[int(index)-1]
-        state.final_topic = f"{sug['title']}: {sug['description']}"
+    if custom_topic and custom_topic.strip():
+        state["final_topic"] = custom_topic
+    elif index is not None and 0 <= int(index)-1 < len(state.get("suggestions", [])):
+        sug = state["suggestions"][int(index)-1]
+        state["final_topic"] = f"{sug['title']}: {sug['description']}"
     else:
-        return gr.update(), "Please select an option or enter a custom topic.", state
+        return gr.update(), "Please select an option or enter a custom topic.", state, gr.update()
     
-    return gr.update(visible=True), f"Targeting: **{state.final_topic}**", state, gr.update(visible=False)
+    return gr.update(visible=True), f"Targeting: **{state['final_topic']}**", state, gr.update(visible=False)
 
 def generate_strategy(hf_key, state):
     planner = Planner(model_name=PLANNER_MODEL, hf_key=hf_key)
     try:
-        plan = planner.plan(state.final_topic)
-        state.research_plan = plan
+        plan = planner.plan(state["final_topic"])
+        state["research_plan"] = plan
         return gr.update(visible=True), plan, state, gr.update(visible=False)
     except Exception as e:
         return gr.update(), f"Error: {e}", state, gr.update()
@@ -100,8 +85,8 @@ def generate_strategy(hf_key, state):
 def decompose_tasks(hf_key, state):
     splitter = Splitter(model_name=SPLITTER_MODEL, hf_key=hf_key)
     try:
-        subtasks = splitter.split(state.research_plan)
-        state.subtasks = subtasks
+        subtasks = splitter.split(state["research_plan"])
+        state["subtasks"] = subtasks
         
         tasks_md = "### 📋 Generated Subtasks\n\n"
         for task in subtasks:
@@ -144,7 +129,7 @@ def run_research(hf_key, tavily_key, state):
     current_findings = []
     log_content = "### 🔍 Agentic Research Progress\n\n"
     
-    for i, task in enumerate(state.subtasks):
+    for i, task in enumerate(state.get("subtasks", [])):
         t_id = task['id']
         t_title = task['title']
         t_desc = task['description']
@@ -159,8 +144,8 @@ def run_research(hf_key, tavily_key, state):
             max_steps=2
         )
         prompt = SUBAGENT_DIRECTION.format(
-            user_query=state.final_topic,
-            research_plan=state.research_plan,
+            user_query=state["final_topic"],
+            research_plan=state["research_plan"],
             subtask_id=t_id,
             subtask_title=t_title,
             subtask_description=t_desc
@@ -179,9 +164,9 @@ def run_research(hf_key, tavily_key, state):
     yield log_content, state, gr.update(), gr.update()
     
     sys_prompt = COORDINATOR_DIRECTION.format(
-        user_query=state.final_topic,
-        research_plan=state.research_plan,
-        subtasks_json=json.dumps(state.subtasks, indent=2)
+        user_query=state["final_topic"],
+        research_plan=state["research_plan"],
+        subtasks_json=json.dumps(state["subtasks"], indent=2)
     )
     user_prompt = f"Synthesize these findings:\n\n" + "\n\n".join(current_findings)
     
@@ -202,16 +187,19 @@ def run_research(hf_key, tavily_key, state):
         
         reviewer = Reviewer(model_name=REVIEWER_MODEL, hf_key=hf_key)
         polished_report = reviewer.review(final_report)
-        state.final_report = polished_report
+        state["final_report"] = polished_report
         
         os.makedirs("temp_outputs", exist_ok=True)
-        pdf_path = f"temp_outputs/research_{int(time.time())}.pdf"
+        pdf_name = f"research_{int(time.time())}.pdf"
+        pdf_path = os.path.join("temp_outputs", pdf_name)
         if reviewer.generate_pdf(polished_report, pdf_path):
-            state.pdf_path = pdf_path
+            state["pdf_path"] = pdf_path
             
-        md_path = f"temp_outputs/research_{int(time.time())}.md"
+        md_name = f"research_{int(time.time())}.md"
+        md_path = os.path.join("temp_outputs", md_name)
         with open(md_path, "w", encoding="utf-8") as f:
             f.write(polished_report)
+        state["md_path"] = md_path
             
         yield log_content + "✅ Research mission accomplished!", state, gr.update(visible=True), polished_report
     except Exception as e:
@@ -219,18 +207,18 @@ def run_research(hf_key, tavily_key, state):
 
 def reset_all():
     return (
-        gr.update(value="", visible=True), # topic_input
-        gr.update(visible=False), # step2_col
-        gr.update(visible=False), # step3_col
-        gr.update(visible=False), # step4_col
-        gr.update(visible=False), # step5_col
-        ResearchState(), # state
-        "", # log_output
-        "" # final_md_display
+        gr.update(value="", visible=True), 
+        gr.update(visible=False), 
+        gr.update(visible=False), 
+        gr.update(visible=False), 
+        gr.update(visible=False), 
+        {}, 
+        "Ready to research.", 
+        ""
     )
 
 with gr.Blocks(theme=theme, title="Deep Research Agent") as demo:
-    state = gr.State(ResearchState())
+    state = gr.State({})
     
     gr.Markdown("# 🧬 Deep Research Agent")
     gr.Markdown("### The ultimate AI research pipeline that browses the web for you.")
@@ -246,14 +234,11 @@ with gr.Blocks(theme=theme, title="Deep Research Agent") as demo:
             log_output = gr.Markdown("Ready to research.")
 
         with gr.Column(scale=4):
-            # --- STEP 1 ---
             with gr.Column(visible=True) as step1_col:
                 gr.Markdown("## 1️⃣ What are you researching?")
                 topic_input = gr.Textbox(label="Enter a broad topic or research question:", placeholder="e.g., The impact of AI on specialized legal services")
                 start_btn = gr.Button("Start Clarification", variant="primary")
-                step1_msg = gr.Markdown()
 
-            # --- STEP 2 ---
             with gr.Column(visible=False) as step2_col:
                 gr.Markdown("## 2️⃣ Refine Your Topic")
                 suggestion_display = gr.Markdown()
@@ -261,9 +246,7 @@ with gr.Blocks(theme=theme, title="Deep Research Agent") as demo:
                     opt_index = gr.Dropdown(choices=["1", "2", "3"], label="Select Option")
                     custom_topic_input = gr.Textbox(label="Or type your own refined topic:")
                 refine_btn = gr.Button("Use Selected/Custom Topic", variant="primary")
-                step2_msg = gr.Markdown()
 
-            # --- STEP 3 ---
             with gr.Column(visible=False) as step3_col:
                 gr.Markdown("## 3️⃣ Strategy & Task Splitting")
                 target_display = gr.Markdown()
@@ -273,13 +256,10 @@ with gr.Blocks(theme=theme, title="Deep Research Agent") as demo:
                 tasks_display = gr.Markdown(visible=False)
                 execute_btn = gr.Button("🚀 Execute Research Agents", variant="primary", visible=False)
 
-            # --- STEP 4 ---
             with gr.Column(visible=False) as step4_col:
                 gr.Markdown("## 4️⃣ Agentic Research in Progress")
-                # Removed log_output from here since it's move to the sidebar
                 gr.Markdown("Agents are browsing the web... check the status panel!")
 
-            # --- STEP 5 ---
             with gr.Column(visible=False) as step5_col:
                 gr.Markdown("## 🏁 Final Research Report")
                 final_md_display = gr.Markdown()
@@ -288,7 +268,6 @@ with gr.Blocks(theme=theme, title="Deep Research Agent") as demo:
                     download_pdf = gr.File(label="Download PDF")
                 new_research_btn = gr.Button("Start New Research", variant="primary")
 
-    # --- CALLBACKS ---
     start_btn.click(
         start_clarification, 
         inputs=[topic_input, hf_key_input, state], 
@@ -324,9 +303,9 @@ with gr.Blocks(theme=theme, title="Deep Research Agent") as demo:
         inputs=[hf_key_input, tavily_key_input, state],
         outputs=[log_output, state, step5_col, final_md_display]
     ).then(
-        lambda s: [s.pdf_path, s.pdf_path.replace(".pdf", ".md") if s.pdf_path else None],
+        lambda s: (s.get("md_path"), s.get("pdf_path")),
         inputs=[state],
-        outputs=[download_pdf, download_md]
+        outputs=[download_md, download_pdf]
     )
 
     reset_btn.click(
@@ -340,4 +319,5 @@ with gr.Blocks(theme=theme, title="Deep Research Agent") as demo:
     )
 
 if __name__ == "__main__":
+    demo.queue()
     demo.launch()