Update app.py

app.py

@@ -4,6 +4,11 @@
 
 import gradio as gr
 import asyncio
+import time
+import csv
+import io
+import zipfile
+from collections import defaultdict
 from typing import Dict, List
 
 from ingest.registry import get_all_adapters
@@ -11,10 +16,6 @@ from policy.kill_switch import KillSwitch
 from analytics import log_event
 
 
-# ============================================================
-# App Metadata
-# ============================================================
-
 APP_TITLE = "Federal FOIA Intelligence Search"
 APP_SUBTITLE = "Public Electronic Reading Rooms Only"
 
@@ -23,12 +24,28 @@ ALL_ADAPTERS = get_all_adapters()
 
 
 # ============================================================
-# Adapter Utilities
+# Utility: FOIA Exemption Heuristic
 # ============================================================
 
-def get_active_adapters(include_stubs: bool, enable_extended: bool) -> Dict[str, object]:
-    active = {}
+def classify_foia_exemption(text: str) -> str:
+    text = text.lower()
+    if "classified" in text or "national security" in text:
+        return "b(1)"
+    if "statute" in text:
+        return "b(3)"
+    if "privacy" in text or "personnel" in text:
+        return "b(6)"
+    if "law enforcement" in text or "investigation" in text:
+        return "b(7)"
+    return "—"
+
+
+# ============================================================
+# Adapter Filtering
+# ============================================================
 
+def get_active_adapters(include_stubs, enable_extended):
+    active = {}
     for name, adapter in ALL_ADAPTERS.items():
         if kill_switch.is_disabled(name):
             continue
