Update app.py

app.py

@@ -4,19 +4,20 @@
 import gradio as gr
 import requests
 import pandas as pd
-import json
+import re
+from typing import Dict, Any, Tuple, List
 from tavily import TavilyClient
 
+# ---------- constants ----------
 KYS_SAMPLE = "https://kys.udiseplus.gov.in/webapp/api/search-schools?searchType=3&searchParam={udise}"
 STATES = [
     "Arunachal_pradesh",
     "Assam",
     "Bihar",
     "Chhattisgarh",
-    "Delhi",
     "Jharkhand",
     "Karnataka",
-    "Madhya pradesh",
+    "Madhya_pradesh",
     "Manipur",
     "Meghalaya",
     "Mizoram",
@@ -25,16 +26,14 @@ STATES = [
     "Puducherry",
     "Rajasthan",
     "Sikkim",
-    "Tamil nadu",
     "Telangana",
     "Tripura",
-    "Uttar pradesh",
-    "Uttarakhand"
+    "Uttar_pradesh",
+    "Uttarakhand",
 ]
 
-
-
-def call_kys_by_udise(udise_code):
+# ---------- network helpers ----------
+def call_kys_by_udise(udise_code: str) -> Dict[str, Any]:
     url = KYS_SAMPLE.format(udise=udise_code)
     try:
         resp = requests.get(url, timeout=10)
@@ -45,7 +44,7 @@ def call_kys_by_udise(udise_code):
         return {"ok": False, "error": str(e), "url": url}
 
 
-def call_search_sdk(api_key, payload_text):
+def call_search_sdk(api_key: str, payload_text: str) -> Dict[str, Any]:
     try:
         client = TavilyClient(api_key)
         resp = client.search(query=payload_text)
@@ -54,31 +53,58 @@ def call_search_sdk(api_key, payload_text):
         return {"ok": False, "error": str(e)}
 
 
-def extract_udise_candidates_from_search(search_json):
-    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)
-    return sorted(list(found))
-
-
-def json_to_table(obj):
+# ---------- UDISE extraction ----------
+UDISE_RE = re.compile(r"\b(?:UDISE|Udise|udise|Udise Code|UDISE Code)[:\s]*([0-9]{6,12})\b")
+GENERIC_NUM_RE = re.compile(r"\b([0-9]{5,12})\b")
+
+def normalize_candidate(c: str) -> str:
+    return c.strip()
+
+
+def extract_udise_code_from_tavily_results(tavily_json: Dict[str, Any]) -> Tuple[List[str], str]:
+    results = tavily_json.get("data", {}).get("results", []) if isinstance(tavily_json, dict) else []
+    candidates: List[str] = []
+    for r in results:
+        content = " ".join([str(r.get(k, "")) for k in ("content", "title")])
+        for m in UDISE_RE.finditer(content):
+            candidates.append(normalize_candidate(m.group(1)))
+        for m in GENERIC_NUM_RE.finditer(content):
+            candidates.append(normalize_candidate(m.group(1)))
+        url = r.get("url", "")
+        for m in GENERIC_NUM_RE.finditer(url):
+            candidates.append(normalize_candidate(m.group(1)))
+
+    # dedupe while preserving order
+    seen = set()
+    unique = []
+    for c in candidates:
+        if c and c not in seen:
+            seen.add(c)
+            unique.append(c)
+
+    # heuristics for best
+    best = ""
+    for c in unique:
+        if len(c) >= 10:
+            best = c
+            break
+    if not best and unique:
+        best = unique[0]
+    return unique, best
+
+
+def extract_udise_candidates_from_search(search_data: Dict[str, Any]) -> List[str]:
+    candidates, _ = extract_udise_code_from_tavily_results(search_data)
+    return candidates
+
+
+# ---------- helpers: JSON -> table ----------
+def json_to_table(obj: Any) -> pd.DataFrame:
     try:
         if isinstance(obj, list):
             return pd.json_normalize(obj)
         if isinstance(obj, dict):
-            for k in ("results", "data", "hits", "items"):
+            for k in ("results", "data", "hits", "items", "content"):
                 if k in obj and isinstance(obj[k], list):
                     return pd.json_normalize(obj[k])
             return pd.json_normalize([obj])
@@ -87,40 +113,36 @@ def json_to_table(obj):
     return pd.DataFrame()
 
 
-def to_table_from_kys(kys_json):
-    """
-    Convert KYS JSON wrapper into a simplified pandas DataFrame showing only
-    selected fields from the `content` list.
-    """
+def to_table_from_kys(kys_json: Dict[str, Any]) -> pd.DataFrame:
     try:
         content = None
         if isinstance(kys_json, dict):
-            inner = kys_json.get("data") if kys_json.get("data") is not None else None
-            if isinstance(inner, dict) and isinstance(inner.get("content"), list):
-                content = inner.get("content")
-            elif isinstance(inner, dict) and isinstance(inner.get("data"), dict) and isinstance(inner.get("data").get("content"), list):
-                content = inner.get("data").get("content")
+            if isinstance(kys_json.get("data"), dict) and isinstance(kys_json["data"].get("content"), list):
+                content = kys_json["data"]["content"]
             elif isinstance(kys_json.get("content"), list):
