app.py

# search_kys_space.py
# Gradio app for Hugging Face Spaces (uses Search API via Tavily SDK internally)

import gradio as gr
import requests
import pandas as pd
import json
from tavily import TavilyClient

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",
    "Manipur",
    "Meghalaya",
    "Mizoram",
    "Nagaland",
    "Odisha",
    "Puducherry",
    "Rajasthan",
    "Sikkim",
    "Tamil nadu",
    "Telangana",
    "Tripura",
    "Uttar pradesh",
    "Uttarakhand"
]



def call_kys_by_udise(udise_code):
    url = KYS_SAMPLE.format(udise=udise_code)
    try:
        resp = requests.get(url, timeout=10)
        resp.raise_for_status()
        data = resp.json()
        return {"ok": True, "url": url, "data": data}
    except Exception as e:
        return {"ok": False, "error": str(e), "url": url}


def call_search_sdk(api_key, payload_text):
    try:
        client = TavilyClient(api_key)
        resp = client.search(query=payload_text)
        return {"ok": True, "data": resp}
    except Exception as e:
        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):
    try:
        if isinstance(obj, list):
            return pd.json_normalize(obj)
        if isinstance(obj, dict):
            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])
    except Exception:
        pass
    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.
    """
    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")
            elif isinstance(kys_json.get("content"), list):
                content = kys_json.get("content")
        if not content:
            return pd.DataFrame()
        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")
            })
        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):
    out = {"kys": None, "search": None, "suggestions": []}
    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"):
            candidates = extract_udise_candidates_from_search(search_res["data"])
            out["suggestions"] = candidates
    else:
        out["search"] = {"ok": False, "error": "Search disabled or SDK not used"}

    if use_kys and school_name and school_name.strip().isdigit() and 6 <= len(school_name.strip()) <= 14:
        kys_res = call_kys_by_udise(school_name.strip())
        out["kys"] = kys_res

    return out


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.
"""
    )

    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)

    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)")
    clear_key_btn = gr.Button("Clear saved key")

    with gr.Row():
        use_search = gr.Checkbox(value=True, label="Call Search API")
        use_kys = gr.Checkbox(value=True, label="Call KYS by UDISE")

    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)
    lookup_btn = gr.Button("Lookup UDISE (Call KYS)")

    kys_output_json = gr.JSON(label="KYS Raw Output", visible=False)
    kys_table = gr.DataFrame(headers=None, label="KYS results (table)")

    saved_key_state = gr.State("")

    def on_run(school, state, key, save_key, use_s, use_k, saved_key):
        effective_key = saved_key if saved_key else key
        res = search_workflow(school, state, effective_key, use_s, use_k)
        tbl = pd.DataFrame()
        if res.get("search") and res["search"].get("ok"):
            tbl = json_to_table(res["search"]["data"])
        suggestions = res.get("suggestions", [])
        new_saved_key = saved_key
        textbox_value = key
        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])

    def on_select_suggestion(choice):
        return choice or ""

    suggestions_dropdown.change(on_select_suggestion, inputs=[suggestions_dropdown], outputs=[udise_input])

    def on_lookup_udise(udise_code):
        if not udise_code or not udise_code.strip().isdigit():
            return {"ok": False, "error": "Provide a numeric UDISE code (6-14 digits)."}, pd.DataFrame()
        kys_res = call_kys_by_udise(udise_code.strip())
        df = pd.DataFrame()
        if kys_res.get("ok"):
            df = to_table_from_kys(kys_res["data"]) if kys_res.get("data") else pd.DataFrame()
        return kys_res, df

    lookup_btn.click(on_lookup_udise, inputs=[udise_input], outputs=[kys_output_json, kys_table])

    def on_clear_key(_):
        return "", ""

    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()