Update app.py

app.py

@@ -1,7 +1,6 @@
 import os
 import re
 import json
-import gradio as gr
 import requests
 from typing import List, Dict
 from googlesearch import search
@@ -31,67 +30,82 @@ def initialize_gemini(api_key: str):
     )
     return model
 
-def google_search_naics(company_name: str, company_description: str = "") -> List[str]:
-    """Find potential NAICS codes for a company using Google search with enhanced context"""
-    # Create a more effective search query with company description if available
-    if company_description:
-        query = f"2022 NAICS code for \"{company_name}\" {company_description} industry classification"
-    else:
-        query = f"2022 NAICS code for \"{company_name}\" company industry classification"
-    
+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"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"
+    ]
+    
     try:
-        search_results = search(query, stop=5, pause=2)
+        print(f"🔎 Searching Google for NAICS codes for '{company_name}'...")
         
-        for result_url in search_results:
+        for query in queries:
+            print(f"  Query: {query}")
             try:
-                response = requests.get(result_url, timeout=5)
-                if response.status_code == 200:
-                    # Extract 6-digit NAICS codes
-                    found_codes = re.findall(r'\b\d{6}\b', response.text)
-                    naics_codes.update(found_codes)
+                # Search with each query, limiting to 3 results per 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
+                            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:
+                                print(f"  Found codes in {result_url}: {found_codes}")
+                    except Exception as e:
+                        print(f"  ⚠️ Error fetching {result_url}: {e}")
             except Exception as e:
-                print(f"Error fetching {result_url}: {e}")
+                print(f"  ⚠️ Error with query '{query}': {e}")
+                continue
         
-        return list(naics_codes)[:5]  # Return up to 5 extracted NAICS codes
+        # Return unique codes, limited to 10 most common
+        return list(naics_codes)[:10]
     except Exception as e:
-        print(f"Error performing Google search: {str(e)}")
+        print(f"❌ Error performing Google search: {str(e)}")
         return []
 
-def get_naics_classification(model, company_name: str, context: str, candidates: List[str], search_query: str) -> dict:
+def get_naics_classification(model, company_name: str, context: str, candidates: List[str]) -> dict:
     """
     Use Gemini AI to determine the most appropriate NAICS code from candidates
-    First provides reasoning, then multiple possibilities with confidence levels
+    First provides reasoning, then returns the NAICS code and explanation
     """
     try:
+        print("🤖 AI is analyzing NAICS classification...")
+        
         # If we have candidate codes from Google search
         if candidates:
+            # Create a prompt that asks for research on the 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.
 
 Company Name: {company_name}
 Context Information: {context}
-Google Search Query Used: {search_query}
-NAICS Code Candidates from Google Search: {candidates}
 
-First, start with a section titled "GOOGLE_FINDINGS:" where you describe what the Google search results suggest about this company based on the NAICS codes found.
-
-Then, in a section titled "REASONING:", explain your reasoning for which industry this company belongs to.
+NAICS Code Candidates from Google Search: {candidates}
 
-Then list 3 potential NAICS classifications with confidence percentages (must add up to 100%).
-Finally, provide your final conclusion.
+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.
 
 Your response should be in this format:
