Update app.py

app.py

@@ -21,18 +21,22 @@ from huggingface_hub import InferenceClient
 
 TASK_REGISTRY: Dict[str, Any] = {}
 
+
 def _register_tasks():
     """
     Import tasks.* modules dynamically and pull their run() functions.
+    Missing modules are ignored so the UI can still boot.
     """
-    def _try(name, module):
+
+    def _try(name: str, module: str):
         try:
             m = __import__(f"tasks.{module}", fromlist=["*"])
             fn = getattr(m, "run", None)
             if callable(fn):
                 TASK_REGISTRY[name] = fn
         except Exception:
-            pass  # In HF Spaces, missing tasks won't break startup.
+            # In Spaces, you might not have all tasks yet; that's fine.
+            pass
 
     _try("lookup_ip", "lookup_ip")
     _try("lookup_domain", "lookup_domain")
@@ -43,14 +47,15 @@ def _register_tasks():
     _try("mitre_map", "mitre_map")
     _try("quickscan", "quickscan")
 
-_register_tasks()
 
+_register_tasks()
 
 # ---------------------------------------------------------------------
 # Task Execution + Normalization
 # ---------------------------------------------------------------------
 
-def call_task(name: str, payload: Dict[str, Any]):
+
+def call_task(name: str, payload: Dict[str, Any]) -> Dict[str, Any]:
     fn = TASK_REGISTRY.get(name)
     if not fn:
         return {"error": f"Unknown tool '{name}'."}
@@ -64,7 +69,7 @@ def call_task(name: str, payload: Dict[str, Any]):
         return {"error": str(e), "traceback": traceback.format_exc()}
 
 
-def normalize_result(res: Dict[str, Any]):
+def normalize_result(res: Dict[str, Any]) -> Dict[str, str]:
     """Ensures consistent UI formatting."""
     pretty = json.dumps(res, indent=2, default=str)
     summary = res.get("summary", "")
@@ -73,7 +78,8 @@ def normalize_result(res: Dict[str, Any]):
     if not markdown and summary:
         markdown = f"## Summary\n\n{summary}"
 
-    safe_json = lambda v: json.dumps(v, indent=2, default=str) if v else ""
+    def safe_json(value: Any) -> str:
+        return json.dumps(value, indent=2, default=str) if value else ""
 
     return {
         "summary": summary,
@@ -89,6 +95,7 @@ def normalize_result(res: Dict[str, Any]):
 # Analyst Copilot LLM
 # ---------------------------------------------------------------------
 
+
 def respond(
     message,
     history,
@@ -101,11 +108,11 @@ def respond(
 ):
     """
     Streaming LLM output using WhiteRabbit Neo or Cybertron.
-    hf_token is a raw string entered by user.
+    `hf_token` is a raw string entered by the user.
     """
     client = InferenceClient(
         model=model_name,
-        token=hf_token,  # Direct string, no OAuthToken object
+        token=hf_token,  # Direct string token
     )
 
     msgs = [{"role": "system", "content": system_prompt}]
@@ -128,12 +135,17 @@ def respond(
 
 
 def inject_osint(history, osint_obj):
-    """Inject OSINT result JSON into copilot context."""
+    """
+    Inject OSINT result JSON into the copilot's chat history as a system message.
+    `history` is the ChatInterface state (list of messages).
+    """
     pretty = json.dumps(osint_obj, indent=2, default=str)
-    history.append({
-        "role": "system",
-        "content": f"### Injected OSINT Result\n```\n{pretty}\n```"
-    })
+    history.append(
+        {
+            "role": "system",
+            "content": f"### Injected OSINT Result\n```\n{pretty}\n```",
+        }
+    )
     return history
 
 
@@ -141,22 +153,49 @@ def inject_osint(history, osint_obj):
 # 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
+    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})
+    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
+    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})
+    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
+    return (
+        norm["summary"],
+        norm["markdown"],
+        norm["json"],
+        norm["mitre"],
+        norm["stix"],
+        raw,
+    )
 
 
 def ui_correlate_iocs(iocs):
@@ -176,6 +215,7 @@ def ui_quickscan(target):
 # MCP Bridge
 # ---------------------------------------------------------------------
 
+
 def ui_bridge(tool, args_json):
     try:
         payload = json.loads(args_json)
@@ -191,10 +231,15 @@ def ui_bridge(tool, args_json):
 # UI Layout
 # ---------------------------------------------------------------------
 
+
 def build_interface():
     with gr.Blocks(title="Parrot OSINT MCP Console") as demo:
-        gr.Markdown("# 🦜 Parrot OSINT MCP Console\nMulti-mode Intelligence Workstation.")
+        gr.Markdown(
+            "# 🦜 Parrot OSINT MCP Console\n"
+            "Multi-mode Intelligence Workstation."
+        )
 
+        # Holds the last OSINT result (dict) to inject into the copilot
         osint_state = gr.State({})
 
         # -------------------------
@@ -218,9 +263,11 @@ def build_interface():
                 run.click(
                     ui_lookup_ip,
                     [ip, enrich, mitre],
-                    [out_s, out_md, out_json, out_mitre, out_stix, osint_state]
+                    [out_s, out_md, out_json, out_mitre, out_stix, osint_state],
                 )
 
+            # You can add more tabs here: Domain Lookup, Hash Lookup, Correlate IOCs, Quickscan
+
         # -------------------------
         # MCP Bridge
         # -------------------------
@@ -235,14 +282,16 @@ def build_interface():
             btn.click(
                 ui_bridge,
                 [tool, args],
-                [out_bridge_json, out_bridge_md, osint_state]
+                [out_bridge_json, out_bridge_md, osint_state],
             )
 
         # -------------------------
         # Analyst Copilot
         # -------------------------
         with gr.Tab("Analyst Copilot"):
-            gr.Markdown("### WhiteRabbit Neo + Cybertron Threat Intelligence Assistant")
+            gr.Markdown(
+                "### WhiteRabbit Neo + Cybertron Threat Intelligence Assistant"
+            )
 
             system_prompt = gr.Textbox(
                 label="System Prompt",
@@ -259,18 +308,23 @@ def build_interface():
                     "berkeley-nest/WhiteRabbitNeo-8B",
                     "cybertronai/cybertron-1.1-1b",
                     "cybertronai/cybertron-1.1-7b",
-                    "cybertronai/cybertron-1.1-32b"
+                    "cybertronai/cybertron-1.1-32b",
                 ],
                 value="berkeley-nest/WhiteRabbitNeo-8B",
             )
 
-            gr.Markdown("### HuggingFace API Token (required for LLM inference)")
+            gr.Markdown(
+                "### HuggingFace API Token (required for LLM inference)"
+            )
             hf_token = gr.Textbox(
                 label="HF Token",
                 type="password",
                 placeholder="hf_xxx...",
             )
 
+            # Chat history state for the copilot (list of messages)
+            chat_state = gr.State([])
+
             chatbot = gr.ChatInterface(
                 respond,
                 type="messages",
@@ -278,17 +332,36 @@ def build_interface():
                     system_prompt,
                     model_select,
                     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"),
+                    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",
+                    ),
                 ],
+                state=chat_state,
             )
 
             inject_btn = gr.Button("Inject Last OSINT Result into Copilot")
             inject_btn.click(
                 inject_osint,
-                [chatbot._chatbot_state, osint_state],
-                [chatbot._chatbot_state],
+                inputs=[chat_state, osint_state],
+                outputs=[chat_state],
             )
 
     return demo