Update app.py

app.py

@@ -6,58 +6,90 @@ import requests
 import pandas as pd
 import json
 import re
+from difflib import get_close_matches
 from tavily import TavilyClient
 
 KYS_SAMPLE = "https://kys.udiseplus.gov.in/webapp/api/search-schools?searchType=3&searchParam={udise}"
 
-# List of valid state codes for UDISE (first 2 digits of UDISE code)
-# List of valid state codes for UDISE (first 2 digits of UDISE code)
-VALID_UDISE_STATE_CODES = [
-    '01',  # Jammu & Kashmir
-    '02',  # Himachal Pradesh
-    '03',  # Punjab
-    '04',  # Chandigarh
-    '05',  # Uttarakhand
-    '06',  # Haryana
-    '07',  # Delhi
-    '08',  # Rajasthan
-    '09',  # Uttar Pradesh
-    '10',  # Bihar
-    '11',  # Sikkim
-    '12',  # Arunachal Pradesh
-    '13',  # Nagaland
-    '14',  # Manipur
-    '15',  # Mizoram
-    '16',  # Tripura
-    '17',  # Meghalaya
-    '18',  # Assam
-    '19',  # West Bengal
-    '20',  # Jharkhand
-    '21',  # Odisha
-    '22',  # Chhattisgarh
-    '23',  # Madhya Pradesh
-    '24',  # Gujarat
-    '25',  # Daman & Diu
-    '26',  # Dadra & Nagar Haveli
-    '27',  # Maharashtra
-    '28',  # Andhra Pradesh
-    '29',  # Karnataka
-    '30',  # Goa
-    '31',  # Lakshadweep
-    '32',  # Kerala
-    '33',  # Tamil Nadu
-    '34',  # Puducherry
-    '35',  # Andaman & Nicobar Islands
-    '36',  # Telangana
-    '37'   # Ladakh (added in newer datasets)
-]
-
-def is_valid_udise(code):
-    """Check if a string is a valid UDISE code."""
+# Mapping of state names to UDISE state codes
+STATE_TO_UDISE_CODE = {
+    'Jammu & Kashmir': '01',
+    'Himachal Pradesh': '02',
+    'Punjab': '03',
+    'Chandigarh': '04',
+    'Uttarakhand': '05',
+    'Haryana': '06',
+    'Delhi': '07',
+    'Rajasthan': '08',
+    'Uttar Pradesh': '09',
+    'Bihar': '10',
+    'Sikkim': '11',
+    'Arunachal Pradesh': '12',
+    'Nagaland': '13',
+    'Manipur': '14',
+    'Mizoram': '15',
+    'Tripura': '16',
+    'Meghalaya': '17',
+    'Assam': '18',
+    'West Bengal': '19',
+    'Jharkhand': '20',
+    'Odisha': '21',
+    'Chhattisgarh': '22',
+    'Madhya Pradesh': '23',
+    'Gujarat': '24',
+    'Daman & Diu': '25',
+    'Dadra & Nagar Haveli': '26',
+    'Maharashtra': '27',
+    'Andhra Pradesh': '28',
+    'Karnataka': '29',
+    'Goa': '30',
+    'Lakshadweep': '31',
+    'Kerala': '32',
+    'Tamil Nadu': '33',
+    'Puducherry': '34',
+    'Andaman & Nicobar Islands': '35',
+    'Telangana': '36',
+    'Ladakh': '37'
+}
+
+# For backward compatibility
+VALID_UDISE_STATE_CODES = list(STATE_TO_UDISE_CODE.values())
+
+def is_valid_udise(code, state_name=None):
+    """
+    Check if a string is a valid UDISE code.
+    
+    Args:
+        code: The UDISE code to validate
+        state_name: Optional state name to validate against the UDISE state code
+        
+    Returns:
+        bool: True if the code is valid, False otherwise
+    """
+    # Basic validation
     if not (code and code.isdigit() and len(code) == 11):
         return False
