350 lines
16 KiB
Python
350 lines
16 KiB
Python
from coagent.connector.schema import Memory, Message
|
|
import random
|
|
from textwrap import dedent
|
|
import copy
|
|
from loguru import logger
|
|
|
|
from coagent.connector.utils import extract_section, parse_section
|
|
|
|
|
|
class PromptManager:
|
|
def __init__(self, role_prompt="", prompt_config=None, monitored_agents=[], monitored_fields=[]):
|
|
self.role_prompt = role_prompt
|
|
self.monitored_agents = monitored_agents
|
|
self.monitored_fields = monitored_fields
|
|
self.field_handlers = {}
|
|
self.context_handlers = {}
|
|
self.field_order = [] # 用于普通字段的顺序
|
|
self.context_order = [] # 单独维护上下文字段的顺序
|
|
self.field_descriptions = {}
|
|
self.omit_if_empty_flags = {}
|
|
self.context_title = "### Context Data\n\n"
|
|
|
|
self.prompt_config = prompt_config
|
|
if self.prompt_config:
|
|
self.register_fields_from_config()
|
|
|
|
def register_field(self, field_name, function=None, title=None, description=None, is_context=True, omit_if_empty=True):
|
|
"""
|
|
注册一个新的字段及其处理函数。
|
|
Args:
|
|
field_name (str): 字段名称。
|
|
function (callable): 处理字段数据的函数。
|
|
title (str, optional): 字段的自定义标题（可选）。
|
|
description (str, optional): 字段的描述（可选，可以是几句话）。
|
|
is_context (bool, optional): 指示该字段是否为上下文字段。
|
|
omit_if_empty (bool, optional): 如果数据为空，是否省略该字段。
|
|
"""
|
|
if not function:
|
|
function = self.handle_custom_data
|
|
|
|
# Register the handler function based on context flag
|
|
if is_context:
|
|
self.context_handlers[field_name] = function
|
|
else:
|
|
self.field_handlers[field_name] = function
|
|
|
|
# Store the custom title if provided and adjust the title prefix based on context
|
|
title_prefix = "####" if is_context else "###"
|
|
if title is not None:
|
|
self.field_descriptions[field_name] = f"{title_prefix} {title}\n\n"
|
|
elif description is not None:
|
|
# If title is not provided but description is, use description as title
|
|
self.field_descriptions[field_name] = f"{title_prefix} {field_name.replace('_', ' ').title()}\n\n{description}\n\n"
|
|
else:
|
|
# If neither title nor description is provided, use the field name as title
|
|
self.field_descriptions[field_name] = f"{title_prefix} {field_name.replace('_', ' ').title()}\n\n"
|
|
|
|
# Store the omit_if_empty flag for this field
|
|
self.omit_if_empty_flags[field_name] = omit_if_empty
|
|
|
|
if is_context and field_name != 'context_placeholder':
|
|
self.context_handlers[field_name] = function
|
|
self.context_order.append(field_name)
|
|
else:
|
|
self.field_handlers[field_name] = function
|
|
self.field_order.append(field_name)
|
|
|
|
def generate_full_prompt(self, **kwargs):
|
|
full_prompt = []
|
|
context_prompts = [] # 用于收集上下文内容
|
|
is_pre_print = kwargs.get("is_pre_print", False) # 用于强制打印所有prompt 字段信息，不管有没有空
|
|
|
|
# 先处理上下文字段
|
|
for field_name in self.context_order:
|
|
handler = self.context_handlers[field_name]
|
|
processed_prompt = handler(**kwargs)
|
|
# Check if the field should be omitted when empty
|
|
if self.omit_if_empty_flags.get(field_name, False) and not processed_prompt and not is_pre_print:
|
|
continue # Skip this field
|
|
title_or_description = self.field_descriptions.get(field_name, f"#### {field_name.replace('_', ' ').title()}\n\n")
|
|
context_prompts.append(title_or_description + processed_prompt + '\n\n')
|
|
|
|
# 处理普通字段，同时查找 context_placeholder 的位置
|
|
for field_name in self.field_order:
|
|
if field_name == 'context_placeholder':
|
|
# 在 context_placeholder 的位置插入上下文数据
|
|
full_prompt.append(self.context_title) # 添加上下文部分的大标题
|
|
full_prompt.extend(context_prompts) # 添加收集的上下文内容
|
|
else:
|
|
handler = self.field_handlers[field_name]
|
|
processed_prompt = handler(**kwargs)
|
|
# Check if the field should be omitted when empty
|
|
if self.omit_if_empty_flags.get(field_name, False) and not processed_prompt and not is_pre_print:
|
|
continue # Skip this field
|
|
title_or_description = self.field_descriptions.get(field_name, f"### {field_name.replace('_', ' ').title()}\n\n")
|
|
full_prompt.append(title_or_description + processed_prompt + '\n\n')
|
|
|
|
# 返回完整的提示，移除尾部的空行
|
|
return ''.join(full_prompt).rstrip('\n')
|
|
|
|
def pre_print(self, **kwargs):
|
|
kwargs.update({"is_pre_print": True})
|
|
prompt = self.generate_full_prompt(**kwargs)
|
|
|
|
input_keys = parse_section(self.role_prompt, 'Response Output Format')
|
|
llm_predict = "\n".join([f"**{k}:**" for k in input_keys])
|
|
return prompt + "\n\n" + "#"*19 + "\n<<<<LLM PREDICT>>>>\n" + "#"*19 + f"\n\n{llm_predict}\n"
|
|
|
|
def handle_custom_data(self, **kwargs):
|
|
return ""
|
|
|
|
def handle_tool_data(self, **kwargs):
|
|
if 'previous_agent_message' not in kwargs:
|
|
return ""
|
|
|
|
previous_agent_message = kwargs.get('previous_agent_message')
|
|
tools = previous_agent_message.tools
|
|
|
|
if not tools:
|
|
return ""
|
|
|
|
tool_strings = []
|
|
for tool in tools:
|
|
args_schema = str(tool.args)
|
|
tool_strings.append(f"{tool.name}: {tool.description}, args: {args_schema}")
|
|
formatted_tools = "\n".join(tool_strings)
|
|
|
|
tool_names = ", ".join([tool.name for tool in tools])
|
|
|
|
tool_prompt = dedent(f"""
|
|
Below is a list of tools that are available for your use:
|
|
{formatted_tools}
|
|
|
|
valid "tool_name" value is:
|
|
{tool_names}
|
|
""")
|
|
|
|
return tool_prompt
|
|
|
|
def handle_agent_data(self, **kwargs):
|
|
if 'agents' not in kwargs:
|
|
return ""
|
|
|
|
agents = kwargs.get('agents')
|
|
random.shuffle(agents)
|
|
agent_names = ", ".join([f'{agent.role.role_name}' for agent in agents])
|
|
agent_descs = []
|
|
for agent in agents:
|
|
role_desc = agent.role.role_prompt.split("####")[1]
|
|
while "\n\n" in role_desc:
|
|
role_desc = role_desc.replace("\n\n", "\n")
|
|
role_desc = role_desc.replace("\n", ",")
|
|
|
|
agent_descs.append(f'"role name: {agent.role.role_name}\nrole description: {role_desc}"')
|
|
|
|
agents = "\n".join(agent_descs)
|
|
agent_prompt = f'''
|
|
Please ensure your selection is one of the listed roles. Available roles for selection:
|
|
{agents}
|
|
Please ensure select the Role from agent names, such as {agent_names}'''
|
|
|
|
return dedent(agent_prompt)
|
|
|
|
def handle_doc_info(self, **kwargs) -> str:
|
|
if 'previous_agent_message' not in kwargs:
|
|
return ""
|
|
previous_agent_message: Message = kwargs.get('previous_agent_message')
|
|
db_docs = previous_agent_message.db_docs
|
|
search_docs = previous_agent_message.search_docs
|
|
code_cocs = previous_agent_message.code_docs
|
|
doc_infos = "\n".join([doc.get_snippet() for doc in db_docs] + [doc.get_snippet() for doc in search_docs] +
|
|
[doc.get_code() for doc in code_cocs])
|
|
return doc_infos
|
|
|
|
def handle_session_records(self, **kwargs) -> str:
|
|
|
|
memory_pool: Memory = kwargs.get('memory_pool', Memory(messages=[]))
|
|
memory_pool = self.select_memory_by_agent_name(memory_pool)
|
|
memory_pool = self.select_memory_by_parsed_key(memory_pool)
|
|
|
|
return memory_pool.to_str_messages(content_key="parsed_output_list", with_tag=True)
|
|
|
|
def handle_current_plan(self, **kwargs) -> str:
|
|
if 'previous_agent_message' not in kwargs:
|
|
return ""
|
|
previous_agent_message = kwargs['previous_agent_message']
|
|
return previous_agent_message.parsed_output.get("CURRENT_STEP", "")
|
|
|
|
def handle_agent_profile(self, **kwargs) -> str:
|
|
return extract_section(self.role_prompt, 'Agent Profile')
|
|
|
|
def handle_output_format(self, **kwargs) -> str:
|
|
return extract_section(self.role_prompt, 'Response Output Format')
|
|
|
|
def handle_response(self, **kwargs) -> str:
|
|
if 'react_memory' not in kwargs:
|
|
return ""
|
|
|
|
react_memory = kwargs.get('react_memory', Memory(messages=[]))
|
|
if react_memory is None:
|
|
return ""
|
|
|
|
return "\n".join(["\n".join([f"**{k}:**\n{v}" for k,v in _dict.items()]) for _dict in react_memory.get_parserd_output()])
|
|
|
|
def handle_task_records(self, **kwargs) -> str:
|
|
if 'task_memory' not in kwargs:
|
|
return ""
|
|
|
|
task_memory: Memory = kwargs.get('task_memory', Memory(messages=[]))
|
|
if task_memory is None:
|
|
return ""
|
|
|
|
return "\n".join(["\n".join([f"**{k}:**\n{v}" for k,v in _dict.items() if k not in ["CURRENT_STEP"]]) for _dict in task_memory.get_parserd_output()])
|
|
|
|
def handle_previous_message(self, message: Message) -> str:
|
|
pass
|
|
|
|
def handle_message_by_role_name(self, message: Message) -> str:
|
|
pass
|
|
|
|
def handle_message_by_role_type(self, message: Message) -> str:
|
|
pass
|
|
|
|
def handle_current_agent_react_message(self, message: Message) -> str:
|
|
pass
|
|
|
|
def extract_codedoc_info_for_prompt(self, message: Message) -> str:
|
|
code_docs = message.code_docs
|
|
doc_infos = "\n".join([doc.get_code() for doc in code_docs])
|
|
return doc_infos
|
|
|
|
def select_memory_by_parsed_key(self, memory: Memory) -> Memory:
|
|
return Memory(
|
|
messages=[self.select_message_by_parsed_key(message) for message in memory.messages
|
|
if self.select_message_by_parsed_key(message) is not None]
|
|
)
|
|
|
|
def select_memory_by_agent_name(self, memory: Memory) -> Memory:
|
|
return Memory(
|
|
messages=[self.select_message_by_agent_name(message) for message in memory.messages
|
|
if self.select_message_by_agent_name(message) is not None]
|
|
)
|
|
|
|
def select_message_by_agent_name(self, message: Message) -> Message:
|
|
# assume we focus all agents
|
|
if self.monitored_agents == []:
|
|
return message
|
|
return None if message is None or message.role_name not in self.monitored_agents else self.select_message_by_parsed_key(message)
|
|
|
|
def select_message_by_parsed_key(self, message: Message) -> Message:
|
|
# assume we focus all key contents
|
|
if message is None:
|
|
return message
|
|
|
|
if self.monitored_fields == []:
|
|
return message
|
|
|
|
message_c = copy.deepcopy(message)
|
|
message_c.parsed_output = {k: v for k,v in message_c.parsed_output.items() if k in self.monitored_fields}
|
|
message_c.parsed_output_list = [{k: v for k,v in parsed_output.items() if k in self.monitored_fields} for parsed_output in message_c.parsed_output_list]
|
|
return message_c
|
|
|
|
def get_memory(self, content_key="role_content"):
|
|
return self.memory.to_tuple_messages(content_key="step_content")
|
|
|
|
def get_memory_str(self, content_key="role_content"):
|
|
return "\n".join([": ".join(i) for i in self.memory.to_tuple_messages(content_key="step_content")])
|
|
|
|
def register_fields_from_config(self):
|
|
|
|
for prompt_field in self.prompt_config:
|
|
|
|
function_name = prompt_field.function_name
|
|
# 检查function_name是否是self的一个方法
|
|
if function_name and hasattr(self, function_name):
|
|
function = getattr(self, function_name)
|
|
else:
|
|
function = self.handle_custom_data
|
|
|
|
self.register_field(prompt_field.field_name,
|
|
function=function,
|
|
title=prompt_field.title,
|
|
description=prompt_field.description,
|
|
is_context=prompt_field.is_context,
|
|
omit_if_empty=prompt_field.omit_if_empty)
|
|
|
|
def register_standard_fields(self):
|
|
self.register_field('agent_profile', function=self.handle_agent_profile, is_context=False)
|
|
self.register_field('tool_information', function=self.handle_tool_data, is_context=False)
|
|
self.register_field('context_placeholder', is_context=True) # 用于标记上下文数据部分的位置
|
|
self.register_field('reference_documents', function=self.handle_doc_info, is_context=True)
|
|
self.register_field('session_records', function=self.handle_session_records, is_context=True)
|
|
self.register_field('output_format', function=self.handle_output_format, title='Response Output Format', is_context=False)
|
|
self.register_field('response', function=self.handle_response, is_context=False, omit_if_empty=False)
|
|
|
|
def register_executor_fields(self):
|
|
self.register_field('agent_profile', function=self.handle_agent_profile, is_context=False)
|
|
self.register_field('tool_information', function=self.handle_tool_data, is_context=False)
|
|
self.register_field('context_placeholder', is_context=True) # 用于标记上下文数据部分的位置
|
|
self.register_field('reference_documents', function=self.handle_doc_info, is_context=True)
|
|
self.register_field('session_records', function=self.handle_session_records, is_context=True)
|
|
self.register_field('current_plan', function=self.handle_current_plan, is_context=True)
|
|
self.register_field('output_format', function=self.handle_output_format, title='Response Output Format', is_context=False)
|
|
self.register_field('response', function=self.handle_response, is_context=False, omit_if_empty=False)
|
|
|
|
def register_fields_from_dict(self, fields_dict):
|
|
# 使用字典注册字段的函数
|
|
for field_name, field_config in fields_dict.items():
|
|
function_name = field_config.get('function', None)
|
|
title = field_config.get('title', None)
|
|
description = field_config.get('description', None)
|
|
is_context = field_config.get('is_context', True)
|
|
omit_if_empty = field_config.get('omit_if_empty', True)
|
|
|
|
# 检查function_name是否是self的一个方法
|
|
if function_name and hasattr(self, function_name):
|
|
function = getattr(self, function_name)
|
|
else:
|
|
function = self.handle_custom_data
|
|
|
|
# 调用已存在的register_field方法注册字段
|
|
self.register_field(field_name, function=function, title=title, description=description, is_context=is_context, omit_if_empty=omit_if_empty)
|
|
|
|
|
|
|
|
def main():
|
|
manager = PromptManager()
|
|
manager.register_standard_fields()
|
|
|
|
manager.register_field('agents_work_progress', title=f"Agents' Work Progress", is_context=True)
|
|
|
|
# 创建数据字典
|
|
data_dict = {
|
|
"agent_profile": "这是代理配置文件...",
|
|
# "tool_list": "这是工具列表...",
|
|
"reference_documents": "这是参考文档...",
|
|
"session_records": "这是会话记录...",
|
|
"agents_work_progress": "这是代理工作进展...",
|
|
"output_format": "这是预期的输出格式...",
|
|
# "response": "这是生成或继续回应的指令...",
|
|
"response": "",
|
|
"test": 'xxxxx'
|
|
}
|
|
|
|
# 组合完整的提示
|
|
full_prompt = manager.generate_full_prompt(data_dict)
|
|
print(full_prompt)
|
|
|
|
if __name__ == "__main__":
|
|
main() |