Update app.py

app.py

@@ -50,72 +50,78 @@ Text: {search_results}"""
         return f"[ERROR] Extraction failed: {str(e)}"
 
 
-# === Gradio interface ===
+# === Gradio interface functions ===
 
-def process_name_with_progress(name: str, progress=gr.Progress()):
-    """Process name """
+def search_only(name: str):
+    """Perform search only and return results"""
     if not name.strip():
-        yield "No name provided", "", ""
-        return
-
-    search_results = ""
-    entities = ""
-    full_names = ""
-
+        return "No name provided", ""
+    
     try:
-        # Step 1: Search
-        progress(0.1, desc="Searching for articles...")
-        search_results += f"Starting search for: {name}\n"
-        yield search_results, "", ""
-
         search_start = time.time()
         articles_output = search_articles(name.strip())
         search_time = time.time() - search_start
+        
+        search_results = f"Search completed for: {name} in {search_time:.2f}s\n\n"
+        search_results += articles_output
+        
+        return search_results, articles_output  # Return both display and raw results
+    except Exception as e:
+        error_msg = f"[ERROR] Search failed: {str(e)}"
+        return error_msg, ""
 
-        search_results += f"Search completed in {search_time:.2f}s\n"
-        search_results += f"{articles_output}\n"
-        yield search_results, "", ""
-
-        # Step 2: Extract entities
-        progress(0.5, desc="Extracting entities...")
-        search_results += "Starting entity extraction...\n"
-        yield search_results, "Extracting entities...", ""
 
+def extract_only(stored_search_results: str):
+    """Extract entities from stored search results"""
+    if not stored_search_results.strip():
+        return "No search results available. Please search first."
+    
+    try:
         extract_start = time.time()
-        entities = extract_entities(articles_output)
+        entities = extract_entities(stored_search_results)
         extract_time = time.time() - extract_start
-        search_results += f"Entity extraction completed in {extract_time:.2f}s\n"
-        yield search_results, entities, ""
-
-
+        
+        extraction_results = f"Entity extraction completed in {extract_time:.2f}s\n\n"
+        extraction_results += entities
+        
+        return extraction_results
     except Exception as e:
-        error_msg = f"[ERROR] {str(e)}"
-        yield search_results + error_msg, entities or error_msg, full_names or error_msg
-
+        return f"[ERROR] Extraction failed: {str(e)}"
 
 
 # === Gradio UI ===
 
-with gr.Blocks(title="Related Entitites Finder") as demo:
+with gr.Blocks(title="Related Entities Finder") as demo:
     gr.Markdown("# 🔎 Related Entities Finder")
     gr.Markdown("Enter a business or project name to search for related articles and extract key entities.")
 
+    # State to store search results between operations
+    search_state = gr.State("")
+
     with gr.Row():
         name_input = gr.Textbox(label="Name", placeholder="Enter business or project name")
         with gr.Column():
-            search_btn = gr.Button("Search (Real-time)", variant="primary")
+            search_btn = gr.Button("Search Articles", variant="primary")
+            extract_btn = gr.Button("Extract Entities", variant="secondary")
     
-
     with gr.Column():
         output1 = gr.Textbox(label="Search Results", lines=50, max_lines=100)
         output2 = gr.Textbox(label="Extracted Entities and Relationships", lines=5, max_lines=10)
-       
+    
+    # Search button click
     search_btn.click(
-        fn=process_name_with_progress,
+        fn=search_only,
         inputs=[name_input],
-        outputs=[output1, output2]
+        outputs=[output1, search_state]
+    )
+    
+    # Extract button click
+    extract_btn.click(
+        fn=extract_only,
+        inputs=[search_state],
+        outputs=[output2]
     )
 
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()