156 lines
5.2 KiB
Python
156 lines
5.2 KiB
Python
from logging import getLogger
|
|
from random import choice
|
|
from typing import Any
|
|
|
|
from taleweave.context import (
|
|
get_agent_for_character,
|
|
get_dungeon_master,
|
|
get_game_systems,
|
|
get_system_data,
|
|
)
|
|
from taleweave.game_system import GameSystem, SystemData
|
|
from taleweave.generate import generate_room
|
|
from taleweave.models.base import dataclass
|
|
from taleweave.models.entity import Room, World
|
|
from taleweave.utils.search import find_character_in_room, find_room
|
|
from taleweave.utils.systems import load_system_data, save_system_data
|
|
from taleweave.utils.template import format_prompt
|
|
|
|
logger = getLogger(__name__)
|
|
|
|
ABDUCTION_INTERVAL = 5 # turns
|
|
|
|
# set up data classes for the data that will be saved and loaded for this system
|
|
|
|
|
|
@dataclass
|
|
class AlienAbduction:
|
|
abductee: str
|
|
source: str
|
|
|
|
|
|
@dataclass
|
|
class AlienData:
|
|
abduction: AlienAbduction | None = None
|
|
laboratory: str | None = None
|
|
|
|
|
|
# system functions
|
|
|
|
|
|
def initialize_aliens(world: World):
|
|
"""
|
|
Initialize the aliens system for the given world and create an alien laboratory room if one does not already exist.
|
|
"""
|
|
data = get_system_data("aliens") or AlienData()
|
|
|
|
if data.laboratory is None:
|
|
logger.info("generating alien laboratory")
|
|
dungeon_master = get_dungeon_master()
|
|
systems = get_game_systems() # TODO: other systems are not set up yet, this is always empty
|
|
laboratory = generate_room(
|
|
dungeon_master,
|
|
world,
|
|
systems,
|
|
additional_prompt="make sure this room is an alien laboratory to which characters could be abducted"
|
|
)
|
|
world.rooms.append(laboratory)
|
|
data.laboratory = laboratory.name
|
|
|
|
return data
|
|
|
|
|
|
def abduct_character(world: World, laboratory_room: Room) -> AlienAbduction | None:
|
|
"""
|
|
Abduct a random character from a random room and place them in the alien laboratory.
|
|
"""
|
|
logger.info("abducting a character")
|
|
valid_rooms = [room for room in world.rooms if room != laboratory_room and room.characters]
|
|
if not valid_rooms:
|
|
logger.info("no valid rooms to abduct from")
|
|
return None
|
|
|
|
source_room = choice(valid_rooms)
|
|
abductee = choice(source_room.characters)
|
|
|
|
logger.info(f"abducting {abductee.name}")
|
|
source_room.characters.remove(abductee)
|
|
laboratory_room.characters.append(abductee)
|
|
|
|
# give the abductee a memory of their abduction
|
|
abduction_memory = format_prompt("system_aliens_memory_abduction", abductee=abductee, laboratory_room=laboratory_room, source_room=source_room)
|
|
abducted_agent = get_agent_for_character(abductee)
|
|
if abducted_agent and abducted_agent.memory:
|
|
abducted_agent.memory.append(abduction_memory)
|
|
|
|
# TODO: give the abductee a chance to escape
|
|
return AlienAbduction(abductee=abductee.name, source=source_room.name)
|
|
|
|
|
|
def return_abductee(world: World, laboratory_room: Room, abduction: AlienAbduction) -> AlienAbduction | None:
|
|
"""
|
|
Return an abductee to their source room.
|
|
"""
|
|
logger.info(f"returning {abduction.abductee} to their room")
|
|
source_room = find_room(world, abduction.source)
|
|
if not source_room:
|
|
raise ValueError("Abductee's source room could not be found")
|
|
|
|
abductee = find_character_in_room(laboratory_room, abduction.abductee)
|
|
if not abductee:
|
|
logger.info(f"abductee {abduction.abductee} is not in the laboratory")
|
|
return
|
|
|
|
# move the abductee back to their source room
|
|
laboratory_room.characters.remove(abductee)
|
|
source_room.characters.append(abductee)
|
|
|
|
# give the abductee a memory of their return
|
|
return_memory = format_prompt("system_aliens_memory_return", abductee=abductee, laboratory_room=laboratory_room, source_room=source_room)
|
|
abducted_agent = get_agent_for_character(abductee)
|
|
if abducted_agent and abducted_agent.memory:
|
|
abducted_agent.memory.append(return_memory)
|
|
|
|
return None
|
|
|
|
|
|
def simulate_aliens(world: World, turn: int, data: AlienData | None = None):
|
|
"""
|
|
Simulate the alien abduction system. Every interval, abduct a character and place them in the alien laboratory or
|
|
return them to their room.
|
|
"""
|
|
data = data or get_system_data("aliens")
|
|
if not data:
|
|
raise ValueError("Alien data is required for simulation")
|
|
|
|
# every N turns, randomly abduct a character and place them in the alien laboratory or return them to their room
|
|
if data.laboratory and turn % ABDUCTION_INTERVAL == 0:
|
|
laboratory_room = find_room(world, data.laboratory)
|
|
if not laboratory_room:
|
|
raise ValueError("Alien laboratory room not found")
|
|
|
|
if data.abduction:
|
|
data.abduction = return_abductee(world, laboratory_room, data.abduction)
|
|
else:
|
|
data.abduction = abduct_character(world, laboratory_room)
|
|
|
|
|
|
def load_aliens_data(file: str) -> Any:
|
|
return load_system_data(AlienData, file)
|
|
|
|
|
|
def save_aliens_data(file: str, data: Any) -> None:
|
|
return save_system_data(AlienData, file, data)
|
|
|
|
|
|
def init():
|
|
return [GameSystem(
|
|
name="alien_abduction",
|
|
data=SystemData(
|
|
load=load_aliens_data,
|
|
save=save_aliens_data,
|
|
),
|
|
initialize=initialize_aliens,
|
|
simulate=simulate_aliens,
|
|
)]
|