use action context managers, sort modules

Makefile

@@ -56,4 +56,4 @@ lint-fix:
 style: lint-fix
 
 typecheck:
-	mypy feedme
+	mypy adventure

adventure/actions.py

@@ -3,8 +3,13 @@ from logging import getLogger
 
 from packit.utils import could_be_json
 
-from adventure.context import broadcast, get_actor_agent_for_name, get_current_context
-from adventure.search import (
+from adventure.context import (
+    action_context,
+    broadcast,
+    get_actor_agent_for_name,
+    world_context,
+)
+from adventure.utils.search import (
     find_actor_in_room,
     find_item_in_actor,
     find_item_in_room,
@@ -22,7 +27,8 @@ def action_look(target: str) -> str:
     Args:
         target: The name of the target to look at.
     """
-    _, action_room, action_actor = get_current_context()
+
+    with action_context() as (action_room, action_actor):
         broadcast(f"{action_actor.name} looks at {target}")
 
         if target.lower() == action_room.name.lower():
@@ -60,8 +66,8 @@ def action_move(direction: str) -> str:
     Args:
         direction: The direction to move in.
     """
-    action_world, action_room, action_actor = get_current_context()
 
+    with world_context() as (action_world, action_room, action_actor):
         destination_name = action_room.portals.get(direction.lower())
         if not destination_name:
             return f"You cannot move {direction} from here."
@@ -84,16 +90,15 @@ def action_take(item_name: str) -> str:
     Args:
         item_name: The name of the item to take.
     """
-    _, action_room, action_actor = get_current_context()
-
+    with action_context() as (action_room, action_actor):
         item = find_item_in_room(action_room, item_name)
-    if item:
+        if not item:
+            return "The {item_name} item is not in the room."
+
         broadcast(f"{action_actor.name} takes the {item_name} item")
         action_room.items.remove(item)
         action_actor.items.append(item)
         return "You take the {item_name} item and put it in your inventory."
-    else:
-        return "The {item_name} item is not in the room."
 
 
 def action_ask(character: str, question: str) -> str:
@@ -105,8 +110,7 @@ def action_ask(character: str, question: str) -> str:
         question: The question to ask them.
     """
     # capture references to the current actor and room, because they will be overwritten
-    _world, _room, action_actor = get_current_context()
-
+    with action_context() as (_, action_actor):
         # sanity checks
         question_actor, question_agent = get_actor_agent_for_name(character)
         if question_actor == action_actor:
@@ -143,8 +147,8 @@ def action_tell(character: str, message: str) -> str:
         message: The message to tell them.
     """
     # capture references to the current actor and room, because they will be overwritten
-    _world, _room, action_actor = get_current_context()
 
+    with action_context() as (_, action_actor):
         # sanity checks
         question_actor, question_agent = get_actor_agent_for_name(character)
         if question_actor == action_actor:
@@ -180,8 +184,7 @@ def action_give(character: str, item_name: str) -> str:
         character: The name of the character to give the item to.
         item_name: The name of the item to give.
     """
-    _, action_room, action_actor = get_current_context()
-
+    with action_context() as (action_room, action_actor):
         destination_actor = find_actor_in_room(action_room, character)
         if not destination_actor:
             return f"The {character} character is not in the room."
@@ -205,8 +208,7 @@ def action_drop(item_name: str) -> str:
         item_name: The name of the item to drop.
     """
 
-    _, action_room, action_actor = get_current_context()
-
+    with action_context() as (action_room, action_actor):
         item = find_item_in_actor(action_actor, item_name)
         if not item:
             return f"You do not have the {item_name} item in your inventory."

adventure/bot/discord.py

@@ -33,14 +33,14 @@ from adventure.player import (
     remove_player,
     set_player,
 )
-from adventure.render_comfy import render_event
+from adventure.render.comfy import render_event
 
 logger = getLogger(__name__)
 client = None
 bot_config: DiscordBotConfig = DEFAULT_CONFIG.bot.discord
 
 active_tasks = set()
-event_messages: Dict[str, str | GameEvent] = {}
+event_messages: Dict[int, str | GameEvent] = {}
 event_queue: Queue[GameEvent] = Queue()
 
 
@@ -81,13 +81,13 @@ class AdventureClient(Client):
         channel = message.channel
         user_name = author.name  # include nick
 
+        if message.content.startswith("!adventure"):
             world = get_current_world()
             if world:
                 active_world = f"Active world: {world.name} (theme: {world.theme})"
             else:
                 active_world = "No active world"
 
-        if message.content.startswith("!adventure"):
             await message.channel.send(f"Hello! Welcome to Adventure! {active_world}")
             return
 
@@ -215,7 +215,14 @@ def stop_bot():
     global client
 
     if client:
-        client.close()
+        close_task = client.loop.create_task(client.close())
+        active_tasks.add(close_task)
+
+        def on_close_task_done(future):
+            logger.info("discord client closed")
+            active_tasks.discard(future)
+
+        close_task.add_done_callback(on_close_task_done)
         client = None
 
 
@@ -299,7 +306,7 @@ async def broadcast_event(message: str | GameEvent):
         event_messages[event_message.id] = message
 
 
-def embed_from_event(event: GameEvent) -> Embed:
+def embed_from_event(event: GameEvent) -> Embed | None:
     if isinstance(event, GenerateEvent):
         return embed_from_generate(event)
     elif isinstance(event, ResultEvent):

adventure/context.py

@@ -37,14 +37,21 @@ game_systems: List[GameSystem] = []
 # TODO: where should this one go?
 actor_agents: Dict[str, Tuple[Actor, Agent]] = {}
 
+STRING_EVENT_TYPE = "message"
+
+
+def get_event_name(event: GameEvent | Type[GameEvent]):
+    return f"event:{event.type}"
+
 
 def broadcast(message: str | GameEvent):
     if isinstance(message, GameEvent):
-        logger.debug(f"broadcasting {message.type}")
-        event_emitter.emit(message.type, message)
+        event_name = get_event_name(message)
+        logger.debug(f"broadcasting {event_name}")
+        event_emitter.emit(event_name, message)
     else:
         logger.warning("broadcasting a string message is deprecated")
-        event_emitter.emit("message", message)
+        event_emitter.emit(STRING_EVENT_TYPE, message)
 
 
 def is_union(type_: Type | UnionType):
@@ -62,10 +69,13 @@ def subscribe(
 
         return
 
+    if event_type is str:
+        event_name = STRING_EVENT_TYPE
+    else:
+        event_name = get_event_name(event_type)
+
     logger.debug(f"subscribing {callback.__name__} to {event_type}")
-    event_emitter.on(
-        event_type.type, callback
-    )  # TODO: should this use str or __name__?
+    event_emitter.on(event_name, callback)
 
 
 def has_dungeon_master():
@@ -74,11 +84,17 @@ def has_dungeon_master():
 
 # region context manager
 @contextmanager
-def with_action_context():
+def action_context():
     room, actor = get_action_context()
     yield room, actor
 
 
+@contextmanager
+def world_context():
+    world, room, actor = get_world_context()
+    yield world, room, actor
+
+
 # endregion
 
 
@@ -94,7 +110,7 @@ def get_action_context() -> Tuple[Room, Actor]:
     return (current_room, current_actor)
 
 
-def get_current_context() -> Tuple[World, Room, Actor]:
+def get_world_context() -> Tuple[World, Room, Actor]:
     if not current_world:
         raise ValueError(
             "The current world must be set before calling action functions"

adventure/main.py

@@ -246,19 +246,19 @@ def main():
     threads = []
 
     if args.render:
-        from adventure.render_comfy import launch_render, render_generated
+        from adventure.render.comfy import launch_render, render_generated
 
         threads.extend(launch_render(config.render))
         if args.render_generated:
             subscribe(GenerateEvent, render_generated)
 
     if args.discord:
-        from adventure.bot_discord import launch_bot
+        from adventure.bot.discord import launch_bot
 
         threads.extend(launch_bot(config.bot.discord))
 
     if args.server:
-        from adventure.server_socket import launch_server, server_system
+        from adventure.server.websocket import launch_server
 
         threads.extend(launch_server(config.server.websocket))
 
@@ -300,6 +300,8 @@ def main():
 
     # make sure the server system runs after any updates
     if args.server:
+        from adventure.server.websocket import server_system
+
         extra_systems.append(GameSystem(simulate=server_system))
 
     # load or generate the world

adventure/optional_actions.py

@@ -4,16 +4,17 @@ from typing import Callable, List
 from packit.agent import Agent, agent_easy_connect
 
 from adventure.context import (
+    action_context,
     broadcast,
     get_agent_for_actor,
-    get_current_context,
     get_dungeon_master,
     has_dungeon_master,
     set_dungeon_master,
+    world_context,
 )
 from adventure.generate import OPPOSITE_DIRECTIONS, generate_item, generate_room
-from adventure.search import find_actor_in_room
 from adventure.utils.effect import apply_effect
+from adventure.utils.search import find_actor_in_room
 from adventure.utils.world import describe_actor, describe_entity
 
 logger = getLogger(__name__)
@@ -39,29 +40,26 @@ def action_explore(direction: str) -> str:
         direction: The direction to explore: north, south, east, or west.
     """
 
-    current_world, current_room, current_actor = get_current_context()
+    with world_context() as (action_world, action_room, action_actor):
         dungeon_master = get_dungeon_master()
 
-    if not current_world:
-        raise ValueError("No world found")
-
-    if direction in current_room.portals:
-        dest_room = current_room.portals[direction]
+        if direction in action_room.portals:
+            dest_room = action_room.portals[direction]
             return f"You cannot explore {direction} from here, that direction already leads to {dest_room}. Please use the move action to go there."
 
-    existing_rooms = [room.name for room in current_world.rooms]
+        existing_rooms = [room.name for room in action_world.rooms]
         try:
             new_room = generate_room(
-            dungeon_master, current_world.theme, existing_rooms=existing_rooms
+                dungeon_master, action_world.theme, existing_rooms=existing_rooms
             )
-        current_world.rooms.append(new_room)
+            action_world.rooms.append(new_room)
 
             # link the rooms together
-        current_room.portals[direction] = new_room.name
-        new_room.portals[OPPOSITE_DIRECTIONS[direction]] = current_room.name
+            action_room.portals[direction] = new_room.name
+            new_room.portals[OPPOSITE_DIRECTIONS[direction]] = action_room.name
 
             broadcast(
-            f"{current_actor.name} explores {direction} of {current_room.name} and finds a new room: {new_room.name}"
+                f"{action_actor.name} explores {direction} of {action_room.name} and finds a new room: {new_room.name}"
             )
             return f"You explore {direction} and find a new room: {new_room.name}"
         except Exception:
@@ -74,11 +72,13 @@ def action_search(unused: bool) -> str:
     Search the room for hidden items.
     """
 
-    action_world, action_room, action_actor = get_current_context()
+    with world_context() as (action_world, action_room, action_actor):
         dungeon_master = get_dungeon_master()
 
         if len(action_room.items) > 2:
-        return "You find nothing hidden in the room. There is no room for more items."
+            return (
+                "You find nothing hidden in the room. There is no room for more items."
+            )
 
         existing_items = [item.name for item in action_room.items]
 
@@ -108,7 +108,7 @@ def action_use(item: str, target: str) -> str:
         item: The name of the item to use.
         target: The name of the character to use the item on, or "self" to use the item on yourself.
     """
-    _, action_room, action_actor = get_current_context()
+    with action_context() as (action_room, action_actor):
         dungeon_master = get_dungeon_master()
 
         action_item = next(
@@ -167,7 +167,7 @@ def action_use(item: str, target: str) -> str:
 
         # make sure both agents remember the outcome
         target_agent = get_agent_for_actor(target_actor)
-    if target_agent:
+        if target_agent and target_agent.memory:
             target_agent.memory.append(outcome)
 
         return outcome

adventure/player.py

@@ -8,7 +8,7 @@ from langchain_core.messages import AIMessage, BaseMessage, HumanMessage
 from packit.agent import Agent
 from packit.utils import could_be_json
 
-from adventure.context import get_current_context
+from adventure.context import action_context
 from adventure.models.event import PromptEvent
 
 logger = getLogger(__name__)
@@ -183,7 +183,7 @@ class RemotePlayer(BasePlayer):
         formatted_prompt = prompt.format(**kwargs)
         self.memory.append(HumanMessage(content=formatted_prompt))
 
-        _, current_room, current_actor = get_current_context()
+        with action_context() as (current_room, current_actor):
             prompt_event = PromptEvent(
                 prompt=formatted_prompt, room=current_room, actor=current_actor
             )

adventure/rpg_systems/crafting_actions.py

@@ -1,6 +1,6 @@
 from random import randint
 
-from adventure.context import broadcast, get_current_context, get_dungeon_master
+from adventure.context import broadcast, get_dungeon_master, world_context
 from adventure.generate import generate_item
 from adventure.models.base import dataclass
 from adventure.models.entity import Item
@@ -26,8 +26,7 @@ def action_craft(item_name: str) -> str:
     Args:
         item_name: The name of the item to craft.
     """
-    action_world, _, action_actor = get_current_context()
-
+    with world_context() as (action_world, _, action_actor):
         if item_name not in recipes:
             return f"There is no recipe to craft a {item_name}."
 
@@ -42,9 +41,13 @@ def action_craft(item_name: str) -> str:
         inventory_items = {item.name for item in action_actor.items}
 
         # Check for sufficient ingredients
-    missing_items = [item for item in recipe.ingredients if item not in inventory_items]
+        missing_items = [
+            item for item in recipe.ingredients if item not in inventory_items
+        ]
         if missing_items:
-        return f"You are missing {' and '.join(missing_items)} to craft {item_name}."
+            return (
+                f"You are missing {' and '.join(missing_items)} to craft {item_name}."
+            )
 
         # Deduct the ingredients from inventory
         for ingredient in recipe.ingredients:

adventure/rpg_systems/language_actions.py

@@ -1,5 +1,5 @@
-from adventure.context import broadcast, get_current_context
-from adventure.search import find_item_in_actor
+from adventure.context import action_context, broadcast
+from adventure.utils.search import find_item_in_actor
 
 
 def action_read(item_name: str) -> str:
@@ -9,8 +9,7 @@ def action_read(item_name: str) -> str:
     Args:
         item_name: The name of the item to read.
     """
-    _, _, action_actor = get_current_context()
-
+    with action_context() as (_, action_actor):
         item = find_item_in_actor(action_actor, item_name)
         if not item:
             return f"You do not have a {item_name} to read."
@@ -18,5 +17,5 @@ def action_read(item_name: str) -> str:
         if "text" in item.attributes:
             broadcast(f"{action_actor.name} reads {item_name}")
             return str(item.attributes["text"])
-    else:
+
         return f"The {item_name} has nothing to read."

adventure/rpg_systems/magic_actions.py

@@ -1,7 +1,7 @@
 from random import randint
 
-from adventure.context import broadcast, get_current_context, get_dungeon_master
-from adventure.search import find_actor_in_room
+from adventure.context import action_context, broadcast, get_dungeon_master
+from adventure.utils.search import find_actor_in_room
 
 
 def action_cast(spell: str, target: str) -> str:
@@ -12,8 +12,7 @@ def action_cast(spell: str, target: str) -> str:
         spell: The name of the spell to cast.
         target: The target of the spell.
     """
-    _, action_room, action_actor = get_current_context()
-
+    with action_context() as (action_room, action_actor):
         target_actor = find_actor_in_room(action_room, target)
         dungeon_master = get_dungeon_master()
 

adventure/rpg_systems/movement_actions.py

@@ -1,7 +1,7 @@
 from random import randint
 
-from adventure.context import broadcast, get_current_context, get_dungeon_master
-from adventure.search import find_item_in_room
+from adventure.context import action_context, broadcast, get_dungeon_master
+from adventure.utils.search import find_item_in_room
 
 
 def action_climb(target: str) -> str:
@@ -11,8 +11,7 @@ def action_climb(target: str) -> str:
     Args:
         target: The object or feature to climb.
     """
-    _, action_room, action_actor = get_current_context()
-
+    with action_context() as (action_room, action_actor):
         dungeon_master = get_dungeon_master()
         # Assume 'climbable' is an attribute that marks climbable targets
         climbable_feature = find_item_in_room(action_room, target)
@@ -31,7 +30,9 @@ def action_climb(target: str) -> str:
                 )
                 return f"You successfully climb the {target}."
             else:
-            broadcast(f"{action_actor.name} fails to climb the {target}. {flavor_text}")
+                broadcast(
+                    f"{action_actor.name} fails to climb the {target}. {flavor_text}"
+                )
                 return f"You fail to climb the {target}."
         else:
             return f"The {target} is not climbable."

adventure/server/websocket.py

@@ -36,9 +36,9 @@ from adventure.player import (
     remove_player,
     set_player,
 )
-from adventure.render_comfy import render_entity, render_event
-from adventure.search import find_actor, find_item, find_room
+from adventure.render.comfy import render_entity, render_event
 from adventure.state import snapshot_world, world_json
+from adventure.utils.search import find_actor, find_item, find_room
 
 logger = getLogger(__name__)
 
@@ -233,7 +233,7 @@ def render_input(data):
     elif "item" in data:
         item_name = data["item"]
         item = find_item(
-            world, item_name, include_actor_inventory=True, include_room_inventory=True
+            world, item_name, include_actor_inventory=True, include_item_inventory=True
         )
         if item:
             render_entity(item)

adventure/sim_systems/combat_actions.py

@@ -1,10 +1,10 @@
 from adventure.context import (
+    action_context,
     broadcast,
     get_agent_for_actor,
-    get_current_context,
     get_dungeon_master,
 )
-from adventure.search import find_actor_in_room, find_item_in_room
+from adventure.utils.search import find_actor_in_room, find_item_in_room
 from adventure.utils.world import describe_entity
 
 
@@ -16,8 +16,7 @@ def action_attack(target: str) -> str:
       target: The name of the character or item to attack.
     """
 
-    _, action_room, action_actor = get_current_context()
-
+    with action_context() as (action_room, action_actor):
         # make sure the target is in the room
         target_actor = find_actor_in_room(action_room, target)
         target_item = find_item_in_room(action_room, target)
@@ -45,7 +44,8 @@ def action_attack(target: str) -> str:
             )
             broadcast(description)
             return description
-    elif target_item:
+
+        if target_item:
             outcome = dungeon_master(
                 f"{action_actor.name} attacks {target} in the {action_room.name}. {describe_entity(action_room)}."
                 f"{describe_entity(action_actor)}. {describe_entity(target_item)}."
@@ -55,7 +55,7 @@ def action_attack(target: str) -> str:
             description = f"{action_actor.name} attacks the {target} in the {action_room.name}. {outcome}"
             broadcast(description)
             return description
-    else:
+
         return f"{target} is not in the {action_room.name}."
 
 
@@ -68,8 +68,7 @@ def action_cast(target: str, spell: str) -> str:
       spell: The name of the spell to cast.
     """
 
-    _, action_room, action_actor = get_current_context()
-
+    with action_context() as (action_room, action_actor):
         # make sure the target is in the room
         target_actor = find_actor_in_room(action_room, target)
         target_item = find_item_in_room(action_room, target)

adventure/sim_systems/hunger_actions.py

@@ -1,5 +1,5 @@
-from adventure.context import get_current_context
-from adventure.search import find_item_in_actor
+from adventure.context import action_context
+from adventure.utils.search import find_item_in_actor
 
 
 def action_cook(item: str) -> str:
@@ -9,8 +9,7 @@ def action_cook(item: str) -> str:
     Args:
       item: The name of the item to cook.
     """
-    _, _, action_actor = get_current_context()
-
+    with action_context() as (_, action_actor):
         target_item = find_item_in_actor(action_actor, item)
         if target_item is None:
             return "You don't have the item to cook."
@@ -37,8 +36,7 @@ def action_eat(item: str) -> str:
     Args:
       item: The name of the item to eat.
     """
-    _, _, action_actor = get_current_context()
-
+    with action_context() as (_, action_actor):
         target_item = find_item_in_actor(action_actor, item)
         if target_item is None:
             return "You don't have the item to eat."

adventure/sim_systems/hygiene_actions.py

@@ -1,4 +1,4 @@
-from adventure.context import get_current_context, get_dungeon_master
+from adventure.context import action_context, get_dungeon_master
 from adventure.utils.world import describe_entity
 
 
@@ -7,7 +7,7 @@ def action_wash(unused: bool) -> str:
     Wash yourself.
     """
 
-    _, action_room, action_actor = get_current_context()
+    with action_context() as (action_room, action_actor):
         hygiene = action_actor.attributes.get("hygiene", "clean")
 
         dungeon_master = get_dungeon_master()

adventure/sim_systems/sleeping_actions.py

@@ -1,4 +1,4 @@
-from adventure.context import get_current_context, get_dungeon_master
+from adventure.context import action_context, get_dungeon_master
 from adventure.utils.world import describe_entity
 
 
@@ -7,8 +7,7 @@ def action_sleep(unused: bool) -> str:
     Sleep until you are rested.
     """
 
-    _, action_room, action_actor = get_current_context()
-
+    with action_context() as (action_room, action_actor):
         dungeon_master = get_dungeon_master()
         outcome = dungeon_master(
             f"{action_actor.name} sleeps in the {action_room.name}. {describe_entity(action_room)}. {describe_entity(action_actor)}"

adventure/simulate.py

@@ -146,6 +146,7 @@ def simulate_world(
             )
 
             logger.debug(f"{actor.name} step result: {result}")
+            if agent.memory:
                 agent.memory.append(result)
 
             result_event = ResultEvent(result=result, room=room, actor=actor)

adventure/state.py

@@ -54,7 +54,7 @@ def snapshot_world(world: World, step: int):
     json_memory = {}
 
     for actor, agent in get_all_actor_agents():
-        json_memory[actor.name] = list(agent.memory)
+        json_memory[actor.name] = list(agent.memory or [])
 
     return {
         "world": json_world,