From fc67fc48337160f5e8dc291982161f004f475d94 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Fri, 15 Mar 2024 14:37:28 +0000 Subject: [PATCH] Fix notebook parsing data --- .../Data-Manipulation-Customising-Red-Agent.ipynb | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/primaite/notebooks/Data-Manipulation-Customising-Red-Agent.ipynb b/src/primaite/notebooks/Data-Manipulation-Customising-Red-Agent.ipynb index 779d89f6..56e9bf5a 100644 --- a/src/primaite/notebooks/Data-Manipulation-Customising-Red-Agent.ipynb +++ b/src/primaite/notebooks/Data-Manipulation-Customising-Red-Agent.ipynb @@ -22,6 +22,7 @@ "# Imports\n", "\n", "from primaite.config.load import data_manipulation_config_path\n", + "from primaite.game.agent.interface import AgentActionHistoryItem\n", "from primaite.session.environment import PrimaiteGymEnv\n", "import yaml\n", "from pprint import pprint" @@ -62,12 +63,12 @@ "source": [ "def friendly_output_red_action(info):\n", " # parse the info dict form step output and write out what the red agent is doing\n", - " red_info = info['agent_actions']['data_manipulation_attacker']\n", - " red_action = red_info[0]\n", + " red_info : AgentActionHistoryItem = info['agent_actions']['data_manipulation_attacker']\n", + " red_action = red_info.action\n", " if red_action == 'DONOTHING':\n", " red_str = 'DO NOTHING'\n", " elif red_action == 'NODE_APPLICATION_EXECUTE':\n", - " client = \"client 1\" if red_info[1]['node_id'] == 0 else \"client 2\"\n", + " client = \"client 1\" if red_info.parameters['node_id'] == 0 else \"client 2\"\n", " red_str = f\"ATTACK from {client}\"\n", " return red_str" ]