taleweave-ai/adventure/simulate.py

from functools import partial
from itertools import count
from logging import getLogger
from math import inf
from typing import Callable, Sequence

from packit.agent import Agent
from packit.conditions import condition_or, condition_threshold
from packit.loops import loop_reduce, loop_retry
from packit.results import multi_function_or_str_result
from packit.toolbox import Toolbox
from packit.utils import could_be_json

from adventure.actions.base import (
    action_ask,
    action_give,
    action_look,
    action_move,
    action_take,
    action_tell,
)
from adventure.actions.planning import (
    check_calendar,
    erase_notes,
    get_recent_notes,
    read_notes,
    replace_note,
    schedule_event,
    take_note,
)
from adventure.context import (
    broadcast,
    get_actor_agent_for_name,
    get_actor_for_agent,
    get_current_step,
    get_current_world,
    set_current_actor,
    set_current_room,
    set_current_step,
    set_current_world,
    set_game_systems,
)
from adventure.game_system import GameSystem
from adventure.models.entity import Actor, Room, World
from adventure.models.event import ActionEvent, ReplyEvent, ResultEvent
from adventure.utils.conversation import make_keyword_condition, summarize_room
from adventure.utils.effect import expire_effects
from adventure.utils.planning import expire_events, get_upcoming_events
from adventure.utils.search import find_room_with_actor
from adventure.utils.world import describe_entity, format_attributes

logger = getLogger(__name__)


def world_result_parser(value, agent, **kwargs):
    current_world = get_current_world()
    if not current_world:
        raise ValueError(
            "The current world must be set before calling world_result_parser"
        )

    logger.debug(f"parsing action for {agent.name}: {value}")

    current_actor = get_actor_for_agent(agent)
    current_room = next(
        (room for room in current_world.rooms if current_actor in room.actors), None
    )

    set_current_room(current_room)
    set_current_actor(current_actor)

    return multi_function_or_str_result(value, agent=agent, **kwargs)


def prompt_actor_action(
    room, actor, agent, action_names, action_toolbox, current_turn
) -> str:
    # collect data for the prompt
    notes_prompt, events_prompt = get_notes_events(actor, current_turn)

    room_actors = [actor.name for actor in room.actors]
    room_items = [item.name for item in room.items]
    room_directions = [portal.name for portal in room.portals]

    actor_attributes = format_attributes(actor)
    # actor_effects = [effect.name for effect in actor.active_effects]
    actor_items = [item.name for item in actor.items]

    # set up a result parser for the agent
    def result_parser(value, agent, **kwargs):
        if not room or not actor:
            raise ValueError("Room and actor must be set before parsing results")

        if could_be_json(value):
            event = ActionEvent.from_json(value, room, actor)
        else:
            event = ReplyEvent.from_text(value, room, actor)

        broadcast(event)

        return world_result_parser(value, agent, **kwargs)

    # prompt and act
    logger.info("starting turn for actor: %s", actor.name)
    result = loop_retry(
        agent,
        (
            "You are currently in the {room_name} room. {room_description}. {attributes}. "
            "The room contains the following characters: {visible_actors}. "
            "The room contains the following items: {visible_items}. "
            "Your inventory contains the following items: {actor_items}."
            "You can take the following actions: {actions}. "
            "You can move in the following directions: {directions}. "
            "{notes_prompt} {events_prompt}"
            "What will you do next? Reply with a JSON function call, calling one of the actions."
            "You can only perform one action per turn. What is your next action?"
        ),
        context={
            "actions": action_names,
            "actor_items": actor_items,
            "attributes": actor_attributes,
            "directions": room_directions,
            "room_name": room.name,
            "room_description": describe_entity(room),
            "visible_actors": room_actors,
            "visible_items": room_items,
            "notes_prompt": notes_prompt,
            "events_prompt": events_prompt,
        },
        result_parser=result_parser,
        toolbox=action_toolbox,
    )

    logger.debug(f"{actor.name} step result: {result}")
    if agent.memory:
        # TODO: make sure this is not duplicating memories and wasting space
        agent.memory.append(result)

    return result


