File size: 12,770 Bytes
f1f9b0c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
import os
import gradio as gr # type: ignore
# Functions for creating the chat interface
from functools import cache
from typing import Literal
import json
from collections import defaultdict
from utils import Environment, Agent, get_context_prompt, dialogue_history_prompt
from sotopia_pi_generate import prepare_model, generate_action
from sotopia_space.constants import MODEL_OPTIONS

DEPLOYED = os.getenv("DEPLOYED", "true").lower() == "true"
DEFAULT_MODEL_SELECTION = "gpt-3.5-turbo"
TEMPERATURE = 0.7
TOP_P = 1
MAX_TOKENS = 1024

ENVIRONMENT_PROFILES = "profiles/environment_profiles.jsonl"
AGENT_PROFILES = "profiles/agent_profiles.jsonl"
RELATIONSHIP_PROFILES = "profiles/relationship_profiles.jsonl"

Action = Literal['none', 'action', 'non-verbal communication', 'speak', 'leave']
ACTION_TYPES: list[Action] = ['none', 'action', 'non-verbal communication', 'speak', 'leave']



@cache
def get_sotopia_profiles(env_file=ENVIRONMENT_PROFILES, agent_file=AGENT_PROFILES, relationship_file=RELATIONSHIP_PROFILES):
    with open(env_file, 'r') as f:
        data = [json.loads(line) for line in f.readlines()]
    
    code_names_count = defaultdict(int)
    environments = []
    environment_dict = {}
    for profile in sorted(data, key=lambda x: x['codename']):
        env_obj = Environment(profile)
        if profile['codename'] in code_names_count:
            environments.append((
                "{}_{:05d}".format(profile['codename'], 
                                   code_names_count[profile['codename']]
                                   ), 
                env_obj._id
                ))
        else:
            environments.append((profile['codename'], env_obj._id))
        environment_dict[env_obj._id] = env_obj
        code_names_count[profile['codename']] += 1
    
    with open(agent_file, 'r') as f:
        data = [json.loads(line) for line in f.readlines()]
    
    agent_dict = {}
    for profile in data:
        agent_obj = Agent(profile)
        agent_dict[agent_obj._id] = agent_obj
        
    with open(relationship_file, 'r') as f:
        data = [json.loads(line) for line in f.readlines()]
    
    relationship_dict = defaultdict(lambda : defaultdict(list))
    for profile in data:
        relationship_dict[profile['relationship']][profile['agent1_id']].append(profile['agent2_id'])
        relationship_dict[profile['relationship']][profile['agent2_id']].append(profile['agent1_id'])
    
    return environments, environment_dict, agent_dict, relationship_dict

def chat_introduction():
    with gr.Column(scale=2):
        gr.Image(
            "images/sotopia.jpg", elem_id="banner-image", show_label=False
        )
    with gr.Column(scale=5):
        gr.Markdown(
            """# Sotopia Space
            **Chat with different social agent models including [sotopia-pi](https://github.com/sotopia-lab/sotopia-pi), GPT and so on in sotopia space!**

            ➡️️ **Intended Use**: Sotopia space is intended to showcase the social intelligence ability of different social agents in interesting social scenarios. 
            
            ✨ **Guidance**: 
            
            Step (1) Select a social scenario that interests you in "Scenario Selection"
            
            Step (2) Select a social agent you want to chat with in "Model Selection"

            Step (3) Select which character you and your social agent will play in the scenario in "User Agent Selection" and "Bot Agent Selection" 
            
            Step (4) Negotiate/debate/cooperate with the social agent to see whether your goal or their social goal can be achieved.

            ⚠️ **Limitations**: The social agent can and will produce factually incorrect information, hallucinating facts and potentially offensive actions. It can produce problematic outputs, especially if prompted to do so.

            🗄️ **Disclaimer**: User prompts and generated replies from the model may be collected solely for the purpose of pure academic research. By using this demo, users implicitly agree to these terms.
            """
        )
    # with gr.Column(scale=1):
    #     toggle_dark = gr.Button(value="Toggle Dark")

