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