Create agent.py

agent.py

@@ -1,284 +1,131 @@
-import os
-import subprocess
-
-import openai
-
-from agent.prompts import (
-    ACTION_PROMPT,
-    ADD_PROMPT,
-    COMPRESS_HISTORY_PROMPT,
-    LOG_PROMPT,
-    LOG_RESPONSE,
-    MODIFY_PROMPT,
-    PREFIX,
-    READ_PROMPT,
-    TASK_PROMPT,
-    UNDERSTAND_TEST_RESULTS_PROMPT,
-)
-from agent.utils import parse_action, parse_file_content, read_python_module_structure
-
-VERBOSE = False
-MAX_HISTORY = 100
-MODEL = "gpt-3.5-turbo"  # "gpt-4"
-
-
-def run_gpt(
-    prompt_template,
-    stop_tokens,
-    max_tokens,
-    module_summary,
-    purpose,
-    **prompt_kwargs,
-):
-    content = PREFIX.format(
-        module_summary=module_summary,
-        purpose=purpose,
-    ) + prompt_template.format(**prompt_kwargs)
-    if VERBOSE:
-        print(LOG_PROMPT.format(content))
-    resp = openai.ChatCompletion.create(
-        model=MODEL,
-        messages=[
-            {"role": "system", "content": content},
-        ],
-        temperature=0.0,
-        max_tokens=max_tokens,
-        stop=stop_tokens if stop_tokens else None,
-    )["choices"][0]["message"]["content"]
-    if VERBOSE:
-        print(LOG_RESPONSE.format(resp))
-    return resp
-
-
-def compress_history(purpose, task, history, directory):
-    module_summary, _, _ = read_python_module_structure(directory)
-    resp = run_gpt(
-        COMPRESS_HISTORY_PROMPT,
-        stop_tokens=["observation:", "task:", "action:", "thought:"],
-        max_tokens=512,
-        module_summary=module_summary,
-        purpose=purpose,
-        task=task,
-        history=history,
-    )
-    history = "observation: {}\n".format(resp)
-    return history
-
-
-def call_main(purpose, task, history, directory, action_input):
-    module_summary, _, _ = read_python_module_structure(directory)
-    resp = run_gpt(
-        ACTION_PROMPT,
-        stop_tokens=["observation:", "task:"],
-        max_tokens=256,
-        module_summary=module_summary,
-        purpose=purpose,
-        task=task,
-        history=history,
-    )
-    lines = resp.strip().strip("\n").split("\n")
-    for line in lines:
-        if line == "":
-            continue
-        if line.startswith("thought: "):
-            history += "{}\n".format(line)
-        elif line.startswith("action: "):
-            action_name, action_input = parse_action(line)
-            history += "{}\n".format(line)
-            return action_name, action_input, history, task
-        else:
-            assert False, "unknown action: {}".format(line)
-    return "MAIN", None, history, task
-
-
-def call_test(purpose, task, history, directory, action_input):
-    result = subprocess.run(
-        ["python", "-m", "pytest", "--collect-only", directory],
-        capture_output=True,
-        text=True,
-    )
-    if result.returncode != 0:
-        history += "observation: there are no tests! Test should be written in a test folder under {}\n".format(
-            directory
-        )
-        return "MAIN", None, history, task
-    result = subprocess.run(
-        ["python", "-m", "pytest", directory], capture_output=True, text=True
-    )
-    if result.returncode == 0:
-        history += "observation: tests pass\n"
-        return "MAIN", None, history, task
-    module_summary, content, _ = read_python_module_structure(directory)
-    resp = run_gpt(
-        UNDERSTAND_TEST_RESULTS_PROMPT,
-        stop_tokens=[],
-        max_tokens=256,
-        module_summary=module_summary,
-        purpose=purpose,
-        task=task,
-        history=history,
-        stdout=result.stdout[:5000],  # limit amount of text
-        stderr=result.stderr[:5000],  # limit amount of text
-    )
-    history += "observation: tests failed: {}\n".format(resp)
-    return "MAIN", None, history, task
-
-
-def call_set_task(purpose, task, history, directory, action_input):
-    module_summary, content, _ = read_python_module_structure(directory)
-    task = run_gpt(
-        TASK_PROMPT,
-        stop_tokens=[],
-        max_tokens=64,
-        module_summary=module_summary,
-        purpose=purpose,
-        task=task,
-        history=history,
-    ).strip("\n")
-    history += "observation: task has been updated to: {}\n".format(task)
-    return "MAIN", None, history, task
-
-
-def call_read(purpose, task, history, directory, action_input):
-    if not os.path.exists(action_input):
-        history += "observation: file does not exist\n"
-        return "MAIN", None, history, task
-    module_summary, content, _ = read_python_module_structure(directory)
-    f_content = (
-        content[action_input] if content[action_input] else "< document is empty >"
-    )
-    resp = run_gpt(
-        READ_PROMPT,
-        stop_tokens=[],
-        max_tokens=256,
-        module_summary=module_summary,
-        purpose=purpose,
-        task=task,
-        history=history,
-        file_path=action_input,
-        file_contents=f_content,
-    ).strip("\n")
-    history += "observation: {}\n".format(resp)
-    return "MAIN", None, history, task
-
-
-def call_modify(purpose, task, history, directory, action_input):
-    if not os.path.exists(action_input):
-        history += "observation: file does not exist\n"
-        return "MAIN", None, history, task
-    (
-        module_summary,
-        content,
-        _,
-    ) = read_python_module_structure(directory)
-    f_content = (
-        content[action_input] if content[action_input] else "< document is empty >"
-    )
-    resp = run_gpt(
-        MODIFY_PROMPT,
-        stop_tokens=["action:", "thought:", "observation:"],
-        max_tokens=2048,
-        module_summary=module_summary,
-        purpose=purpose,
-        task=task,
-        history=history,
-        file_path=action_input,
-        file_contents=f_content,
-    )
-    new_contents, description = parse_file_content(resp)
-    if new_contents is None:
-        history += "observation: failed to modify file\n"
-        return "MAIN", None, history, task
-
-    with open(action_input, "w") as f:
-        f.write(new_contents)
-
-    history += "observation: file successfully modified\n"
-    history += "observation: {}\n".format(description)
-    return "MAIN", None, history, task
-
-
-def call_add(purpose, task, history, directory, action_input):
-    d = os.path.dirname(action_input)
-    if not d.startswith(directory):
-        history += "observation: files must be under directory {}\n".format(directory)
-    elif not action_input.endswith(".py"):
-        history += "observation: can only write .py files\n"
-    else:
-        if d and not os.path.exists(d):
-            os.makedirs(d)
-        if not os.path.exists(action_input):
-            module_summary, _, _ = read_python_module_structure(directory)
-            resp = run_gpt(
-                ADD_PROMPT,
-                stop_tokens=["action:", "thought:", "observation:"],
-                max_tokens=2048,
-                module_summary=module_summary,
-                purpose=purpose,
-                task=task,
-                history=history,
-                file_path=action_input,
-            )
-            new_contents, description = parse_file_content(resp)
-            if new_contents is None:
-                history += "observation: failed to write file\n"
-                return "MAIN", None, history, task
-
-            with open(action_input, "w") as f:
-                f.write(new_contents)
-
-            history += "observation: file successfully written\n"
-            history += "obsertation: {}\n".format(description)
-        else:
-            history += "observation: file already exists\n"
-    return "MAIN", None, history, task
-
-
-NAME_TO_FUNC = {
-    "MAIN": call_main,
-    "UPDATE-TASK": call_set_task,
-    "MODIFY-FILE": call_modify,
-    "READ-FILE": call_read,
-    "ADD-FILE": call_add,
-    "TEST": call_test,
-}
-
-
-def run_action(purpose, task, history, directory, action_name, action_input):
-    if action_name == "COMPLETE":
-        exit(0)
-
-    # compress the history when it is long
-    if len(history.split("\n")) > MAX_HISTORY:
-        if VERBOSE:
-            print("COMPRESSING HISTORY")
-        history = compress_history(purpose, task, history, directory)
-
-    assert action_name in NAME_TO_FUNC
-
-    print("RUN: ", action_name, action_input)
-    return NAME_TO_FUNC[action_name](purpose, task, history, directory, action_input)
-
-
-def run(purpose, directory, task=None):
-    history = ""
-    action_name = "UPDATE-TASK" if task is None else "MAIN"
-    action_input = None
-    while True:
-        print("")
-        print("")
-        print("---")
-        print("purpose:", purpose)
-        print("task:", task)
-        print("---")
-        print(history)
-        print("---")
-
-        action_name, action_input, history, task = run_action(
-            purpose,
-            task,
-            history,
-            directory,
-            action_name,
-            action_input,
-        )
+PREFIX = """You are an expert online researcher
+You are working on the task outlined here
+Purpose:
+{purpose}
+"""
+
+ACTION_PROMPT = """
+You have access to the following tools:
+- action: UPDATE-TASK action_input=NEW_TASK
+- action: SEARCH action_input=PAGE_URL
+- action: COMPLETE
+Instructions
+- Write a response satisfying the request
+- Set action: SEARCH action_input=URL of the page you need to search
+- Complete the current task as best you can
+- When the task is complete, update the task
+- SEARCH to verify that information you are providing
+Use the following format:
+task: the input task you must complete
+thought: you should always think about what to do
+action: the action to take (should be one of [UPDATE-TASK, SEARCH, COMPLETE]) action_input=XXX
+observation: the result of the action
+thought: you should always think after an observation
+action: always perform a SEARCH to verify your observation
+... (thought/action/observation/thought can repeat N times)
+You are attempting to complete the task
+task: {task}
+{history}"""
+
+SEARCH_QUERY = """
+You are attempting to complete the task
+task: {task}
+Progress:
+{history}
+What is the URL for the webpage that I need to search for information on?
+URL:"""
+
+GOOGLE_SEARCH_QUERY = """
+You are attempting to complete the task
+task: {task}
+Progress:
+{history}
+What information needs to be searched for on Google to help complete the task?
+Query:"""
+
+
+TASK_PROMPT = """
+You are attempting to complete the task
+task: {task}
+Progress:
+{history}
+Tasks should be small, isolated, and independent
+What should the task be for us to achieve the purpose?
+task: """
+
+READ_PROMPT = """
+You are attempting to complete the task
+task: {task}
+Progress:
+{history}
+{file_path}
+---
+{file_contents}
+---
+Return your thoughts about the file relevant to completing the task (in a paragraph)
+Mention any specific functions, arguments, or details needed
+"""
+
+ADD_PROMPT = """
+You are attempting to complete the task
+task: {task}
+Progress:
+{history}
+Write a new file called {file_path} with contents between ---
+After the contents write a paragraph on what was inserted with details
+"""
+
+MODIFY_PROMPT = """
+You are attempting to complete the task
+task: {task}
+Progress:
+{history}
+{file_path}
+---
+{file_contents}
+---
+Return the complete modified {file_path} contents between ---
+After the contents write a paragraph on what was changed with details
+"""
+
+
+UNDERSTAND_TEST_RESULTS_PROMPT = """
+You are attempting to complete the task
+task: {task}
+Progress:
+{history}
+Test results:
+STDOUT
+---
+{stdout}
+---
+STDERR
+---
+{stderr}
+---
+Describe why the tests failed and how to fix them (in a paragraph)
+"""
+
+
+COMPRESS_HISTORY_PROMPT = """
+You are attempting to complete the task
+task: {task}
+Progress:
+{history}
+Compress the timeline of progress above into a single summary (as a paragraph)
+Include all important milestones, the current challenges, and implementation details necessary to proceed
+"""
+
+LOG_PROMPT = """
+PROMPT
+**************************************
+{}
+**************************************
+"""
+
+LOG_RESPONSE = """
+RESPONSE
+**************************************
+{}
+**************************************
+"""