def get_notes_events(actor: Actor, current_turn: int):
    recent_notes = get_recent_notes(actor)
    upcoming_events = get_upcoming_events(actor, current_turn)

    if len(recent_notes) > 0:
        notes = "\n".join(recent_notes)
        notes_prompt = f"Your recent notes are: {notes}\n"
    else:
        notes_prompt = "You have no recent notes.\n"

    if len(upcoming_events) > 0:
        current_step = get_current_step()
        events = [
            f"{event.name} in {event.turn - current_step} turns"
            for event in upcoming_events
        ]
        events = "\n".join(events)
        events_prompt = f"Upcoming events are: {events}\n"
    else:
        events_prompt = "You have no upcoming events.\n"

    return notes_prompt, events_prompt


def prompt_actor_think(
    room: Room, actor: Actor, agent: Agent, planner_toolbox: Toolbox, current_turn: int
) -> str:
    notes_prompt, events_prompt = get_notes_events(actor, current_turn)

    event_count = len(actor.planner.calendar.events)
    note_count = len(actor.planner.notes)

    logger.info("starting planning for actor: %s", actor.name)
    _, condition_end, result_parser = make_keyword_condition("You are done planning.")
    stop_condition = condition_or(condition_end, partial(condition_threshold, max=3))

    result = loop_reduce(
        agent,
        "You are about to start your turn. Plan your next action carefully. Take notes and schedule events to help keep track of your goals. "
        "You can check your notes for important facts or check your calendar for upcoming events. You have {note_count} notes. "
        "If you have plans with other characters, schedule them on your calendar. You have {event_count} events on your calendar. "
        "{room_summary}"
        "Think about your goals and any quests that you are working on, and plan your next action accordingly. "
        "Try to keep your notes accurate and up-to-date. Replace or erase old notes when they are no longer accurate or useful. "
        "Do not keeps notes about upcoming events, use your calendar for that. "
        "You can perform up to 3 planning actions in a single turn. When you are done planning, reply with 'END'."
        "{notes_prompt} {events_prompt}",
        context={
            "event_count": event_count,
            "events_prompt": events_prompt,
            "note_count": note_count,
            "notes_prompt": notes_prompt,
            "room_summary": summarize_room(room, actor),
        },
        result_parser=result_parser,
        stop_condition=stop_condition,
        toolbox=planner_toolbox,
    )

    if agent.memory:
        agent.memory.append(result)

    return result


def simulate_world(
    world: World,
    steps: float | int = inf,
    actions: Sequence[Callable[..., str]] = [],
    systems: Sequence[GameSystem] = [],
):
    logger.info("simulating the world")
    set_current_world(world)
    set_game_systems(systems)

    # build a toolbox for the actions
    action_tools = Toolbox(
        [
            action_ask,
            action_give,
            action_look,
            action_move,
            action_take,
            action_tell,
            *actions,
        ]
    )
    action_names = action_tools.list_tools()

    # build a toolbox for the planners
    planner_toolbox = Toolbox(
        [
            take_note,
            read_notes,
            replace_note,
            erase_notes,
            schedule_event,
            check_calendar,
        ]
    )

    # simulate each actor
    for i in count():
        current_step = get_current_step()
        logger.info(f"simulating step {i} of {steps} (world step {current_step})")

        for actor_name in world.order:
            actor, agent = get_actor_agent_for_name(actor_name)
            if not agent or not actor:
                logger.error(f"agent or actor not found for name {actor_name}")
                continue

            room = find_room_with_actor(world, actor)
            if not room:
                logger.error(f"actor {actor_name} is not in a room")
                continue

            # prep context
            set_current_room(room)
            set_current_actor(actor)

            # decrement effects on the actor and remove any that have expired
            expire_effects(actor)
            expire_events(actor, current_step)

            # give the actor a chance to think and check their planner
            if agent.memory and len(agent.memory) > 0:
                try:
                    thoughts = prompt_actor_think(
                        room, actor, agent, planner_toolbox, current_step
                    )
                    logger.debug(f"{actor.name} thinks: {thoughts}")
                except Exception:
                    logger.exception(f"error during planning for actor {actor.name}")

            result = prompt_actor_action(
                room, actor, agent, action_names, action_tools, current_step
            )
            result_event = ResultEvent(result=result, room=room, actor=actor)
            broadcast(result_event)

        for system in systems:
            if system.simulate:
                system.simulate(world, current_step)

        set_current_step(current_step + 1)
        if i >= steps:
            logger.info("reached step limit at world step %s", current_step + 1)
            break