From de98cfcdc9e2e2e9d32553ac27460a48baa45847 Mon Sep 17 00:00:00 2001 From: Owen Colegrove Date: Sun, 23 Apr 2023 22:29:01 -0400 Subject: [PATCH] Move to builder paradigm --- automata/core/agents/automata_agent.py | 136 +++++++++++------- .../core/agents/tests/test_automata_agent.py | 17 ++- automata/core/tests/conftest.py | 22 +-- automata/scripts/docstring_cleanup.py | 56 +++----- automata/scripts/main_automata.py | 20 +-- .../python_indexer_tool_manager.py | 22 +-- .../python_writer_tool_manager.py | 23 +-- 7 files changed, 161 insertions(+), 135 deletions(-) diff --git a/automata/core/agents/automata_agent.py b/automata/core/agents/automata_agent.py index 63aed65b..0e75dc98 100644 --- a/automata/core/agents/automata_agent.py +++ b/automata/core/agents/automata_agent.py @@ -6,28 +6,22 @@ Example: - inputs = {"model": args.model} llm_toolkits = load_llm_toolkits(tools_list, **inputs) initial_payload = { "overview": python_inexer.get_overview(), } - logger.info("Passing in instructions: %s", args.instructions) - logger.info("-" * 100) - agent = AutomataAgent( - initial_payload=initial_payload, - instructions=args.instructions, - llm_toolkits=llm_toolkits, - version=args.version, - model=args.model, - session_id=args.session_id, - stream=args.stream, - ) + agent = (AutomataAgentBuilder() + .with_initial_payload(initial_payload) + .with_llm_toolkits(llm_toolkits) + .with_instructions(instructions) + .with_model(model) + .build()) + + agent.run() - next_instruction = agent.iter_task(instructions) - ... - TODO - Add error checking to ensure that we don't terminate when + TODO - Add error checking to ensure that we don't terminate when our previous result returned an error """ import logging @@ -47,6 +41,55 @@ logger = logging.getLogger(__name__) +class AutomataAgentBuilder: + def __init__(self): + self._instance = AutomataAgent() + + def with_initial_payload(self, initial_payload: Dict[str, str]): + self._instance.initial_payload = initial_payload + return self + + def with_llm_toolkits(self, llm_toolkits: Dict[ToolkitType, Toolkit]): + self._instance.llm_toolkits = llm_toolkits + return self + + def with_instructions(self, instructions: str): + self._instance.instructions = instructions + return self + + def with_version(self, version: AgentConfig): + self._instance.version = version + return self + + def with_model(self, model: str): + self._instance.model = model + return self + + def with_stream(self, stream: bool): + self._instance.stream = stream + return self + + def with_verbose(self, verbose: bool): + self._instance.verbose = verbose + return self + + def with_max_iters(self, max_iters: int): + self._instance.max_iters = max_iters + return self + + def with_temperature(self, temperature: float): + self._instance.temperature = temperature + return self + + def with_session_id(self, session_id: Optional[str]): + self._instance.session_id = session_id + return self + + def build(self): + self._instance._setup() + return self._instance + + class AutomataAgent: """ AutomataAgent is an autonomous agent that performs the actual work of the Automata @@ -57,19 +100,7 @@ class AutomataAgent: CONTINUE_MESSAGE = "Continue, and return a result JSON when finished." NUM_DEFAULT_MESSAGES = 3 - def __init__( - self, - initial_payload: Dict[str, str], - llm_toolkits: Dict[ToolkitType, Toolkit], - instructions: str, - version: AgentConfig = AgentConfig.AUTOMATA_MASTER_V2, - model: str = "gpt-4", - stream: bool = False, - verbose: bool = True, - max_iters: int = 1_000_000, # default to ~infinite iterations - temperature: float = 0.7, - session_id: Optional[str] = None, - ): + def __init__(self): """ Args: initial_payload (Dict[str, str]): Initial payload to send to the agent. @@ -88,34 +119,39 @@ def __init__( replay_messages() -> List[Dict[str, str]]: Replays agent messages buffer. """ - + self.initial_payload = {} + self.llm_toolkits = {} + self.instructions = "" + self.version = AgentConfig.AUTOMATA_MASTER_V2 + self.model = "gpt-4" + self.stream = False + self.verbose = True + self.max_iters = 1_000_000 + self.temperature = 0.7 + self.session_id = None + + def _setup(self): + # Put the setup logic here that was originally in the __init__ method # Initialize OpenAI API Key openai.api_key = OPENAI_API_KEY # noqa F405 # Initialize state variables - self.model = model - self.version = version - self.toolkits = llm_toolkits - self.messages: List[Dict[str, str]] = [] - self.stream = stream - self.verbose = verbose - self.max_iters = max_iters - self.temperature = temperature - - initial_payload["tools"] = "".join( + self.messages = [] + self.tokenizer = GPT2Tokenizer.from_pretrained("gpt2") + + self.initial_payload["tools"] = "".join( [ f"\n{tool.name}: {tool.description}\n" - for toolkit in self.toolkits.values() + for toolkit in self.llm_toolkits.values() for tool in toolkit.tools ] ) - prompt = self._load_prompt(initial_payload) + prompt = self._load_prompt() self._init_database() - if session_id: - self.session_id = session_id + if self.session_id: self._load_previous_interactions() else: self.session_id = str(uuid.uuid4()) @@ -125,12 +161,12 @@ def __init__( "role": "assistant", "content": 'Thought: I will begin by initializing myself. {"tool": "automata-initializer", "input": "Hello, I am Automata, OpenAI\'s most skilled coding system. How may I assit you today?"}', }, - {"role": "user", "content": f'Observation:\n{{"task_0":"{instructions}"}}'}, + {"role": "user", "content": f'Observation:\n{{"task_0":"{self.instructions}"}}'}, ] for message in initial_messages: self._save_interaction(message) - self.tokenizer = GPT2Tokenizer.from_pretrained("gpt2") + if self.verbose: logger.info("Initializing with Prompt:%s\n" % (prompt)) logger.info("-" * 100) @@ -210,9 +246,6 @@ def run(self) -> str: while True: self.iter_task() - - # if AutomataAgent.is_completion_message(self.messages[-1]["content"]): - # return self.messages[-1]["content"] # Check the previous previous agent message to see if it is a completion message if AutomataAgent.is_completion_message(self.messages[-2]["content"]): return self.messages[-2]["content"] @@ -238,13 +271,14 @@ def replay_messages(self) -> str: return "No completion message found." def extend_last_instructions(self, new_message: str) -> None: + """Extend the last instructions with a new message.""" previous_message = self.messages[-1] self.messages[-1] = { "role": previous_message["role"], "content": f"{previous_message}\n{new_message}", } - def _load_prompt(self, initial_payload: Dict[str, str]) -> str: + def _load_prompt(self) -> str: """Load the prompt from a config specified at initialization.""" with open( format_config_path("agent_configs", f"{self.version.value}.yaml"), @@ -254,7 +288,7 @@ def _load_prompt(self, initial_payload: Dict[str, str]) -> str: prompt = loaded_yaml["template"] if loaded_yaml["input_variables"] is not None: for arg in loaded_yaml["input_variables"]: - prompt = prompt.replace(f"{{{arg}}}", initial_payload[arg]) + prompt = prompt.replace(f"{{{arg}}}", self.initial_payload[arg]) return prompt def _process_input(self, response_text: str): @@ -276,7 +310,7 @@ def _process_input(self, response_text: str): outputs.append(requested_tool_input) else: tool_found = False - for toolkit in self.toolkits.values(): + for toolkit in self.llm_toolkits.values(): for tool in toolkit.tools: if tool.name == requested_tool: tool_output = tool.run(requested_tool_input, verbose=False) diff --git a/automata/core/agents/tests/test_automata_agent.py b/automata/core/agents/tests/test_automata_agent.py index bc1813a5..9fe8e943 100644 --- a/automata/core/agents/tests/test_automata_agent.py +++ b/automata/core/agents/tests/test_automata_agent.py @@ -1,7 +1,7 @@ import pytest from automata.core import load_llm_toolkits -from automata.core.agents.automata_agent import AutomataAgent +from automata.core.agents.automata_agent import AutomataAgentBuilder from automata.core.utils import root_py_path from automata.tools.python_tools.python_indexer import PythonIndexer @@ -19,10 +19,15 @@ def automata_agent(): "overview": overview, } - agent = AutomataAgent( - initial_payload=initial_payload, - instructions="Test instruction.", - llm_toolkits=mock_llm_toolkits, + # initial_payload=initial_payload, + instructions = "Test instruction." + + agent = ( + AutomataAgentBuilder() + .with_initial_payload(initial_payload) + .with_instructions(instructions) + .with_llm_toolkits(mock_llm_toolkits) + .build() ) return agent @@ -31,7 +36,7 @@ def test_automata_agent_init(automata_agent): assert automata_agent is not None assert automata_agent.model == "gpt-4" assert automata_agent.session_id is not None - assert len(automata_agent.toolkits.keys()) > 0 + assert len(automata_agent.llm_toolkits.keys()) > 0 def test_automata_agent_iter_task( diff --git a/automata/core/tests/conftest.py b/automata/core/tests/conftest.py index 5ea50d90..500eac45 100644 --- a/automata/core/tests/conftest.py +++ b/automata/core/tests/conftest.py @@ -7,7 +7,7 @@ from automata.configs.agent_configs import AgentConfig from automata.core import load_llm_toolkits -from automata.core.agents.automata_agent import AutomataAgent +from automata.core.agents.automata_agent import AutomataAgentBuilder from automata.core.utils import root_py_path from automata.tools.python_tools.python_indexer import PythonIndexer @@ -23,15 +23,17 @@ def build_agent_with_params( model="gpt-3.5-turbo", ): initial_payload, mock_llm_toolkits = automata_params - agent = AutomataAgent( - initial_payload=initial_payload, - instructions=instructions, - llm_toolkits=mock_llm_toolkits, - verbose=True, - max_iters=max_iters, - version=version, - temperature=temperature, - model=model, + agent = ( + AutomataAgentBuilder() + .with_initial_payload(initial_payload) + .with_instructions(instructions) + .with_llm_toolkits(mock_llm_toolkits) + .with_verbose(True) + .with_max_iters(max_iters) + .with_version(version) + .with_temperature(temperature) + .with_model(model) + .build() ) return agent diff --git a/automata/scripts/docstring_cleanup.py b/automata/scripts/docstring_cleanup.py index ebae1f05..418bd40f 100644 --- a/automata/scripts/docstring_cleanup.py +++ b/automata/scripts/docstring_cleanup.py @@ -6,7 +6,7 @@ from automata.configs.agent_configs import AgentConfig from automata.core import load_llm_toolkits -from automata.core.agents.automata_agent import AutomataAgent +from automata.core.agents.automata_agent import AutomataAgentBuilder from automata.core.utils import get_logging_config, root_py_path from automata.tools.python_tools import PythonIndexer @@ -40,48 +40,26 @@ def update_docstrings(): raw_code = file.read() logger.info("Initial Code:\n%s" % (raw_code)) - agent = AutomataAgent( - initial_payload={}, - instructions=f"BEFORE WRITING CODE, begin with a multi-bullet description of what the following file is responsible for:\n {raw_code}" - f" FOLLOWING that, your objective is to increase the modularity, readability, and maintainability of the file." - f" BE SURE to include module docstrings, and docstrings for every function, class, method, and function." - f" You may introduce comments where applicable," - f" and rename functions and variables if it significantly improves the code." - f" NOTE the key changes you have made. The code may be written dirrectly in your prompt and a developer will paste it into the file.", - llm_toolkits=llm_toolkits, - version=AgentConfig(args.version), - model="gpt-4", - stream=True, - verbose=True, + agent = ( + AutomataAgentBuilder() + .with_instructions( + f"BEFORE WRITING CODE, begin with a multi-bullet description of what the following file is responsible for:\n {raw_code}" + f" FOLLOWING that, your objective is to increase the modularity, readability, and maintainability of the file." + f" BE SURE to include module docstrings, and docstrings for every function, class, method, and function." + f" You may introduce comments where applicable," + f" and rename functions and variables if it significantly improves the code." + f" NOTE the key changes you have made. The code may be written dirrectly in your prompt and a developer will paste it into the file." + ) + .with_llm_toolkits(llm_toolkits) + .with_version(AgentConfig(args.version)) + .with_model("gpt-4") + .with_stream(True) + .with_verbose(True) + .build() ) agent.run() break - # for module_path in python_indexer.module - - # (python_indexer, _) = (tool_payload["python_indexer"], tool_payload["python_writer"]) - # overview = python_indexer.get_overview() - # initial_payload = {"overview": overview} - # for function in python_parser.function_dict.values(): - # path = function.py_path - # raw_code = python_parser.get_raw_code(path) - # docstring = python_parser.get_docstring(path) - # if "No results found." not in docstring: - # continue - # logger.info("Prev Docstring:\n%s" % docstring) - # logger.info("Prev Raw Code:\n%s" % raw_code) - # instructions = f"The following code is located at the path {path}:\n\n{raw_code}\n\nPlease fetch the code from the raw file, then write relevant docstrings for this piece of code, and lastly, use the python-writer to write the result to disk." - # agent = AutomataAgent( - # initial_payload=initial_payload, - # instructions=instructions, - # llm_toolkits=exec_tools, - # version=args.version, - # model=args.model, - # session_id=args.session_id, - # stream=args.stream, - # verbose=True, - # ) - # agent.run() if __name__ == "__main__": diff --git a/automata/scripts/main_automata.py b/automata/scripts/main_automata.py index 4fa5b6eb..b7865b78 100644 --- a/automata/scripts/main_automata.py +++ b/automata/scripts/main_automata.py @@ -5,7 +5,7 @@ from automata.configs.agent_configs import AgentConfig from automata.core import Toolkit, ToolkitType, load_llm_toolkits -from automata.core.agents.automata_agent import AutomataAgent +from automata.core.agents.automata_agent import AutomataAgentBuilder from automata.core.utils import get_logging_config, root_py_path from automata.tools.python_tools.python_indexer import PythonIndexer @@ -64,14 +64,16 @@ def main(): } logger.info("Passing in instructions: %s", args.instructions) logger.info("-" * 100) - agent = AutomataAgent( - initial_payload=initial_payload, - instructions=args.instructions, - llm_toolkits=llm_toolkits, - version=args.version, - model=args.model, - session_id=args.session_id, - stream=args.stream, + agent = ( + AutomataAgentBuilder() + .with_initial_payload(initial_payload) + .with_instructions(args.instructions) + .with_llm_toolkits(llm_toolkits) + .with_version(args.version) + .with_model(args.model) + .with_session_id(args.session_id) + .with_stream(args.stream) + .build() ) logger.info("Running the agent now...") diff --git a/automata/tools/tool_management/python_indexer_tool_manager.py b/automata/tools/tool_management/python_indexer_tool_manager.py index 373385be..7dc32694 100644 --- a/automata/tools/tool_management/python_indexer_tool_manager.py +++ b/automata/tools/tool_management/python_indexer_tool_manager.py @@ -128,20 +128,22 @@ def _run_indexer_retrieve_raw_code(self, input_str: str) -> str: def _run_automata_indexer_retrieve_code(self, path_str: str) -> str: """Automata retrieves the code of the python package, module, standalone function, class, or method at the given python path, without docstrings.""" from automata.core import load_llm_toolkits - from automata.core.agents.automata_agent import AutomataAgent + from automata.core.agents.automata_agent import AutomataAgentBuilder try: initial_payload = {"overview": self.indexer.get_overview()} instructions = f"Retrieve the code for {path_str}" - agent = AutomataAgent( - initial_payload=initial_payload, - instructions=instructions, - llm_toolkits=load_llm_toolkits(["python_indexer"]), - version=self.automata_version, - model=self.model, - stream=self.stream, - verbose=self.verbose, - temperature=self.temperature, + agent = ( + AutomataAgentBuilder() + .with_initial_payload(initial_payload) + .with_instructions(instructions) + .with_llm_toolkits(load_llm_toolkits(["python_indexer"])) + .with_version(self.automata_version) + .with_model(self.model) + .with_stream(self.stream) + .with_verbose(self.verbose) + .with_temperature(self.temperature) + .build() ) result = agent.run() result = clean_agent_result(result) diff --git a/automata/tools/tool_management/python_writer_tool_manager.py b/automata/tools/tool_management/python_writer_tool_manager.py index 3f2dd8c3..de0b721d 100644 --- a/automata/tools/tool_management/python_writer_tool_manager.py +++ b/automata/tools/tool_management/python_writer_tool_manager.py @@ -105,22 +105,25 @@ def _writer_update_module(self, input_str: str) -> str: def _automata_update_module(self, input_str: str) -> str: """Creates an AutomataAgent to write the given task.""" from automata.core import load_llm_toolkits - from automata.core.agents.automata_agent import AutomataAgent + from automata.core.agents.automata_agent import AutomataAgentBuilder try: initial_payload = { "overview": self.writer.indexer.get_overview(), } - agent = AutomataAgent( - initial_payload=initial_payload, - instructions=input_str, - llm_toolkits=load_llm_toolkits(["python_writer"]), - version=self.automata_version, - model=self.model, - stream=self.stream, - verbose=self.verbose, - temperature=self.temperature, + agent = ( + AutomataAgentBuilder() + .with_initial_payload(initial_payload) + .with_instructions(input_str) + .with_llm_toolkits(load_llm_toolkits(["python_writer"])) + .with_version(self.automata_version) + .with_model(self.model) + .with_stream(self.stream) + .with_verbose(self.verbose) + .with_temperature(self.temperature) + .build() ) + agent.run() return "Success"