Update app.py

app.py

@@ -103,54 +103,90 @@ 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, school_name="", state_name=""):
     """
-    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.
+    Extract UDISE codes from Tavily search results with scoring based on relevance.
+    
+    Args:
+        search_json: The raw Tavily API response
+        school_name: Optional school name for better matching
+        state_name: Optional state name for better matching
+        
+    Returns:
+        List of UDISE codes sorted by relevance
     """
     found = set()
+    candidates = []
+    
+    # Normalize search terms for matching
+    school_tokens = [t for t in re.sub(r"\s+", " ", school_name).strip().lower().split() if t]
+    state_norm = re.sub(r"\s+", " ", state_name).strip().lower()
     
     # 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', [])
-    
+    results = search_json.get('data', {}).get('results', [])
     if not isinstance(results, list):
         return []
     
-    # Compile patterns once for better performance
+    # Patterns to match UDISE codes
     patterns = [
         r'UDISE\s*[Cc]ode\s*:?\s*(\d{11})(?![0-9])',  # Matches "UDISE Code: 12345678901"
         r'Udise\s+School\s+Code\s*:?\s*(\d{11})(?![0-9])',  # Matches "Udise School Code: 12345678901"
+        r'udise\s+code:?\s*(\d{11})(?![0-9])',  # Case-insensitive match
         r'(?<![0-9])(\d{11})(?![0-9])'  # Matches exactly 11 digits not part of a longer number
     ]
     
-    # Compile patterns for better performance
-    compiled_patterns = [re.compile(pattern) for pattern in patterns]
-    
     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 compiled_patterns:
-                matches = pattern.finditer(text)
-                for match in matches:
-                    # Get the first group if it exists, otherwise the whole match
-                    udise_code = match.group(1) if len(match.groups()) > 0 else match.group(0)
-                    if udise_code and is_valid_udise(udise_code):
-                        found.add(udise_code)
+        # Combine title and content for better matching
+        title = str(result.get('title', ''))
+        content = str(result.get('content', ''))
+        snippet = f"{title} {content}"
+        snippet_low = snippet.lower()
+        
+        # Skip if no UDISE code pattern found
+        if not any(re.search(p, snippet, re.IGNORECASE) for p in patterns):
+            continue
+            
+        # Calculate relevance score
+        score = 0
+        
+        # Higher score if school name tokens are found
+        if school_tokens and any(tok in snippet_low for tok in school_tokens):
+            score += 2
+            
+        # Higher score if state name is mentioned
+        if state_norm and state_norm in snippet_low:
+            score += 1
+            
+        # Extra points for exact school name match
+        if school_name and school_name.lower() in snippet_low:
+            score += 3
+            
+        # Extract all UDISE codes from this result
+        for pattern in patterns:
+            for match in re.finditer(pattern, snippet, re.IGNORECASE):
+                udise_code = match.group(1) if len(match.groups()) > 0 else match.group(0)
+                if udise_code and is_valid_udise(udise_code):
+                    found.add(udise_code)
+                    # Store with score for sorting
+                    candidates.append((score, udise_code))
+    
+    # Sort by score (descending) and return unique codes
+    candidates.sort(key=lambda x: x[0], reverse=True)
+    seen = set()
+    result = []
+    
+    for score, code in candidates:
+        if code not in seen:
+            seen.add(code)
+            result.append(code)
     
-    return sorted(list(found))
+    return result
 
 
 def json_to_table(obj):
@@ -208,9 +244,14 @@ def search_workflow(school_name, state_name, search_key, use_search=True, use_ky
         search_res = call_search_sdk(search_key, payload_text)
         out["search"] = search_res
         if search_res.get("ok"):
-            candidates = extract_udise_candidates_from_search(search_res["data"])
+            # Pass both school_name and state_name to improve UDISE code extraction
+            candidates = extract_udise_candidates_from_search(
+                search_res["data"],
+                school_name=school_name,
+                state_name=state_name
+            )
             out["suggestions"] = candidates
-            if candidates:  # If we have candidates, set the first one
+            if candidates and candidates[0] != "No UDISE codes found":
                 out["first_candidate"] = candidates[0]
     else:
         out["search"] = {"ok": False, "error": "Search disabled or SDK not used"}