def create_user_agent_dropdown(environment_id):
    _, environment_dict, agent_dict, relationship_dict = get_sotopia_profiles()
    environment = environment_dict[environment_id]
    
    user_agents_list = []
    unique_agent_ids = set()
    for x, _ in relationship_dict[environment.relationship].items():
        unique_agent_ids.add(x)
    
    for agent_id in unique_agent_ids:
        user_agents_list.append((agent_dict[agent_id].name, agent_id))
    return gr.Dropdown(choices=user_agents_list, value=user_agents_list[0][1] if user_agents_list else None, label="User Agent Selection")

def create_bot_agent_dropdown(environment_id, user_agent_id):
    _, environment_dict, agent_dict, relationship_dict = get_sotopia_profiles()
    environment, user_agent = environment_dict[environment_id], agent_dict[user_agent_id]
    
    bot_agent_list = []
    for neighbor_id in relationship_dict[environment.relationship][user_agent.agent_id]:
        bot_agent_list.append((agent_dict[neighbor_id].name, neighbor_id))
        
    return gr.Dropdown(choices=bot_agent_list, value=bot_agent_list[0][1] if bot_agent_list else None,  label="Bot Agent Selection")

def create_environment_info(environment_dropdown):
    _, environment_dict, _, _ = get_sotopia_profiles()
    environment = environment_dict[environment_dropdown]
    text = environment.scenario
    return gr.Textbox(label="Scenario", lines=1, value=text)

def create_user_info(user_agent_dropdown):
    _, _, agent_dict, _ = get_sotopia_profiles()
    user_agent = agent_dict[user_agent_dropdown]
    text = f"{user_agent.background} {user_agent.personality}"
    return gr.Textbox(label="User Agent Profile", lines=4, value=text)

def create_bot_info(bot_agent_dropdown):
    _, _, agent_dict, _ = get_sotopia_profiles()
    bot_agent = agent_dict[bot_agent_dropdown]
    text = f"{bot_agent.background} {bot_agent.personality}"
    return gr.Textbox(label="Bot Agent Profile", lines=4, value=text)

def create_user_goal(environment_dropdown):
    _, environment_dict, _, _ = get_sotopia_profiles()
    text = environment_dict[environment_dropdown].agent_goals[0]
    text = text.replace('(', '').replace(')', '')
    if "<extra_info>" in text:
        text = text.replace("<extra_info>", "\n\n")
        text = text.replace("</extra_info>", "\n")
    if "<strategy_hint>" in text:
        text = text.replace("<strategy_hint>", "\n\n")
        text = text.replace("</strategy_hint>", "\n")
    return gr.Textbox(label="User Agent Goal", lines=4, value=text)

def create_bot_goal(environment_dropdown):
    _, environment_dict, _, _ = get_sotopia_profiles()
    text = environment_dict[environment_dropdown].agent_goals[1]
    text = text.replace('(', '').replace(')', '')
    if "<extra_info>" in text:
        text = text.replace("<extra_info>", "\n\n")
        text = text.replace("</extra_info>", "\n")
    if "<strategy_hint>" in text:
        text = text.replace("<strategy_hint>", "\n\n")
        text = text.replace("</strategy_hint>", "\n")
    return gr.Textbox(label="Bot Agent Goal", lines=4, value=text)

