Update app.py

app.py

@@ -1,12 +1,20 @@
 import gradio as gr
 import requests
+import re
+from typing import List, Dict
+import os
 import time
 from duckduckgo_search import DDGS
 
-# DuckDuckGo search
+# ----------------------------
+# 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))
@@ -22,19 +30,21 @@ def search_articles(name: str) -> str:
     except Exception as e:
         return f"Search failed: {str(e)}"
 
-# Call your Modal Mistral endpoint
 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=30
+            timeout=90  # Increased timeout
         )
         if response.status_code == 200:
             return response.json().get("response", "No entities extracted")
@@ -44,16 +54,19 @@ 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=30
+            timeout=90  # Increased timeout
         )
         if response.status_code == 200:
             return response.json().get("response", "No full names found")
@@ -62,73 +75,81 @@ Provide full names with their roles/titles where mentioned."""
     except Exception as e:
         return f"Full name extraction failed: {str(e)}"
 
-# Pipeline with separate status
-def process_name_with_progress(name: str, progress=gr.Progress()):
+# ----------------------------
+# Gradio Interface Logic
+# ----------------------------
+
+def process_name_with_progress(name: str, progress=gr.Progress(track_tqdm=True)):
     if not name.strip():
-        yield "", "", "", "Please enter a name"
-        return
+        yield "", "", ""
+
+    # Start with all outputs empty
+    search_results = ""
+    entities = ""
+    full_names = ""
 
     try:
-        yield "", "", "", "Searching for articles..."
-        progress(0.1)
+        # Step 1: Search
+        progress(0.1, desc="Searching for articles...")
         search_results = search_articles(name.strip())
+        yield search_results, "", ""
 
-        yield search_results, "", "", "Extracting entities..."
-        progress(0.4)
+        # Step 2: Extract Entities
+        progress(0.4, desc="Extracting entities...")
         entities = extract_entities(search_results)
+        yield search_results, entities, ""
 
-        yield search_results, entities, "", "Finding full names..."
-        progress(0.7)
+        # Step 3: Resolve Full Names
+        progress(0.7, desc="Finding full names...")
         full_names = find_full_names(search_results, entities)
-
-        progress(1.0)
-        yield search_results, entities, full_names, "Done!"
+        progress(1.0, desc="Done.")
+        yield search_results, entities, full_names
 
     except Exception as e:
-        error = f"Error: {str(e)}"
-        yield error, error, error, error
+        error_msg = f"Error: {str(e)}"
+        yield search_results or error_msg, entities or error_msg, full_names or error_msg
 
-# Simple non-streamed debug version
 def process_name_simple(name: str):
+    """Simple debugging version (no progress bars)"""
     if not name.strip():
-        return "", "", "", "Please enter a name"
-    try:
-        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, "Done!"
-    except Exception as e:
-        error = f"Error: {str(e)}"
-        return error, error, error, error
+        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
+# ----------------------------
 
-# UI layout
 with gr.Blocks(title="Name Research Tool") as demo:
     gr.Markdown("# 🔍 Name Research Tool")
-    gr.Markdown("Enter a business or project name to search related articles and extract key entities.")
+    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.Row():
-        with gr.Column():
-            output1 = gr.Textbox(label="Search Results", lines=10)
-            output2 = gr.Textbox(label="Extracted Entities", lines=5)
-            output3 = gr.Textbox(label="Full Names & Roles", lines=5)
-        status = gr.Textbox(label="Status", interactive=False)
+    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)
 
+    # Real-time (streamed) version
     search_btn.click(
         fn=process_name_with_progress,
         inputs=[name_input],
-        outputs=[output1, output2, output3, status]
+        outputs=[output1, output2, output3]
     )
 
+    # Debug version (no progress)
     debug_btn.click(
         fn=process_name_simple,
         inputs=[name_input],
-        outputs=[output1, output2, output3, status]
+        outputs=[output1, output2, output3]
     )
 
 if __name__ == "__main__":