Update app.py

app.py

@@ -81,131 +81,6 @@ def normalize_result(res: Dict[str, Any]):
 # ANALYST COPILOT (LLM)
 # ---------------------------------------------------------------------
 
-def respond(
-    message,
-    history,
-    system_prompt,
-    model_name,
-    hf_token,
-    temperature,
-    top_p,
-    max_tokens,
-):
-    """Streaming response from WhiteRabbit Neo or Cybertron."""
-    client = InferenceClient(model=model_name, token=hf_token.token)
-
-    msgs = [{"role": "system", "content": system_prompt}]
-    msgs.extend(history)
-    msgs.append({"role": "user", "content": message})
-
-    buf = ""
-    for chunk in client.chat_completion(
-        messages=msgs,
-        max_tokens=max_tokens,
-        temperature=temperature,
-        top_p=top_p,
-        stream=True,
-    ):
-        delta = chunk.choices[0].delta.content
-        if delta:
-            buf += delta
-            yield buf
-
-def inject_osint(history, osint_obj):
-    """Inject raw JSON results into the chat context."""
-    pretty = json.dumps(osint_obj, indent=2, default=str)
-    history.append({
-        "role": "system",
-        "content": f"### Injected OSINT Result\n```\n{pretty}\n```"
-    })
-    return history
-
-# ---------------------------------------------------------------------
-# OSINT Dashboard Callbacks
-# ---------------------------------------------------------------------
-
-def ui_lookup_ip(ip, enrich, mitre):
-    raw = call_task("lookup_ip", {"ip": ip, "enrich": enrich, "map_mitre": mitre})
-    norm = normalize_result(raw)
-    return norm["summary"], norm["markdown"], norm["json"], norm["mitre"], norm["stix"], raw
-
-def ui_lookup_domain(domain, enrich, mitre):
-    raw = call_task("lookup_domain", {"domain": domain, "enrich": enrich, "map_mitre": mitre})
-    norm = normalize_result(raw)
-    return norm["summary"], norm["markdown"], norm["json"], norm["mitre"], norm["stix"], raw
-
-def ui_lookup_hash(h, ht, enrich, mitre):
-    raw = call_task("lookup_hash", {"hash": h, "hash_type": ht, "enrich": enrich, "map_mitre": mitre})
-    norm = normalize_result(raw)
-    return norm["summary"], norm["markdown"], norm["json"], norm["mitre"], norm["stix"], raw
-
-def ui_correlate_iocs(iocs):
-    lst = [x.strip() for x in iocs.splitlines() if x.strip()]
-    raw = call_task("correlate_iocs", {"iocs": lst})
-    norm = normalize_result(raw)
-    return norm["summary"], norm["markdown"], norm["json"], norm["mitre"], raw
-
-def ui_quickscan(target):
-    raw = call_task("quickscan", {"target": target})
-    norm = normalize_result(raw)
-    return norm["summary"], norm["markdown"], norm["json"], raw
-
-# ---------------------------------------------------------------------
-# MCP Bridge
-# ---------------------------------------------------------------------
-
-def ui_bridge(tool, args_json):
-    try:
-        payload = json.loads(args_json)
-    except Exception as e:
-        return json.dumps({"error": str(e)}, indent=2), "", {}
-    raw = call_task(tool, payload)
-    norm = normalize_result(raw)
-    return norm["json"], norm["markdown"], raw
-
-# ---------------------------------------------------------------------
-# BUILD UI
-# ---------------------------------------------------------------------
-
-def build_interface():
-    with gr.Blocks(title="Parrot OSINT MCP Console") as demo:
-        gr.Markdown("# Parrot OSINT MCP Console")
-
-        osint_state = gr.State({})
-
-        # -------------------------
-        # OSINT Dashboard
-        # -------------------------
-        with gr.Tab("OSINT Dashboard"):
-            with gr.Tab("IP"):
-                ip = gr.Textbox(label="IP Address")
-                enrich = gr.Checkbox(value=True, label="Enrich")
-                mitre = gr.Checkbox(value=True, label="MITRE Map")
-                run = gr.Button("Run IP Lookup")
-
-                s = gr.Textbox(label="Summary")
-                md = gr.Markdown()
-                js = gr.Code(language="json")
-                mt = gr.Code(language="json")
-                st = gr.Code(language="json")
-
-                run.click(ui_lookup_ip, [ip, enrich, mitre], [s, md, js, mt, st, osint_state])
-
-            # Add other tabs (Domain, Hash, etc.)
-            # Your earlier implementation plugs in cleanly.
-
-        # -------------------------
-        # MCP Bridge
-        # -------------------------
-        with gr.Tab("MCP Bridge"):
-            tool = gr.Dropdown(sorted(TASK_REGISTRY.keys()))
-            args = gr.Code(language="json")
-            btn = gr.Button("Call Tool")
-            out_js = gr.Code(language="json")
-            out_md = gr.Markdown()
-
-            btn.click(ui_bridge, [tool, args], [out_js, out_md, osint_state])
-
         # -------------------------
         # Analyst Copilot
         # -------------------------
@@ -231,13 +106,16 @@ def build_interface():
                 value="berkeley-nest/WhiteRabbitNeo-8B",
             )
 
+            gr.Markdown("### Provide your HuggingFace API Token for the LLM:")
+            hf_token = gr.OAuthToken()
+
             chatbot = gr.ChatInterface(
                 respond,
                 type="messages",
                 additional_inputs=[
                     system_prompt,
                     model_select,
-                    gr.OAuthToken(label="HF Token"),
+                    hf_token,
                     gr.Slider(0.1, 2.0, value=0.7, step=0.1, label="Temperature"),
                     gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top-p"),
                     gr.Slider(32, 4096, value=512, step=32, label="Max Tokens"),
@@ -251,7 +129,7 @@ def build_interface():
                 outputs=[chatbot._chatbot_state],
             )
 
-    return demo
+            return demo
 
 
 if __name__ == "__main__":