-    # Optional: Validate state code (first 2 digits)
-    return code[:2] in VALID_UDISE_STATE_CODES
+        
+    state_code = code[:2]
+    
+    # Check if state code is valid
+    if state_code not in VALID_UDISE_STATE_CODES:
+        return False
+    
+    # If state_name is provided, validate against it
+    if state_name:
+        state_name = state_name.strip().title()
+        # Handle special case for 'Uttar pradesh' vs 'Uttar Pradesh'
+        state_name = state_name.replace('_', ' ')
+        expected_code = STATE_TO_UDISE_CODE.get(state_name)
+        if not expected_code:
+            print(f"Warning: Unknown state name: {state_name}")
+            return False
+        if state_code != expected_code:
+            print(f"UDISE code {code} state code {state_code} does not match expected state {state_name} ({expected_code})")
+            return False
+    
+    return True
 
 STATES = [
     "Arunachal_pradesh",
@@ -103,13 +135,24 @@ def call_search_sdk(api_key, payload_text):
         return {"ok": False, "error": str(e)}
 
 
-def extract_udise_candidates_from_search(search_json):
+def extract_udise_candidates_from_search(search_json, state_name=None, search_query=None):
     """
-    Extract UDISE codes from Tavily search results.
-    Returns a list of unique UDISE codes found in the search results.
+    Extract UDISE codes and school information from Tavily search results.
+    
+    Args:
+        search_json: JSON response from Tavily search
+        state_name: Optional state name to validate UDISE codes against
+        search_query: Original search query to help with fuzzy matching
+        
+    Returns:
+        list: List of dictionaries containing UDISE codes and school information
     """
     print("\n===== Extracting UDISE Codes =====")
-    found = set()
+    if state_name:
+        print(f"Validating UDISE codes against state: {state_name}")
+    
+    found_codes = set()
+    school_info = []  # List to store school information
     
     # Check if we have valid search results
     if not search_json or not isinstance(search_json, dict):
@@ -123,7 +166,7 @@ def extract_udise_candidates_from_search(search_json):
     
     print(f"Found {len(results)} search results")
     
-    # Patterns to match UDISE codes
+    # Patterns to match UDISE codes and school information
     patterns = [
         r'UDISE[^\d]*(?:code|Code|CODE)[^\d]*(\d{11})(?![0-9])',
         r'Udise[^\d]*(?:School[^\d]*Code|Code)[^\d]*(\d{11})(?![0-9])',
@@ -134,26 +177,70 @@ def extract_udise_candidates_from_search(search_json):
         if not isinstance(result, dict):
             continue
             
-        # Combine title and content for searching
-        text = f"{result.get('title', '')} {result.get('content', '')}"
+        # Get title and content
+        title = result.get('title', '')
+        content = result.get('content', '')
+        url = result.get('url', '')
+        text = f"{title} {content}"
         
         # Check for UDISE codes using all patterns
         for pattern in patterns:
             matches = re.finditer(pattern, text, re.IGNORECASE)
             for match in matches:
                 udise_code = match.group(1) if len(match.groups()) > 0 else match.group(0)
-                if udise_code and is_valid_udise(udise_code):
-                    print(f"Found UDISE code: {udise_code}")
-                    found.add(udise_code)
+                if udise_code and is_valid_udise(udise_code, state_name) and udise_code not in found_codes:
+                    print(f"Found valid UDISE code: {udise_code}")
+                    found_codes.add(udise_code)
+                    
+                    # Extract school name - try to find the most relevant text
+                    school_name = title
+                    
+                    # If title is too short or doesn't seem like a school name, try to find a better match
+                    if len(school_name.split()) < 2 or any(word in school_name.lower() for word in ['udise', 'code', 'school']):
+                        # Look for a school-like name in the content
+                        school_matches = re.findall(r'([A-Z][a-z]+(?:\s+[A-Z][a-z]+)+\s*(?:School|School|High School|High School|Vidyalaya|Vidyalaya|Vidyalayam|Vidyalayam|Vidhya|Vidhya|Vidya|Vidya|Public School|Public School|Govt|Government|Kendriya|Jawahar|Navodaya|Sainik|Army|Air Force|Navy|Central School|Central School|CBSE|ICSE|State Board|State Board|EM|EM|TM|TM|Primary|Primary|Upper Primary|Upper Primary|Higher Secondary|Higher Secondary|HSS|HSS|HS|HS|UPS|UPS|PS|PS))', content, re.IGNORECASE)
+                        if school_matches:
+                            school_name = school_matches[0][0].strip()
+                    
+                    school_info.append({
+                        'udise': udise_code,
+                        'name': school_name,
+                        'source': url,
+                        'snippet': content[:200] + '...' if len(content) > 200 else content
+                    })
     
-    # Convert set to list and sort
-    result = sorted(list(found))
+    # If we have a search query, sort results by relevance to the query
+    if search_query and school_info:
+        # Extract just the school names for fuzzy matching
+        school_names = [s['name'] for s in school_info]
+        
+        # Get fuzzy matches and their scores
+        matches = get_close_matches(
+            search_query.lower(),
+            [name.lower() for name in school_names],
+            n=len(school_names),
+            cutoff=0.3  # Lower cutoff to allow more fuzzy matches
+        )
+        
+        # Create a dictionary to map lowercase names to their original objects with scores
+        school_map = {s['name'].lower(): s for s in school_info}
+        
+        # Rebuild the school_info list in order of best match
+        sorted_schools = []
+        for match in matches:
+            if match in school_map:
+                sorted_schools.append(school_map[match])
+                del school_map[match]
+        
+        # Add any remaining schools that didn't match the fuzzy search
+        sorted_schools.extend(school_map.values())
+        school_info = sorted_schools
     
-    if not result:
-        print("No valid UDISE codes found")
+    if not school_info:
+        print("No valid school information found with UDISE codes")
         return []
     
-    return result
+    return school_info
 
 
 def json_to_table(obj):
@@ -204,18 +291,29 @@ def to_table_from_kys(kys_json):
 
 
 def search_workflow(school_name, state_name, search_key, use_search=True, use_kys=True):
-    out = {"kys": None, "search": None, "suggestions": [], "first_candidate": None}
+    out = {"kys": None, "search": None, "suggestions": [], "first_candidate": None, "school_info": []}
     payload_text = f"{school_name or ''} {state_name or ''} UDISE code".strip()
 
     if use_search:
         search_res = call_search_sdk(search_key, payload_text)
         out["search"] = search_res
         if search_res.get("ok"):
-            # Pass both school_name and state_name to improve UDISE code extraction
-            candidates = extract_udise_candidates_from_search(
-                search_res["data"]
+            # Pass school_name for fuzzy matching and state_name for validation
+            school_info = extract_udise_candidates_from_search(
+                search_res["data"],
+                state_name=state_name,
+                search_query=school_name
             )
-            out["suggestions"] = candidates
+            
+            # Extract just the UDISE codes for backward compatibility
+            candidates = [info['udise'] for info in school_info]
+            
+            out["suggestions"] = [
+                f"{info['name']} (UDISE: {info['udise']})" 
+                for info in school_info
+            ]
+            out["school_info"] = school_info
+            
             if candidates and candidates[0] != "No UDISE codes found":
                 out["first_candidate"] = candidates[0]
     else:
@@ -231,10 +329,9 @@ def search_workflow(school_name, state_name, search_key, use_search=True, use_ky
 with gr.Blocks() as demo:
     gr.Markdown(
         """
-# Search + KYS Lookup (Hugging Face Space)
-This version uses the Search SDK internally. Provide your API key in the textbox.
-Enter a school name (or UDISE code) and select the state; the app calls Search SDK
-with the combined query `<school> <state> UDISE code`, then optionally calls KYS.
+# Find School UDISE Code
+Provide your API key in the textbox.
+Enter a school name and select the state
 """
     )
 
@@ -270,11 +367,21 @@ with the combined query `<school> <state> UDISE code`, then optionally calls KYS
         if res.get("search") and res["search"].get("ok"):
             tbl = json_to_table(res["search"]["data"])
         
-        # Ensure suggestions is always a list and handle empty state
-        suggestions = res.get("suggestions", [])
-        if not suggestions:
-            suggestions = ["No UDISE codes found"]
-            
+        # Get school info and format suggestions with school names and UDISE codes
+        school_info = res.get("school_info", [])
+        suggestions = []
+        first_candidate = ""
+        
+        if school_info:
+            # Format suggestions as "School Name (UDISE: 12345678901)"
+            suggestions = [
+                f"{info['name']} (UDISE: {info['udise']})" 
+                for info in school_info
+            ]
+            first_candidate = school_info[0]['udise'] if school_info else ""
+        else:
+            suggestions = ["No matching schools found"]
+        
         # Always save the key to state if a new one is provided
         new_saved_key = key or saved_key
             
@@ -283,7 +390,7 @@ with the combined query `<school> <state> UDISE code`, then optionally calls KYS
             res.get("search"),  # output_json
             tbl,  # search_table
             {"choices": suggestions, "__type__": "update"},  # Update dropdown choices
-            res.get("first_candidate", ""),  # This will update udise_input
+            first_candidate,  # This will update udise_input with the UDISE code
             new_saved_key,  # saved_key_state
             res.get("kys")  # kys_output_json
         )
@@ -302,8 +409,15 @@ with the combined query `<school> <state> UDISE code`, then optionally calls KYS
     )
 
     def on_select_suggestion(choice):
-        # Only update if a valid choice is selected (not the placeholder)
-        return choice if choice and choice != "No UDISE codes found" else ""
+        # Extract UDISE code from the selected choice
+        if not choice or choice in ["No matching schools found", "No UDISE codes found"]:
+            return ""
+        
+        # Extract UDISE code from the format "School Name (UDISE: 12345678901)"
+        match = re.search(r'\(UDISE:\s*(\d+)\)', choice)
+        if match:
+            return match.group(1)
+        return ""
 
     suggestions_dropdown.change(
         on_select_suggestion, 
@@ -328,14 +442,5 @@ with the combined query `<school> <state> UDISE code`, then optionally calls KYS
 
     show_raw_checkbox.change(toggle_raw, inputs=[show_raw_checkbox], outputs=[output_json, kys_output_json])
 
-    gr.Markdown(
-        """
----
-**Notes:**
-- Your API key is stored in-memory for the current session only.
-- It is not written to disk or logged.
-"""
-    )
-
 if __name__ == "__main__":
     demo.launch()