Update app.py

app.py

@@ -1,4 +1,58 @@
-import os
+def google_search_company_info(company_name: str) -> str:
+    """
+    Search for basic company information to help with NAICS classification
+    """
+    company_info = ""
+    
+    # Create search queries focused on company information
+    queries = [
+        f"what is {company_name} company",
+        f"{company_name} company about us",
+        f"{company_name} business description",
+        f"{company_name} company profile",
+        f"what does {company_name} company do"
+    ]
+    
+    try:
+        print(f"🔍 Searching for information about '{company_name}'...")
+        
+        for query in queries[:2]:  # Limit to first 2 queries to save time
+            try:
+                # Search with each query
+                search_results = search(query, stop=2, pause=2)
+                
+                for result_url in search_results:
+                    try:
+                        response = requests.get(result_url, timeout=5)
+                        if response.status_code == 200:
+                            # Extract text from paragraphs
+                            from bs4 import BeautifulSoup
+                            soup = BeautifulSoup(response.text, 'html.parser')
+                            paragraphs = soup.find_all('p')
+                            
+                            # Get text from first 3 substantial paragraphs
+                            for p in paragraphs:
+                                text = p.get_text().strip()
+                                if len(text) > 100 and company_name.lower() in text.lower():
+                                    company_info += text + "\n\n"
+                                    if len(company_info) > 500:
+                                        break
+                            
+                            if len(company_info) > 500:
+                                break
+                    except Exception as e:
+                        print(f"  ⚠️ Error fetching {result_url}: {e}")
+                        
+                if len(company_info) > 500:
+                    break
+            except Exception as e:
+                print(f"  ⚠️ Error with query '{query}': {e}")
+                continue
+        
+        return company_info.strip()
+    except Exception as e:
+        print(f"❌ Error searching for company info: {str(e)}")
+        return ""import os
 import re
 import json
 import requests
@@ -40,11 +94,11 @@ def google_search_naics(company_name: str) -> List[str]:
     
     # Create multiple search queries for better results
     queries = [
-        f"NAICS code for {company_name}",
-        f"what is {company_name} company NAICS code",
-        f"{company_name} business entity NAICS classification",
-        f"{company_name} industry classification NAICS",
-        f"{company_name} company information NAICS"
+        f"2022 NAICS code for {company_name}",
+        f"NAICS 2022 classification for {company_name}",
+        f"{company_name} business NAICS 2022 code",
+        f"{company_name} industry NAICS code 2022",
+        f"what is {company_name} company NAICS code"
     ]
     
     try:
@@ -87,6 +141,16 @@ def get_naics_classification(model, company_name: str, context: str, candidates:
     try:
         print("🤖 AI is analyzing NAICS classification...")
         
+        # Get additional company information from Google
+        company_info = google_search_company_info(company_name)
+        if company_info:
+            print(f"📝 Found additional company information:\n{company_info[:200]}...")
+            # Add the found information to the context
+            if context:
+                context = f"{context}\n\nAdditional information found online:\n{company_info}"
+            else:
+                context = f"Information found online:\n{company_info}"
+        
         # If we have candidate codes from Google search
         if candidates:
             # Create a prompt that asks for research on the candidates
@@ -233,38 +297,106 @@ def classify_company(company_name: str, company_description: str, api_key: str =
 
 # Create the Gradio interface
 def create_gradio_interface():
+    # Check if API key is set in environment
+    has_api_key = bool(os.environ.get('GEMINI_API_KEY'))
+    
     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.")
+        gr.Markdown("Enter a company name to find its appropriate NAICS code. The tool will search for information about the company and relevant NAICS codes online.")
         
         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")
+                company_description = gr.Textbox(label="Additional Context (optional)", placeholder="Any additional information about the company")
+                
+                # Only show API key input if not set in environment
+                if not has_api_key:
+                    api_key = gr.Textbox(
+                        label="Gemini API Key (required)", 
+                        placeholder="Enter your Google Gemini API key",
+                        type="password"
+                    )
+                else:
+                    api_key = gr.Textbox(visible=False, value="")
+                
+                submit_btn = gr.Button("Find NAICS Code", variant="primary")
             
             with gr.Column():
+                status_output = gr.Markdown(label="Status")
                 naics_output = gr.Markdown(label="NAICS Code")
-                research_output = gr.Markdown(label="Research")
-                reasoning_output = gr.Markdown(label="Reasoning")
+                with gr.Accordion("Company Information", open=False):
+                    company_info_output = gr.Markdown()
+                with gr.Accordion("NAICS Codes Research", open=False):
+                    research_output = gr.Markdown()
+                with gr.Accordion("Classification Reasoning", open=True):
+                    reasoning_output = gr.Markdown()
+        
+        # Functions for the interface
+        def process_company(company_name, company_description, api_key):
+            if not company_name:
+                return "Please enter a company name", "", "", "", ""
+            
+            # Use API key from input or environment
+            key_to_use = api_key if api_key else os.environ.get('GEMINI_API_KEY')
+            if not key_to_use:
+                return "No API key provided. Please enter your Gemini API key.", "", "", "", ""
+            
+            status_md = "🔍 Searching for company information...\n\n"
+            yield status_md, "", "", "", ""
+            
+            # Get company info first
+            company_info = google_search_company_info(company_name)
+            if company_info:
+                company_info_md = f"## Information found about {company_name}\n\n{company_info}"
+                status_md += "✅ Found company information\n\n"
+            else:
+                company_info_md = f"No detailed information found for {company_name}"
+                status_md += "⚠️ No company information found\n\n"
+            
+            yield status_md, "", company_info_md, "", ""
+                
+            # Get NAICS candidates
+            status_md += "🔍 Searching for NAICS codes...\n\n"
+            yield status_md, "", company_info_md, "", ""
+            
+            # Run the core functionality
+            result = find_naics_code(company_name, company_description, key_to_use)
+            
+            if "candidates" in result and result["candidates"]:
+                status_md += f"✅ Found {len(result['candidates'])} potential NAICS codes\n\n"
+            else:
+                status_md += "⚠️ No specific NAICS codes found in search results\n\n"
+                
+            status_md += "🤖 Analyzing classification...\n\n"
+            yield status_md, "", company_info_md, "", ""
+            
+            # Format the NAICS code output
+            naics_code_md = f"## NAICS Code: {result['naics_code']}"
+            
+            # Format the research output
+            research_md = ""
+            if "research" in result and result["research"]:
+                research_md = f"## Research on NAICS Codes\n\n{result['research']}"
+                
+            # Format the reasoning output
+            reasoning_md = f"## Analysis\n\n{result['reasoning']}"
+            
+            status_md += "✅ Classification complete!"
+            
+            return status_md, naics_code_md, company_info_md, research_md, reasoning_md
         
         submit_btn.click(
-            classify_company, 
+            process_company, 
             inputs=[company_name, company_description, api_key], 
-            outputs=[naics_output, research_output, reasoning_output]
+            outputs=[status_output, naics_output, company_info_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"]
+                ["Apple Inc", "Tech company"],
+                ["Walmart", "Retail store"],
+                ["Goldman Sachs", "Investment bank"],
+                ["Ford Motor Company", "Automobile manufacturer"]
             ],
             inputs=[company_name, company_description]
         )