Update app.py

app.py

@@ -5,45 +5,53 @@ from typing import List, Dict
 from ingest.cia_reading_room import CIAAdapter
 from ingest.fbi_vault import FBIAdapter
 
-# Register adapters here
-ADAPTERS = [
-    CIAAdapter(),        # stub (safe)
-    FBIAdapter(),        # live
+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
+)
+
+BASE_ADAPTERS = [
+    CIAAdapter(),      # stub
+    FBIAdapter(),      # live
 ]
 
-# -------------------------
-# Async search core
-# -------------------------
-async def federated_search_async(query: str) -> List[Dict]:
-    tasks = [adapter.search(query) for adapter in ADAPTERS]
+EXTENDED_ADAPTERS = [
+    CIAExtendedAdapter(),
+    NSAExtendedAdapter(),
+    NROExtendedAdapter(),
+    AATIPExtendedAdapter(),
+    SAPExtendedAdapter(),
+    TENCAPExtendedAdapter(),
+    SpecialActivitiesExtendedAdapter()
+]
+
+
+async def federated_search_async(query: str, extended_mode: bool) -> List[Dict]:
+    adapters = list(BASE_ADAPTERS)
+    if extended_mode:
+        adapters.extend(EXTENDED_ADAPTERS)
+
+    tasks = [a.search(query) for a in adapters]
     results = await asyncio.gather(*tasks, return_exceptions=True)
 
-    flat: List[Dict] = []
+    flat = []
     for r in results:
-        if isinstance(r, Exception):
-            continue
-        flat.extend(r)
-
+        if isinstance(r, list):
+            flat.extend(r)
     return flat
 
 
-# -------------------------
-# Sync wrapper (HF-safe)
-# -------------------------
-def federated_search(query: str) -> List[Dict]:
-    loop = asyncio.get_event_loop()
-    if loop.is_running():
-        # HF / Gradio case
-        return asyncio.create_task(federated_search_async(query))
+def format_result(r: Dict) -> str:
+    if r.get("extended"):
+        badge = "🟣 LIVE (Extended)"
     else:
-        return loop.run_until_complete(federated_search_async(query))
-
+        badge = "🟢 LIVE" if r.get("live") else "⚪ STUB"
 
-# -------------------------
-# UI helpers
-# -------------------------
-def format_result(r: Dict) -> str:
-    badge = "🟢 LIVE" if r.get("live") else "⚪ STUB"
     return f"""
 ### {r['title']}
 **{badge} · {r['source']}**
@@ -54,12 +62,15 @@ def format_result(r: Dict) -> str:
 """
 
 
-def run_search(query: str, hide_stubs: bool):
-    results = federated_search(query)
-
-    # If async task returned
-    if asyncio.isfuture(results):
-        results = asyncio.get_event_loop().run_until_complete(results)
+def run_search(query: str, hide_stubs: bool, extended_mode: bool):
+    loop = asyncio.get_event_loop()
+    if loop.is_running():
+        task = asyncio.create_task(federated_search_async(query, extended_mode))
+        results = loop.run_until_complete(task)
+    else:
+        results = loop.run_until_complete(
+            federated_search_async(query, extended_mode)
+        )
 
     if hide_stubs:
         results = [r for r in results if r.get("live")]
@@ -70,26 +81,26 @@ def run_search(query: str, hide_stubs: bool):
     return "\n\n---\n\n".join(format_result(r) for r in results)
 
 
-# -------------------------
-# Gradio UI
-# -------------------------
 with gr.Blocks() as demo:
     gr.Markdown("# 📜 FOIA Federated Search")
     gr.Markdown(
-        "Search across public government FOIA reading rooms. "
-        "Stub sources are clearly labeled."
+        "Search across public government FOIA reading rooms.\n\n"
+        "- ⚪ Stub = placeholder\n"
+        "- 🟢 Live = public indexed source\n"
+        "- 🟣 Extended = advanced public sources (opt-in)"
     )
 
-    query = gr.Textbox(label="Search term", placeholder="e.g. UAP radar")
-    hide_stubs = gr.Checkbox(
-        label="Hide stub (placeholder) sources",
+    query = gr.Textbox(label="Search term")
+    hide_stubs = gr.Checkbox(label="Hide stub sources", value=False)
+    extended_mode = gr.Checkbox(
+        label="Enable Extended Features (advanced public agencies)",
         value=False
     )
     output = gr.Markdown()
 
     gr.Button("Search").click(
         run_search,
-        inputs=[query, hide_stubs],
+        inputs=[query, hide_stubs, extended_mode],
         outputs=output
     )