-GOOGLE_FINDINGS: [Describe what the Google search results suggest about this company based on the NAICS codes found]
-
-REASONING: [Your detailed reasoning about the company's industry classification]
-
-POSSIBILITY_1: [Industry name] - NAICS Code [6-digit code] - [XX]% confidence
-POSSIBILITY_2: [Industry name] - NAICS Code [6-digit code] - [XX]% confidence  
-POSSIBILITY_3: [Industry name] - NAICS Code [6-digit code] - [XX]% confidence
-
-CONCLUSION: I am [XX]% confident this company is [industry description] which is NAICS code [6-digit code]
+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]
+NAICS_CODE: [6-digit NAICS code]
 """
         # If no candidates were found from Google search
         else:
@@ -100,157 +114,86 @@ You are a NAICS code classification expert. Based on the company information pro
 
 Company Name: {company_name}
 Context Information: {context}
-Google Search Query Used: {search_query}
-
-First, start with a section titled "GOOGLE_FINDINGS:" where you acknowledge that the Google search did not return any specific NAICS codes for this company.
 
-Then, in a section titled "REASONING:", explain your reasoning for which industry this company belongs to based on the limited information available.
-
-Then list 3 potential NAICS classifications with confidence percentages (must add up to 100%).
-Finally, provide your final conclusion.
+First, analyze what industry this company likely belongs to based on its name and the provided context.
+Consider standard business classifications and determine the most appropriate category.
+Then provide the single most appropriate 6-digit NAICS code.
 
 Your response should be in this format:
-GOOGLE_FINDINGS: No specific NAICS codes were found in the Google search results using the query "{search_query}".
-
-REASONING: [Your detailed reasoning about the company's industry classification based on the limited information available]
-
-POSSIBILITY_1: [Industry name] - NAICS Code [6-digit code] - [XX]% confidence
-POSSIBILITY_2: [Industry name] - NAICS Code [6-digit code] - [XX]% confidence  
-POSSIBILITY_3: [Industry name] - NAICS Code [6-digit code] - [XX]% confidence
-
-CONCLUSION: I am [XX]% confident this company is [industry description] which is NAICS code [6-digit code]
+REASONING: [Your detailed reasoning about the company's industry classification, including what business activities it likely performs]
+NAICS_CODE: [6-digit NAICS code]
 """
         response = model.generate_content(prompt)
         response_text = response.text.strip()
         
-        # Extract Google findings
-        google_findings_match = re.search(r'GOOGLE_FINDINGS:(.*?)REASONING:', response_text, re.DOTALL | re.IGNORECASE)
-        google_findings = google_findings_match.group(1).strip() if google_findings_match else "No Google findings provided."
-        
-        # Extract reasoning
-        reasoning_match = re.search(r'REASONING:(.*?)POSSIBILITY_1:', response_text, re.DOTALL | re.IGNORECASE)
-        reasoning = reasoning_match.group(1).strip() if reasoning_match else "No reasoning provided."
-        
-        # Extract possibilities
-        possibilities = []
+        # Create result dictionary
+        result = {}
         
-        # Try to extract possibility 1
-        poss1_match = re.search(r'POSSIBILITY_1:(.*?)POSSIBILITY_2:', response_text, re.DOTALL | re.IGNORECASE)
-        if poss1_match:
-            possibilities.append(poss1_match.group(1).strip())
-            
-        # Try to extract possibility 2  
-        poss2_match = re.search(r'POSSIBILITY_2:(.*?)POSSIBILITY_3:', response_text, re.DOTALL | re.IGNORECASE)
-        if poss2_match:
-            possibilities.append(poss2_match.group(1).strip())
-            
-        # Try to extract possibility 3
-        poss3_match = re.search(r'POSSIBILITY_3:(.*?)CONCLUSION:', response_text, re.DOTALL | re.IGNORECASE)
-        if poss3_match:
-            possibilities.append(poss3_match.group(1).strip())
+        # 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 conclusion
-        conclusion_match = re.search(r'CONCLUSION:(.*?)$', response_text, re.DOTALL | re.IGNORECASE)
-        conclusion = conclusion_match.group(1).strip() if conclusion_match else "No conclusion provided."
+        # 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."
         
-        # Extract final NAICS code from conclusion
-        naics_match = re.search(r'NAICS code (\d{6})', conclusion)
+        # Extract NAICS code
+        naics_match = re.search(r'NAICS_CODE:(.*?)(\d{6})', response_text, re.DOTALL)
         if naics_match:
-            naics_code = naics_match.group(1)
+            result["naics_code"] = naics_match.group(2)
         else:
-            # Try to find any 6-digit code in the conclusion
-            code_match = re.search(r'\b(\d{6})\b', conclusion)
-            naics_code = code_match.group(1) if code_match else "000000"
-        
-        return {
-            "naics_code": naics_code,
-            "google_findings": google_findings,
-            "reasoning": reasoning,
-            "possibilities": possibilities,
-            "conclusion": conclusion
-        }
+            # Try to find any 6-digit code in the response
+            code_match = re.search(r'\b(\d{6})\b', response_text)
+            result["naics_code"] = code_match.group(1) if code_match else "000000"
+            
+        return result
     except Exception as e:
