Update app.py

app.py

@@ -2,7 +2,8 @@ import os
 import re
 import json
 import requests
-from typing import List, Dict
+from typing import List, Dict, Optional, Tuple
+import gradio as gr
 from googlesearch import search
 import google.generativeai as genai
 from google.generativeai.types import HarmCategory, HarmBlockThreshold
@@ -156,44 +157,123 @@ NAICS_CODE: [6-digit NAICS code]
             "reasoning": f"Error analyzing company: {str(e)}"
         }
 
-def main():
-    """Main function to run the NAICS classifier"""
-    print("🚀 NAICS Code Finder\n")
-
-    # Step 1: Get API Key
-    api_key = input("Enter your Google Gemini API Key: ")
-    model = initialize_gemini(api_key)
-
-    while True:
-        # Step 2: Get Company Info
-        company_name = input("\nEnter the company name (or 'exit' to quit): ")
-        if company_name.lower() == 'exit':
-            break
-            
-        context = input("Enter a brief description of the company (or press Enter for none): ")
+def find_naics_code(company_name: str, context: str = "", api_key: Optional[str] = None) -> Dict:
+    """
+    Core function to find NAICS code for a company that can be called from different interfaces
+    
+    Args:
+        company_name: Name of the company
+        context: Brief description of the company (optional)
+        api_key: Google Gemini API key (if None, will try to get from environment variable)
         
-        # Step 3: Find NAICS Code Candidates
-        naics_candidates = google_search_naics(company_name)
+    Returns:
+        Dictionary with NAICS code, reasoning, and optional research
+    """
+    # Get API key from environment if not provided
+    if not api_key:
+        api_key = os.environ.get('GEMINI_API_KEY')
+        if not api_key:
+            return {
+                "error": "No API key provided. Set GEMINI_API_KEY environment variable or pass as parameter.",
+                "naics_code": "000000",
+                "reasoning": "Error: API key missing"
+            }
+    
+    # Initialize Gemini model
+    try:
+        model = initialize_gemini(api_key)
+    except Exception as e:
+        return {
+            "error": f"Failed to initialize Gemini API: {str(e)}",
+            "naics_code": "000000",
+            "reasoning": f"Error: {str(e)}"
+        }
+    
+    # Find NAICS Code Candidates via Google search
+    naics_candidates = google_search_naics(company_name)
 
-        if not naics_candidates:
-            print("❌ No NAICS codes found from Google search.")
-            # Ask Gemini to suggest a code even without candidates
-            result = get_naics_classification(model, company_name, context, [])
-        else:
-            print(f"✅ Found {len(naics_candidates)} NAICS candidates: {naics_candidates}")
-            # Use Gemini to select the best code
-            result = get_naics_classification(model, company_name, context, naics_candidates)
+    # Get classification from Gemini
+    if not naics_candidates:
+        print("No NAICS codes found from Google search.")
+        result = get_naics_classification(model, company_name, context, [])
+    else:
+        print(f"Found {len(naics_candidates)} NAICS candidates: {naics_candidates}")
+        result = get_naics_classification(model, company_name, context, naics_candidates)
+    
+    # Add metadata
+    result["company_name"] = company_name
+    result["context"] = context
+    result["candidates"] = naics_candidates
+    
+    return result
 
-        # Display research findings if available
-        if "research" in result:
-            print(f"\n📊 NAICS Code Research:\n{result['research']}")
+# Gradio interface function
+def classify_company(company_name: str, company_description: str, api_key: str = None) -> Tuple[str, str, str]:
+    """Process inputs from Gradio and return formatted results"""
+    if not api_key:
+        api_key = os.environ.get('GEMINI_API_KEY')
+    
+    if not company_name:
+        return "Error: Company name is required", "", ""
+    
+    result = find_naics_code(company_name, company_description, api_key)
+    
+    # Format the NAICS code output
+    naics_code = f"**NAICS Code: {result['naics_code']}**"
+    
+    # Format the research output
+    research = ""
+    if "research" in result and result["research"]:
+        research = f"## Research on NAICS Codes\n\n{result['research']}"
+    
+    # Format the reasoning output
+    reasoning = f"## Analysis\n\n{result['reasoning']}"
+    
+    return naics_code, research, reasoning
 
-        # Display reasoning
-        print(f"\n🧠 Reasoning:\n{result['reasoning']}")
+# Create the Gradio interface
+def create_gradio_interface():
+    with gr.Blocks(title="NAICS Code Finder") as demo:
+        gr.Markdown("# NAICS Code Finder")
+        gr.Markdown("Enter a company name and optional description to find the most appropriate NAICS code.")
+        
+        with gr.Row():
+            with gr.Column():
+                company_name = gr.Textbox(label="Company Name", placeholder="Enter company name")
+                company_description = gr.Textbox(label="Company Description (optional)", placeholder="Brief description of the company")
+                api_key = gr.Textbox(
+                    label="Gemini API Key (optional)", 
+                    placeholder="Enter your API key or set GEMINI_API_KEY env variable",
+                    visible=not bool(os.environ.get('GEMINI_API_KEY'))
+                )
+                submit_btn = gr.Button("Find NAICS Code")
             
-        # Output the NAICS code
-        print(f"\n🏆 NAICS Code: {result['naics_code']}")
-        print("-" * 80)
+            with gr.Column():
+                naics_output = gr.Markdown(label="NAICS Code")
+                research_output = gr.Markdown(label="Research")
+                reasoning_output = gr.Markdown(label="Reasoning")
+        
+        submit_btn.click(
+            classify_company, 
+            inputs=[company_name, company_description, api_key], 
+            outputs=[naics_output, research_output, reasoning_output]
+        )
+        
+        gr.Examples(
+            [
+                ["Apple Inc", "Tech company that makes iPhones and computers"],
+                ["Starbucks", "Coffee shop chain"],
+                ["Bank of America", "Banking and financial services"],
+                ["Tesla", "Electric vehicle manufacturer"]
+            ],
+            inputs=[company_name, company_description]
+        )
+        
+    return demo
+
+# Create and launch the interface
+demo = create_gradio_interface()
 
+# For Spaces deployment
 if __name__ == "__main__":
-    main()
+    demo.launch()