From bd42c115b490efa9a99e2123ca1bdf7fd8ed5b87 Mon Sep 17 00:00:00 2001
From: Sean Sube <seansube@gmail.com>
Date: Sun, 26 May 2024 22:27:13 -0500
Subject: [PATCH] add option to generate new rooms

---
 adventure/main.py          | 47 ++++++++++++++++++++++++++++----------
 adventure/models/config.py |  9 --------
 2 files changed, 35 insertions(+), 21 deletions(-)

diff --git a/adventure/main.py b/adventure/main.py
index 48524e8..ca20f55 100644
--- a/adventure/main.py
+++ b/adventure/main.py
@@ -7,7 +7,6 @@ from dotenv import load_dotenv
 from packit.agent import Agent, agent_easy_connect
 from packit.utils import logger_with_colors
 
-from adventure.context import get_system_data, set_system_data
 from adventure.utils.file import load_yaml
 
 # configure logging
@@ -29,9 +28,15 @@ logger = logger_with_colors(__name__)  # , level="DEBUG")
 load_dotenv(environ.get("ADVENTURE_ENV", ".env"), override=True)
 
 if True:
-    from adventure.context import set_current_step, set_dungeon_master, subscribe
+    from adventure.context import (
+        get_system_data,
+        set_current_step,
+        set_dungeon_master,
+        set_system_data,
+        subscribe,
+    )
     from adventure.game_system import GameSystem
-    from adventure.generate import generate_world
+    from adventure.generate import generate_room, generate_world, link_rooms
     from adventure.models.config import DEFAULT_CONFIG, Config
     from adventure.models.entity import World, WorldState
     from adventure.models.event import GenerateEvent
@@ -70,6 +75,12 @@ def parse_args():
         nargs="*",
         help="Extra actions to include in the simulation",
     )
+    parser.add_argument(
+        "--add-rooms",
+        default=0,
+        type=int,
+        help="The number of new rooms to generate before starting the simulation",
+    )
     parser.add_argument(
         "--config",
         type=str,
@@ -200,10 +211,21 @@ def load_or_generate_world(
 ):
     world_file = args.world + ".json"
     world_state_file = args.state or (args.world + ".state.json")
+    add_rooms = args.add_rooms
 
     memory = {}
     step = 0
 
+    # prepare an agent for the world builder
+    llm = agent_easy_connect()
+    world_builder = Agent(
+        "World Builder",
+        f"You are an experienced game master creating a visually detailed world for a new adventure. "
+        f"{world_prompt.flavor}. The theme is: {world_prompt.theme}.",
+        {},
+        llm,
+    )
+
     if path.exists(world_state_file):
         logger.info(f"loading world state from {world_state_file}")
         with open(world_state_file, "r") as f:
@@ -223,15 +245,6 @@ def load_or_generate_world(
         load_or_initialize_system_data(args, systems, world)
     else:
         logger.info(f"generating a new world using theme: {world_prompt.theme}")
-        llm = agent_easy_connect()
-        world_builder = Agent(
-            "World Builder",
-            f"You are an experienced game master creating a visually detailed world for a new adventure. "
-            f"{world_prompt.flavor}. The theme is: {world_prompt.theme}.",
-            {},
-            llm,
-        )
-
         world = generate_world(
             world_builder,
             args.world,
@@ -242,6 +255,16 @@ def load_or_generate_world(
         save_world(world, world_file)
         save_system_data(args, systems)
 
+    new_rooms = []
+    for i in range(add_rooms):
+        logger.info(f"generating room {i + 1} of {add_rooms}")
+        room = generate_room(world_builder, world, systems)
+        new_rooms.append(room)
+        world.rooms.append(room)
+
+    if new_rooms:
+        link_rooms(world_builder, world, systems, new_rooms)
+
     create_agents(world, memory=memory, players=players)
     return (world, world_state_file, step)
 
diff --git a/adventure/models/config.py b/adventure/models/config.py
index 83cdb8b..044d832 100644
--- a/adventure/models/config.py
+++ b/adventure/models/config.py
@@ -20,11 +20,6 @@ class BotConfig:
     discord: DiscordBotConfig
 
 
-@dataclass
-class PromptConfig:
-    prompts: Dict[str, str]
-
-
 @dataclass
 class RenderConfig:
     cfg: int | IntRange
@@ -79,7 +74,6 @@ class WorldConfig:
 @dataclass
 class Config:
     bot: BotConfig
-    prompt: PromptConfig
     render: RenderConfig
     server: ServerConfig
     world: WorldConfig
@@ -87,9 +81,6 @@ class Config:
 
 DEFAULT_CONFIG = Config(
     bot=BotConfig(discord=DiscordBotConfig(channels=["adventure"])),
-    prompt=PromptConfig(
-        prompts={},
-    ),
     render=RenderConfig(
         cfg=IntRange(min=5, max=8),
         checkpoints=[