File size: 6,497 Bytes
4bdab37
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
from prompt_toolkit import prompt
from prompt_toolkit.completion import WordCompleter
from prompt_toolkit.styles import Style
from rich.console import Console
from rich.text import Text
from rich.color import ANSI_COLOR_NAMES
import random

from ..arena import Arena, TooManyInvalidActions
from ..backends.human import HumanBackendError

ASCII_ART = r"""
_________  .__               __      _____                                   
\_   ___ \ |  |__  _____   _/  |_   /  _  \  _______   ____    ____  _____   
/    \  \/ |  |  \ \__  \  \   __\ /  /_\  \ \_  __ \W/ __ \  /    \ \__  \  
\     \____|   Y  \ / __ \_ |  |  /    |    \ |  | \/\  ___/ |   |  \ / __ \_
 \______  /|___|  /(____  / |__|  \____|__  / |__|    \___  >|___|  /(____  /
        \/      \/      \/                \/              \/      \/      \/ 
"""

visible_colors = [color for color in ANSI_COLOR_NAMES.keys() if
                  color not in ["black", "white", "red", "green"] and "grey" not in color]

MAX_STEPS = 5

import logging

# Set logging level to ERROR
logging.getLogger().setLevel(logging.ERROR)


class ArenaCLI:
    """
    The CLI user interface for ChatArena.
    """

    def __init__(self, arena: Arena):
        self.arena = arena

    def launch(self, max_steps: int = None, interactive: bool = True):
        """
        Run the CLI
        """
        if not interactive and max_steps is None:
            max_steps = MAX_STEPS

        console = Console()
        # Print ascii art
        console.print(ASCII_ART, style="bold dark_orange3")
        timestep = self.arena.reset()
        console.print("🏟 Chat Arena Initialized!", style="bold green")

        env = self.arena.environment
        players = self.arena.players

        env_desc = self.arena.global_prompt
        num_players = env.num_players
        player_colors = random.sample(visible_colors, num_players)  # sample different colors for players
        name_to_color = dict(zip(env.player_names, player_colors))
        # System and Moderator messages are printed in red
        name_to_color["System"] = "red"
        name_to_color["Moderator"] = "red"

        console.print(f"[bold green underline]Environment ({env.type_name}) description:[/]\n{env_desc}")

        # Print the player name, role_desc and backend_type
        for i, player in enumerate(players):
            player_name = Text(f"[{player.name} ({player.backend.type_name})] Role Description:")
            player_name.stylize(f"bold {name_to_color[player.name]} underline")
            console.print(player_name)
            console.print(player.role_desc)

        console.print("\n========= Arena Start! ==========\n", style="bold green")

        step = 0
        while not timestep.terminal:
            if interactive:
                command = prompt([('class:command', "command (n/r/q/s/h) > ")],
                                 style=Style.from_dict({'command': 'blue'}),
                                 completer=WordCompleter(
                                     ['next', 'n', 'reset', 'r', 'exit', 'quit', 'q', 'help', 'h', 'save', 's']))
                command = command.strip()

                if command == "help" or command == "h":
                    console.print("Available commands:")
                    console.print("    [bold]next or n or <Enter>[/]: next step")
                    console.print("    [bold]exit or quit or q[/]: exit the game")
                    console.print("    [bold]help or h[/]: print this message")
                    console.print("    [bold]reset or r[/]: reset the game")
                    console.print("    [bold]save or s[/]: save the history to file")
                    continue
                elif command == "exit" or command == "quit" or command == "q":
                    break
                elif command == "reset" or command == "r":
                    timestep = self.arena.reset()
                    console.print("\n========= Arena Reset! ==========\n", style="bold green")
                    continue
                elif command == "next" or command == "n" or command == "":
                    pass
                elif command == "save" or command == "s":
                    # Prompt to get the file path
                    file_path = prompt([('class:command', "save file path > ")],
                                       style=Style.from_dict({'command': 'blue'}))
                    file_path = file_path.strip()
                    # Save the history to file
                    self.arena.save_history(file_path)
                    # Print the save success message
                    console.print(f"History saved to {file_path}", style="bold green")
                else:
                    console.print(f"Invalid command: {command}", style="bold red")
                    continue

            try:
                timestep = self.arena.step()
            except HumanBackendError as e:
                # Handle human input and recover with the game update
                human_player_name = env.get_next_player()
                if interactive:
                    human_input = prompt(
                        [('class:user_prompt', f"Type your input for {human_player_name}: ")],
                        style=Style.from_dict({'user_prompt': 'ansicyan underline'})
                    )
                    # If not, the conversation does not stop
                    timestep = env.step(human_player_name, human_input)
                else:
                    raise e  # cannot recover from this error in non-interactive mode
            except TooManyInvalidActions as e:
                # Print the error message
                console.print(f"Too many invalid actions: {e}", style="bold red")
                break

            # The messages that are not yet logged
            messages = [msg for msg in env.get_observation() if not msg.logged]
            # Print the new messages
            for msg in messages:
                message_text = Text(f"[{msg.agent_name}->{msg.visible_to}]: {msg.content}")
                message_text.stylize(f"bold {name_to_color[msg.agent_name]}", 0,
                                     len(f"[{msg.agent_name}->{msg.visible_to}]:"))
                console.print(message_text)
                msg.logged = True

            step += 1
            if max_steps is not None and step >= max_steps:
                break

        console.print("\n========= Arena Ended! ==========\n", style="bold red")