Update app.py

app.py

@@ -4,20 +4,19 @@
 import gradio as gr
 import requests
 import pandas as pd
-import re
-from typing import Dict, Any, Tuple, List
+import json
 from tavily import TavilyClient
 
-# ---------- constants ----------
 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",
+    "Madhya pradesh",
     "Manipur",
     "Meghalaya",
     "Mizoram",
@@ -26,14 +25,16 @@ STATES = [
     "Puducherry",
     "Rajasthan",
     "Sikkim",
+    "Tamil nadu",
     "Telangana",
     "Tripura",
-    "Uttar_pradesh",
-    "Uttarakhand",
+    "Uttar pradesh",
+    "Uttarakhand"
 ]
 
-# ---------- network helpers ----------
-def call_kys_by_udise(udise_code: str) -> Dict[str, Any]:
+
+
+def call_kys_by_udise(udise_code):
     url = KYS_SAMPLE.format(udise=udise_code)
     try:
         resp = requests.get(url, timeout=10)
@@ -44,7 +45,7 @@ def call_kys_by_udise(udise_code: str) -> Dict[str, Any]:
         return {"ok": False, "error": str(e), "url": url}
 
 
-def call_search_sdk(api_key: str, payload_text: str) -> Dict[str, Any]:
+def call_search_sdk(api_key, payload_text):
     try:
         client = TavilyClient(api_key)
         resp = client.search(query=payload_text)
@@ -53,58 +54,31 @@ def call_search_sdk(api_key: str, payload_text: str) -> Dict[str, Any]:
         return {"ok": False, "error": str(e)}
 
 
-# ---------- 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:
+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):
     try:
         if isinstance(obj, list):
             return pd.json_normalize(obj)
         if isinstance(obj, dict):
-            for k in ("results", "data", "hits", "items", "content"):
+            for k in ("results", "data", "hits", "items"):
                 if k in obj and isinstance(obj[k], list):
                     return pd.json_normalize(obj[k])
             return pd.json_normalize([obj])
@@ -113,36 +87,40 @@ def json_to_table(obj: Any) -> pd.DataFrame:
     return pd.DataFrame()
 
 
-def to_table_from_kys(kys_json: Dict[str, Any]) -> 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.
+    """
     try:
         content = None
         if isinstance(kys_json, dict):
-            if isinstance(kys_json.get("data"), dict) and isinstance(kys_json["data"].get("content"), list):
-                content = kys_json["data"]["content"]
+            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")
             elif isinstance(kys_json.get("content"), list):
-                content = kys_json["content"]
+                content = kys_json.get("content")
         if not content:
             return pd.DataFrame()
-
-        rows = [
-            {
+        rows = []
+        for r in content:
+            rows.append({
                 "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"),
-            }
-            for r in content
-        ]
+                "Management Type": r.get("schMgmtType")
+            })
         return pd.DataFrame(rows)
     except Exception as e:
         print("to_table_from_kys error:", e)
         return pd.DataFrame()
 
 
-# ---------- main workflow ----------
-def search_workflow(school_name: str, state_name: str, search_key: str, use_search: bool, use_kys: bool):
+def search_workflow(school_name, state_name, search_key, use_search, use_kys):
     out = {"kys": None, "search": None, "suggestions": []}
     payload_text = f"{school_name or ''} {state_name or ''} UDISE code".strip()
 
@@ -162,23 +140,19 @@ def search_workflow(school_name: str, state_name: str, search_key: str, use_sear
     return out
 
 
-# ---------- Gradio UI ----------
 with gr.Blocks() as demo:
-    gr.Markdown("# Search + KYS Lookup (Hugging Face Space)")
+    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.
+"""
+    )
 
     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)")
@@ -190,17 +164,15 @@ with gr.Blocks() as demo:
 
     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)
@@ -220,24 +192,10 @@ with gr.Blocks() as demo:
         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")
 
-        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],
-    )
+    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 ""
@@ -260,11 +218,21 @@ with gr.Blocks() as demo:
 
     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()
+    demo.launch()