def sotopia_info_accordion(accordion_visible=True):
    environments, _, _, _ = get_sotopia_profiles()
    
    with gr.Accordion("Create your sotopia space!", open=accordion_visible):
        with gr.Row():
            environment_dropdown = gr.Dropdown(
                choices=environments,
                label="Scenario Selection",
                value=environments[0][1] if environments else None,
                interactive=True,
            )
            model_name_dropdown = gr.Dropdown(
                choices=MODEL_OPTIONS,
                value=DEFAULT_MODEL_SELECTION,
                interactive=True,
                label="Model Selection"
            )

        with gr.Row():
            user_agent_dropdown = create_user_agent_dropdown(environment_dropdown.value)
            bot_agent_dropdown = create_bot_agent_dropdown(environment_dropdown.value, user_agent_dropdown.value)

    with gr.Accordion("Check your social task!", open=accordion_visible):

        scenario_info_display = create_environment_info(environment_dropdown.value)
            
        with gr.Row():
            bot_goal_display = create_bot_goal(environment_dropdown.value)
            user_goal_display = create_user_goal(environment_dropdown.value)
            

        
        with gr.Row():
            bot_agent_info_display = create_bot_info(bot_agent_dropdown.value)
            user_agent_info_display = create_user_info(user_agent_dropdown.value)

    # Update user dropdown when scenario changes
    environment_dropdown.change(fn=create_user_agent_dropdown, inputs=[environment_dropdown], outputs=[user_agent_dropdown])
    # Update bot dropdown when user or scenario changes
    user_agent_dropdown.change(fn=create_bot_agent_dropdown, inputs=[environment_dropdown, user_agent_dropdown], outputs=[bot_agent_dropdown])
    # Update scenario information when scenario changes
    environment_dropdown.change(fn=create_environment_info, inputs=[environment_dropdown], outputs=[scenario_info_display])
    # Update user agent profile when user changes
    user_agent_dropdown.change(fn=create_user_info, inputs=[user_agent_dropdown], outputs=[user_agent_info_display])
    # Update bot agent profile when bot changes
    bot_agent_dropdown.change(fn=create_bot_info, inputs=[bot_agent_dropdown], outputs=[bot_agent_info_display])
    # Update user goal when scenario changes
    environment_dropdown.change(fn=create_user_goal, inputs=[environment_dropdown], outputs=[user_goal_display])
    # Update bot goal when scenario changes
    environment_dropdown.change(fn=create_bot_goal, inputs=[environment_dropdown], outputs=[bot_goal_display])

    return model_name_dropdown, environment_dropdown, user_agent_dropdown, bot_agent_dropdown

def instructions_accordion(instructions, according_visible=False):
    with gr.Accordion("Instructions", open=False, visible=according_visible):
        instructions = gr.Textbox(
            lines=10,
            value=instructions,
            interactive=False,
            placeholder="Instructions",
            show_label=False,
            max_lines=10,
            visible=False,
        )
    return instructions

def chat_tab():
   # history are input output pairs
    _, environment_dict, agent_dict, _ = get_sotopia_profiles()
    def run_chat(
        message,
        history,
        environment_selection,
        user_agent_dropdown,
        bot_agent_dropdown,
        model_selection:str
    ):
        environment = environment_dict[environment_selection]
        user_agent = agent_dict[user_agent_dropdown]
        bot_agent = agent_dict[bot_agent_dropdown]
        
        context = get_context_prompt(bot_agent, user_agent, environment)
        dialogue_history, next_turn_idx = dialogue_history_prompt(message, history, user_agent, bot_agent)
        prompt_history = f"{context}{dialogue_history}"
        agent_action = generate_action(model_selection, prompt_history, next_turn_idx, ACTION_TYPES, bot_agent.name, TEMPERATURE)
        return agent_action.to_natural_language()
    
    with gr.Column():
        with gr.Blocks():
            model_name_dropdown, scenario_dropdown, user_agent_dropdown, bot_agent_dropdown = sotopia_info_accordion()
            
        with gr.Column():
            with gr.Accordion("Start the conversation to achieve your goal!", open=True):
                gr.ChatInterface(
                    fn=run_chat,
                    chatbot=gr.Chatbot(
                        height=620,
                        render=False,
                        show_label=False,
                        rtl=False,
                        avatar_images=(
                            "images/profile1.jpg",
                            "images/profile2.jpg",
                        ),
                    ),
                    textbox=gr.Textbox(
                        placeholder="Write your message here...",
                        render=False,
                        scale=7,
                        rtl=False,
                    ),
                    additional_inputs=[
                        scenario_dropdown,
                        user_agent_dropdown,
                        bot_agent_dropdown,
                        model_name_dropdown,
                    ],
                    submit_btn="Send",
                    stop_btn="Stop",
                    retry_btn="🔄 Retry",
                    undo_btn="↩️ Delete",
                    clear_btn="🗑️ Clear",
                )