Files
MetaGPT/PR/action_node.py
T

352 lines
13 KiB
Python

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Time : 2023/12/11 18:45
@Author : alexanderwu
@File : action_node.py
NOTE: You should use typing.List instead of list to do type annotation. Because in the markdown extraction process,
we can use typing to extract the type of the node, but we cannot use built-in list to extract.
"""
import json
from typing import Any, Dict, List, Optional, Tuple, Type
from pydantic import BaseModel, create_model, model_validator
from tenacity import retry, stop_after_attempt, wait_random_exponential
from metagpt.config import CONFIG
from metagpt.llm import BaseLLM
from metagpt.logs import logger
from metagpt.provider.postprocess.llm_output_postprocess import llm_output_postprocess
from metagpt.utils.common import OutputParser, general_after_log
TAG = "CONTENT"
LANGUAGE_CONSTRAINT = "Language: Please use the same language as Human INPUT."
FORMAT_CONSTRAINT = f"Format: output wrapped inside [{TAG}][/{TAG}] like format example, nothing else."
# Delete comments in json
# If you don't want JSONDecodeError to occur, you can add Delete comments in json after FORMAT_CONSTRAINT
SIMPLE_TEMPLATE = """
## context
{context}
-----
## format example
{example}
## nodes: "<node>: <type> # <instruction>"
{instruction}
## constraint
{constraint}
## action
Follow instructions of nodes, generate output and make sure it follows the format example.
"""
def dict_to_markdown(d, prefix="- ", kv_sep="\n", postfix="\n"):
markdown_str = ""
for key, value in d.items():
markdown_str += f"{prefix}{key}{kv_sep}{value}{postfix}"
return markdown_str
class ActionNode:
"""ActionNode is a tree of nodes."""
schema: str # raw/json/markdown, default: ""
# Action Context
context: str # all the context, including all necessary info
llm: BaseLLM # LLM with aask interface
children: dict[str, "ActionNode"]
# Action Input
key: str # Product Requirement / File list / Code
expected_type: Type # such as str / int / float etc.
# context: str # everything in the history.
instruction: str # the instructions should be followed.
example: Any # example for In Context-Learning.
# Action Output
content: str
instruct_content: BaseModel
def __init__(
self,
key: str,
expected_type: Type,
instruction: str,
example: Any,
content: str = "",
children: dict[str, "ActionNode"] = None,
schema: str = "",
):
self.key = key
self.expected_type = expected_type
self.instruction = instruction
self.example = example
self.content = content
self.children = children if children is not None else {}
self.schema = schema
def __str__(self):
return (
f"{self.key}, {repr(self.expected_type)}, {self.instruction}, {self.example}"
f", {self.content}, {self.children}"
)
def __repr__(self):
return self.__str__()
def add_child(self, node: "ActionNode"):
"""增加子ActionNode"""
self.children[node.key] = node
def add_children(self, nodes: List["ActionNode"]):
"""批量增加子ActionNode"""
for node in nodes:
self.add_child(node)
@classmethod
def from_children(cls, key, nodes: List["ActionNode"]):
"""直接从一系列的子nodes初始化"""
obj = cls(key, str, "", "")
obj.add_children(nodes)
return obj
def get_children_mapping(self, exclude=None) -> Dict[str, Tuple[Type, Any]]:
"""获得子ActionNode的字典,以key索引"""
exclude = exclude or []
return {k: (v.expected_type, ...) for k, v in self.children.items() if k not in exclude}
def get_self_mapping(self) -> Dict[str, Tuple[Type, Any]]:
"""get self key: type mapping"""
return {self.key: (self.expected_type, ...)}
def get_mapping(self, mode="children", exclude=None) -> Dict[str, Tuple[Type, Any]]:
"""get key: type mapping under mode"""
if mode == "children" or (mode == "auto" and self.children):
return self.get_children_mapping(exclude=exclude)
return {} if exclude and self.key in exclude else self.get_self_mapping()
@classmethod
def create_model_class(cls, class_name: str, mapping: Dict[str, Tuple[Type, Any]]):
"""基于pydantic v1的模型动态生成,用来检验结果类型正确性"""
def check_fields(cls, values):
required_fields = set(mapping.keys())
missing_fields = required_fields - set(values.keys())
if missing_fields:
raise ValueError(f"Missing fields: {missing_fields}")
unrecognized_fields = set(values.keys()) - required_fields
if unrecognized_fields:
logger.warning(f"Unrecognized fields: {unrecognized_fields}")
return values
validators = {"check_missing_fields_validator": model_validator(mode="before")(check_fields)}
new_class = create_model(class_name, __validators__=validators, **mapping)
return new_class
def create_children_class(self, exclude=None):
"""使用object内有的字段直接生成model_class"""
class_name = f"{self.key}_AN"
mapping = self.get_children_mapping(exclude=exclude)
return self.create_model_class(class_name, mapping)
def to_dict(self, format_func=None, mode="auto", exclude=None) -> Dict:
"""将当前节点与子节点都按照node: format的格式组织成字典"""
# 如果没有提供格式化函数,使用默认的格式化方式
if format_func is None:
format_func = lambda node: f"{node.instruction}"
# 使用提供的格式化函数来格式化当前节点的值
formatted_value = format_func(self)
# 创建当前节点的键值对
if mode == "children" or (mode == "auto" and self.children):
node_dict = {}
else:
node_dict = {self.key: formatted_value}
if mode == "root":
return node_dict
# 遍历子节点并递归调用 to_dict 方法
exclude = exclude or []
for _, child_node in self.children.items():
if child_node.key in exclude:
continue
node_dict.update(child_node.to_dict(format_func))
return node_dict
def compile_to(self, i: Dict, schema, kv_sep) -> str:
if schema == "json":
return json.dumps(i, indent=4)
elif schema == "markdown":
return dict_to_markdown(i, kv_sep=kv_sep)
else:
return str(i)
def tagging(self, text, schema, tag="") -> str:
if not tag:
return text
if schema == "json":
return f"[{tag}]\n" + text + f"\n[/{tag}]"
else: # markdown
return f"[{tag}]\n" + text + f"\n[/{tag}]"
def _compile_f(self, schema, mode, tag, format_func, kv_sep, exclude=None) -> str:
nodes = self.to_dict(format_func=format_func, mode=mode, exclude=exclude)
text = self.compile_to(nodes, schema, kv_sep)
return self.tagging(text, schema, tag)
def compile_instruction(self, schema="markdown", mode="children", tag="", exclude=None) -> str:
"""compile to raw/json/markdown template with all/root/children nodes"""
format_func = lambda i: f"{i.expected_type} # {i.instruction}"
return self._compile_f(schema, mode, tag, format_func, kv_sep=": ", exclude=exclude)
def compile_example(self, schema="json", mode="children", tag="", exclude=None) -> str:
"""compile to raw/json/markdown examples with all/root/children nodes"""
# 这里不能使用f-string,因为转译为str后再json.dumps会额外加上引号,无法作为有效的example
# 错误示例:"File list": "['main.py', 'const.py', 'game.py']", 注意这里值不是list,而是str
format_func = lambda i: i.example
return self._compile_f(schema, mode, tag, format_func, kv_sep="\n", exclude=exclude)
def compile(self, context, schema="json", mode="children", template=SIMPLE_TEMPLATE, exclude=[]) -> str:
"""
mode: all/root/children
mode="children": 编译所有子节点为一个统一模板,包括instruction与example
mode="all": NotImplemented
mode="root": NotImplemented
schmea: raw/json/markdown
schema="raw": 不编译,context, lang_constaint, instruction
schema="json":编译context, example(json), instruction(markdown), constraint, action
schema="markdown": 编译context, example(markdown), instruction(markdown), constraint, action
"""
if schema == "raw":
return context + "\n\n## Actions\n" + LANGUAGE_CONSTRAINT + "\n" + self.instruction
# FIXME: json instruction会带来格式问题,如:"Project name": "web_2048 # 项目名称使用下划线",
# compile example暂时不支持markdown
instruction = self.compile_instruction(schema="markdown", mode=mode, exclude=exclude)
example = self.compile_example(schema=schema, tag=TAG, mode=mode, exclude=exclude)
# nodes = ", ".join(self.to_dict(mode=mode).keys())
constraints = [LANGUAGE_CONSTRAINT, FORMAT_CONSTRAINT]
constraint = "\n".join(constraints)
prompt = template.format(
context=context,
example=example,
instruction=instruction,
constraint=constraint,
)
return prompt
@retry(
wait=wait_random_exponential(min=1, max=20),
stop=stop_after_attempt(6),
after=general_after_log(logger),
)
async def _aask_v1(
self,
prompt: str,
output_class_name: str,
output_data_mapping: dict,
system_msgs: Optional[list[str]] = None,
schema="markdown", # compatible to original format
timeout=CONFIG.timeout,
) -> (str, BaseModel):
"""Use ActionOutput to wrap the output of aask"""
content = await self.llm.aask(prompt, system_msgs, timeout=timeout)
logger.debug(f"llm raw output:\n{content}")
output_class = self.create_model_class(output_class_name, output_data_mapping)
if schema == "json":
parsed_data = llm_output_postprocess(
output=content, schema=output_class.model_json_schema(), req_key=f"[/{TAG}]"
)
else: # using markdown parser
parsed_data = OutputParser.parse_data_with_mapping(content, output_data_mapping)
logger.debug(f"parsed_data:\n{parsed_data}")
instruct_content = output_class(**parsed_data)
return content, instruct_content
def get(self, key):
return self.instruct_content.model_dump()[key]
def set_recursive(self, name, value):
setattr(self, name, value)
for _, i in self.children.items():
i.set_recursive(name, value)
def set_llm(self, llm):
self.set_recursive("llm", llm)
def set_context(self, context):
self.set_recursive("context", context)
async def simple_fill(self, schema, mode, timeout=CONFIG.timeout, exclude=None):
prompt = self.compile(context=self.context, schema=schema, mode=mode, exclude=exclude)
if schema != "raw":
mapping = self.get_mapping(mode, exclude=exclude)
class_name = f"{self.key}_AN"
content, scontent = await self._aask_v1(prompt, class_name, mapping, schema=schema, timeout=timeout)
self.content = content
self.instruct_content = scontent
else:
self.content = await self.llm.aask(prompt)
self.instruct_content = None
return self
async def fill(self, context, llm, schema="json", mode="auto", strgy="simple", timeout=CONFIG.timeout, exclude=[]):
"""Fill the node(s) with mode.
:param context: Everything we should know when filling node.
:param llm: Large Language Model with pre-defined system message.
:param schema: json/markdown, determine example and output format.
- raw: free form text
- json: it's easy to open source LLM with json format
- markdown: when generating code, markdown is always better
:param mode: auto/children/root
- auto: automated fill children's nodes and gather outputs, if no children, fill itself
- children: fill children's nodes and gather outputs
- root: fill root's node and gather output
:param strgy: simple/complex
- simple: run only once
- complex: run each node
:param timeout: Timeout for llm invocation.
:param exclude: The keys of ActionNode to exclude.
:return: self
"""
self.set_llm(llm)
self.set_context(context)
if self.schema:
schema = self.schema
if strgy == "simple":
return await self.simple_fill(schema=schema, mode=mode, timeout=timeout, exclude=exclude)
elif strgy == "complex":
# 这里隐式假设了拥有children
tmp = {}
for _, i in self.children.items():
if exclude and i.key in exclude:
continue
child = await i.simple_fill(schema=schema, mode=mode, timeout=timeout, exclude=exclude)
tmp.update(child.instruct_content.dict())
cls = self.create_children_class()
self.instruct_content = cls(**tmp)
return self