Update app.py

app.py

@@ -1,19 +1,163 @@
 import gradio as gr
-from model import search_articles, extract_entities, find_full_names
+import requests
+import re
+from typing import List, Dict
+import os
+import time
+from duckduckgo_search import DDGS
 
-def process_name(name: str):
-    """Process name through search and entity extraction pipeline"""
+# 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))
+        if not results:
+            return f"No articles found for {name}"
+        articles = []
+        for i, result in enumerate(results, 1):
+            article = f"**{i}. {result['title']}**\n"
+            article += f"Source: {result['href']}\n"
+            article += f"{result['body']}\n"
+            articles.append(article)
+        return "\n\n".join(articles)
+    except Exception as e:
+        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
+            }
+        )
+        if response.status_code == 200:
+            return response.json().get("response", "No entities extracted")
+        else:
+            return f"API Error: {response.status_code}"
+    except Exception as e:
+        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
+            }
+        )
+        if response.status_code == 200:
+            return response.json().get("response", "No full names found")
+        else:
+            return f"API Error: {response.status_code}"
+    except Exception as e:
+        return f"Full name extraction failed: {str(e)}"
+
+# Gradio interface functions
+def process_name_with_progress(name: str, progress=gr.Progress()):
+    """Process name through search and entity extraction pipeline with real-time updates"""
     if not name.strip():
         return "", "", ""
     
+    # Initialize outputs
+    search_results = ""
+    entities = ""
+    full_names = ""
+    
+    try:
+        # Step 1: Search for articles
+        progress(0.1, desc="Searching for articles...")
+        yield "Searching for articles...", "", ""
+        
+        search_start = time.time()
+        search_results = search_articles(name.strip())
+        search_time = time.time() - search_start
+        print(f"Search took: {search_time:.2f} seconds")
+        
+        progress(0.4, desc="Articles found! Extracting entities...")
+        yield search_results, "Extracting entities from articles...", ""
+        
+        # Step 2: Extract entities from search results
+        extract_start = time.time()
+        entities = extract_entities(search_results)
+        extract_time = time.time() - extract_start
+        print(f"Entity extraction took: {extract_time:.2f} seconds")
+        
+        progress(0.7, desc="Entities extracted! Finding full names...")
+        yield search_results, entities, "Finding full names and roles..."
+        
+        # Step 3: Find full names of entities
+        names_start = time.time()
+        full_names = find_full_names(search_results, entities)
+        names_time = time.time() - names_start
+        print(f"Full name extraction took: {names_time:.2f} seconds")
+        
+        progress(1.0, desc="Complete!")
+        yield search_results, entities, full_names
+        
+    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
+
+def process_name_simple(name: str):
+    """Simple version without progress - for debugging bottlenecks"""
+    if not name.strip():
+        return "", "", ""
+    
+    print(f"Starting process for: {name}")
+    total_start = time.time()
+    
     # Search for articles
+    print("Step 1: Searching articles...")
+    search_start = time.time()
     search_results = search_articles(name.strip())
+    search_time = time.time() - search_start
+    print(f"Search completed in: {search_time:.2f} seconds")
     
     # Extract entities from search results
+    print("Step 2: Extracting entities...")
+    extract_start = time.time()
     entities = extract_entities(search_results)
+    extract_time = time.time() - extract_start
+    print(f"Entity extraction completed in: {extract_time:.2f} seconds")
     
     # Find full names of entities
+    print("Step 3: Finding full names...")
+    names_start = time.time()
     full_names = find_full_names(search_results, entities)
+    names_time = time.time() - names_start
+    print(f"Full name extraction completed in: {names_time:.2f} seconds")
+    
+    total_time = time.time() - total_start
+    print(f"Total process time: {total_time:.2f} seconds")
     
     return search_results, entities, full_names
 
@@ -24,20 +168,28 @@ with gr.Blocks(title="Name Research Tool") as demo:
     
     with gr.Row():
         name_input = gr.Textbox(label="Name", placeholder="Enter business or project name")
-        search_btn = gr.Button("Search", variant="primary")
+        with gr.Column():
+            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)
     
+    # Real-time search with progress
     search_btn.click(
-        fn=process_name,
+        fn=process_name_with_progress,
         inputs=[name_input],
         outputs=[output1, output2, output3]
     )
     
-    # TODO: Add CSV upload functionality here
+    # Debug search with timing info
+    debug_btn.click(
+        fn=process_name_simple,
+        inputs=[name_input],
+        outputs=[output1, output2, output3]
+    )
 
 if __name__ == "__main__":
     demo.launch()