Update app.py

app.py

@@ -1,20 +1,12 @@
 import gradio as gr
 import requests
-import re
-from typing import List, Dict
-import os
 import time
 from duckduckgo_search import DDGS
 
-# ----------------------------
-# Model functions
-# ----------------------------
-
+# --- Model Functions ---
 def search_articles(name: str) -> str:
-    """Search for 3 newspaper articles containing the name and keywords using DuckDuckGo"""
     keywords = ['founders', 'partners', 'funders', 'owners']
     search_query = f'"{name}" ({" OR ".join(keywords)}) site:news'
-
     try:
         with DDGS() as ddgs:
             results = list(ddgs.text(search_query, max_results=3))
@@ -31,20 +23,17 @@ def search_articles(name: str) -> str:
         return f"Search failed: {str(e)}"
 
 def extract_entities(search_results: str) -> str:
-    """Extract entities using Mistral 7B endpoint"""
     modal_endpoint = "https://msoaresdiego--mistral-llm-endpoint-fastapi-app.modal.run/generate"
-
     prompt = f"""Extract all person names and organization names from the following text.
 Format as:
 PERSON: [name]
 ORG: [organization name]
 Text: {search_results}"""
-
     try:
         response = requests.post(
             modal_endpoint,
             json={"prompt": prompt, "max_tokens": 500, "temperature": 0.1},
-            timeout=90  # Increased timeout
+            timeout=180  # Increased timeout
         )
         if response.status_code == 200:
             return response.json().get("response", "No entities extracted")
@@ -54,19 +43,16 @@ Text: {search_results}"""
         return f"Extraction failed: {str(e)}"
 
 def find_full_names(search_results: str, entities: str) -> str:
-    """Find full names using Mistral 7B endpoint"""
     modal_endpoint = "https://msoaresdiego--mistral-llm-endpoint-fastapi-app.modal.run/generate"
-
     prompt = f"""Based on the search results, find the full names and titles/roles for these entities:
 Entities: {entities}
 Search Results: {search_results}
 Provide full names with their roles/titles where mentioned."""
-
     try:
         response = requests.post(
             modal_endpoint,
             json={"prompt": prompt, "max_tokens": 300, "temperature": 0.1},
-            timeout=90  # Increased timeout
+            timeout=180  # Increased timeout
         )
         if response.status_code == 200:
             return response.json().get("response", "No full names found")
@@ -75,81 +61,60 @@ Provide full names with their roles/titles where mentioned."""
     except Exception as e:
         return f"Full name extraction failed: {str(e)}"
 
-# ----------------------------
-# Gradio Interface Logic
-# ----------------------------
-
+# --- Pipeline Function with Progress & Status ---
 def process_name_with_progress(name: str, progress=gr.Progress(track_tqdm=True)):
     if not name.strip():
-        yield "", "", ""
+        yield "", "", "", "Please enter a name."
 
-    # Start with all outputs empty
     search_results = ""
     entities = ""
     full_names = ""
 
     try:
-        # Step 1: Search
-        progress(0.1, desc="Searching for articles...")
+        progress(0.1)
+        status = "🔍 Searching for articles..."
         search_results = search_articles(name.strip())
-        yield search_results, "", ""
+        yield search_results, "", "", status
 
-        # Step 2: Extract Entities
-        progress(0.4, desc="Extracting entities...")
+        progress(0.4)
+        status = "📄 Extracting entities from articles..."
         entities = extract_entities(search_results)
-        yield search_results, entities, ""
+        yield search_results, entities, "", status
 
-        # Step 3: Resolve Full Names
-        progress(0.7, desc="Finding full names...")
+        progress(0.7)
+        status = "🧠 Finding full names and roles..."
         full_names = find_full_names(search_results, entities)
-        progress(1.0, desc="Done.")
-        yield search_results, entities, full_names
+        yield search_results, entities, full_names, status
 
-    except Exception as e:
-        error_msg = f"Error: {str(e)}"
-        yield search_results or error_msg, entities or error_msg, full_names or error_msg
+        progress(1.0)
+        yield search_results, entities, full_names, "✅ Complete!"
 
-def process_name_simple(name: str):
-    """Simple debugging version (no progress bars)"""
-    if not name.strip():
-        return "", "", ""
-
-    search_results = search_articles(name.strip())
-    entities = extract_entities(search_results)
-    full_names = find_full_names(search_results, entities)
-    return search_results, entities, full_names
-
-# ----------------------------
-# Gradio UI
-# ----------------------------
+    except Exception as e:
+        err = f"❌ Error: {str(e)}"
+        yield search_results or err, entities or err, full_names or err, err
 
+# --- Gradio Interface ---
 with gr.Blocks(title="Name Research Tool") as demo:
-    gr.Markdown("# 🔍 Name Research Tool")
+    gr.Markdown("# Name Research Tool")
     gr.Markdown("Enter a business or project name to search for related articles and extract key entities.")
 
     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")
-            debug_btn = gr.Button("🐞 Search (Debug Mode)", variant="secondary")
+            search_btn = gr.Button("Search (Real-time)", variant="primary")
+            debug_btn = gr.Button("Search (Debug Mode)", variant="secondary")
 
     with gr.Column():
         output1 = gr.Textbox(label="Search Results", lines=10, max_lines=20)
         output2 = gr.Textbox(label="Extracted Entities", lines=5, max_lines=10)
         output3 = gr.Textbox(label="Full Names", lines=5, max_lines=10)
+        status_output = gr.Textbox(label="Status / Progress", lines=1, interactive=False)
 
-    # Real-time (streamed) version
+    # Real-time search with progress
     search_btn.click(
         fn=process_name_with_progress,
         inputs=[name_input],
-        outputs=[output1, output2, output3]
-    )
-
-    # Debug version (no progress)
-    debug_btn.click(
-        fn=process_name_simple,
-        inputs=[name_input],
-        outputs=[output1, output2, output3]
+        outputs=[output1, output2, output3, status_output]
     )
 
 if __name__ == "__main__":