|
|
import re |
|
|
from typing import Dict, List, Tuple |
|
|
|
|
|
|
|
|
class Agent: |
|
|
def __init__(self, agent_profile: Dict[str, str]): |
|
|
self._id = agent_profile["agent_id"] |
|
|
|
|
|
self.agent_profile = agent_profile |
|
|
self.agent_id = agent_profile["agent_id"] |
|
|
self.name = self.get_name(agent_profile) |
|
|
self.background = self.get_background(agent_profile) |
|
|
self.secret = agent_profile["secret"] |
|
|
self.personality = agent_profile["personality_and_values"] |
|
|
self.goal = "" |
|
|
|
|
|
def get_name(self, agent_profile: Dict[str, str]) -> str: |
|
|
return agent_profile["first_name"] + " " + agent_profile["last_name"] |
|
|
|
|
|
def get_background(self, agent_profile: Dict[str, str]) -> str: |
|
|
name = self.name |
|
|
return f"{name} is a {agent_profile['age']}-year-old {agent_profile['gender'].lower()} {agent_profile['occupation']}. {agent_profile['public_info']}" |
|
|
|
|
|
|
|
|
class Environment: |
|
|
def __init__(self, env_profile: Dict[str, str]): |
|
|
self._id = env_profile["env_id"] |
|
|
|
|
|
self.environment_profile = env_profile |
|
|
self.codename = env_profile["codename"] |
|
|
self.scenario = env_profile["scenario"] |
|
|
self.agent_goals = env_profile["agent_goals"] |
|
|
self.relationship = env_profile["relationship"] |
|
|
|
|
|
def to_dict(self) -> Dict[str, str]: |
|
|
return self.environment_profile |
|
|
|
|
|
|
|
|
def get_context_prompt( |
|
|
machine_agent: Agent, human_agent: Agent, environment: Environment |
|
|
) -> str: |
|
|
return f"Here is the context of this interaction:\n Scenario: {environment.scenario}\nParticipants: {human_agent.name} and {machine_agent.name}\n{human_agent.name}'s background: {human_agent.background} Personality and values description: {human_agent.personality} \n{machine_agent.name}'s background: {machine_agent.background} Personality and values description: {machine_agent.personality} {machine_agent.name}'s secrets: {machine_agent.secret}\n{human_agent.name}'s goal: Unknown\n{machine_agent.name}'s goal: {environment.agent_goals[1]}\nConversation Starts:" |
|
|
|
|
|
|
|
|
def dialogue_history_prompt( |
|
|
message: str, history: List[List[str]], user_agent: Agent, bot_agent: Agent |
|
|
) -> Tuple[str, int]: |
|
|
dialogue_history = "" |
|
|
for idx, turn in enumerate(history): |
|
|
user_message, bot_message = turn |
|
|
|
|
|
user_turn_idx = idx * 2 |
|
|
bot_turn_idx = idx * 2 + 1 |
|
|
dialogue_history = f"""{dialogue_history}\n\nTurn #{user_turn_idx} {user_message}"\n\nTurn #{bot_turn_idx} {bot_message}""" |
|
|
curr_turn_idx = len(history) * 2 |
|
|
dialogue_history = ( |
|
|
f"""{dialogue_history}\n\nTurn #{curr_turn_idx} {message}\n""" |
|
|
) |
|
|
return dialogue_history, curr_turn_idx + 1 |
|
|
|
|
|
|
|
|
def format_docstring(docstring: str) -> str: |
|
|
"""Format a docstring for use in a prompt template.""" |
|
|
return re.sub("\n +", "\n", docstring).strip() |
|
|
|
