Update app.py

app.py

@@ -1,55 +1,62 @@
 """
 Federal FOIA Intelligence Search
 Public Electronic Reading Rooms Only
-HF Spaces compatible (no asyncio.run)
+HF Spaces compatible
 """
 
-import time
 import asyncio
+import time
 import gradio as gr
 from typing import List, Dict
+
 from ingest.registry import get_enabled_adapters
 from ingest.health import get_health_snapshot
-from ingest.cluster import semantic_cluster
-from ingest.export import export_zip
-from ingest.utils import highlight_terms
+from ingest.cluster import build_embeddings, cluster_embeddings
+from ingest.export import export_journalist_zip
+from ingest.utils import highlight_terms, classify_foia_exemptions, format_citation
+from ingest.coverage import build_coverage_heatmap
 
 TITLE = "🏛️ Federal FOIA Intelligence Search"
 SUBTITLE = "Public Electronic Reading Rooms Only"
 
-# ---------- Async-safe search wrapper ----------
-
-async def run_search(query: str, enable_live: bool, hide_stub: bool):
-    adapters = get_enabled_adapters(enable_live=enable_live)
-    results = []
-    per_agency_counts = {}
 
-    tasks = []
-    for adapter in adapters:
-        tasks.append(adapter.search(query))
+# ---------------- Async-safe search ----------------
 
+async def _run_search(query: str, enable_live: bool, hide_stub: bool):
+    adapters = get_enabled_adapters(enable_live=enable_live)
+    tasks = [a.search(query) for a in adapters]
     responses = await asyncio.gather(*tasks, return_exceptions=True)
 
-    for adapter, response in zip(adapters, responses):
-        if isinstance(response, Exception):
+    results: List[Dict] = []
+    per_agency_counts: Dict[str, int] = {}
+
+    for adapter, resp in zip(adapters, responses):
+        if isinstance(resp, Exception):
             continue
         if hide_stub and adapter.is_stub:
             continue
-        per_agency_counts[adapter.source_name] = len(response)
-        for r in response:
+        per_agency_counts[adapter.source_name] = len(resp)
+        for r in resp:
             r["source"] = adapter.source_name
             r["latency"] = adapter.last_latency
+            r["exemptions"] = classify_foia_exemptions(r.get("snippet", ""))
+            r["citation"] = format_citation(r)
             results.append(r)
 
-    clustered = semantic_cluster(results)
-    return results, clustered, per_agency_counts
+    # Embeddings + clustering
+    embeddings = build_embeddings(results)
+    clusters = cluster_embeddings(results, embeddings)
+
+    coverage = build_coverage_heatmap(per_agency_counts)
+
+    return results, clusters, per_agency_counts, coverage
 
 
 def search_ui(query, enable_live, hide_stub):
     loop = asyncio.new_event_loop()
     asyncio.set_event_loop(loop)
-    results, clusters, counts = loop.run_until_complete(
-        run_search(query, enable_live, hide_stub)
+    results, clusters, counts, coverage = loop.run_until_complete(
+        _run_search(query, enable_live, hide_stub)
     )
 
     table = []
@@ -59,20 +66,21 @@ def search_ui(query, enable_live, hide_stub):
             highlight_terms(r["title"], query),
             r["url"],
             r.get("snippet", ""),
+            ",".join(r["exemptions"]),
             f"{r.get('latency', 0):.2f}s"
         ])
 
-    return table, clusters, counts
+    return table, clusters, counts, coverage
 
 
-# ---------- UI ----------
+# ---------------- UI ----------------
 
 with gr.Blocks(title="Federal FOIA Intelligence Search") as demo:
     gr.Markdown(f"# {TITLE}")
     gr.Markdown(f"### {SUBTITLE}")
 
     gr.Markdown(
-        "**Terms of Use:** This tool indexes only documents already released "
+        "**Terms of Use:** This application indexes only documents already released "
         "to the public via official FOIA Electronic Reading Rooms."
     )
 
@@ -85,30 +93,31 @@ with gr.Blocks(title="Federal FOIA Intelligence Search") as demo:
         hide_stub = gr.Checkbox(label="Hide Stub Sources", value=True)
 
     results_table = gr.Dataframe(
-        headers=["Agency", "Title", "URL", "Snippet", "Latency"],
+        headers=["Agency", "Title", "URL", "Snippet", "FOIA Exemptions", "Latency"],
         interactive=False,
         wrap=True
     )
 
-    clusters_box = gr.JSON(label="Semantic Clusters (Explainable)")
+    clusters_box = gr.JSON(label="Semantic Clusters (Interactive Graph Data)")
     counts_box = gr.JSON(label="Per-Agency Result Counts")
+    coverage_box = gr.JSON(label="Coverage Heatmap Data")
 
     health_box = gr.JSON(
         value=get_health_snapshot(),
         label="Agency Health / Latency"
     )
 
-    export_btn = gr.Button("Export Selected → Journalist ZIP")
+    export_btn = gr.Button("Export → Journalist ZIP")
     export_status = gr.Textbox(label="Export Status")
 
     search_btn.click(
         fn=search_ui,
         inputs=[query, enable_live, hide_stub],
-        outputs=[results_table, clusters_box, counts_box]
+        outputs=[results_table, clusters_box, counts_box, coverage_box]
     )
 
     export_btn.click(
-        fn=export_zip,
+        fn=export_journalist_zip,
         inputs=[results_table],
         outputs=[export_status]
     )