Update app.py

app.py

@@ -1,8 +1,9 @@
+# tavily_kys_space.py
+# Gradio app for Hugging Face Spaces (Tavily API key textbox + Save key in session)
+
 import gradio as gr
 import requests
-import json
 import pandas as pd
-from difflib import get_close_matches
 
 KYS_SAMPLE = "https://kys.udiseplus.gov.in/webapp/api/search-schools?searchType=3&searchParam={udise}"
 
@@ -23,12 +24,12 @@ def call_tavily(tavily_url, api_key, payload_text):
     if api_key:
         headers["Authorization"] = api_key
     body = {"query": payload_text}
+
     try:
         resp = requests.post(tavily_url, json=body, headers=headers, timeout=12)
         resp.raise_for_status()
         return {"ok": True, "url": tavily_url, "status_code": resp.status_code, "data": resp.json()}
     except Exception as e:
-        # fallback to GET
         try:
             params = {"q": payload_text}
             resp = requests.get(tavily_url, params=params, headers=headers, timeout=12)
@@ -39,7 +40,6 @@ def call_tavily(tavily_url, api_key, payload_text):
 
 
 def extract_udise_candidates_from_tavily(tavily_json):
-    # Heuristic: walk the JSON and collect strings that look like UDISE (digits, length 6-14)
     found = set()
 
     def walk(obj):
@@ -61,20 +61,13 @@ def extract_udise_candidates_from_tavily(tavily_json):
 
 
 def json_to_table(obj):
-    # Try to convert a JSON object/array to pandas DataFrame for display
     try:
         if isinstance(obj, list):
-            df = pd.json_normalize(obj)
-            return df
+            return pd.json_normalize(obj)
         if isinstance(obj, dict):
-            # Often tavily returns {"results": [...]}
-            if "results" in obj and isinstance(obj["results"], list):
-                return pd.json_normalize(obj["results"])
-            # common patterns: top-level list wrapped in data or hits
-            for k in ("data", "hits", "items", "results"):
+            for k in ("results", "data", "hits", "items"):
                 if k in obj and isinstance(obj[k], list):
                     return pd.json_normalize(obj[k])
-            # fallback: flatten dict to single-row table
             return pd.json_normalize([obj])
     except Exception:
         pass
@@ -88,8 +81,7 @@ def to_table_from_kys(kys_json):
             items = items["data"]
         if not items:
             return pd.DataFrame()
-        df = pd.json_normalize(items)
-        return df
+        return pd.json_normalize(items)
     except Exception:
         return pd.DataFrame()
 