-                content = kys_json.get("content")
+                content = kys_json["content"]
         if not content:
             return pd.DataFrame()
-        rows = []
-        for r in content:
-            rows.append({
+
+        rows = [
+            {
                 "School Name": r.get("schoolName"),
                 "School ID": r.get("schoolId"),
                 "Pincode": r.get("pincode"),
                 "State": r.get("stateName"),
                 "District": r.get("districtName"),
-                "Management Type": r.get("schMgmtType")
-            })
+                "Management Type": r.get("schMgmtType"),
+            }
+            for r in content
+        ]
         return pd.DataFrame(rows)
     except Exception as e:
         print("to_table_from_kys error:", e)
         return pd.DataFrame()
 
 
-def search_workflow(school_name, state_name, search_key, use_search, use_kys):
+# ---------- main workflow ----------
+def search_workflow(school_name: str, state_name: str, search_key: str, use_search: bool, use_kys: bool):
     out = {"kys": None, "search": None, "suggestions": []}
     payload_text = f"{school_name or ''} {state_name or ''} UDISE code".strip()
 
@@ -140,21 +162,23 @@ def search_workflow(school_name, state_name, search_key, use_search, use_kys):
     return out
 
 
+# ---------- Gradio UI ----------
 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.
-"""
-    )
+    gr.Markdown("# Search + KYS Lookup (Hugging Face Space)")
 
     with gr.Row():
-        inp = gr.Textbox(label="School name or UDISE code", placeholder="e.g. GOVT SEC SCHOOL DARLONG or 12345678901", lines=1)
-        state_dropdown = gr.Dropdown(choices=STATES, label="State", value=STATES[0] if STATES else "", interactive=True, allow_custom_value=True)
+        inp = gr.Textbox(
+            label="School name or UDISE code",
+            placeholder="e.g. GOVT SEC SCHOOL DARLONG or 12345678901",
+            lines=1,
+        )
+        state_dropdown = gr.Dropdown(
+            choices=STATES,
+            label="State",
+            value=STATES[0] if STATES else "",
+            interactive=True,
+            allow_custom_value=True,
+        )
 
     search_key = gr.Textbox(label="Search API Key (required)", placeholder="api-key...", lines=1)
     save_key_toggle = gr.Checkbox(value=False, label="Save key in session (keeps key between interactions)")
@@ -166,15 +190,17 @@ with the combined query `<school> <state> UDISE code`, then optionally calls KYS
 
     run = gr.Button("Search")
 
-    # By default hide raw JSON outputs; users can toggle visibility with `show_raw_checkbox`
     show_raw_checkbox = gr.Checkbox(value=False, label="Show raw JSON outputs")
-
     output_json = gr.JSON(label="Raw Search Output (JSON)", visible=False)
     search_table = gr.DataFrame(headers=None, label="Search results (table)")
 
     gr.Markdown("### UDISE candidates found in Search results")
     suggestions_dropdown = gr.Dropdown(choices=[], label="UDISE candidates (from Search)")
-    udise_input = gr.Textbox(label="UDISE to lookup (editable)", placeholder="Pick a candidate or type a UDISE code...", lines=1)
+    udise_input = gr.Textbox(
+        label="UDISE to lookup (editable)",
+        placeholder="Pick a candidate or type a UDISE code...",
+        lines=1,
+    )
     lookup_btn = gr.Button("Lookup UDISE (Call KYS)")
 
     kys_output_json = gr.JSON(label="KYS Raw Output", visible=False)
@@ -194,10 +220,24 @@ 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])
+        suggestions_update = gr.Dropdown.update(
+            choices=suggestions, value=suggestions[0] if suggestions else None
+        )
+        return (
+            res.get("search"),
+            tbl,
+            suggestions_update,
+            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],
+    )
 
     def on_select_suggestion(choice):
         return choice or ""
@@ -220,21 +260,11 @@ with the combined query `<school> <state> UDISE code`, then optionally calls KYS
 
     clear_key_btn.click(on_clear_key, inputs=[saved_key_state], outputs=[saved_key_state, search_key])
 
-    # Toggle visibility handler for raw JSON outputs
     def toggle_raw(visible: bool):
         return gr.update(visible=visible), gr.update(visible=visible)
 
     show_raw_checkbox.change(toggle_raw, inputs=[show_raw_checkbox], outputs=[output_json, kys_output_json])
 
-    gr.Markdown(
-        """
----
-**Notes:**
-- The 'Save key in session' toggle keeps the key active for the current session.
-- The key is stored only in-memory (not written to disk) and is not logged.
-- Use 'Clear saved key' to remove it from the session.
-"""
-    )
 
 if __name__ == "__main__":
     demo.launch()