File size: 9,944 Bytes
962d35d 343bc6e 962d35d 343bc6e 962d35d 343bc6e 962d35d 343bc6e 962d35d 343bc6e 962d35d 343bc6e 962d35d 343bc6e 962d35d |
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 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 |
import json
from copy import deepcopy
from typing import Any, Dict, List
from flow_modules.aiflows.ChatFlowModule import ChatAtomicFlow
from dataclasses import dataclass
@dataclass
class Command:
name: str
description: str
input_args: List[str]
class Controller_JarvisFlow(ChatAtomicFlow):
"""This class is a controller for JarvisFlow, it takes the plan generated by the planner, logs of previous executions,
depending on the initial goal or the subsequent feedback from the branching executors (and the human), to decide which
executor to call next (or to exit by calling finish).
*Configuration Parameters*:
- `commands` (dict): a dictionary of commands that the controller can call, each command has a name, a description, and a list of input arguments.
The commands will be injected into the system message prompt template.
- `system_message_prompt_template` (str): the template for the system message prompt, there are several components needs to be injected into the
template, including the commands, plan, plan_file_location, logs, and the goal. The injection of commands is done then initalizing the flow,
the rest of the components are injected at the beginning of each run.
- `previous_messages` (int): a sliding window of previous messages that will be passed to the model. This is the central part of short-term memory management.
*Input Interface Non Initialized*:
- `goal` (str): the initial goal of the conversation, this is the input to the model.
- `memory_files` (dict): a dictionary of file locations that contains the plan, logs.
- `plan` (str): the plan generated by the planner, the plan will change (marked as done, or re-plan) as execution preceeds.
- `logs` (str): the logs of previous executions, the logs will be appended as execution preceeds.
*Input Interface Initialized*:
- `result` (str): the result of the previous execution, this is the input to the model.
- `memory_files` (dict): a dictionary of file locations that contains the plan, logs.
- `plan` (str): the plan generated by the planner, the plan will change (marked as done, or re-plan) as execution preceeds.
- `logs` (str): the logs of previous executions, the logs will be appended as execution preceeds.
- `goal` (str): the initial goal, this is kept because the goal is also injected into the system prompts so that Jarvis does not
forget what the goal is, when the memory sliding window is implemented.
*Output Interface*:
- `command` (str): the command to be executed by the executor.
- `command_args` (dict): the arguments of the command to be executed by the executor.
"""
def __init__(
self,
commands: List[Command],
**kwargs):
"""Initialize the flow, inject the commands into the system message prompt template.
:param commands: a list of commands that the controller can call.
:type commands: List[Command]
:param kwargs: other parameters.
:type kwargs: Dict[str, Any]
"""
super().__init__(**kwargs)
self.system_message_prompt_template = self.system_message_prompt_template.partial(
commands=self._build_commands_manual(commands),
plan="no plans yet",
plan_file_location="no plan file location yet",
logs="no logs yet",
)
self.hint_for_model = """
Make sure your response is in the following format:
Response Format:
{
"command": "call one of the subordinates",
"command_args": {
"arg name": "value"
}
}
"""
def _get_content_file_location(self, input_data, content_name):
"""
Get the location of the file that contains the content: plan, logs, code_library
:param input_data: the input data to the flow
:type input_data: Dict[str, Any]
:param content_name: the name of the content
:type content_name: str
:raises AssertionError: if the content is not in the memory_files
:raises AssertionError: if memory_files is not passed to the flow
:return: the location of the file that contains the content
"""
# get the location of the file that contains the content: plan, logs, code_library
assert "memory_files" in input_data, "memory_files not passed to Jarvis/Controller"
assert content_name in input_data["memory_files"], f"{content_name} not in memory files"
return input_data["memory_files"][content_name]
def _get_content(self, input_data, content_name):
"""
Get the content of the file that contains the content: plan, logs, code_library
:param input_data: the input data to the flow
:type input_data: Dict[str, Any]
:param content_name: the name of the content
:type content_name: str
:raises AssertionError: if the content is not in the input_data
:return: the content of the file that contains the content
"""
# get the content of the file that contains the content: plan, logs, code_library
assert content_name in input_data, f"{content_name} not passed to Jarvis/Controller"
content = input_data[content_name]
if len(content) == 0:
content = f'No {content_name} yet'
return content
@staticmethod
def _build_commands_manual(commands: List[Command]) -> str:
"""Build the manual for the commands.
:param commands: a list of commands that the controller can call.
:type commands: List[Command]
:return: the manual for the commands.
:rtype: str
"""
ret = ""
for i, command in enumerate(commands):
command_input_json_schema = json.dumps(
{input_arg: f"YOUR_{input_arg.upper()}" for input_arg in command.input_args})
ret += f"{i + 1}. {command.name}: {command.description} Input arguments (given in the JSON schema): {command_input_json_schema}\n"
return ret
@classmethod
def instantiate_from_config(cls, config):
"""Setting up the flow from the config file. In particular, setting up the prompts, backend, and commands.
:param config: the config file.
:type config: Dict[str, Any]
:return: the instantiated flow.
:rtype: Controller_JarvisFlow
"""
flow_config = deepcopy(config)
kwargs = {"flow_config": flow_config}
# ~~~ Set up prompts ~~~
kwargs.update(cls._set_up_prompts(flow_config))
# ~~~Set up backend ~~~
kwargs.update(cls._set_up_backend(flow_config))
# ~~~ Set up commands ~~~
commands = flow_config["commands"]
commands = [
Command(name, command_conf["description"], command_conf["input_args"]) for name, command_conf in
commands.items()
]
kwargs.update({"commands": commands})
# ~~~ Instantiate flow ~~~
return cls(**kwargs)
def _update_prompts_and_input(self, input_data: Dict[str, Any]):
"""Hinting the model to output in json format, updating the plan, logs to the system prompts.
:param input_data: the input data to the flow.
:type input_data: Dict[str, Any]
"""
if 'goal' in input_data:
input_data['goal'] += self.hint_for_model
if 'result' in input_data:
input_data['result'] += self.hint_for_model
plan_file_location = self._get_content_file_location(input_data, "plan")
plan_content = self._get_content(input_data, "plan")
logs_content = self._get_content(input_data, "logs")
self.system_message_prompt_template = self.system_message_prompt_template.partial(
plan_file_location=plan_file_location,
plan=plan_content,
logs=logs_content
)
def run(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
"""Run the flow, update the system prompts, and run the model.
:param input_data: the input data to the flow.
:type input_data: Dict[str, Any]
:return: the output of the flow.
:rtype: Dict[str, Any]
"""
self._update_prompts_and_input(input_data)
# ~~~when conversation is initialized, append the updated system prompts to the chat history ~~~
if self._is_conversation_initialized():
updated_system_message_content = self._get_message(self.system_message_prompt_template, input_data)
self._state_update_add_chat_message(content=updated_system_message_content,
role=self.flow_config["system_name"])
# ~~~run the model, special mechanism to deal with situations where the output is not in json format. ~~~
while True:
api_output = super().run(input_data)["api_output"].strip()
try:
start = api_output.index("{")
end = api_output.rindex("}") + 1
json_str = api_output[start:end]
return json.loads(json_str)
except (ValueError, json.decoder.JSONDecodeError, json.JSONDecodeError):
updated_system_message_content = self._get_message(self.system_message_prompt_template, input_data)
self._state_update_add_chat_message(content=updated_system_message_content,
role=self.flow_config["system_name"])
new_goal = "The previous respond cannot be parsed with json.loads. Next time, do not provide any comments or code blocks. Make sure your next response is purely json parsable."
new_input_data = input_data.copy()
new_input_data['result'] = new_goal
input_data = new_input_data |