227 lines
8.4 KiB
Python
227 lines
8.4 KiB
Python
"""
|
||
提示词组装器 (Prompt Assembler)
|
||
|
||
负责根据 SillyTavern 规范将角色卡、世界书、聊天历史等组件
|
||
拼装成最终的 LLM 消息列表。
|
||
"""
|
||
import re
|
||
from typing import List, Dict, Optional
|
||
from langchain_core.messages import SystemMessage, HumanMessage, AIMessage, BaseMessage
|
||
|
||
from models.internal import CharacterCard, ChatMessage, WorldInfoEntry
|
||
|
||
|
||
class PromptConfig:
|
||
"""提示词组装配置"""
|
||
def __init__(
|
||
self,
|
||
an_position: str = "after_history", # "before_history" or "after_history"
|
||
an_depth: int = 4,
|
||
post_history_instructions: Optional[str] = None
|
||
):
|
||
self.an_position = an_position
|
||
self.an_depth = an_depth
|
||
self.post_history_instructions = post_history_instructions
|
||
|
||
|
||
class PromptAssembler:
|
||
"""
|
||
轻量级提示词组装核心
|
||
|
||
不依赖复杂的框架，只负责纯粹的文本拼接和位置插入。
|
||
"""
|
||
|
||
# SillyTavern 的位置枚举映射
|
||
POS_WI_BEFORE = 0
|
||
POS_WI_AFTER = 1
|
||
POS_EXAMPLES_BEFORE = 2
|
||
POS_EXAMPLES_AFTER = 3
|
||
POS_AN_TOP = 4
|
||
POS_AN_BOTTOM = 5
|
||
POS_DEPTH = 6
|
||
POS_OUTLET = 7
|
||
|
||
def assemble(
|
||
self,
|
||
character: CharacterCard,
|
||
chat_history: List[ChatMessage],
|
||
user_input: str,
|
||
active_entries: List[WorldInfoEntry],
|
||
config: PromptConfig = PromptConfig()
|
||
) -> List[BaseMessage]:
|
||
"""
|
||
执行完整的提示词组装流程
|
||
|
||
Returns:
|
||
List[BaseMessage]: 准备好发送给 LLM 的消息列表
|
||
"""
|
||
# 1. 按位置分组世界书条目
|
||
grouped_entries = self._group_entries_by_position(active_entries)
|
||
|
||
# 2. 组装 Story String (包含 Pos 0-3)
|
||
story_string = self._build_story_string(character, grouped_entries)
|
||
|
||
# 3. 组装 Author's Note (包含 Pos 4-5)
|
||
authors_note_content = self._build_authors_note(grouped_entries, config.an_depth)
|
||
|
||
# 4. 处理 Chat History 并注入 Depth 条目 (Pos 6)
|
||
processed_history = self._inject_depth_entries(chat_history, grouped_entries.get(self.POS_DEPTH, []))
|
||
|
||
# 5. 准备 Outlet 替换字典 (Pos 7)
|
||
outlet_map = {entry.uid: entry.content for entry in grouped_entries.get(self.POS_OUTLET, [])}
|
||
|
||
# 6. 最终封装为 Messages
|
||
return self._wrap_to_messages(
|
||
story_string,
|
||
authors_note_content,
|
||
processed_history,
|
||
user_input,
|
||
outlet_map,
|
||
config
|
||
)
|
||
|
||
def _group_entries_by_position(self, entries: List[WorldInfoEntry]) -> Dict[int, List[WorldInfoEntry]]:
|
||
"""将激活的条目按 position 分组"""
|
||
grouped = {}
|
||
for entry in entries:
|
||
# 这里假设 entry.position 存储的是我们定义的 0-7 整数
|
||
pos = entry.position if isinstance(entry.position, int) else 1 # 默认为 wiAfter
|
||
if pos not in grouped:
|
||
grouped[pos] = []
|
||
grouped[pos].append(entry)
|
||
|
||
# 对每个组内的条目按 order 排序
|
||
for pos in grouped:
|
||
grouped[pos].sort(key=lambda x: x.order)
|
||
return grouped
|
||
|
||
def _build_story_string(self, character: CharacterCard, grouped: Dict) -> str:
|
||
"""组装故事字符串 (Story String)"""
|
||
parts = []
|
||
|
||
# Pos 0: wiBefore
|
||
for entry in grouped.get(self.POS_WI_BEFORE, []):
|
||
parts.append(entry.content)
|
||
|
||
# 角色核心信息
|
||
parts.append(f"[Character('{character.name}')]\n{character.description}\n")
|
||
parts.append(f"Personality: {character.personality}\n")
|
||
parts.append(f"Scenario: {character.scenario}\n")
|
||
|
||
# Pos 1: wiAfter
|
||
for entry in grouped.get(self.POS_WI_AFTER, []):
|
||
parts.append(entry.content)
|
||
|
||
# Pos 2: Examples Before
|
||
for entry in grouped.get(self.POS_EXAMPLES_BEFORE, []):
|
||
parts.append(entry.content)
|
||
|
||
# 示例对话
|
||
if character.mes_example:
|
||
parts.append(f"<START>\n{character.mes_example}")
|
||
|
||
# Pos 3: Examples After
|
||
for entry in grouped.get(self.POS_EXAMPLES_AFTER, []):
|
||
parts.append(entry.content)
|
||
|
||
return "\n".join(parts)
|
||
|
||
def _build_authors_note(self, grouped: Dict, depth: int) -> str:
|
||
"""组装作者笔记 (Author's Note)"""
|
||
parts = []
|
||
|
||
# Pos 4: AN Top
|
||
for entry in grouped.get(self.POS_AN_TOP, []):
|
||
parts.append(str(entry.content) if entry.content else "")
|
||
|
||
# AN 核心内容 (这里简化为一个占位，实际应从角色卡或设置获取)
|
||
parts.append(f"[Author's note at depth {depth}]")
|
||
|
||
# Pos 5: AN Bottom
|
||
for entry in grouped.get(self.POS_AN_BOTTOM, []):
|
||
parts.append(str(entry.content) if entry.content else "")
|
||
|
||
return "\n".join(parts)
|
||
|
||
def _inject_depth_entries(self, history: List[ChatMessage], depth_entries: List[WorldInfoEntry]) -> List[Dict]:
|
||
"""
|
||
在聊天历史的指定深度插入条目 (Pos 6)
|
||
返回一个包含 role 和 content 的字典列表，方便后续转换
|
||
|
||
✅ 过滤已被总结的消息（is_summarized=True 且 mes=""）
|
||
"""
|
||
# 先将历史转换为中间格式，过滤掉空消息（已被总结）
|
||
msg_list = []
|
||
for msg in history:
|
||
# ✅ 跳过已被总结的空消息
|
||
if msg.is_summarized and (msg.mes == "" or msg.mes.strip() == ""):
|
||
continue
|
||
msg_list.append({"role": "user" if msg.is_user else "assistant", "content": msg.mes})
|
||
|
||
# 按 depth 分组插入
|
||
# d0 通常指最新用户输入之前，即列表末尾
|
||
for entry in depth_entries:
|
||
depth = entry.depth if entry.depth is not None else 0
|
||
# 计算插入索引 (从后往前数)
|
||
insert_index = max(0, len(msg_list) - depth)
|
||
|
||
# 确定角色
|
||
role_map = {"system": "system", "user": "user", "assistant": "assistant"}
|
||
role = role_map.get(str(entry.position).split('_')[-1] if '_' in str(entry.position) else "system", "system")
|
||
|
||
msg_list.insert(insert_index, {"role": "system", "content": entry.content})
|
||
|
||
return msg_list
|
||
|
||
def _replace_outlets(self, text: str, outlet_map: Dict[str, str]) -> str:
|
||
"""执行 Outlet 宏替换 (Pos 7)"""
|
||
def replacer(match):
|
||
uid = match.group(1)
|
||
return outlet_map.get(uid, "")
|
||
|
||
# 匹配 {{outlet::UID}}
|
||
return re.sub(r"\{\{outlet::([^}]+)\}\}", replacer, text)
|
||
|
||
def _wrap_to_messages(
|
||
self,
|
||
story_string: str,
|
||
an_content: str,
|
||
history: List[Dict],
|
||
user_input: str,
|
||
outlet_map: Dict[str, str],
|
||
config: PromptConfig
|
||
) -> List[BaseMessage]:
|
||
"""将组装好的文本块封装为 LangChain Messages"""
|
||
messages = []
|
||
|
||
# 1. System Message (Story String + Outlet 替换)
|
||
final_story = self._replace_outlets(story_string, outlet_map)
|
||
if final_story:
|
||
messages.append(SystemMessage(content=final_story))
|
||
|
||
# 2. Author's Note (根据配置位置插入)
|
||
if an_content and config.an_position == "before_history":
|
||
messages.append(SystemMessage(content=self._replace_outlets(an_content, outlet_map)))
|
||
|
||
# 3. Chat History
|
||
for msg_data in history:
|
||
if msg_data["role"] == "user":
|
||
messages.append(HumanMessage(content=msg_data["content"]))
|
||
elif msg_data["role"] == "assistant":
|
||
messages.append(AIMessage(content=msg_data["content"]))
|
||
else:
|
||
messages.append(SystemMessage(content=msg_data["content"]))
|
||
|
||
# 4. Author's Note (如果在 History 之后)
|
||
if an_content and config.an_position == "after_history":
|
||
messages.append(SystemMessage(content=self._replace_outlets(an_content, outlet_map)))
|
||
|
||
# 5. Post-History Instructions & User Input
|
||
final_input = user_input
|
||
if config.post_history_instructions:
|
||
final_input = f"{config.post_history_instructions}\n\n{user_input}"
|
||
|
||
messages.append(HumanMessage(content=final_input))
|
||
|
||
return messages
|