Update app.py

app.py

@@ -8,12 +8,11 @@ import json
 from tavily import TavilyClient
 
 KYS_SAMPLE = "https://kys.udiseplus.gov.in/webapp/api/search-schools?searchType=3&searchParam={udise}"
-STATES = [
+states = [
     "Arunachal_pradesh",
     "Assam",
     "Bihar",
     "Chhattisgarh",
-    "Delhi",
     "Jharkhand",
     "Karnataka",
     "Madhya pradesh",
@@ -25,7 +24,6 @@ STATES = [
     "Puducherry",
     "Rajasthan",
     "Sikkim",
-    "Tamil nadu",
     "Telangana",
     "Tripura",
     "Uttar pradesh",
@@ -55,21 +53,50 @@ def call_search_sdk(api_key, payload_text):
 
 
 def extract_udise_candidates_from_search(search_json):
+    """
+    Extract UDISE codes from Tavily search results.
+    Looks for patterns like "UDISE Code: 12345678901" or "Udise School Code: 12345678901"
+    in the content and title fields of search results.
+    """
     found = set()
-    def walk(obj):
-        if isinstance(obj, dict):
-            for v in obj.values():
-                walk(v)
-        elif isinstance(obj, list):
-            for item in obj:
-                walk(item)
-        elif isinstance(obj, str):
-            s = obj.strip()
-            tokens = s.replace('-', ' ').split()
-            for t in tokens:
-                if t.isdigit() and 6 <= len(t) <= 14:
-                    found.add(t)
-    walk(search_json)
+    
+    # Check if we have valid search results
+    if not search_json or not isinstance(search_json, dict):
+        return []
+        
+    data = search_json.get('data', {})
+    results = data.get('results', [])
+    
+    if not isinstance(results, list):
+        return []
+    
+    # Patterns to match UDISE codes in text
+    patterns = [
+        r'UDISE\s*[Cc]ode\s*:?\s*(\d{11})',  # Matches "UDISE Code: 12345678901"
+        r'Udise\s+School\s+Code\s*:?\s*(\d{11})',  # Matches "Udise School Code: 12345678901"
+        r'\b(\d{11})\b'  # Matches any 11-digit number as fallback
+    ]
+    
+    for result in results:
+        if not isinstance(result, dict):
+            continue
+            
+        # Check both content and title fields
+        for field in ['content', 'title']:
+            text = result.get(field, '')
+            if not isinstance(text, str):
+                continue
+                
+            # Try each pattern
+            for pattern in patterns:
+                import re
+                matches = re.findall(pattern, text)
+                for match in matches:
+                    if isinstance(match, tuple):  # If pattern has groups
+                        match = match[0]  # Take the first group
+                    if match.isdigit() and len(match) == 11:  # UDISE codes are 11 digits
+                        found.add(match)
+    
     return sorted(list(found))
 
 
@@ -121,7 +148,7 @@ def to_table_from_kys(kys_json):
 
 
 def search_workflow(school_name, state_name, search_key, use_search, use_kys):
-    out = {"kys": None, "search": None, "suggestions": []}
+    out = {"kys": None, "search": None, "suggestions": [], "first_candidate": None}
     payload_text = f"{school_name or ''} {state_name or ''} UDISE code".strip()
 
     if use_search:
@@ -130,6 +157,8 @@ def search_workflow(school_name, state_name, search_key, use_search, use_kys):
         if search_res.get("ok"):
             candidates = extract_udise_candidates_from_search(search_res["data"])
             out["suggestions"] = candidates
+            if candidates:  # If we have candidates, set the first one
+                out["first_candidate"] = candidates[0]
     else:
         out["search"] = {"ok": False, "error": "Search disabled or SDK not used"}
 
@@ -192,10 +221,30 @@ with the combined query `<school> <state> UDISE code`, then optionally calls KYS
         if save_key:
             new_saved_key = key or saved_key
             textbox_value = new_saved_key
-        # return raw JSON objects (even if hidden) and the table + suggestions
-        return res.get("search"), tbl, suggestions, new_saved_key, textbox_value, res.get("kys")
-
-    run.click(on_run, inputs=[inp, state_dropdown, search_key, save_key_toggle, use_search, use_kys, saved_key_state], outputs=[output_json, search_table, suggestions_dropdown, saved_key_state, search_key, kys_output_json])
+        # Return the first candidate along with other values
+        return (
+            res.get("search"),  # output_json
+            tbl,  # search_table
+            suggestions,  # suggestions_dropdown
+            res.get("first_candidate", ""),  # This will update udise_input
+            new_saved_key,  # saved_key_state
+            textbox_value,  # search_key
+            res.get("kys")  # kys_output_json
+        )
+
+    run.click(
+        on_run, 
+        inputs=[inp, state_dropdown, search_key, save_key_toggle, use_search, use_kys, saved_key_state], 
+        outputs=[
+            output_json, 
+            search_table, 
+            suggestions_dropdown, 
+            udise_input,  # This will be updated with first_candidate
+            saved_key_state, 
+            search_key, 
+            kys_output_json
+        ]
+    )
 
     def on_select_suggestion(choice):
         return choice or ""