Merge pull request #9 from kevinbackhouse/yaml-headers

Add header to the yaml files

Author: kevinbackhouse

available_tools.py

@@ -1,9 +1,49 @@
+import logging
+
+class VersionException(Exception):
+    pass
+
+class FileTypeException(Exception):
+    pass
+
 class AvailableTools:
     """
     This class is used for storing dictionaries of all the available
     personalities, taskflows, and prompts.
     """
-    def __init__(self, personalities: dict, taskflows: dict, prompts: dict):
-        self.personalities = personalities
-        self.taskflows = taskflows
-        self.prompts = prompts
+    def __init__(self, yamls: dict):
+        self.personalities = {}
+        self.taskflows = {}
+        self.prompts = {}
+        self.toolboxes = {}
+
+        # Iterate through all the yaml files and divide them into categories.
+        # Each file should contain a header like this:
+        #
+        #   seclab-taskflow-agent:
+        #     type: taskflow
+        #     version: 1
+        #
+        for path, yaml in yamls.items():
+            try:
+                header = yaml['seclab-taskflow-agent']
+                version = header['version']
+                if version != 1:
+                    raise VersionException(str(version))
+                filetype = header['type']
+                if filetype == 'personality':
+                    self.personalities.update({path: yaml})
+                elif filetype == 'taskflow':
+                    self.taskflows.update({path: yaml})
+                elif filetype == 'prompt':
+                    self.prompts.update({path: yaml})
+                elif filetype == 'toolbox':
+                    self.toolboxes.update({path: yaml})
+                else:
+                    raise FileTypeException(str(filetype))
+            except KeyError as err:
+                logging.error(f'{path} does not contain the key {err.args[0]}')
+            except VersionException as err:
+                logging.error(f'{path}: seclab-taskflow-agent version {err.args[0]} is not supported')
+            except FileTypeException as err:
+                logging.error(f'{path}: seclab-taskflow-agent file type {err.args[0]} is not supported')

main.py

@@ -79,7 +79,8 @@ def parse_prompt_args(available_tools: AvailableTools,
     l = args[0].l
     return p, t, l, ' '.join(args[0].prompt), help_msg
 
-async def deploy_task_agents(agents: dict,
+async def deploy_task_agents(available_tools: AvailableTools,
+                             agents: dict,
                              prompt: str,
                              async_task: bool = False,
                              toolboxes_override: list = [],
@@ -120,7 +121,7 @@ async def deploy_task_agents(agents: dict,
     tool_filter = create_static_tool_filter(blocked_tool_names=blocked_tools) if blocked_tools else None
 
     # fetch mcp params
-    mcp_params = mcp_client_params(YamlParser('toolboxes').get_yaml_dict(recurse=True), toolboxes)
+    mcp_params = mcp_client_params(available_tools.toolboxes, toolboxes)
     for tb, (params, confirms, server_prompt, client_session_timeout) in mcp_params.items():
         server_prompts.append(server_prompt)
         # https://openai.github.io/openai-agents-python/mcp/
@@ -401,6 +402,7 @@ async def on_handoff_hook(
             raise ValueError("No such personality!")
 
         await deploy_task_agents(
+            available_tools,
             { p:personality },
             prompt,
             run_hooks=TaskRunHooks(
@@ -575,6 +577,7 @@ async def run_prompts(async_task=False, max_concurrent_tasks=5):
                         async def _deploy_task_agents(resolved_agents, prompt):
                             async with semaphore:
                                 result = await deploy_task_agents(
+                                    available_tools,
                                     # pass agents and prompt by assignment, they change in-loop
                                     resolved_agents,
                                     prompt,
@@ -626,9 +629,10 @@ async def _deploy_task_agents(resolved_agents, prompt):
 
 if __name__ == '__main__':
     available_tools = AvailableTools(
-        personalities = YamlParser('personalities').get_yaml_dict(),
-        taskflows = YamlParser('taskflows').get_yaml_dict(),
-        prompts = YamlParser('prompts').get_yaml_dict(dir_namespace=True))
+        YamlParser('personalities').get_yaml_dict() |
+        YamlParser('taskflows').get_yaml_dict() |
+        YamlParser('prompts').get_yaml_dict(dir_namespace=True) |
+        YamlParser('toolboxes').get_yaml_dict(recurse=True))
 
     p, t, l, user_prompt, help_msg = parse_prompt_args(available_tools)
 

personalities/assistant.yaml

@@ -1,3 +1,7 @@
+seclab-taskflow-agent:
+  type: personality
+  version: 1
+
 personality: |
   You are a helpful assistant.
 task: |

personalities/c_auditer.yaml

@@ -1,3 +1,7 @@
+seclab-taskflow-agent:
+  type: personality
+  version: 1
+
 personality: |
   Your name is Ronald. You are a C programming language security expert.
   You have the ability to call tools to aid you in your security reviews.

personalities/examples/apple_expert.yaml

@@ -1,5 +1,9 @@
+seclab-taskflow-agent:
+  type: personality
+  version: 1
+
 personality: |
-  You are an an apples expert.
+  You are an apples expert.
 
 task: |
   Answer any questions about apples.

personalities/examples/banana_expert.yaml

@@ -1,3 +1,7 @@
+seclab-taskflow-agent:
+  type: personality
+  version: 1
+
 personality: |
   You are a bananas expert.
 

personalities/examples/echo.yaml

@@ -1,3 +1,7 @@
+seclab-taskflow-agent:
+  type: personality
+  version: 1
+
 personality: |
   You are a simple echo bot. You use echo tools to echo things.
   

personalities/examples/example_triage_agent.yaml

@@ -1,3 +1,7 @@
+seclab-taskflow-agent:
+  type: personality
+  version: 1
+
 personality: |
   You are a triage agent. You route tasks to other agents.
 

personalities/examples/fruit_expert.yaml

@@ -1,3 +1,7 @@
+seclab-taskflow-agent:
+  type: personality
+  version: 1
+
 personality: |
   Your name is Bob. You are a fruit expert.
 

personalities/examples/orange_expert.yaml

@@ -1,3 +1,7 @@
+seclab-taskflow-agent:
+  type: personality
+  version: 1
+
 personality: |
   You are an oranges expert.