Update app.py

app.py

@@ -32,157 +32,121 @@ def initialize_gemini(api_key: str):
     )
     return model
 
-def google_search_company_info(company_name: str) -> str:
+def combined_google_search(company_name: str) -> Tuple[str, List[str]]:
     """
-    Search for basic company information to help with NAICS classification
+    Combined search function that finds both company information and NAICS codes
+    
+    Returns:
+        Tuple containing (company_info, naics_code_candidates)
     """
     company_info = ""
+    naics_codes = set()
     
-    # 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"
+    # Create comprehensive search queries
+    info_queries = [
+        f"what is {company_name} company business industry sector",
+        f"{company_name} company about us business description",
+        f"{company_name} company profile what they 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
-                            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 ""
-
-def google_search_naics(company_name: str) -> List[str]:
-    """
-    Find potential NAICS codes for a company using multiple targeted Google searches
-    Uses more specific search queries to improve results
-    """
-    naics_codes = set()
-    
-    # Create multiple search queries for better results
-    queries = [
-        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"
+    naics_queries = [
+        f"2022 NAICS code for {company_name} company",
+        f"{company_name} NAICS 2022 classification",
+        f"what is {company_name} industry NAICS code 2022"
     ]
     
+    all_queries = info_queries + naics_queries
+    
     try:
-        print(f"🔎 Searching Google for NAICS codes for '{company_name}'...")
+        print(f"🔍 Searching for information about '{company_name}'...")
         
-        for query in queries:
+        for query in all_queries:
             print(f"  Query: {query}")
             try:
-                # Search with each query, limiting to 3 results per query
+                # Search with each query
                 search_results = search(query, stop=3, pause=2)
                 
                 for result_url in search_results:
                     try:
                         response = requests.get(result_url, timeout=5)
                         if response.status_code == 200:
-                            # Extract 6-digit NAICS codes
+                            # Extract NAICS codes
                             found_codes = re.findall(r'\b\d{6}\b', response.text)
-                            naics_codes.update(found_codes)
-                            
-                            # If we find codes, print them
                             if found_codes:
+                                naics_codes.update(found_codes)
                                 print(f"  Found codes in {result_url}: {found_codes}")
+                            
+                            # Extract company information
+                            if len(company_info) < 1000:  # Only if we need more info
+                                soup = BeautifulSoup(response.text, 'html.parser')
+                                paragraphs = soup.find_all('p')
+                                
+                                # Get text from paragraphs that mention the company
+                                for p in paragraphs:
+                                    text = p.get_text().strip()
+                                    if len(text) > 80 and company_name.lower() in text.lower():
+                                        company_info += text + "\n\n"
+                                        if len(company_info) > 1000:
+                                            break
+                            
                     except Exception as e:
                         print(f"  ⚠️ Error fetching {result_url}: {e}")
+                        
+                # If we have enough information, move to the next query
+                if len(company_info) > 1000 and len(naics_codes) > 0:
+                    break
+                    
             except Exception as e:
                 print(f"  ⚠️ Error with query '{query}': {e}")
                 continue
         
-        # Return unique codes, limited to 10 most common
-        return list(naics_codes)[:10]
+        # Return company info and NAICS codes
+        return company_info.strip(), list(naics_codes)[:10]
     except Exception as e:
-        print(f"❌ Error performing Google search: {str(e)}")
-        return []
+        print(f"❌ Error during Google search: {str(e)}")
+        return "", []
 
-def get_naics_classification(model, company_name: str, context: str, candidates: List[str]) -> dict:
+def analyze_naics_code(model, company_name: str, context: str, company_info: str, naics_candidates: List[str]) -> dict:
     """
-    Use Gemini AI to determine the most appropriate NAICS code from candidates
-    First provides reasoning, then returns the NAICS code and explanation
+    Use Gemini AI to determine the most appropriate NAICS code
     """
     try:
         print("🤖 AI is analyzing NAICS classification...")
         
-        # Get additional company information from Google
-        company_info = google_search_company_info(company_name)
+        # Combine provided context with discovered company info
         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}"
+                combined_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
+                combined_context = f"Information found online:\n{company_info}"
+        else:
+            combined_context = context
+            
+        # Create the prompt based on whether we have candidate codes
+        if naics_candidates:
             prompt = f"""
-You are a NAICS code classification expert. Based on the company information provided and the NAICS code candidates found from Google search, determine the most appropriate NAICS code.
+You are a NAICS code classification expert. Based on the company information provided and any NAICS code candidates found from online research, determine the most appropriate NAICS code.
 
 Company Name: {company_name}
-Context Information: {context}
+Information about the company: {combined_context}
 
-NAICS Code Candidates from Google Search: {candidates}
+NAICS Code Candidates found in research: {naics_candidates}
 
-First, research what these NAICS codes represent:
-1. For each NAICS code candidate, briefly explain what industry or business activity it corresponds to.
-2. Then explain which industry classification best matches this company based on the name and context provided.
-3. Finally, select the single most appropriate NAICS code from the candidates, or suggest a different one if none match.
+First, analyze what these NAICS codes represent and which industry this company belongs to based on the information provided.
+Then select the single most appropriate 6-digit NAICS code.
 
 Your response should be in this format:
-RESEARCH: [Brief explanation of what each NAICS code represents]
-REASONING: [Your detailed reasoning about why the chosen industry classification is most appropriate for this company]
+REASONING: [Your detailed reasoning about why the chosen industry classification is most appropriate for this company, including what business activities it performs]
 NAICS_CODE: [6-digit NAICS code]
 """
-        # If no candidates were found from Google search
         else:
             prompt = f"""
 You are a NAICS code classification expert. Based on the company information provided, determine the most appropriate NAICS code.
 
 Company Name: {company_name}
-Context Information: {context}
+Information about the company: {combined_context}
 
-First, analyze what industry this company likely belongs to based on its name and the provided context.
+Analyze what industry this company likely belongs to based on its name and the provided information.
 Consider standard business classifications and determine the most appropriate category.
 Then provide the single most appropriate 6-digit NAICS code.
 
@@ -196,12 +160,6 @@ NAICS_CODE: [6-digit NAICS code]
         # Create result dictionary
         result = {}
         
-        # Extract research if available
-        if "RESEARCH:" in response_text:
-            research_match = re.search(r'RESEARCH:(.*?)REASONING:', response_text, re.DOTALL | re.IGNORECASE)
-            if research_match:
-                result["research"] = research_match.group(1).strip()
-        
         # Extract reasoning
         reasoning_match = re.search(r'REASONING:(.*?)NAICS_CODE:', response_text, re.DOTALL | re.IGNORECASE)
         result["reasoning"] = reasoning_match.group(1).strip() if reasoning_match else "No reasoning provided."
@@ -225,15 +183,7 @@ NAICS_CODE: [6-digit NAICS code]
 
 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)
-        
-    Returns:
-        Dictionary with NAICS code, reasoning, and optional research
+    Core function to find NAICS code for a company
     """
     # Get API key from environment if not provided
     if not api_key:
@@ -255,20 +205,16 @@ def find_naics_code(company_name: str, context: str = "", api_key: Optional[str]
             "reasoning": f"Error: {str(e)}"
         }
     
-    # Find NAICS Code Candidates via Google search
-    naics_candidates = google_search_naics(company_name)
+    # Run the combined search
+    company_info, naics_candidates = combined_google_search(company_name)
 
-    # 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)
+    # Get AI analysis
+    result = analyze_naics_code(model, company_name, context, company_info, naics_candidates)
     
     # Add metadata
     result["company_name"] = company_name
     result["context"] = context
+    result["company_info"] = company_info
     result["candidates"] = naics_candidates
     
     return result
@@ -280,7 +226,7 @@ def create_gradio_interface():
     
     with gr.Blocks(title="NAICS Code Finder") as demo:
         gr.Markdown("# NAICS Code Finder")
-        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.")
+        gr.Markdown("Enter a company name to find its appropriate NAICS code. The tool will search for information about the company and find the most appropriate classification.")
         
         with gr.Row():
             with gr.Column():
@@ -304,69 +250,55 @@ def create_gradio_interface():
                 naics_output = gr.Markdown(label="NAICS Code")
                 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", "", "", "", ""
+                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.", "", "", "", ""
+                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, "", ""
+            status_md = "🔍 Searching for company information and NAICS codes...\n\n"
+            yield status_md, "", "", ""
             
             # Run the core functionality
             result = find_naics_code(company_name, company_description, key_to_use)
             
+            # Update status based on results
+            if "company_info" in result and result["company_info"]:
+                status_md += "✅ Found company information\n\n"
+                company_info_md = f"## Information found about {company_name}\n\n{result['company_info']}"
+            else:
+                status_md += "⚠️ Limited company information found\n\n"
+                company_info_md = f"Limited information found for {company_name}"
+            
             if "candidates" in result and result["candidates"]:
-                status_md += f"✅ Found {len(result['candidates'])} potential NAICS codes\n\n"
+                status_md += f"✅ Found {len(result['candidates'])} potential NAICS codes: {', '.join(result['candidates'])}\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, "", ""
+            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
+            return status_md, naics_code_md, company_info_md, reasoning_md
         
         submit_btn.click(
             process_company, 
             inputs=[company_name, company_description, api_key], 
-            outputs=[status_output, naics_output, company_info_output, research_output, reasoning_output]
+            outputs=[status_output, naics_output, company_info_output, reasoning_output]
         )
         
         gr.Examples(