@@ -106,7 +98,6 @@ def search_workflow(school_name_or_udise, tavily_url, tavily_key, use_tavily, us
     else:
         out["tavily"] = {"ok": False, "error": "Tavily disabled or no URL provided"}
 
-    # If user provided a numeric UDISE directly and KYS requested, do KYS
     if use_kys and school_name_or_udise and school_name_or_udise.strip().isdigit() and 6 <= len(school_name_or_udise.strip()) <= 14:
         kys_res = call_kys_by_udise(school_name_or_udise.strip())
         out["kys"] = kys_res
@@ -115,17 +106,20 @@ def search_workflow(school_name_or_udise, tavily_url, tavily_key, use_tavily, us
 
 
 with gr.Blocks() as demo:
-    gr.Markdown("""
-    # Tavily + KYS Search (Hugging Face Space)
-    
-    Enter a school name or a UDISE code and call KYS (UDISE) or your Tavily search endpoint.
-    """)
+    # Fixed: properly terminated Markdown strings
+    gr.Markdown('''# Tavily + KYS Search (Hugging Face Space)
 
+Enter a school name or a UDISE code, Tavily endpoint, and optionally an API key.
+Use the 'Save key in session' toggle to keep the key for the current browser session.''')
 
     with gr.Row():
-        inp = gr.Textbox(label="School name or UDISE code", placeholder="e.g. 123456789012 or 'Govt High School...'", lines=1)
+        inp = gr.Textbox(label="School name", placeholder="'Govt High School...'", lines=1)
         tavily_url = gr.Textbox(label="Tavily endpoint URL (POST/GET)", placeholder="https://your-tavily.example/api/search")
-        tavily_key = gr.Textbox(label="Tavily API Key / Authorization (optional)", placeholder="Bearer ... or APIKEY ...")
+
+    # Tavily API key textbox + save toggle
+    tavily_key = gr.Textbox(label="Tavily API Key (optional)", placeholder="Bearer ... or APIKEY ...", 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_tavily = gr.Checkbox(value=True, label="Call Tavily")
@@ -134,36 +128,43 @@ with gr.Blocks() as demo:
     run = gr.Button("Search Tavily")
 
     output_json = gr.JSON(label="Raw Tavily Output (JSON)")
-    tavily_table = gr.DataFrame(headers=None, label="Tavily results (table)")
+    tavily_table = gr.DataFrame(headers=None, label="Search Results (table)")
 
-    gr.Markdown('''### UDISE candidates found in Tavily results
-The dropdown will list all numeric tokens resembling UDISE codes found in Tavily's response. Choose any candidate (or edit the UDISE field) and click 'Lookup UDISE' to call KYS.")
-    suggestions_dropdown = gr.Dropdown(choices=[], label="UDISE candidates (from Tavily)''')
+    gr.Markdown("### UDISE candidates found in search results")
+    suggestions_dropdown = gr.Dropdown(choices=[], label="UDISE candidates")
     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")
     kys_table = gr.DataFrame(headers=None, label="KYS results (table)")
 
-    def on_run(school, turl, tkey, utav, ukys):
-        res = search_workflow(school, turl, tkey, utav, ukys)
-        # Prepare tavily table
+    # A small gr.State to persist the saved key value inside the session
+    saved_key_state = gr.State("")
+
+    def on_run(school, turl, tkey, save_key, utav, ukys, saved_key):
+        # Determine effective key: if save_key and saved_key present use saved_key, else use tkey
+        effective_key = saved_key if saved_key else tkey
+        res = search_workflow(school, turl, effective_key, utav, ukys)
         tav_df = pd.DataFrame()
         if res.get("tavily") and res["tavily"].get("ok"):
             tav_df = json_to_table(res["tavily"]["data"])
-        # suggestions for dropdown
         suggestions = res.get("suggestions", [])
-        return res.get("tavily"), tav_df, suggestions
+        # If user asked to save the key for session, return it into saved_key_state and keep the textbox value
+        new_saved_key = saved_key
+        textbox_value = tkey
+        if save_key:
+            new_saved_key = tkey or saved_key
+            textbox_value = new_saved_key
+        return res.get("tavily"), tav_df, suggestions, new_saved_key, textbox_value
 
-    run.click(on_run, inputs=[inp, tavily_url, tavily_key, use_tavily, use_kys], outputs=[output_json, tavily_table, suggestions_dropdown])
+    # run.click outputs: tavily json, tavily table, suggestions dropdown, saved_key_state, tavily_key (textbox)
+    run.click(on_run, inputs=[inp, tavily_url, tavily_key, save_key_toggle, use_tavily, use_kys, saved_key_state], outputs=[output_json, tavily_table, suggestions_dropdown, saved_key_state, tavily_key])
 
-    # When user selects a suggestion, populate the editable UDISE textbox
     def on_select_suggestion(choice):
         return choice or ""
 
     suggestions_dropdown.change(on_select_suggestion, inputs=[suggestions_dropdown], outputs=[udise_input])
 
-    # Lookup button handler: call KYS for whatever is in 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()
@@ -175,12 +176,18 @@ The dropdown will list all numeric tokens resembling UDISE codes found in Tavily
 
     lookup_btn.click(on_lookup_udise, inputs=[udise_input], outputs=[kys_output_json, kys_table])
 
+    # Clear saved key handler: empties both the saved_key_state and the textbox
+    def on_clear_key(_):
+        return "", ""
+
+    clear_key_btn.click(on_clear_key, inputs=[saved_key_state], outputs=[saved_key_state, tavily_key])
+
     gr.Markdown('''---
-    **Notes:**
-    - Tavily results are attempted to be flattened into a table for easy scanning; complex nested structures will be shown as JSON above.
-    - The UDISE candidates dropdown contains numeric tokens heuristically detected in the Tavily response; you can select one to fill the UDISE field, edit it, or type a different code.
-    - The 'Search Tavily' button only queries Tavily and fills the table + suggestions; use 'Lookup UDISE' to call KYS for the chosen code.
-    ''')
+**Notes:**
+- The 'Save key in session' toggle will keep the key active for the current browser session and across interactions in the Space UI.
+- The key is stored only in-memory for the session (not written to disk) and is not logged.
+- Use the 'Clear saved key' button to remove it from the session.''')
+
 
 if __name__ == "__main__":
     demo.launch()