File size: 2,876 Bytes
0c51b93 |
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 |
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
# TODOTODO (haofeiyu): we first assume that human talks first
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()
|