Update app.py

app.py

@@ -1,51 +1,235 @@
 import asyncio
+import time
+from typing import List, Dict
+
 import gradio as gr
-from typing import Dict, List
 
-from ingest.registry import get_all_adapters, BASE_AGENCIES, EXTENDED_AGENCIES
+from ingest.registry import (
+    get_all_adapters,
+    BASE_AGENCIES,
+    EXTENDED_AGENCIES,
+)
+from citations import cite
+import coverage
 import saved_searches
 
 
+# -----------------------------
+# Adapter Registry
+# -----------------------------
+
 ALL_ADAPTERS = get_all_adapters()
 
 
+# -----------------------------
+# Async Search Engine
+# -----------------------------
+
+async def _search_adapter(adapter, query: str) -> Dict:
+    start = time.time()
+    try:
+        results = await adapter.search(query)
+        latency = round(time.time() - start, 2)
+        return {
+            "ok": True,
+            "results": results,
+            "latency": latency,
+            "source": adapter.source_name,
+        }
+    except Exception as e:
+        return {
+            "ok": False,
+            "results": [],
+            "latency": None,
+            "source": adapter.source_name,
+            "error": str(e),
+        }
+
+
 async def federated_search_async(
     query: str,
     enabled_agencies: List[str],
 ):
     adapters = [
-        ALL_ADAPTERS[k]
-        for k in enabled_agencies
-        if k in ALL_ADAPTERS
+        ALL_ADAPTERS[a]
+        for a in enabled_agencies
+        if a in ALL_ADAPTERS
     ]
 
-    tasks = [a.search(query) for a in adapters]
-    results = await asyncio.gather(*tasks, return_exceptions=True)
+    tasks = [_search_adapter(a, query) for a in adapters]
+    responses = await asyncio.gather(*tasks)
 
-    flat = []
-    for r in results:
-        if isinstance(r, list):
-            flat.extend(r)
-    return flat
+    flat_results = []
+    health = {}
 
+    for r in responses:
+        health[r["source"]] = {
+            "ok": r["ok"],
+            "latency": r["latency"],
+        }
+        if r["ok"]:
+            flat_results.extend(r["results"])
 
-def run_search(query, hide_stubs, extended, enabled_agencies):
-    loop = asyncio.get_event_loop()
+    return flat_results, health
+
+
+# -----------------------------
+# Rendering Helpers
+# -----------------------------
+
+def badge(result: Dict) -> str:
+    badges = []
+    badges.append("🟢 LIVE" if result.get("live") else "🟡 STUB")
+    if result.get("extended"):
+        badges.append("⚠️ EXTENDED")
+    return " · ".join(badges)
+
+
+def render_result(r: Dict) -> str:
+    return f"""
+### {r['title']}
+**{r['source']} · {badge(r)}**
+
+{r['snippet']}
+
+🔗 {r['url']}
+
+<details>
+<summary>📑 Citation</summary>
+
+{cite(r)}
+
+</details>
+"""
+
+
+def render_health(health: Dict) -> str:
+    lines = ["### 🩺 Agency Health\n"]
+    for agency, h in health.items():
+        if h["ok"]:
+            lines.append(f"- **{agency}**: 🟢 {h['latency']}s")
+        else:
+            lines.append(f"- **{agency}**: 🔴 unavailable")
+    return "\n".join(lines)
+
+
+# -----------------------------
+# Sync Wrapper (HF Safe)
+# -----------------------------
+
+def run_search(
+    query: str,
+    hide_stubs: bool,
+    enable_extended: bool,
+    enabled_extended_agencies: List[str],
+):
+    if not query.strip():
+        return "Enter a search term.", "", ""
 
     agencies = list(BASE_AGENCIES)
-    if extended:
-        agencies.extend([a for a in enabled_agencies if a in EXTENDED_AGENCIES])
 
-    results = loop.run_until_complete(
+    if enable_extended:
+        agencies.extend(
+            [a for a in enabled_extended_agencies if a in EXTENDED_AGENCIES]
+        )
+
+    loop = asyncio.get_event_loop()
+    results, health = loop.run_until_complete(
         federated_search_async(query, agencies)
     )
 
-    saved_searches.save(query, extended)
+    saved_searches.save(query, enable_extended)
 
     if hide_stubs:
         results = [r for r in results if r.get("live")]
 
     if not results:
-        return "No results found."
+        return "No results found.", render_health(health), coverage.render_coverage()
+
+    rendered = "\n\n---\n\n".join(render_result(r) for r in results)
+    return rendered, render_health(health), coverage.render_coverage()
+
+
+# -----------------------------
+# Gradio UI
+# -----------------------------
+
+with gr.Blocks(title="Federated FOIA Search") as demo:
+    gr.Markdown(
+        """
+# 🏛️ Federated FOIA Document Search
+Search **public government FOIA reading rooms** across multiple agencies.
+
+- Default: **Safe, public-only sources**
+- Extended mode: **Additional public releases (opt-in)**
+- No authentication
+- No classified or restricted systems
+"""
+    )
+
+    with gr.Row():
+        query = gr.Textbox(
+            label="Search term",
+            placeholder="e.g. UAP radar incident",
+        )
+
+    with gr.Row():
+        hide_stubs = gr.Checkbox(
+            label="Hide stub (non-live) sources",
+            value=False,
+        )
+        enable_extended = gr.Checkbox(
+            label="Enable Extended Features (live but sensitive public sources)",
+            value=False,
+        )
+
+    with gr.Accordion("Extended Agency Kill Switches", open=False):
+        extended_agencies = gr.CheckboxGroup(
+            choices=EXTENDED_AGENCIES,
+            label="Enable specific extended agencies",
+        )
+
+    search_btn = gr.Button("🔍 Search")
+
+    results_md = gr.Markdown()
+    health_md = gr.Markdown()
+    coverage_md = gr.Markdown()
+
+    search_btn.click(
+        fn=run_search,
+        inputs=[
+            query,
+            hide_stubs,
+            enable_extended,
+            extended_agencies,
+        ],
+        outputs=[
+            results_md,
+            health_md,
+            coverage_md,
+        ],
+    )
+
+    with gr.Accordion("Saved Searches", open=False):
+        gr.Markdown(saved_searches.render())
+
+    with gr.Accordion("About & Legal", open=False):
+        gr.Markdown(
+            """
+**Legal Notice**
+
+This tool:
+- Accesses **only publicly available FOIA Electronic Reading Rooms**
+- Respects robots.txt and rate limits
+- Does not bypass access controls
+- Is intended for journalism, research, and public-interest use
+"""
+        )
+
+
+# -----------------------------
+# HF Entry Point
+# -----------------------------
 
-    return "\n\n---\n\n".join(render(r) for r in results)
+if __name__ == "__main__":
+    demo.launch()