两人副本与龙与地下城¶

from typing import Callable, List

from langchain.schema import (
    HumanMessage,
    SystemMessage,
)
from langchain_openai import ChatOpenAI

class DialogueAgent:
    def __init__(
        self,
        name: str,
        system_message: SystemMessage,
        model: ChatOpenAI,
    ) -> None:
        self.name = name
        self.system_message = system_message
        self.model = model
        self.prefix = f"{self.name}: "  # 设置对话代理的前缀
        self.reset()  # 重置对话历史

    def reset(self):
        self.message_history = ["Here is the conversation so far."]  # 初始化对话历史，包含一条初始消息

    def send(self) -> str:
        """
        将对话模型应用于对话历史并返回消息字符串
        """
        message = self.model.invoke(
            [
                self.system_message,
                HumanMessage(content="\n".join(self.message_history + [self.prefix])),
            ]
        )
        return message.content

    def receive(self, name: str, message: str) -> None:
        """
        将{name}说的{message}连接到对话历史中
        """
        self.message_history.append(f"{name}: {message}")

class DialogueSimulator:
    def __init__(
        self,
        agents: List[DialogueAgent],
        selection_function: Callable[[int, List[DialogueAgent]], int],
    ) -> None:
        self.agents = agents
        self._step = 0
        self.select_next_speaker = selection_function

    def reset(self):
        for agent in self.agents:
            agent.reset()

    def inject(self, name: str, message: str):
        """
        用{name}发送的{message}开始对话
        """
        for agent in self.agents:
            agent.receive(name, message)

        # 增加时间步
        self._step += 1

    def step(self) -> tuple[str, str]:
        # 1. 选择下一个发言者
        speaker_idx = self.select_next_speaker(self._step, self.agents)
        speaker = self.agents[speaker_idx]

        # 2. 下一个发言者发送消息
        message = speaker.send()

        # 3. 所有人接收消息
        for receiver in self.agents:
            receiver.receive(speaker.name, message)

        # 4. 增加时间步
        self._step += 1

        return speaker.name, message

# 定义主角名字为 "Harry Potter"
protagonist_name = "Harry Potter"

# 定义讲述者名字为 "Dungeon Master"
storyteller_name = "Dungeon Master"

# 定义任务为 "寻找所有伏地魔的七个魂器"
quest = "Find all of Lord Voldemort's seven horcruxes."

# 定义任务头脑风暴的字数限制为 50
word_limit = 50

# 游戏描述
game_description = f"""这里是一个龙与地下城游戏的主题：{quest}。
        这个游戏中有一个玩家：主角，{protagonist_name}。
        故事由讲述者，{storyteller_name}，叙述。"""

# 系统消息，描述玩家细节
player_descriptor_system_message = SystemMessage(
    content="你可以为龙与地下城玩家的描述添加细节。"
)

# 主角描述提示信息
protagonist_specifier_prompt = [
    player_descriptor_system_message,
    HumanMessage(
        content=f"""{game_description}
        请回复一个关于主角{protagonist_name}的富有创意的描述，限制在{word_limit}个字以内。
        直接对{protagonist_name}讲话。
        不要添加其他内容。"""
    ),
]

# 获取主角描述
protagonist_description = ChatOpenAI(temperature=1.0)(
    protagonist_specifier_prompt
).content

# 叙述者描述提示信息
storyteller_specifier_prompt = [
    player_descriptor_system_message,
    HumanMessage(
        content=f"""{game_description}
        请回复一个关于叙述者{storyteller_name}的富有创意的描述，限制在{word_limit}个字以内。
        直接对{storyteller_name}讲话。
        不要添加其他内容。"""
    ),
]

# 获取叙述者描述
storyteller_description = ChatOpenAI(temperature=1.0)(
    storyteller_specifier_prompt
).content

# 创建主角系统消息
protagonist_system_message = SystemMessage(
    content=(
        f"""{game_description}
永远不要忘记你是主角，{protagonist_name}，而我是叙事者，{storyteller_name}。
你的角色描述如下：{protagonist_description}。
你将提出你计划采取的行动，我将解释当你采取这些行动时会发生什么。
以第一人称从{protagonist_name}的角度说话。
描述自己的身体动作时，请用'*'括起你的描述。
不要改变角色！
不要从{storyteller_name}的角度说话。
不要忘记以“轮到你了，{storyteller_name}”结束发言。
不要添加其他内容。
记住你是主角，{protagonist_name}。
当你从你的角度说话结束时立即停止说话。
"""
    )
)

# 创建叙事者系统消息
storyteller_system_message = SystemMessage(
    content=(
        f"""{game_description}
永远不要忘记你是叙事者，{storyteller_name}，而我是主角，{protagonist_name}。
你的角色描述如下：{storyteller_description}。
我将提出我计划采取的行动，你将解释当我采取这些行动时会发生什么。
以第一人称从{storyteller_name}的角度说话。
描述自己的身体动作时，请用'*'括起你的描述。
不要改变角色！
不要从{protagonist_name}的角度说话。
不要忘记以“轮到你了，{protagonist_name}”结束发言。
不要添加其他内容。
记住你是叙事者，{storyteller_name}。
当你从你的角度说话结束时立即停止说话。
"""
    )
)

# 创建主角对话代理
protagonist = DialogueAgent(
    name=protagonist_name,  # 主角名称
    system_message=protagonist_system_message,  # 主角系统消息
    model=ChatOpenAI(temperature=0.2),  # 主角对话模型
)

# 创建讲故事者对话代理
storyteller = DialogueAgent(
    name=storyteller_name,  # 讲故事者名称
    system_message=storyteller_system_message,  # 讲故事者系统消息
    model=ChatOpenAI(temperature=0.2),  # 讲故事者对话模型
)

def select_next_speaker(step: int, agents: List[DialogueAgent]) -> int:
    # 计算下一个发言者的索引，取余操作可以循环选择发言者
    idx = step % len(agents)
    return idx