Update agent.py

agent.py

@@ -1,207 +1,284 @@
-from typing import List, Dict, Optional
-from custom_types import (
-    Code,
-    Prompt,
-    AppType,
-    File,
-    Space,
-    Tutorial,
-    App,
-    WebApp,
-    GradioApp,
-    StreamlitApp,
-    ReactApp,
-    Code,
-)
-from prompts import (
-    createLlamaPrompt,
-    createSpace,
-    isPythonOrGradioAppPrompt,
-    isReactAppPrompt,
-    isStreamlitAppPrompt,
-    getWebApp,
-    getGradioApp,
-    getReactApp,
-    getStreamlitApp,
-    parseTutorial,
-    generateFiles,
+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 huggingface_hub import InferenceClient
+from agent.utils import parse_action, parse_file_content, read_python_module_structure
 
-class Agent:
-    def __init__(self, prompts: Dict[str, any]):
-        self.prompts = prompts
-        self.client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
+VERBOSE = False
+MAX_HISTORY = 100
+MODEL = "gpt-3.5-turbo"  # "gpt-4"
 
-    def process(self, user_input: str) -> str:
-        """ Processes the user's input and generates code. """
-        # Parse the user's input
-        app_type, app_name, app_description, app_features, app_dependencies, app_space, app_tutorial = self.parse_input(user_input)
 
-        # Generate a prompt for the Llama model
-        prompt = self.prompts["createLlamaPrompt"](
-            app_type, app_name, app_description, app_features, app_dependencies, app_space, app_tutorial
-        )
+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
 
-        # Generate code using the Llama model
-        code = self.generate_code(prompt)
 
-        # Generate files for the application
-        files = self.prompts["generateFiles"](
-            app_type, app_name, app_description, app_features, app_dependencies, app_space, app_tutorial
-        )
+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
 
-        # Return the generated code and files
-        return f"Code: {code}\nFiles: {files}"
-
-    def parse_input(self, user_input: str) -> tuple:
-        """ Parses the user's input and extracts the relevant information. """
-        # Extract the app type
-        app_type = self.extract_app_type(user_input)
-        # Extract the app name
-        app_name = self.extract_app_name(user_input)
-        # Extract the app description
-        app_description = self.extract_app_description(user_input)
-        # Extract the app features
-        app_features = self.extract_app_features(user_input)
-        # Extract the app dependencies
-        app_dependencies = self.extract_app_dependencies(user_input)
-        # Extract the app space
-        app_space = self.extract_app_space(user_input)
-        # Extract the app tutorial
-        app_tutorial = self.extract_app_tutorial(user_input)
-
-        return app_type, app_name, app_description, app_features, app_dependencies, app_space, app_tutorial
-
-    def extract_app_type(self, user_input: str) -> AppType:
-        """ Extracts the app type from the user's input. """
-        # Check if the user specified a specific app type
-        if "web app" in user_input:
-            return AppType.WEB_APP
-        elif "gradio app" in user_input:
-            return AppType.GRADIO_APP
-        elif "streamlit app" in user_input:
-            return AppType.STREAMLIT_APP
-        elif "react app" in user_input:
-            return AppType.REACT_APP
-        # Otherwise, assume the user wants a web app
-        return AppType.WEB_APP
-
-    def extract_app_name(self, user_input: str) -> str:
-        """ Extracts the app name from the user's input. """
-        # Find the substring "app name is:"
-        start_index = user_input.find("app name is:") + len("app name is:")
-        # Find the end of the app name
-        end_index = user_input.find(".", start_index)
-        # Extract the app name
-        app_name = user_input[start_index:end_index].strip()
-        return app_name
-
-    def extract_app_description(self, user_input: str) -> str:
-        """ Extracts the app description from the user's input. """
-        # Find the substring "app description is:"
-        start_index = user_input.find("app description is:") + len("app description is:")
-        # Find the end of the app description
-        end_index = user_input.find(".", start_index)
-        # Extract the app description
-        app_description = user_input[start_index:end_index].strip()
-        return app_description
-
-    def extract_app_features(self, user_input: str) -> List[str]:
-        """ Extracts the app features from the user's input. """
-        # Find the substring "app features are:"
-        start_index = user_input.find("app features are:") + len("app features are:")
-        # Find the end of the app features
-        end_index = user_input.find(".", start_index)
-        # Extract the app features
-        app_features_str = user_input[start_index:end_index].strip()
-        # Split the app features string into a list
-        app_features = app_features_str.split(", ")
-        return app_features
-
-    def extract_app_dependencies(self, user_input: str) -> List[str]:
-        """ Extracts the app dependencies from the user's input. """
-        # Find the substring "app dependencies are:"
-        start_index = user_input.find("app dependencies are:") + len("app dependencies are:")
-        # Find the end of the app dependencies
-        end_index = user_input.find(".", start_index)
-        # Extract the app dependencies
-        app_dependencies_str = user_input[start_index:end_index].strip()
-        # Split the app dependencies string into a list
-        app_dependencies = app_dependencies_str.split(", ")
-        return app_dependencies
-
-    def extract_app_space(self, user_input: str) -> Optional[Space]:
-        """ Extracts the app space from the user's input. """
-        # Find the substring "app space is:"
-        start_index = user_input.find("app space is:") + len("app space is:")
-        # Find the end of the app space
-        end_index = user_input.find(".", start_index)
-        # Extract the app space
-        app_space_str = user_input[start_index:end_index].strip()
-        # Create a Space object
-        app_space = Space(space=app_space_str)
-        return app_space
-
-    def extract_app_tutorial(self, user_input: str) -> Optional[Tutorial]:
-        """ Extracts the app tutorial from the user's input. """
-        # Find the substring "app tutorial is:"
-        start_index = user_input.find("app tutorial is:") + len("app tutorial is:")
-        # Find the end of the app tutorial
-        end_index = user_input.find(".", start_index)
-        # Extract the app tutorial
-        app_tutorial_str = user_input[start_index:end_index].strip()
-        # Create a Tutorial object
-        app_tutorial = Tutorial(tutorial=app_tutorial_str)
-        return app_tutorial
-
-    def generate_code(self, prompt: Prompt) -> Code:
-        """ Generates code using the Llama model. """
-        # Send the prompt to the Llama model
-        response = self.client(prompt.prompt)
-        # Extract the generated code
-        code = response["generated_text"]
-        code = code.replace("```", "")
-        code = code.replace("```", "")
-        # Create a Code object
-        code = Code(code=code, language="python")
-        return code
-
-    def generate_files(self, app_type: AppType, app_name: str, app_description: str, app_features: List[str], app_dependencies: List[str], app_space: Optional[Space] = None, app_tutorial: Optional[Tutorial] = None) -> List[File]:
-        """ Generates files for the application. """
-        # Generate files based on the app type
-        files = self.prompts["generateFiles"](
-            app_type, app_name, app_description, app_features, app_dependencies, app_space, app_tutorial
+
+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 files
-
-def main():
-    """ Main function for the application. """
-    # Create an agent
-    agent = Agent(
-        prompts={
-            "createLlamaPrompt": createLlamaPrompt,
-            "createSpace": createSpace,
-            "isPythonOrGradioAppPrompt": isPythonOrGradioAppPrompt,
-            "isReactAppPrompt": isReactAppPrompt,
-            "isStreamlitAppPrompt": isStreamlitAppPrompt,
-            "getWebApp": getWebApp,
-            "getGradioApp": getGradioApp,
-            "getReactApp": getReactApp,
-            "getStreamlitApp": getStreamlitApp,
-            "parseTutorial": parseTutorial,
-            "generateFiles": generateFiles,
-        }
+        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
 
-    # Get user input
-    user_input = input("Enter your request: ")
 
-    # Process the user's input
-    response = agent.process(user_input)
+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
+
 
-    # Print the response
-    print(response)
+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,
+}
 
-if __name__ == "__main__":
-    main()
+
+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,
+        )