Update app.py

app.py

@@ -1,35 +1,24 @@
 import asyncio
 import gradio as gr
-from typing import List, Dict
+from typing import Dict, List
 
-from ingest.cia_reading_room import CIAAdapter
-from ingest.fbi_vault import FBIAdapter
+from ingest.registry import get_all_adapters, BASE_AGENCIES, EXTENDED_AGENCIES
+import saved_searches
 
-from ingest.cia_extended import CIAExtendedAdapter
-from ingest.nsa_extended import NSAExtendedAdapter
-from ingest.nro_extended import NROExtendedAdapter
-from ingest.dod_special_extended import (
-    AATIPExtendedAdapter,
-    SAPExtendedAdapter,
-    TENCAPExtendedAdapter,
-    SpecialActivitiesExtendedAdapter
-)
 
-import saved_searches
+ALL_ADAPTERS = get_all_adapters()
 
-BASE_ADAPTERS = [CIAAdapter(), FBIAdapter()]
-EXTENDED_ADAPTERS = [
-    CIAExtendedAdapter(),
-    NSAExtendedAdapter(),
-    NROExtendedAdapter(),
-    AATIPExtendedAdapter(),
-    SAPExtendedAdapter(),
-    TENCAPExtendedAdapter(),
-    SpecialActivitiesExtendedAdapter()
-]
 
-async def federated_search_async(query: str, extended: bool):
-    adapters = BASE_ADAPTERS + (EXTENDED_ADAPTERS if extended else [])
+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
+    ]
+
     tasks = [a.search(query) for a in adapters]
     results = await asyncio.gather(*tasks, return_exceptions=True)
 
@@ -39,27 +28,17 @@ async def federated_search_async(query: str, extended: bool):
             flat.extend(r)
     return flat
 
-def badge(r):
-    h = r.get("health", {})
-    latency = f"{h.get('latency_ms','?')}ms"
-    status = "🟢" if h.get("ok") else "🔴"
-    if r.get("extended"):
-        return f"🟣 EXT · {status} {latency}"
-    return "🟢 LIVE" if r.get("live") else "⚪ STUB"
-
-def render(r):
-    return f"""
-### {r['title']}
-**{badge(r)} · {r['source']}**
 
-{r['snippet']}
+def run_search(query, hide_stubs, extended, enabled_agencies):
+    loop = asyncio.get_event_loop()
 
-🔗 {r['url']}
-"""
+    agencies = list(BASE_AGENCIES)
+    if extended:
+        agencies.extend([a for a in enabled_agencies if a in EXTENDED_AGENCIES])
 
-def run_search(query, hide_stubs, extended):
-    loop = asyncio.get_event_loop()
-    results = loop.run_until_complete(federated_search_async(query, extended))
+    results = loop.run_until_complete(
+        federated_search_async(query, agencies)
+    )
 
     saved_searches.save(query, extended)
 
@@ -69,44 +48,4 @@ def run_search(query, hide_stubs, extended):
     if not results:
         return "No results found."
 
-    return "\n\n---\n\n".join(render(r) for r in results)
-
-def show_saved():
-    rows = saved_searches.list_saved()
-    if not rows:
-        return "No saved searches yet."
-    return "\n".join(
-        f"- **{r['query']}** (extended={r['extended']})"
-        for r in rows
-    )
-
-with gr.Blocks(css="""
-.gradio-container { max-width: 1100px }
-""") as demo:
-    gr.Markdown("# 📜 FOIA Federated Search")
-    gr.Markdown(
-        "Transparent federated search across public FOIA reading rooms.\n\n"
-        "**Badges**: ⚪ Stub · 🟢 Live · 🟣 Extended · 🔴 Unhealthy"
-    )
-
-    with gr.Row():
-        query = gr.Textbox(label="Search term", scale=3)
-        hide_stubs = gr.Checkbox(label="Hide stub sources", value=False)
-        extended = gr.Checkbox(
-            label="Enable Extended Features",
-            value=False,
-            info="Advanced public sources. Disabled by default."
-        )
-
-    output = gr.Markdown()
-    saved = gr.Markdown()
-
-    with gr.Row():
-        gr.Button("Search").click(
-            run_search,
-            inputs=[query, hide_stubs, extended],
-            outputs=output
-        )
-        gr.Button("Saved Searches").click(show_saved, outputs=saved)
-
-demo.launch()
+    return "\n\n---\n\n".join(render(r) for r in results)