Update app.py

app.py

@@ -81,6 +81,129 @@ 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\nA multi-mode intelligence workstation.")
+
+        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])
+
+        # -------------------------
+        # 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
         # -------------------------
@@ -129,8 +252,11 @@ def normalize_result(res: Dict[str, Any]):
                 outputs=[chatbot._chatbot_state],
             )
 
-            return demo
+    return demo
 
+# ---------------------------------------------------------------------
+# MAIN
+# ---------------------------------------------------------------------
 
 if __name__ == "__main__":
     demo = build_interface()