-        print(f"Error getting NAICS classification: {str(e)}")
+        print(f"❌ Error getting NAICS classification: {str(e)}")
         return {
             "naics_code": "000000",
-            "google_findings": "Error occurred during Google search.",
-            "reasoning": f"Error analyzing company: {str(e)}",
-            "possibilities": [],
-            "conclusion": "Error in analysis"
+            "reasoning": f"Error analyzing company: {str(e)}"
         }
 
-def find_naics_code(api_key, company_name, company_description):
-    """Main function to find NAICS code that will be called by Gradio"""
-    if not api_key or not company_name:
-        return "Please provide both API key and company name."
-    
-    try:
-        # Initialize Gemini API
-        model = initialize_gemini(api_key)
-        
-        # Search for NAICS candidates with company description for better context
-        naics_candidates = google_search_naics(company_name, company_description)
-        # Store search query for reporting
-        search_query = f"2022 NAICS code for \"{company_name}\" {company_description} industry classification" if company_description else f"2022 NAICS code for \"{company_name}\" company industry classification"
+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): ")
         
-        # Get classification
+        # Step 3: Find NAICS Code Candidates
+        naics_candidates = google_search_naics(company_name)
+
         if not naics_candidates:
-            result = get_naics_classification(model, company_name, company_description, [], search_query)
+            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:
-            result = get_naics_classification(model, company_name, company_description, naics_candidates, search_query)
-        
-        # Format the output with NAICS code at the end
-        output = f"## Analysis for {company_name}\n\n"
-        
-        # Display search query prominently at the top
-        output += f"**Google Search Query Used:**\n`{search_query}`\n\n"
-        
-        # Add Google findings first
-        if 'google_findings' in result and result['google_findings']:
-            output += f"**Google Search Findings:**\n{result['google_findings']}\n\n"
-        
-        # Then reasoning
-        output += f"**Reasoning:**\n{result['reasoning']}\n\n"
-        
-        # Add possibilities section
-        if 'possibilities' in result and result['possibilities']:
-            output += f"**Possible Classifications:**\n\n"
-            for i, possibility in enumerate(result['possibilities'], 1):
-                output += f"{i}. {possibility}\n\n"
-        
-        # Add conclusion
-        if 'conclusion' in result and result['conclusion']:
-            output += f"**Conclusion:**\n{result['conclusion']}\n\n"
-        
-        # Add final NAICS code at the very end
-        output += f"**FINAL NAICS CODE: {result['naics_code']}**"
-        
-        return output
-    
-    except Exception as e:
-        return f"Error: {str(e)}"
+            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)
 
-# Create Gradio Interface
-with gr.Blocks(title="NAICS Code Finder") as app:
-    gr.Markdown("# NAICS Code Finder")
-    gr.Markdown("This app helps you find the appropriate NAICS code for a company based on its name and description.")
-    
-    with gr.Row():
-        with gr.Column():
-            api_key = gr.Textbox(label="Google Gemini API Key", placeholder="Enter your Gemini API key here", type="password")
-            company_name = gr.Textbox(label="Company Name", placeholder="Enter the company name")
-            company_description = gr.Textbox(label="Company Description", placeholder="Enter a brief description of the company", lines=5)
+        # Display research findings if available
+        if "research" in result:
+            print(f"\n📊 NAICS Code Research:\n{result['research']}")
+
+        # Display reasoning
+        print(f"\n🧠 Reasoning:\n{result['reasoning']}")
             
-            submit_btn = gr.Button("Find NAICS Code")
-        
-        with gr.Column():
-            output = gr.Markdown(label="Result")
-    
-    submit_btn.click(
-        fn=find_naics_code,
-        inputs=[api_key, company_name, company_description],
-        outputs=output
-    )
+        # Output the NAICS code
+        print(f"\n🏆 NAICS Code: {result['naics_code']}")
+        print("-" * 80)
 
 if __name__ == "__main__":
-    app.launch()
+    main()