@@ -37,81 +54,120 @@ def get_active_adapters(include_stubs: bool, enable_extended: bool) -> Dict[str,
         if getattr(adapter, "extended", False) and not enable_extended:
             continue
         active[name] = adapter
-
     return active
 
 
-def build_coverage_heatmap(include_stubs: bool, enable_extended: bool):
-    """
-    Returns a table describing agency coverage state.
-    """
-    rows = []
+# ============================================================
+# Async Search + Metrics
+# ============================================================
 
-    for name, adapter in ALL_ADAPTERS.items():
-        if kill_switch.is_disabled(name):
-            status = "🔴 Disabled"
-        elif getattr(adapter, "is_stub", False):
-            status = "🟡 Stub"
-        else:
-            status = "🟢 Live"
+async def timed_search(adapter, query):
+    start = time.time()
+    try:
+        results = await adapter.search(query)
+        latency = round(time.time() - start, 2)
+        return results, latency, "🟢 Healthy"
+    except Exception:
+        return [], None, "🔴 Error"
 
-        scope = "Extended" if getattr(adapter, "extended", False) else "Default"
 
-        if scope == "Extended" and not enable_extended:
-            status = "⚪ Hidden (Extended Disabled)"
+async def federated_search(query, adapters):
+    tasks = {
+        name: timed_search(adapter, query)
+        for name, adapter in adapters.items()
+    }
 
-        if getattr(adapter, "is_stub", False) and not include_stubs:
-            status = "⚪ Hidden (Stubs Disabled)"
+    responses = await asyncio.gather(*tasks.values())
 
-        rows.append([
-            name,
-            scope,
-            status
-        ])
+    results = []
+    metrics = {}
 
-    return rows
+    for (name, _), (docs, latency, health) in zip(tasks.items(), responses):
+        metrics[name] = {
+            "count": len(docs),
+            "latency": latency,
+            "health": health
+        }
+        for d in docs:
+            d["agency"] = name
+            d["foia_exemption"] = classify_foia_exemption(
+                d.get("snippet", "")
+            )
+            results.append(d)
+
+    return results, metrics
 
 
 # ============================================================
-# Async Search
+# Search Handler
 # ============================================================
 
-async def federated_search(query: str, adapters: Dict[str, object]) -> List[dict]:
-    tasks = [adapter.search(query) for adapter in adapters.values()]
-    results = await asyncio.gather(*tasks, return_exceptions=True)
-
-    flattened = []
-    for r in results:
-        if isinstance(r, Exception):
-            continue
-        flattened.extend(r)
-
-    return flattened
-
-
 def search_handler(query, include_stubs, enable_extended):
     if not query.strip():
-        return [], "⚠️ Please enter a search term."
+        return [], [], "⚠️ Enter a search term."
 
     log_event("search", {"query": query})
 
     adapters = get_active_adapters(include_stubs, enable_extended)
 
-    if not adapters:
-        return [], "⚠️ No active sources enabled."
-
     try:
-        results = asyncio.run(federated_search(query, adapters))
+        results, metrics = asyncio.run(
+            federated_search(query, adapters)
+        )
     except RuntimeError:
         loop = asyncio.get_event_loop()
-        results = loop.run_until_complete(
+        results, metrics = loop.run_until_complete(
             federated_search(query, adapters)
         )
 
-    if not results:
-        return [], "No public FOIA documents found."
+    table = [
+        [
+            r["agency"],
+            r.get("title"),
+            r.get("url"),
+            r.get("snippet"),
+            r["foia_exemption"]
+        ]
+        for r in results
+    ]
+
+    heatmap = [
+        [
+            agency,
+            m["count"],
+            m["latency"] or "—",
+            m["health"]
+        ]
+        for agency, m in metrics.items()
+    ]
+
+    status = f"Found {len(results)} documents across {len(metrics)} agencies."
+    return table, heatmap, status
+
+
+# ============================================================
+# Journalist ZIP Export
+# ============================================================
+
+def export_selected(rows):
+    if not rows:
+        return None
+
+    buf = io.BytesIO()
+    with zipfile.ZipFile(buf, "w") as z:
+        csv_buf = io.StringIO()
+        writer = csv.writer(csv_buf)
+        writer.writerow(["Agency", "Title", "URL", "Snippet", "FOIA Exemption"])
+        for r in rows:
+            writer.writerow(r)
+        z.writestr("index.csv", csv_buf.getvalue())
+        z.writestr(
+            "README.txt",
+            "Federal FOIA Intelligence Search\nPublic Electronic Reading Rooms Only"
+        )
 
-    return results, f"Found {len(results)} public FOIA documents."
+    buf.seek(0)
+    return buf
 
 
 # ============================================================
@@ -119,108 +175,52 @@ def search_handler(query, include_stubs, enable_extended):
 # ============================================================
 
 with gr.Blocks(title=APP_TITLE) as demo:
-    gr.Markdown(
-        f"""
+    gr.Markdown(f"""
 # 🏛️ **{APP_TITLE}**
 ### *{APP_SUBTITLE}*
 
-Federated discovery across **lawfully published U.S. Government FOIA
-Electronic Reading Rooms**.
-
-**Safeguards**
-- Public documents only
-- Robots.txt enforced
-- No restricted systems
-- No classified access
-- Per-agency kill switches
-        """
-    )
-
-    query = gr.Textbox(
-        label="Search FOIA Reading Rooms",
-        placeholder="e.g. UAP, MKULTRA, satellite reconnaissance"
-    )
+Public-interest discovery across U.S. Government FOIA Electronic Reading Rooms.
+""")
 
-    with gr.Row():
-        include_stubs = gr.Checkbox(
-            label="Include Stub Sources (Coverage Only)",
-            value=True
-        )
-        enable_extended = gr.Checkbox(
-            label="Enable Extended Features (Live Intelligence-Related Rooms)",
-            value=False
-        )
+    query = gr.Textbox(label="Search FOIA Reading Rooms")
+    include_stubs = gr.Checkbox(label="Include Stub Sources", value=True)
+    enable_extended = gr.Checkbox(label="Enable Extended Features", value=False)
 
     search_btn = gr.Button("Search")
-
     status = gr.Markdown()
 
     results = gr.Dataframe(
-        headers=["source", "title", "url", "snippet"],
-        datatype=["str", "str", "str", "str"],
-        interactive=False,
+        headers=["Agency", "Title", "URL", "Snippet", "FOIA Exemption"],
+        interactive=True,
         wrap=True
     )
 
+    coverage = gr.Dataframe(
+        headers=["Agency", "Result Count", "Latency (s)", "Health"],
+        interactive=False
+    )
+
     search_btn.click(
-        fn=search_handler,
+        search_handler,
         inputs=[query, include_stubs, enable_extended],
-        outputs=[results, status]
+        outputs=[results, coverage, status]
     )
 
-    # ========================================================
-    # Coverage Heatmap Panel
-    # ========================================================
-
-    with gr.Accordion("🗺️ Agency Coverage Map", open=False):
-        gr.Markdown(
-            """
-This panel shows **search coverage status** by agency.
+    export_btn = gr.Button("🧾 Export Selected Results (ZIP)")
+    export_file = gr.File()
 
-- 🟢 Live = searchable public FOIA reading room
-- 🟡 Stub = declared coverage (not yet live)
-- 🔴 Disabled = kill switch / unavailable
-            """
-        )
-
-        coverage_table = gr.Dataframe(
-            headers=["Agency", "Scope", "Coverage Status"],
-            datatype=["str", "str", "str"],
-            interactive=False
-        )
-
-        def refresh_coverage(include_stubs, enable_extended):
-            return build_coverage_heatmap(include_stubs, enable_extended)
-
-        include_stubs.change(
-            refresh_coverage,
-            inputs=[include_stubs, enable_extended],
-            outputs=[coverage_table]
-        )
-
-        enable_extended.change(
-            refresh_coverage,
-            inputs=[include_stubs, enable_extended],
-            outputs=[coverage_table]
-        )
-
-        # Initial render
-        coverage_table.value = build_coverage_heatmap(True, False)
+    export_btn.click(
+        export_selected,
+        inputs=[results],
+        outputs=[export_file]
+    )
 
-    gr.Markdown(
-        """
+    gr.Markdown("""
 ---
-### ⚖️ Legal & Ethical Notice
-
-This system indexes **only documents already released to the public**
-by U.S. Government agencies under the Freedom of Information Act (FOIA).
-
-It does **not**:
-- Circumvent access controls
-- Access classified systems
-- Declassify or infer restricted information
-        """
-    )
+### ⚖️ Legal Notice
+FOIA exemption labels are **heuristic indicators only** and do not replace
+official agency determinations.
+""")
 
 if __name__ == "__main__":
     demo.launch()