app.py

import gradio as gr
import requests
import time
from duckduckgo_search import DDGS

# === Model functions ===

def search_articles(name: str) -> str:
    """Search for 2 newspaper articles containing the name and keywords using DuckDuckGo"""
    keywords = ['founders', 'partners', 'funders', 'owners']
    search_query = f'"{name}" ({" OR ".join(keywords)}) site:news'
    try:
        print(f"[DEBUG] Search query: {search_query}")
        with DDGS() as ddgs:
            results = list(ddgs.text(search_query, max_results=2))
            print(f"[DEBUG] Raw results: {results}")
        if not results:
            return f"No articles found for {name}"

        articles = []
        for i, result in enumerate(results, 1):
            article = f"**{i}. {result.get('title', 'No Title')}**\n"
            article += f"Source: {result.get('href', 'No URL')}\n"
            article += f"{result.get('body', 'No Body')}\n"
            articles.append(article)

        return "\n\n".join(articles)
    except Exception as e:
        return f"[ERROR] 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.Do not extract products and service names. Only individuals and organizations.Bring the full details of the name in the newspaper article. For example, if only ACME is mentioned as company name, bring only ACME. IF ACME Inc is mentioned as company name, then you have to extract ACME Inc. In addition, define the relationship between the entity and the company that is being searched. For example, is ACME Inc an owner of the company being searched? Then write 'owner'. Is ACME Inc. a funder of the company being searched? Then write 'funder' 
Format as:
PERSON: [name]
ORG: [organization name]
Text: {search_results}"""
    try:
        response = requests.post(
            modal_endpoint,
            json={"prompt": prompt, "max_tokens": 2500, "temperature": 0.15}
        )
        if response.status_code == 200:
            return response.json().get("response", "No entities extracted")
        else:
            return f"[ERROR] API Error: {response.status_code}"
    except Exception as e:
        return f"[ERROR] Extraction failed: {str(e)}"


# === Gradio interface functions ===

def search_only(name: str):
    """Perform search only and return results"""
    if not name.strip():
        return "No name provided", ""
    
    try:
        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, ""


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(stored_search_results)
        extract_time = time.time() - extract_start
        
        extraction_results = f"Entity extraction completed in {extract_time:.2f}s\n\n"
        extraction_results += entities
        
        return extraction_results
    except Exception as e:
        return f"[ERROR] Extraction failed: {str(e)}"


# === Gradio UI ===

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 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=search_only,
        inputs=[name_input],
        outputs=[output1, search_state]
    )
    
    # Extract button click
    extract_btn.click(
        fn=extract_only,
        inputs=[search_state],
        outputs=[output2]
    )


if __name__ == "__main__":
    demo.launch()