def _get_init_data(self, packet, channel_info: ChannelInfo):
"""Get the initialization data for the assigned agent's task"""
task_runner = channel_info.job.task_runner
agent_id = packet.data["provider_data"]["agent_id"]
agent_info = self.agents[agent_id]
assert isinstance(
agent_info.agent, Agent
), f"Can only get init unit data for Agents, not OnboardingAgents, got {agent_info}"
unit_data = task_runner.get_init_data_for_agent(agent_info.agent)
agent_data_packet = Packet(
packet_type=PACKET_TYPE_INIT_DATA,
sender_id=SYSTEM_CHANNEL_ID,
receiver_id=channel_info.channel_id,
data={
"request_id": packet.data["request_id"],
"init_data": unit_data
},
)
self.message_queue.append(agent_data_packet)
if isinstance(unit_data,
dict) and unit_data.get("raw_messages") is not None:
# TODO bring these into constants somehow
for message in unit_data["raw_messages"]:
packet = Packet.from_dict(message)
packet.receiver_id = agent_id
agent_info.agent.pending_observations.append(packet)
def on_message(msg_json):
"""Incoming message handler defers to the internal handler"""
try:
packet_dict = json.loads(msg_json)
packet = Packet.from_dict(packet_dict)
self.on_message(self.channel_id, packet)
except Exception as e:
# TODO(CLEAN) properly handle only failed from_dict calls
logger.exception(repr(e), exc_info=True)
raise