Merge pull request #6 from samshapley/feature/tools_

feature / Tools
samshapley · Sep 1, 2023 · 2ff0df3 · 2ff0df3
2 parents 1c3c570 + 8354ff0
commit 2ff0df3
Show file tree

Hide file tree

Showing 24 changed files with 314 additions and 32 deletions.
diff --git a/README.md b/README.md
@@ -25,6 +25,7 @@ To run OrchestrAI, you'll need:
 - OpenAI Python library  
 - networkx
 - PyYAML
+
 - wandb
 - matplotlib
 
@@ -107,9 +108,23 @@ pipeline:
 ```
 Once you have defined your pipeline, you can ensure it will be run by specifying the pipeline name in the config.yml file.
 
+### Tools
+*Please note, tools have just been added, default tools_enabled = False*
+*Tools currently decreases agent performance, as the system prompt becomes too large and confusing for the model to understand. This is being worked on.*
+*Tools definitely doesn't work for GPT-3.5, it's just not powerful enough to handle the extended system prompt.*
+
+OrchestrAI supports giving modules access to tools. A tool is an additional function the module can use during its operation. Currently, the only tool support is GENERATE_IMAGE. 
+
+The model is instructed to use tools when enabled via `tool_prompt.txt`. It will return a tag to activate the tool, which contains the JSON object required for the tool. For example, to generate an image with DALLE and store it to the file, the model will use the following tag.
+
+<b><@ { "tool_name": "GENERATE_IMAGE", "filename": choose , "prompt" : "Descriptive image prompt" } @></b>
+
+Create a new one within the `tools` folder, and add the selection logic to `tool_manager.py`. Multiple tools can be specified by the AI in the response, to be executed in order. I recommend trying tools with the `story_pipeline.yml` to see how they work in illustrating a story. 
+
 ### Running the Script
 
 Ensure that you've added your OpenAI API key to `config.yml`.
+You can toggle wandb logging, tools, and openai defaults in the config.yml file.
 
 To run OrchestrAI, execute `agent.py`:
 
@@ -200,13 +215,15 @@ The modularity of OrchestrAI makes it easy to add new AI capabilities as needed
 
 ## Logging with Weights and Biases
 
+![Screenshot 2023-09-01 at 12 23 50](https://github.com/samshapley/OrchestrAI/assets/93387313/682f1edc-19fb-48c4-b4fe-e1b76cd9f8cc)
+
 By default, interactions are logged using `log_action' during the process to a file created at the start of the agent. This file is then renamed and moved to the logs folder at the termination of the agent. This allows you to see the full history of the agent's interactions.
 
 However, we can leverage the power of Wandb Prompts (https://docs.wandb.ai/guides/prompts?_gl=1)
 
 Provided you've set up wandb, you can enable wandb logging in the config.yml file.
 
-This allows you to log the agent as a run, with modules as child runs of the chain. This allows you to see the full history of the agent's interactions, and the full history of each module's interactions. This is useful for debugging, and for understanding the agent's behaviour as you explore different pipelines and modules.
+This allows you to log the agent as a run, with modules as child runs of the chain. This allows you to see the full history of the agent's interactions, and the full history of each module's interactions. This is useful for debugging, and for understanding the agent's behaviour as you explore different pipelines and modules. With tools enabled, we also log the tools used in each module as child runs of the LLM.
 
 ------------------
 

diff --git a/__pycache__/ai.cpython-311.pyc b/__pycache__/ai.cpython-311.pyc
diff --git a/__pycache__/codebase_manager.cpython-311.pyc b/__pycache__/codebase_manager.cpython-311.pyc
diff --git a/__pycache__/helpers.cpython-311.pyc b/__pycache__/helpers.cpython-311.pyc
diff --git a/__pycache__/local_logging.cpython-311.pyc b/__pycache__/local_logging.cpython-311.pyc
diff --git a/__pycache__/modules.cpython-311.pyc b/__pycache__/modules.cpython-311.pyc
diff --git a/__pycache__/orchestrate.cpython-311.pyc b/__pycache__/orchestrate.cpython-311.pyc
diff --git a/__pycache__/wandb_logging.cpython-311.pyc b/__pycache__/wandb_logging.cpython-311.pyc
diff --git a/agent.py b/agent.py
@@ -20,12 +20,15 @@
 
 # Obtain the config variables
 wandb_enabled = config['wandb_enabled']
+tools_enabled = config['tools_enabled']
 pipeline_name = config['pipeline'] 
 pipeline_path = "pipelines/" + pipeline_name + ".yml"
 
 if wandb_enabled: # Initialize wandb if it's enabled
     wandb.init(project="OrchestrAI")
     wandb.config.wandb_enabled = wandb_enabled
+if tools_enabled:
+    print("\033[93mTools are enabled.\033[00m")
 
 def main():
     print("\033[95m ------- Welcome to OrchestrAI ------\033[00m")

diff --git a/ai.py b/ai.py
@@ -39,7 +39,7 @@ def __init__(self, module_name, model_config=None, openai=openai):
         self.frequency_penalty = model_config.get('frequency_penalty', default_frequency_penalty) if model_config else default_frequency_penalty
         self.presence_penalty = model_config.get('presence_penalty', default_presence_penalty) if model_config else default_presence_penalty
         self.module_name = module_name
-        self.system = h.load_system_prompt(module_name)
+        self.system, self.component_prompts = h.load_system_prompt(module_name)
         self.messages = [{"role": "system", "content": self.system}]
 
 
@@ -99,6 +99,7 @@ def generate_response(self, prompt):
         "status_code": status_code,
         "status_message": status_message,
         "system_prompt": self.system,
+        "component_prompts": self.component_prompts,
         "prompt": prompt,
         "max_tokens": self.max_tokens,
         "top_p": self.top_p,

diff --git a/codebase_manager.py b/codebase_manager.py
@@ -3,13 +3,25 @@
 import sys
 import subprocess
 import re
+import wandb_logging as wb
+import globals
+from datetime import datetime
+import time
+
+# open the config file
+import yaml
+with open('config.yml', 'r') as f:
+    config = yaml.safe_load(f)
+
+wandb_enabled = config['wandb_enabled'] # Set the wandb_enabled flag
 
 ### Methods used by engineer, debugger and modify_codebase to build repositories
 class CodebaseManager:
 
     def __init__(self, directory):
-        self.directory = directory
-        self.req_file_path = os.path.join(self.directory, "requirements.txt")
+        self.parent_directory = 'generated_outputs'
+        self.directory = os.path.join(self.parent_directory, directory)
+        self.req_file_path = os.path.join(self.directory, 'requirements.txt')
 
     @staticmethod
     def extract_code(chat):
@@ -31,8 +43,9 @@ def extract_code(chat):
 
     def update_codebase(self, files):
         """Create or replace any script in the updates list."""
+
         if not os.path.exists(self.directory):
-            os.mkdir(self.directory)
+            os.makedirs(self.directory)
 
         # Create an empty requirements.txt file only if it doesn't exist
         if not os.path.exists(self.req_file_path):
@@ -45,6 +58,7 @@ def update_codebase(self, files):
 
             with open(os.path.join(self.directory, file_name), "w") as file:
                 file.write(file_content)
+
 
     def compress_codebase(self):
         """Extracts the existing codebase from the directory.
@@ -72,6 +86,7 @@ def compress_codebase(self):
                             # Get the relative path of the file
                             relative_filepath = os.path.relpath(filepath, self.directory)
                             result_content.append(f"----{relative_filepath}----\n{content}")
+                            e= None
                     except Exception as e:
                         print(f"Error processing file {filepath}: {e}")
 
@@ -97,6 +112,7 @@ def run_main(self):
                 - The return code (exit status) of the subprocess. A value of 0 typically indicates successful execution, while any other value suggests an error.
                 - The error message (if any) captured from the stderr of the subprocess.
         """
+        start_time_ms = round(datetime.now().timestamp() * 1000)
 
         # Start a new process to run the main.py script. 
         # stdout=None ensures the standard output (e.g., print statements, prompts for input) 
@@ -120,6 +136,18 @@ def run_main(self):
         # Wait for the subprocess to finish and capture any error message from the stderr.
         _, stderr = process.communicate()
 
+        if wandb_enabled:
+            wb.wandb_log_tool(
+                tool_name="run_main",
+                start_time_ms=start_time_ms,
+                inputs={"python_command": python_command,
+                        "run_script_path": os.path.join(self.directory, "main.py")},
+                outputs={"stdout": process.stdout,
+                        "stderr": stderr},
+                status= "success" if process.returncode == 0 else "error",
+                parent=globals.llm_span   
+            )
+
         return process.returncode, stderr
 
 

diff --git a/config.yml b/config.yml
@@ -1,9 +1,8 @@
 ### Pipeline ###
-pipeline: 'engineering_pipeline'
+pipeline: 'advertising_pipeline'
 
 ### API Keys ###
-openai_api_key: 'YOUR_API_KEY_HERE' 
-
+openai_api_key: 'YOUR API KEY HERE'
 
 ### OpenAI Default Configs ###
 ### All calls to OpenAI will use these configs unless otherwise specified in pipeline ###
@@ -18,4 +17,11 @@ default_presence_penalty: 0
 ### Logging ###
 wandb_enabled: false
 
-working_codebase_dirname: 'generated_code' # A folder that exists 
+### Tools ### -- Early days, currently only an image generator tool.
+# Current version of tools leads to decreased agent performance (uses up context), so is disabled by default.
+tools_enabled: false
+
+
+### Files to save outputs to generate_output folder ###
+working_codebase_dirname: 'code' 
+working_image_dirname: 'images' 
diff --git a/helpers.py b/helpers.py
@@ -1,5 +1,12 @@
 import yaml
 import matplotlib.pyplot as plt
+from tool_manager import compress_tool_prompt
+
+# Load the configuration
+with open('config.yml', 'r') as f:
+    config = yaml.safe_load(f)
+
+tools_enabled = config['tools_enabled']
 
 def load_pipeline(file_path):
     """Load a pipeline configuration from a YAML file."""
@@ -11,14 +18,31 @@ def load_pipeline(file_path):
 def load_system_prompt(module_name):
     # Load the generic system prompt
     with open('general_system.txt', 'r') as file:
-        system_prompt = file.read().replace('\n', '')
+        general_system_prompt = file.read().replace('\n', '')
+
+    system_prompt = general_system_prompt
+
+    if tools_enabled:
+        with open(f'tools/tool_prompt.txt', 'r') as file:
+            tool_prompt = file.read().replace('\n', '')
+
+        tool_prompt = compress_tool_prompt(tool_prompt) # Compress the tool prompt
+        system_prompt += '\n\n' + tool_prompt + '\n'
+    else:
+        tool_prompt = None
 
     with open(f'system_prompts/{module_name}.txt', 'r') as file:
         module_prompt = file.read().replace('\n', '')
 
     system_prompt += '\n\n --- ' + module_name.upper() + ' ---\n\n' +  module_prompt + '\n'
 
-    return system_prompt
+    component_prompts = {
+        'general_system_prompt': general_system_prompt,
+        'tool_prompt': tool_prompt,
+        'module_prompt': module_prompt,
+    }
+
+    return system_prompt, component_prompts
 
 def visualize_pipeline(nx, G):
     # Increase distance between nodes by setting k parameter

diff --git a/local_logging.py b/local_logging.py
@@ -1,11 +1,12 @@
 import json
-import os
 import time
 
 def log_action(action):
     """Super simple, log any json object you want during pipeline execution."""
     with open('memory_log.json', 'r+') as file:
         data = json.load(file)
+        timestamp = round(time.time() * 1000)
+        action['action_timestamp'] = timestamp
         data['actions'].append(action)
         file.seek(0)
         json.dump(data, file)
diff --git a/modules.py b/modules.py
@@ -1,13 +1,13 @@
 # modules.py
 import os
 import yaml
-import helpers as h
 from ai import AI
 from codebase_manager import CodebaseManager
 import wandb_logging as wb
 import local_logging as ll
 import globals
-
+import tool_manager as tool
+from datetime import datetime
 
 # Load the configuration
 with open('config.yml', 'r') as f:
@@ -16,10 +16,15 @@
 wandb_enabled = config['wandb_enabled'] # Set the wandb_enabled flag
 working_codebase_dirname = config['working_codebase_dirname'] # Where code is accessed during a pipeline run
 
+# Communicate with working codebase
 codebase_portal = CodebaseManager(working_codebase_dirname)
 
+wandb_enabled = config['wandb_enabled'] # Set the wandb_enabled flag
+tools_enabled = config['tools_enabled'] # Set the tools_enabled flag
+
 def start_module(module_input, dummy=None):   
     """This function is used to invoke the start module, to accept user input into the pipeline""" 
+    start_time_ms = round(datetime.now().timestamp() * 1000) 
     module_name = "start_module"
     print("\033[92mPlease specify the task you want to perform:\033[00m")
     start = input()
@@ -29,14 +34,16 @@ def start_module(module_input, dummy=None):
     ll.log_action({"module": module_name, "input": module_input, "output": output})
     if wandb_enabled:
         wb.wandb_log_tool(tool_name = module_name,
+                    start_time_ms = start_time_ms,
                     inputs    = {},
                     outputs   = {"original_input": output},
                     parent    = globals.chain_span,
                     status    = "success")
-
+        
     return output
 
 def human_intervention(module_input, dummy=None):
+    start_time_ms = round(datetime.now().timestamp() * 1000) 
     module_name = "human_intervention"
     print("Please provide additional information to guide the agent:")
     additional_info = input()
@@ -46,6 +53,7 @@ def human_intervention(module_input, dummy=None):
     ll.log_action({"module": module_name, "input": module_input, "output": output})
     if wandb_enabled:
         wb.wandb_log_tool(tool_name = module_name,
+                    start_time_ms= start_time_ms,
                     inputs    = {},
                     outputs   = {"human_intervention": output},
                     parent    = globals.chain_span,
@@ -68,6 +76,8 @@ def chameleon(prompt, module_name, model_config=None):
     ll.log_action({"module": module_name, "input": prompt, "output": response})
     if wandb_enabled:
         wb.wandb_log_llm(response, ai.model, ai.temperature, parent = globals.chain_span)
+    if tools_enabled:
+        tool.use_tools(response["response_text"])
 
     return response["response_text"]
 
@@ -85,6 +95,8 @@ def engineer(prompt, model_config=None):
     ll.log_action({"module": module_name, "input": prompt, "output": response})
     if wandb_enabled:
         wb.wandb_log_llm(response, ai.model, ai.temperature, parent = globals.chain_span)
+    if tools_enabled:
+        tool.use_tools(response["response_text"])
 
     response_text = response["response_text"]
 
@@ -110,7 +122,7 @@ def debugger(codebase, model_config=None):
         # Install dependencies before starting the debugging process
     print("\033[94mInstalling dependencies...\033[00m")
     print(codebase_portal.list_dependencies())
-    os.system(f"pip3 install -r {working_codebase_dirname}/requirements.txt")
+    os.system(f"pip3 install -r generated_outputs/{working_codebase_dirname}/requirements.txt")
     print("\033[92mDependencies installed!\033[00m")
 
     human_intervention_provided = False
@@ -161,13 +173,15 @@ def debugger(codebase, model_config=None):
                 ll.log_action({"module": module_name, "input": prompt, "output": debug_response})
                 if wandb_enabled:
                     wb.wandb_log_llm(debug_response, ai.model, ai.temperature, parent = globals.chain_span)
+                if tools_enabled:
+                    tool.use_tools(debug_response["response_text"])
 
                 debugged_code = codebase_portal.extract_code(debug_response["response_text"])
                 codebase_portal.update_codebase(debugged_code)
 
                 if any("requirements.txt" in file_name for file_name, _ in debugged_code):
                     print("\033[94mReinstalling updated dependencies...\033[00m")
-                    os.system(f"pip3 install -r {working_codebase_dirname}/requirements.txt")
+                    os.system(f"pip3 install -r generated_outputs/{working_codebase_dirname}/requirements.txt")
                     print("\033[92mUpdated dependencies installed!\033[00m")
 
                 print("\033[93mDebugger module has made an attempt to fix. Rerunning main.py...\033[00m")
@@ -205,6 +219,8 @@ def modify_codebase(codebase, model_config=None):
         ll.log_action({"module": module_name, "input": codebase, "output": response})
         if wandb_enabled:
             wb.wandb_log_llm(response, ai.model, ai.temperature, parent = globals.chain_span)
+        if tools_enabled:
+            tool.use_tools(response["response_text"])
 
         # Parse the chat and extract files
         files = codebase_portal.extract_code(response["response_text"])
@@ -230,6 +246,8 @@ def create_readme(codebase, model_config=None):
     ll.log_action({"module": module_name, "input": codebase, "output": response})
     if wandb_enabled:
         wb.wandb_log_llm(response, ai.model, ai.temperature, parent = globals.chain_span)
+    if tools_enabled:
+        tool.use_tools(response["response_text"])
 
     # Save the response to a README.md file to the working codebase
     codebase_portal.update_codebase([("README.md", response["response_text"])])

diff --git a/pipelines/advertising_pipeline.yml b/pipelines/advertising_pipeline.yml
@@ -12,9 +12,10 @@ pipeline:
     inputs: [request, advert_plan]
     output_name: slogan
   - module: script_writer
-    supplement: "This is for a TV spot, 30 seconds long."
-    model_config:
-      model: 'gpt-3.5-turbo'
-      temperature: 1.2
+    supplement: "TV spot, 30 seconds"
     inputs: [request, advert_plan, slogan]
-    output_name: script
+    output_name: script
+  - module: illustrator
+    supplement: "Create the storyboard"
+    inputs: [request, script]
+    output_name: illustrations