diff --git a/.github/workflows/smoketest.yaml b/.github/workflows/smoketest.yaml index ecf136b..4b2779d 100644 --- a/.github/workflows/smoketest.yaml +++ b/.github/workflows/smoketest.yaml @@ -49,19 +49,19 @@ jobs: run: | source .venv/bin/activate - python main.py -p GitHubSecurityLab/seclab-taskflow-agent/personalities/assistant 'explain modems to me please' - python main.py -p GitHubSecurityLab/seclab-taskflow-agent/personalities/c_auditer 'explain modems to me please' - python main.py -p GitHubSecurityLab/seclab-taskflow-agent/personalities/examples/echo 'explain modems to me please' - python main.py -t GitHubSecurityLab/seclab-taskflow-agent/taskflows/CVE-2023-2283/CVE-2023-2283 - python main.py -t GitHubSecurityLab/seclab-taskflow-agent/taskflows/examples/echo - python main.py -t GitHubSecurityLab/seclab-taskflow-agent/taskflows/examples/example - python main.py -t GitHubSecurityLab/seclab-taskflow-agent/taskflows/examples/example_globals - python main.py -t GitHubSecurityLab/seclab-taskflow-agent/taskflows/examples/example_inputs - python main.py -t GitHubSecurityLab/seclab-taskflow-agent/taskflows/examples/example_large_list_result_iter - python main.py -t GitHubSecurityLab/seclab-taskflow-agent/taskflows/examples/example_repeat_prompt - python main.py -t GitHubSecurityLab/seclab-taskflow-agent/taskflows/examples/example_repeat_prompt_async - python main.py -t GitHubSecurityLab/seclab-taskflow-agent/taskflows/examples/example_repeat_prompt_dictionary - python main.py -t GitHubSecurityLab/seclab-taskflow-agent/taskflows/examples/example_reusable_prompt - python main.py -t GitHubSecurityLab/seclab-taskflow-agent/taskflows/examples/example_reusable_taskflows - python main.py -t GitHubSecurityLab/seclab-taskflow-agent/taskflows/examples/example_triage_taskflow - python main.py -t GitHubSecurityLab/seclab-taskflow-agent/taskflows/examples/single_step_taskflow + python main.py -p personalities.assistant 'explain modems to me please' + python main.py -p personalities.c_auditer 'explain modems to me please' + python main.py -p personalities.examples.echo 'explain modems to me please' + python main.py -t taskflows.CVE-2023-2283.CVE-2023-2283 + python main.py -t taskflows.examples.echo + python main.py -t taskflows.examples.example + python main.py -t taskflows.examples.example_globals + python main.py -t taskflows.examples.example_inputs + python main.py -t taskflows.examples.example_large_list_result_iter + python main.py -t taskflows.examples.example_repeat_prompt + python main.py -t taskflows.examples.example_repeat_prompt_async + python main.py -t taskflows.examples.example_repeat_prompt_dictionary + python main.py -t taskflows.examples.example_reusable_prompt + python main.py -t taskflows.examples.example_reusable_taskflows + python main.py -t taskflows.examples.example_triage_taskflow + python main.py -t taskflows.examples.single_step_taskflow diff --git a/README.md b/README.md index 4c870f2..6f7e0e2 100644 --- a/README.md +++ b/README.md @@ -147,43 +147,23 @@ Every YAML files used by the Seclab Taskflow Agent must include a header like th seclab-taskflow-agent: version: 1 filetype: taskflow - filekey: GitHubSecurityLab/seclab-taskflow-agent/taskflows/CVE-2023-2283/CVE-2023-2283 ``` -The `filetype` determines whether the file defines a personality, toolbox, or -taskflow. This means that different types of files can be stored in the same directory. -A `filetype` can be one of the followings: - - taskflow - - personality - - toolbox - - prompt - - model_config - -We'll explain these file types in more detail in the following sections. - -The `filekey` is a unique name for the file. It is used to allow -cross-referencing between files. For example, a taskflow can reference -a personality by its filekey. Because filekeys are used for -cross-referencing (rather than file paths), it means that you can move -a file to a different directory without breaking the links. This also -means that you can easily import new files by dropping them into a sub-directory. -We recommend including something like your -GitHub `/` in your filekeys to make them globally unique. - -In the above example, it is a `taskflow` file with `filekey` `GitHubSecurityLab/seclab-taskflow-agent/taskflows/CVE-2023-2283/CVE-2023-2283`. The `filekey` is needed to run the taskflow from command line, e.g.: - -``` -python3 main.py -t GitHubSecurityLab/seclab-taskflow-agent/taskflows/CVE-2023-2283/CVE-2023-2283 -``` - -will run the taskflow. - The `version` number in the header should always be 1. It means that the file uses version 1 of the seclab-taskflow-agent syntax. If we ever need to make a major change to the syntax, then we'll update the version number. This will hopefully enable us to make changes without breaking backwards compatibility. +The `filetype` determines whether the file defines a personality, toolbox, etc. +This means that different types of files can be stored in the same directory. +A `filetype` can be one of the following: + - taskflow + - personality + - toolbox + - prompt + - model_config + We'll now explain the role of different types of files and functionalities available to them. ## Personalities @@ -199,7 +179,6 @@ Example: seclab-taskflow-agent: version: 1 filetype: personality - filekey: GitHubSecurityLab/seclab-taskflow-agent/personalities/examples/echo personality: | You are a simple echo bot. You use echo tools to echo things. @@ -209,7 +188,7 @@ task: | # personality toolboxes map to mcp servers made available to this Agent toolboxes: - - GitHubSecurityLab/seclab-taskflow-agent/toolboxes/echo + - toolboxes.echo ``` In the above, the `personality` and `task` field specifies the system prompt to be used whenever this `personality` is used. @@ -219,7 +198,7 @@ files of the `filetype` `toolbox`. Personalities can be used in two ways. First it can be used standalone with a prompt input from the command line: ``` -python3 main.py -p GitHubSecurityLab/seclab-taskflow-agent/personalities/examples/echo "echo this message" +python3 main.py -p personalities.examples.echo "echo this message" ``` In this case, `personality` and `task` from `GitHubSecurityLab/seclab-taskflow-agent/personalities/examples/echo` are used as the @@ -233,29 +212,29 @@ taskflow: - task: ... agents: - - GitHubSecurityLab/seclab-taskflow-agent/personalities/assistant + - personalities.assistant user_prompt: | Fetch all the open pull requests from `github/codeql` github repository. You do not need to provide a summary of the results. toolboxes: - - GitHubSecurityLab/seclab-taskflow-agent/toolboxes/github_official + - toolboxes.github_official ``` -In this case, the `personality` specified in `agents` provides the system prompt and the user prompt is specified in `user_prompt` field of the task. A big difference in this case is that the `toolboxes` specified in the `task` will overwrite the `toolboxes` that the `personality` has access to. So in the above example, the `GitHubSecurityLab/seclab-taskflow-agent/personalities/assistant` will have access to the `GitHubSecurityLab/seclab-taskflow-agent/toolboxes/github_official` toolbox instead of its own toolbox. It is important to note that the `personalities` toolboxes get *overwritten* in this case, so whenever a `toolboxes` field is provided in a `task`, it'll use the provided toolboxes and `personality` loses access to its own toolboxes. e.g. +In this case, the `personality` specified in `agents` provides the system prompt and the user prompt is specified in `user_prompt` field of the task. A big difference in this case is that the `toolboxes` specified in the `task` will overwrite the `toolboxes` that the `personality` has access to. So in the above example, the `personalities.assistant` will have access to the `toolboxes.github_official` toolbox instead of its own toolbox. It is important to note that the `personalities` toolboxes get *overwritten* in this case, so whenever a `toolboxes` field is provided in a `task`, it'll use the provided toolboxes and `personality` loses access to its own toolboxes. e.g. ```yaml taskflow: - task: ... agents: - - GitHubSecurityLab/seclab-taskflow-agent/personalities/examples/echo + - personalities.examples.echo user_prompt: | echo this toolboxes: - - GitHubSecurityLab/seclab-taskflow-agent/toolboxes/github_official + - toolboxes.github_official ``` -In the above `task`, `GitHubSecurityLab/seclab-taskflow-agent/personalities/examples/echo` will only have access to the `GitHubSecurityLab/seclab-taskflow-agent/toolboxes/github_official` and can no longer access the `GitHubSecurityLab/seclab-taskflow-agent/toolboxes/echo` `toolbox`. (Unless it is added also in the `task` `toolboxes`) +In the above `task`, `personalities.examples.echo` will only have access to the `toolboxes.github_official` and can no longer access the `toolboxes.echo` `toolbox`. (Unless it is added also in the `task` `toolboxes`) ## Toolboxes @@ -269,7 +248,6 @@ For example, to start a stdio MCP server that are implemented in a python file: seclab-taskflow-agent: version: 1 filetype: toolbox - filekey: toolboxes/echo server_params: kind: stdio @@ -319,7 +297,6 @@ Example: seclab-taskflow-agent: version: 1 filetype: taskflow - filekey: taskflows/examples/example.yaml taskflow: - task: @@ -391,6 +368,13 @@ Taskflows support [Agent handoffs](https://openai.github.io/openai-agents-python See the [taskflow examples](taskflows/examples) for other useful Taskflow patterns such as repeatable and asynchronous templated prompts. + +You can run a taskflow from the command line like this: + +``` +python3 main.py -t taskflows.CVE-2023-2283.CVE-2023-2283 +``` + ## Prompt Prompts are configured through YAML files of `filetype` `prompt`. They define a reusable prompt that can be referenced in `taskflow` files. @@ -401,13 +385,12 @@ They contain only one field, the `prompt` field, which is used to replace any `{ seclab-taskflow-agent: version: 1 filetype: prompt - filekey: GitHubSecurityLab/seclab-taskflow-agent/prompts/examples/example_prompt prompt: | Tell me more about bananas as well. ``` -would replace any `{{ PROMPT_GitHubSecurityLab/seclab-taskflow-agent/prompts/examples/example_prompt }}` template parameter found in the `user_prompt` section in a taskflow: +would replace any `{{ PROMPT_prompts.examples.example_prompt }}` template parameter found in the `user_prompt` section in a taskflow: ```yaml - task: @@ -416,7 +399,7 @@ would replace any `{{ PROMPT_GitHubSecurityLab/seclab-taskflow-agent/prompts/exa user_prompt: | Tell me more about apples. - {{ PROMPTS_GitHubSecurityLab/seclab-taskflow-agent/prompts/examples/example_prompt }} + {{ PROMPTS_prompts.examples.example_prompt }} ``` becomes: @@ -439,7 +422,6 @@ Model configs are configured through YAML files of `filetype` `model_config`. Th seclab-taskflow-agent: version: 1 filetype: model_config - filekey: GitHubSecurityLab/seclab-taskflow-agent/configs/model_config models: gpt_latest: gpt-5 ``` @@ -447,7 +429,7 @@ models: A `model_config` file can be used in a `taskflow` and the values defined in `models` can then be used throughout. ```yaml -model_config: GitHubSecurityLab/seclab-taskflow-agent/configs/model_config +model_config: configs.model_config taskflow: - task: @@ -478,7 +460,7 @@ taskflow: - task: must_complete: true agents: - - GitHubSecurityLab/seclab-taskflow-agent/personalities/assistant + - personalities.assistant user_prompt: | Store the json array ["apples", "oranges", "bananas"] in the `fruits` memory key. env: diff --git a/available_tools.py b/available_tools.py index 2df1437..4a8276e 100644 --- a/available_tools.py +++ b/available_tools.py @@ -1,64 +1,88 @@ +from enum import Enum import logging +import importlib.resources +import yaml -class VersionException(Exception): +class BadToolNameError(Exception): pass -class FileIDException(Exception): +class VersionException(Exception): pass class FileTypeException(Exception): pass -def add_yaml_to_dict(table, key, yaml): - """Add the yaml to the table, but raise an error if the id isn't unique """ - if key in table: - raise FileIDException(str(key)) - table.update({key: yaml}) +class AvailableToolType(Enum): + Personality = "personality" + Taskflow = "taskflow" + Prompt = "prompt" + Toolbox = "toolbox" + ModelConfig = "model_config" class AvailableTools: """ This class is used for storing dictionaries of all the available personalities, taskflows, and prompts. """ - def __init__(self, yamls: dict): - self.personalities = {} - self.taskflows = {} - self.prompts = {} - self.toolboxes = {} - self.model_config = {} + def __init__(self): + self.__yamlcache = {} + + def get_personality(self, name: str): + return self.get_tool(AvailableToolType.Personality, name) + + def get_taskflow(self, name: str): + return self.get_tool(AvailableToolType.Taskflow, name) + + def get_prompt(self, name: str): + return self.get_tool(AvailableToolType.Prompt, name) + + def get_toolbox(self, name: str): + return self.get_tool(AvailableToolType.Toolbox, name) + + def get_model_config(self, name: str): + return self.get_tool(AvailableToolType.ModelConfig, name) - # 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'] + def get_tool(self, tooltype: AvailableToolType, toolname: str): + """for example: available_tools.get_tool("personality", "personalities/fruit_expert") + This method first checks whether the tool has already been loaded. If not, it + finds the yaml file and parses it. It also checks that the filetype in the header + matches the expected tooltype. + """ + try: + return self.__yamlcache[tooltype][toolname] + except KeyError: + pass + # Split the string to get the path and filename. + components = toolname.rsplit('.', 1) + if len(components) == 2: + path = components[0] + filename = components[1] + else: + path = '' + filename = toolname + try: + d = importlib.resources.files(path) + if not d.is_dir(): + raise BadToolNameError(f'Cannot load {toolname} because {d} is not a valid directory.') + f = d.joinpath(filename + ".yaml") + with open(f) as s: + y = yaml.safe_load(s) + header = y['seclab-taskflow-agent'] version = header['version'] if version != 1: raise VersionException(str(version)) - filekey = header['filekey'] - filetype = header['filetype'] - if filetype == 'personality': - add_yaml_to_dict(self.personalities, filekey, yaml) - elif filetype == 'taskflow': - add_yaml_to_dict(self.taskflows, filekey, yaml) - elif filetype == 'prompt': - add_yaml_to_dict(self.prompts, filekey, yaml) - elif filetype == 'toolbox': - add_yaml_to_dict(self.toolboxes, filekey, yaml) - elif filetype == 'model_config': - add_yaml_to_dict(self.model_config, filekey, 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 FileIDException as err: - logging.error(f'{path}: file ID {err.args[0]} is not unique') - except FileTypeException as err: - logging.error(f'{path}: seclab-taskflow-agent file type {err.args[0]} is not supported') + filetype = header['filetype'] + if filetype != tooltype.value: + raise FileTypeException( + f'Error in {f}: expected filetype to be {tooltype}, but it\'s {filetype}.') + if tooltype not in self.__yamlcache: + self.__yamlcache[tooltype] = {} + self.__yamlcache[tooltype][toolname] = y + return y + except ModuleNotFoundError as e: + raise BadToolNameError(f'Cannot load {toolname}: {e}') + except FileNotFoundError: + # deal with editor temp files etc. that might have disappeared + raise BadToolNameError(f'Cannot load {toolname} because {f} is not a valid file.') + except ValueError as e: + raise BadToolNameError(f'Cannot load {toolname}: {e}') diff --git a/configs/model_config.yaml b/configs/model_config.yaml index a037365..e236b4c 100644 --- a/configs/model_config.yaml +++ b/configs/model_config.yaml @@ -1,7 +1,6 @@ seclab-taskflow-agent: version: 1 filetype: model_config - filekey: GitHubSecurityLab/seclab-taskflow-agent/configs/model_config models: sonnet_default: claude-sonnet-4 sonnet_latest: claude-sonnet-4.5 diff --git a/main.py b/main.py index 02bdfcf..b45362b 100644 --- a/main.py +++ b/main.py @@ -21,13 +21,12 @@ from agents import Tool, RunContextWrapper, TContext, Agent from openai import BadRequestError, APITimeoutError, RateLimitError from openai.types.responses import ResponseTextDeltaEvent -from typing import Any +from typing import Callable from shell_utils import shell_tool_call from mcp_utils import DEFAULT_MCP_CLIENT_SESSION_TIMEOUT, ReconnectingMCPServerStdio, AsyncDebugMCPServerStdio, MCPNamespaceWrap, mcp_client_params, mcp_system_prompt, StreamableMCPThread, compress_name from render_utils import render_model_output, flush_async_output from env_utils import TmpEnv -from yaml_parser import YamlParser from agent import TaskAgent from capi import list_tool_call_models from available_tools import AvailableTools @@ -66,12 +65,6 @@ def parse_prompt_args(available_tools: AvailableTools, parser.add_argument('prompt', nargs=argparse.REMAINDER) #parser.add_argument('remainder', nargs=argparse.REMAINDER, help="Remaining args") help_msg = parser.format_help() - help_msg += "\nAvailable Personalities:\n\n" - for k in available_tools.personalities: - help_msg += f"`{k}`\n" - help_msg += "\nAvailable Taskflows:\n\n" - for k in available_tools.taskflows: - help_msg += f"`{k}`\n" help_msg += "\nExamples:\n\n" help_msg += "`-p assistant explain modems to me please`\n" try: @@ -127,7 +120,7 @@ async def deploy_task_agents(available_tools: AvailableTools, tool_filter = create_static_tool_filter(blocked_tool_names=blocked_tools) if blocked_tools else None # fetch mcp params - mcp_params = mcp_client_params(available_tools.toolboxes, toolboxes) + mcp_params = mcp_client_params(available_tools, 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/ @@ -403,9 +396,7 @@ async def on_handoff_hook( await render_model_output(f"\n** 🤖🤝 Agent Handoff: {source.name} -> {agent.name}\n") if p: - personality = available_tools.personalities.get(p) - if personality is None: - raise ValueError(f"No such personality: {p}") + personality = available_tools.get_personality(p) await deploy_task_agents( available_tools, @@ -417,13 +408,7 @@ async def on_handoff_hook( if t: - taskflow = available_tools.taskflows.get(t) - if taskflow is None: - taskflow_list = '\n'.join(available_tools.taskflows.keys()) - await render_model_output( - f"** 🤖❗ Input Error: No such taskflow: {t}. Available taskflows are:\n{taskflow_list}" - ) - raise ValueError(f"No such taskflow: {t}") + taskflow = available_tools.get_taskflow(t) await render_model_output(f"** 🤖💪 Running Task Flow: {t}\n") @@ -432,9 +417,7 @@ async def on_handoff_hook( model_config = taskflow.get('model_config', {}) model_keys = [] if model_config: - model_dict = available_tools.model_config.get(model_config, {}) - if not model_dict: - raise ValueError(f"No such model config: {model_config}") + model_dict = available_tools.get_model_config(model_config) model_dict = model_dict.get('models', {}) if model_dict: if not isinstance(model_dict, dict): @@ -451,7 +434,7 @@ async def on_handoff_hook( # can tweak reusable task configurations as they see fit uses = task_body.get('uses', '') if uses: - reusable_taskflow = available_tools.taskflows.get(uses) + reusable_taskflow = available_tools.get_taskflow(uses) if reusable_taskflow is None: raise ValueError(f"No such reusable taskflow: {uses}") if len(reusable_taskflow['taskflow']) > 1: @@ -484,31 +467,35 @@ async def on_handoff_hook( async_task = task_body.get('async', False) max_concurrent_tasks = task_body.get('async_limit', 5) - def preprocess_prompt(prompt: str, tag: str, kv: dict, kv_subkey=None): + def preprocess_prompt(prompt: str, tag: str, kv: Callable[[str], dict], kv_subkey=None): _prompt = prompt for full_match in re.findall(r"\{\{\s+" + tag + r"_(?:.*?)\s+\}\}", prompt): _m = re.search(r"\{\{\s+" + tag + r"_(.*?)\s+\}\}", full_match) if _m: key = _m.group(1) - if key in kv: - _prompt = _prompt.replace( - full_match, - str(kv.get(key)[kv_subkey]) if kv_subkey else str(kv.get(key))) - else: + v = kv(key) + if not v: raise KeyError(f"No such prompt key available: {key}") + _prompt = _prompt.replace( + full_match, + str(v[kv_subkey]) if kv_subkey else str(v)) return _prompt # pre-process the prompt for any prompts if prompt: - prompt = preprocess_prompt(prompt, 'PROMPTS', available_tools.prompts, 'prompt') + prompt = preprocess_prompt(prompt, 'PROMPTS', + lambda key: available_tools.get_prompt(key), + 'prompt') # pre-process the prompt for any inputs if prompt and inputs: - prompt = preprocess_prompt(prompt, 'INPUTS', inputs) + prompt = preprocess_prompt(prompt, 'INPUTS', + lambda key: inputs.get(key)) # pre-process the prompt for any globals if prompt and global_variables: - prompt = preprocess_prompt(prompt, 'GLOBALS', global_variables) + prompt = preprocess_prompt(prompt, 'GLOBALS', + lambda key: global_variables.get(key)) with TmpEnv(env): prompts_to_run = [] @@ -590,7 +577,7 @@ async def run_prompts(async_task=False, max_concurrent_tasks=5): p, _, _, prompt, _ = parse_prompt_args(available_tools, prompt) agents.append(p) for p in agents: - personality = available_tools.personalities.get(p) + personality = available_tools.get_personality(p) if personality is None: raise ValueError(f"No such personality: {p}") resolved_agents[p] = personality @@ -652,12 +639,7 @@ async def _deploy_task_agents(resolved_agents, prompt): if __name__ == '__main__': cwd = pathlib.Path.cwd() - available_tools = AvailableTools( - YamlParser(cwd).get_yaml_dict((cwd/'personalities').rglob('*')) | - YamlParser(cwd).get_yaml_dict((cwd/'taskflows').rglob('*')) | - YamlParser(cwd).get_yaml_dict((cwd/'prompts').rglob('*')) | - YamlParser(cwd).get_yaml_dict((cwd/'toolboxes').rglob('*')) | - YamlParser(cwd).get_yaml_dict((cwd/'configs').rglob('*'))) + available_tools = AvailableTools() p, t, l, user_prompt, help_msg = parse_prompt_args(available_tools) diff --git a/mcp_utils.py b/mcp_utils.py index 26bf4a5..6dca905 100644 --- a/mcp_utils.py +++ b/mcp_utils.py @@ -16,14 +16,14 @@ from agents.mcp import MCPServerStdio from env_utils import swap_env +from available_tools import AvailableToolType, AvailableTools DEFAULT_MCP_CLIENT_SESSION_TIMEOUT = 120 # The openai API complains if the name of a tool is longer than 64 -# chars. But we're encouraging people to use long descriptive -# filekeys to avoid accidental collisions, so it's very easy to go -# over the limit. So this function converts a name to a 12 character -# hash. +# chars. But it's easy to go over the limit if the yaml file is in a +# nested sub-directory, so this function converts a name to a 12 +# character hash. def compress_name(name): m = hashlib.sha256() m.update(name.encode('utf-8')) @@ -284,24 +284,106 @@ async def call_tool(self, *args, **kwargs): result = await self._obj.call_tool(*args, **kwargs) return result -def mcp_client_params(available_toolboxes: dict, requested_toolboxes: list): +def mcp_client_params(available_tools: AvailableTools, requested_toolboxes: list): """Return all the data needed to initialize an mcp server client""" client_params = {} for tb in requested_toolboxes: - if tb not in available_toolboxes: - e = f"Task requested non-existent toolbox {tb}!" - logging.critical(e) - raise ValueError(e) - else: - kind = available_toolboxes[tb]['server_params'].get('kind') - reconnecting = available_toolboxes[tb]['server_params'].get('reconnecting', False) - server_params = {'kind': kind, 'reconnecting': reconnecting} - match kind: - case 'stdio': - env = available_toolboxes[tb]['server_params'].get('env') - args = available_toolboxes[tb]['server_params'].get('args') - logging.debug(f"Initializing toolbox: {tb}\nargs:\n{args }\nenv:\n{env}\n") - if env and isinstance(env, dict): + toolbox = available_tools.get_tool(AvailableToolType.Toolbox, tb) + kind = toolbox['server_params'].get('kind') + reconnecting = toolbox['server_params'].get('reconnecting', False) + server_params = {'kind': kind, 'reconnecting': reconnecting} + match kind: + case 'stdio': + env = toolbox['server_params'].get('env') + args = toolbox['server_params'].get('args') + logging.debug(f"Initializing toolbox: {tb}\nargs:\n{args }\nenv:\n{env}\n") + if env and isinstance(env, dict): + for k, v in dict(env).items(): + try: + env[k] = swap_env(v) + except LookupError as e: + logging.critical(e) + logging.info("Assuming toolbox has default configuration available") + del env[k] + logging.debug(f"Tool call environment: {env}") + if args and isinstance(args, list): + for i, v in enumerate(args): + args[i] = swap_env(v) + logging.debug(f"Tool call args: {args}") + server_params['command'] = toolbox['server_params'].get('command') + server_params['args'] = args + server_params['env'] = env + # XXX: SSE is deprecated in the MCP spec, but keep it around for now + case 'sse': + headers = toolbox['server_params'].get('headers') + # support {{ env SOMETHING }} for header values as well for e.g. tokens + if headers and isinstance(headers, dict): + for k, v in headers.items(): + headers[k] = swap_env(v) + optional_headers = toolbox['server_params'].get('optional_headers') + # support {{ env SOMETHING }} for header values as well for e.g. tokens + if optional_headers and isinstance(optional_headers, dict): + for k, v in dict(optional_headers).items(): + try: + optional_headers[k] = swap_env(v) + except LookupError as e: + del optional_headers[k] + if isinstance(headers, dict): + if isinstance(optional_headers, dict): + headers.update(optional_headers) + elif isinstance(optional_headers, dict): + headers = optional_headers + # if None will default to float(5) in client code + timeout = toolbox['server_params'].get('timeout') + server_params['url'] = toolbox['server_params'].get('url') + server_params['headers'] = headers + server_params['timeout'] = timeout + # for more involved local MCP servers, jsonrpc over stdio seems less than reliable + # as an alternative you can configure local toolboxes to use the streamable transport + # but still be started/stopped on demand similar to stdio mcp servers + # all it requires is a streamable config that also has cmd/args/env set + case 'streamable': + headers = toolbox['server_params'].get('headers') + # support {{ env SOMETHING }} for header values as well for e.g. tokens + if headers and isinstance(headers, dict): + for k, v in headers.items(): + headers[k] = swap_env(v) + optional_headers = toolbox['server_params'].get('optional_headers') + # support {{ env SOMETHING }} for header values as well for e.g. tokens + if optional_headers and isinstance(optional_headers, dict): + for k, v in dict(optional_headers).items(): + try: + optional_headers[k] = swap_env(v) + except LookupError as e: + del optional_headers[k] + if isinstance(headers, dict): + if isinstance(optional_headers, dict): + headers.update(optional_headers) + elif isinstance(optional_headers, dict): + headers = optional_headers + # if None will default to float(5) in client code + timeout = toolbox['server_params'].get('timeout') + server_params['url'] = toolbox['server_params'].get('url') + server_params['headers'] = headers + server_params['timeout'] = timeout + # if command/args/env is set, we also need to start this MCP server ourselves + # this way we can use the streamable transport for MCP servers that get fussy + # over stdio jsonrpc polling + env = toolbox['server_params'].get('env') + args = toolbox['server_params'].get('args') + cmd = toolbox['server_params'].get('command') + if cmd is not None: + logging.debug(f"Initializing streamable toolbox: {tb}\nargs:\n{args }\nenv:\n{env}\n") + exe = shutil.which(cmd) + if exe is None: + raise FileNotFoundError(f"Could not resolve path to {cmd}") + start_cmd = [exe] + if args is not None and isinstance(args, list): + for i, v in enumerate(args): + args[i] = swap_env(v) + start_cmd += args + server_params['command'] = start_cmd + if env is not None and isinstance(env, dict): for k, v in dict(env).items(): try: env[k] = swap_env(v) @@ -309,99 +391,13 @@ def mcp_client_params(available_toolboxes: dict, requested_toolboxes: list): logging.critical(e) logging.info("Assuming toolbox has default configuration available") del env[k] - logging.debug(f"Tool call environment: {env}") - if args and isinstance(args, list): - for i, v in enumerate(args): - args[i] = swap_env(v) - logging.debug(f"Tool call args: {args}") - server_params['command'] = available_toolboxes[tb]['server_params'].get('command') - server_params['args'] = args server_params['env'] = env - # XXX: SSE is deprecated in the MCP spec, but keep it around for now - case 'sse': - headers = available_toolboxes[tb]['server_params'].get('headers') - # support {{ env SOMETHING }} for header values as well for e.g. tokens - if headers and isinstance(headers, dict): - for k, v in headers.items(): - headers[k] = swap_env(v) - optional_headers = available_toolboxes[tb]['server_params'].get('optional_headers') - # support {{ env SOMETHING }} for header values as well for e.g. tokens - if optional_headers and isinstance(optional_headers, dict): - for k, v in dict(optional_headers).items(): - try: - optional_headers[k] = swap_env(v) - except LookupError as e: - del optional_headers[k] - if isinstance(headers, dict): - if isinstance(optional_headers, dict): - headers.update(optional_headers) - elif isinstance(optional_headers, dict): - headers = optional_headers - # if None will default to float(5) in client code - timeout = available_toolboxes[tb]['server_params'].get('timeout') - server_params['url'] = available_toolboxes[tb]['server_params'].get('url') - server_params['headers'] = headers - server_params['timeout'] = timeout - # for more involved local MCP servers, jsonrpc over stdio seems less than reliable - # as an alternative you can configure local toolboxes to use the streamable transport - # but still be started/stopped on demand similar to stdio mcp servers - # all it requires is a streamable config that also has cmd/args/env set - case 'streamable': - headers = available_toolboxes[tb]['server_params'].get('headers') - # support {{ env SOMETHING }} for header values as well for e.g. tokens - if headers and isinstance(headers, dict): - for k, v in headers.items(): - headers[k] = swap_env(v) - optional_headers = available_toolboxes[tb]['server_params'].get('optional_headers') - # support {{ env SOMETHING }} for header values as well for e.g. tokens - if optional_headers and isinstance(optional_headers, dict): - for k, v in dict(optional_headers).items(): - try: - optional_headers[k] = swap_env(v) - except LookupError as e: - del optional_headers[k] - if isinstance(headers, dict): - if isinstance(optional_headers, dict): - headers.update(optional_headers) - elif isinstance(optional_headers, dict): - headers = optional_headers - # if None will default to float(5) in client code - timeout = available_toolboxes[tb]['server_params'].get('timeout') - server_params['url'] = available_toolboxes[tb]['server_params'].get('url') - server_params['headers'] = headers - server_params['timeout'] = timeout - # if command/args/env is set, we also need to start this MCP server ourselves - # this way we can use the streamable transport for MCP servers that get fussy - # over stdio jsonrpc polling - env = available_toolboxes[tb]['server_params'].get('env') - args = available_toolboxes[tb]['server_params'].get('args') - cmd = available_toolboxes[tb]['server_params'].get('command') - if cmd is not None: - logging.debug(f"Initializing streamable toolbox: {tb}\nargs:\n{args }\nenv:\n{env}\n") - exe = shutil.which(cmd) - if exe is None: - raise FileNotFoundError(f"Could not resolve path to {cmd}") - start_cmd = [exe] - if args is not None and isinstance(args, list): - for i, v in enumerate(args): - args[i] = swap_env(v) - start_cmd += args - server_params['command'] = start_cmd - if env is not None and isinstance(env, dict): - for k, v in dict(env).items(): - try: - env[k] = swap_env(v) - except LookupError as e: - logging.critical(e) - logging.info("Assuming toolbox has default configuration available") - del env[k] - server_params['env'] = env - case _: - raise ValueError(f"Unsupported MCP transport {kind}") - confirms = available_toolboxes[tb].get('confirm', []) - server_prompt = available_toolboxes[tb].get('server_prompt', '') - client_session_timeout = float(available_toolboxes[tb].get('client_session_timeout', 0)) - client_params[tb] = (server_params, confirms, server_prompt, client_session_timeout) + case _: + raise ValueError(f"Unsupported MCP transport {kind}") + confirms = toolbox.get('confirm', []) + server_prompt = toolbox.get('server_prompt', '') + client_session_timeout = float(toolbox.get('client_session_timeout', 0)) + client_params[tb] = (server_params, confirms, server_prompt, client_session_timeout) return client_params def mcp_system_prompt(system_prompt: str, task: str, diff --git a/personalities/assistant.yaml b/personalities/assistant.yaml index 6d3bb35..036b5c7 100644 --- a/personalities/assistant.yaml +++ b/personalities/assistant.yaml @@ -1,7 +1,6 @@ seclab-taskflow-agent: version: 1 filetype: personality - filekey: GitHubSecurityLab/seclab-taskflow-agent/personalities/assistant personality: | You are a helpful assistant. diff --git a/personalities/c_auditer.yaml b/personalities/c_auditer.yaml index 556b73b..004e53b 100644 --- a/personalities/c_auditer.yaml +++ b/personalities/c_auditer.yaml @@ -1,7 +1,6 @@ seclab-taskflow-agent: version: 1 filetype: personality - filekey: GitHubSecurityLab/seclab-taskflow-agent/personalities/c_auditer personality: | Your name is Ronald. You are a C programming language security expert. @@ -15,5 +14,5 @@ task: | your findings where possible. toolboxes: - - GitHubSecurityLab/seclab-taskflow-agent/toolboxes/memcache - - GitHubSecurityLab/seclab-taskflow-agent/toolboxes/codeql + - toolboxes.memcache + - toolboxes.codeql diff --git a/personalities/examples/apple_expert.yaml b/personalities/examples/apple_expert.yaml index 16683fe..c178d67 100644 --- a/personalities/examples/apple_expert.yaml +++ b/personalities/examples/apple_expert.yaml @@ -1,7 +1,6 @@ seclab-taskflow-agent: version: 1 filetype: personality - filekey: GitHubSecurityLab/seclab-taskflow-agent/personalities/examples/apple_expert personality: | You are an apples expert. diff --git a/personalities/examples/banana_expert.yaml b/personalities/examples/banana_expert.yaml index edb09de..d879458 100644 --- a/personalities/examples/banana_expert.yaml +++ b/personalities/examples/banana_expert.yaml @@ -1,7 +1,6 @@ seclab-taskflow-agent: version: 1 filetype: personality - filekey: GitHubSecurityLab/seclab-taskflow-agent/personalities/examples/banana_expert personality: | You are a bananas expert. diff --git a/personalities/examples/echo.yaml b/personalities/examples/echo.yaml index 25152c9..770a752 100644 --- a/personalities/examples/echo.yaml +++ b/personalities/examples/echo.yaml @@ -1,7 +1,6 @@ seclab-taskflow-agent: version: 1 filetype: personality - filekey: GitHubSecurityLab/seclab-taskflow-agent/personalities/examples/echo personality: | You are a simple echo bot. You use echo tools to echo things. @@ -10,5 +9,5 @@ task: | Echo user inputs using the echo tools. toolboxes: - - GitHubSecurityLab/seclab-taskflow-agent/toolboxes/echo + - toolboxes.echo diff --git a/personalities/examples/example_triage_agent.yaml b/personalities/examples/example_triage_agent.yaml index 246a414..74c764c 100644 --- a/personalities/examples/example_triage_agent.yaml +++ b/personalities/examples/example_triage_agent.yaml @@ -1,7 +1,6 @@ seclab-taskflow-agent: version: 1 filetype: personality - filekey: GitHubSecurityLab/seclab-taskflow-agent/personalities/examples/example_triage_agent personality: | You are a triage agent. You route tasks to other agents. diff --git a/personalities/examples/fruit_expert.yaml b/personalities/examples/fruit_expert.yaml index e987904..31597c8 100644 --- a/personalities/examples/fruit_expert.yaml +++ b/personalities/examples/fruit_expert.yaml @@ -1,7 +1,6 @@ seclab-taskflow-agent: version: 1 filetype: personality - filekey: GitHubSecurityLab/seclab-taskflow-agent/personalities/examples/fruit_expert personality: | Your name is Bob. You are a fruit expert. diff --git a/personalities/examples/orange_expert.yaml b/personalities/examples/orange_expert.yaml index aa4b953..4581ca5 100644 --- a/personalities/examples/orange_expert.yaml +++ b/personalities/examples/orange_expert.yaml @@ -1,7 +1,6 @@ seclab-taskflow-agent: version: 1 filetype: personality - filekey: GitHubSecurityLab/seclab-taskflow-agent/personalities/examples/orange_expert personality: | You are an oranges expert. diff --git a/prompts/examples/example_prompt.yaml b/prompts/examples/example_prompt.yaml index 1e9f38b..0800a4c 100644 --- a/prompts/examples/example_prompt.yaml +++ b/prompts/examples/example_prompt.yaml @@ -1,7 +1,6 @@ seclab-taskflow-agent: version: 1 filetype: prompt - filekey: GitHubSecurityLab/seclab-taskflow-agent/prompts/examples/example_prompt prompt: | Tell me more about bananas as well. diff --git a/taskflows/CVE-2023-2283/CVE-2023-2283.yaml b/taskflows/CVE-2023-2283/CVE-2023-2283.yaml index c6617ff..4b3922d 100644 --- a/taskflows/CVE-2023-2283/CVE-2023-2283.yaml +++ b/taskflows/CVE-2023-2283/CVE-2023-2283.yaml @@ -1,25 +1,24 @@ seclab-taskflow-agent: version: 1 filetype: taskflow - filekey: GitHubSecurityLab/seclab-taskflow-agent/taskflows/CVE-2023-2283/CVE-2023-2283 -model_config: GitHubSecurityLab/seclab-taskflow-agent/configs/model_config +model_config: configs.model_config taskflow: - task: must_complete: true headless: true agents: - - GitHubSecurityLab/seclab-taskflow-agent/personalities/assistant + - personalities.assistant user_prompt: | Clear the memory cache. toolboxes: - - GitHubSecurityLab/seclab-taskflow-agent/toolboxes/memcache + - toolboxes.memcache - task: model: gpt_latest must_complete: false agents: - - GitHubSecurityLab/seclab-taskflow-agent/personalities/c_auditer + - personalities.c_auditer user_prompt: | You are auditing code using the `libssh-mirror/libssh-codeql` CodeQL database. @@ -79,15 +78,15 @@ taskflow: 6. Make small and concise single line notes while you work. Update the existing value for `notes` in memory as you work. toolboxes: - - GitHubSecurityLab/seclab-taskflow-agent/toolboxes/codeql - - GitHubSecurityLab/seclab-taskflow-agent/toolboxes/memcache + - toolboxes.codeql + - toolboxes.memcache - task: must_complete: true agents: - - GitHubSecurityLab/seclab-taskflow-agent/personalities/c_auditer + - personalities.c_auditer user_prompt: | Fetch your audit notes from memory using the `notes` key. Do not perform any additional security review, only show me your notes. toolboxes: - - GitHubSecurityLab/seclab-taskflow-agent/toolboxes/memcache + - toolboxes.memcache diff --git a/taskflows/GRAMMAR.md b/taskflows/GRAMMAR.md index 409d111..75a1cd3 100644 --- a/taskflows/GRAMMAR.md +++ b/taskflows/GRAMMAR.md @@ -38,7 +38,6 @@ For example, to use the `personality` defined in the following: seclab-taskflow-agent: version: 1 filetype: personality - filekey: GitHubSecurityLab/seclab-taskflow-agent/personalities/assistant personality: | You are a helpful assistant. @@ -58,7 +57,7 @@ The task should include the `filekey` in its list of `agents`: ```yaml - task: agents: - - GitHubSecurityLab/seclab-taskflow-agent/personalities/assistant + - personalities.assistant ... ``` @@ -148,7 +147,7 @@ The iterable can also contain a list of primitives like string or number, in whi max_steps: 5 must_complete: true agents: - - GitHubSecurityLab/seclab-taskflow-agent/personalities/assistant + - personalities.assistant user_prompt: | Store the json array [1, 2, 3] in memory under the `test_repeat_prompt` key as a json object, then retrieve @@ -160,7 +159,7 @@ The iterable can also contain a list of primitives like string or number, in whi must_complete: true repeat_prompt: true agents: - - GitHubSecurityLab/seclab-taskflow-agent/personalities/assistant + - personalities.assistant user_prompt: | What is the integer value of {{ RESULT }}? ``` @@ -256,7 +255,7 @@ Toolboxes are MCP server configurations. They can be defined at the Agent level - task: ... toolboxes: - - GitHubSecurityLab/seclab-taskflow-agent/toolboxes/codeql + - toolboxes.codeql ``` If no `toolboxes` is specified, then the `toolboxes` defined in the `personality` of the `agent` is used: @@ -264,28 +263,28 @@ If no `toolboxes` is specified, then the `toolboxes` defined in the `personality ```yaml - task: agents: - - GitHubSecurityLab/seclab-taskflow-agent/personalities/c_auditer + - personalities.c_auditer user_prompt: | List all the files in the codeql database `some/codeql/db`. - task: ``` -In the above `task`, as no `toolboxes` is specified, the `toolboxes` defined in the `personality` of `GitHubSecurityLab/seclab-taskflow-agent/personalities/c_auditer` is used. +In the above `task`, as no `toolboxes` is specified, the `toolboxes` defined in the `personality` of `personalities.c_auditer` is used. Note that when `toolboxes` is defined for a task, it *overwrites* the `toolboxes` that are available. For example, in the following `task`: ```yaml - task: agents: - - GitHubSecurityLab/seclab-taskflow-agent/personalities/c_auditer + - personalities.c_auditer user_prompt: | List all the files in the codeql database `some/codeql/db`. toolboxes: - - GitHubSecurityLab/seclab-taskflow-agent/toolboxes/echo + - toolboxes.echo ``` -For this task, the `agent` `GitHubSecurityLab/seclab-taskflow-agent/personalities/c_auditer` will have access to the `GitHubSecurityLab/seclab-taskflow-agent/toolboxes/echo` tool. +For this task, the `agent` `personalities.c_auditer` will have access to the `toolboxes.echo` tool. ### Headless Runs @@ -427,15 +426,14 @@ Example: user_prompt: | Tell me more about apples. - {{ PROMPTS_GitHubSecurityLab/seclab-taskflow-agent/prompts/examples/example_prompt }} + {{ PROMPTS_prompts.examples.example_prompt }} ``` -and `GitHubSecurityLab/seclab-taskflow-agent/prompts/examples/example_prompt` is the following: +and `prompts.examples.example_prompt` is the following: ```yaml seclab-taskflow-agent: version: 1 filetype: prompt - filekey: GitHubSecurityLab/seclab-taskflow-agent/prompts/examples/example_prompt prompt: | Tell me more about bananas as well. @@ -461,9 +459,8 @@ LLM models can be configured in a taskflow by setting the `model_config` field t seclab-taskflow-agent: version: 1 filetype: taskflow - filekey: GitHubSecurityLab/seclab-taskflow-agent/taskflows/CVE-2023-2283/CVE-2023-2283 -model_config: GitHubSecurityLab/seclab-taskflow-agent/configs/model_config +model_config: configs.model_config ``` @@ -473,7 +470,6 @@ The variables defined in the `model_config` file can then be used throughout the seclab-taskflow-agent: version: 1 filetype: model_config - filekey: GitHubSecurityLab/seclab-taskflow-agent/configs/model_config models: gpt_latest: gpt-5 ``` @@ -485,7 +481,7 @@ When `gpt_latest` is used in the taskflow to specify a model, the value `gpt-5` model: gpt_latest must_complete: false agents: - - GitHubSecurityLab/seclab-taskflow-agent/personalities/c_auditer + - personalities.c_auditer user_prompt: | ``` diff --git a/taskflows/examples/echo.yaml b/taskflows/examples/echo.yaml index c01334a..3e56208 100644 --- a/taskflows/examples/echo.yaml +++ b/taskflows/examples/echo.yaml @@ -1,7 +1,6 @@ seclab-taskflow-agent: version: 1 filetype: taskflow - filekey: GitHubSecurityLab/seclab-taskflow-agent/taskflows/examples/echo taskflow: - task: @@ -9,13 +8,13 @@ taskflow: max_steps: 5 must_complete: true agents: - - GitHubSecurityLab/seclab-taskflow-agent/personalities/examples/echo + - personalities.examples.echo user_prompt: | Hello - task: must_complete: true agents: - - GitHubSecurityLab/seclab-taskflow-agent/personalities/examples/echo + - personalities.examples.echo user_prompt: | Goodbye env: diff --git a/taskflows/examples/example.yaml b/taskflows/examples/example.yaml index 15203dd..267527c 100644 --- a/taskflows/examples/example.yaml +++ b/taskflows/examples/example.yaml @@ -1,7 +1,6 @@ seclab-taskflow-agent: version: 1 filetype: taskflow - filekey: GitHubSecurityLab/seclab-taskflow-agent/taskflows/examples/example taskflow: - task: @@ -13,8 +12,8 @@ taskflow: must_complete: true # taskflows can set a primary (first entry) and handoff (additional entries) agent agents: - - GitHubSecurityLab/seclab-taskflow-agent/personalities/c_auditer - - GitHubSecurityLab/seclab-taskflow-agent/personalities/examples/fruit_expert + - personalities.c_auditer + - personalities.examples.fruit_expert user_prompt: | Store an example vulnerable C program that uses `strcpy` in the `vulnerable_c_example` memory key and explain why `strcpy` @@ -35,13 +34,13 @@ taskflow: # this normally only has the memcache toolbox, but we extend it here with # the GHSA toolbox toolboxes: - - GitHubSecurityLab/seclab-taskflow-agent/toolboxes/memcache - - GitHubSecurityLab/seclab-taskflow-agent/toolboxes/codeql + - toolboxes.memcache + - toolboxes.codeql - task: must_complete: true model: gpt-4.1 agents: - - GitHubSecurityLab/seclab-taskflow-agent/personalities/c_auditer + - personalities.c_auditer user_prompt: | Retrieve C code for security review from the `vulnerable_c_example` memory key and perform a review. @@ -51,7 +50,7 @@ taskflow: MEMCACHE_STATE_DIR: "example_taskflow/" MEMCACHE_BACKEND: "dictionary_file" toolboxes: - - GitHubSecurityLab/seclab-taskflow-agent/toolboxes/memcache + - toolboxes.memcache # headless mode does not prompt for tool call confirms configured for a server # note: this will auto-allow, if you want control over potentially dangerous # tool calls, then you should NOT run a task in headless mode (default: false) @@ -64,6 +63,6 @@ taskflow: - task: repeat_prompt: true agents: - - GitHubSecurityLab/seclab-taskflow-agent/personalities/assistant + - personalities.assistant user_prompt: | What kind of fruit is {{ RESULT }}? diff --git a/taskflows/examples/example_globals.yaml b/taskflows/examples/example_globals.yaml index 24c9864..cfaf310 100644 --- a/taskflows/examples/example_globals.yaml +++ b/taskflows/examples/example_globals.yaml @@ -1,13 +1,12 @@ seclab-taskflow-agent: version: 1 filetype: taskflow - filekey: GitHubSecurityLab/seclab-taskflow-agent/taskflows/examples/example_globals globals: fruit: bananas taskflow: - task: agents: - - GitHubSecurityLab/seclab-taskflow-agent/personalities/examples/fruit_expert + - personalities.examples.fruit_expert user_prompt: | Tell me more about {{ GLOBALS_fruit }}. diff --git a/taskflows/examples/example_inputs.yaml b/taskflows/examples/example_inputs.yaml index 5c6d255..322f7a7 100644 --- a/taskflows/examples/example_inputs.yaml +++ b/taskflows/examples/example_inputs.yaml @@ -1,12 +1,11 @@ seclab-taskflow-agent: version: 1 filetype: taskflow - filekey: GitHubSecurityLab/seclab-taskflow-agent/taskflows/examples/example_inputs taskflow: - task: agents: - - GitHubSecurityLab/seclab-taskflow-agent/personalities/examples/fruit_expert + - personalities.examples.fruit_expert inputs: fruit: apples user_prompt: | diff --git a/taskflows/examples/example_large_list_result_iter.yaml b/taskflows/examples/example_large_list_result_iter.yaml index 8a58253..651c788 100644 --- a/taskflows/examples/example_large_list_result_iter.yaml +++ b/taskflows/examples/example_large_list_result_iter.yaml @@ -1,25 +1,24 @@ seclab-taskflow-agent: version: 1 filetype: taskflow - filekey: GitHubSecurityLab/seclab-taskflow-agent/taskflows/examples/example_large_list_result_iter taskflow: - task: exclude_from_context: true must_complete: true agents: - - GitHubSecurityLab/seclab-taskflow-agent/personalities/assistant + - personalities.assistant user_prompt: | Fetch all the open pull requests from `github/codeql` github repository. You do not need to provide a summary of the results. toolboxes: - - GitHubSecurityLab/seclab-taskflow-agent/toolboxes/github_official + - toolboxes.github_official env: GITHUB_MCP_TOOLSETS: pull_requests - task: must_complete: true repeat_prompt: true agents: - - GitHubSecurityLab/seclab-taskflow-agent/personalities/assistant + - personalities.assistant user_prompt: | Echo this: The title is {{ RESULT_title }} and the url is {{ RESULT_url }}. diff --git a/taskflows/examples/example_repeat_prompt.yaml b/taskflows/examples/example_repeat_prompt.yaml index 470dc48..1fc02b2 100644 --- a/taskflows/examples/example_repeat_prompt.yaml +++ b/taskflows/examples/example_repeat_prompt.yaml @@ -1,14 +1,13 @@ seclab-taskflow-agent: version: 1 filetype: taskflow - filekey: GitHubSecurityLab/seclab-taskflow-agent/taskflows/examples/example_repeat_prompt taskflow: - task: max_steps: 5 must_complete: true agents: - - GitHubSecurityLab/seclab-taskflow-agent/personalities/assistant + - personalities.assistant user_prompt: | Store the json array [1, 2, 3] in memory under the `test_repeat_prompt` key as a json object, then retrieve @@ -17,13 +16,13 @@ taskflow: MEMCACHE_STATE_DIR: "example_repeat_prompt_taskflow/" MEMCACHE_BACKEND: "dictionary_file" toolboxes: - - GitHubSecurityLab/seclab-taskflow-agent/toolboxes/memcache + - toolboxes.memcache - task: # if the last mcp tool result is iterable # repeat_prompt can iter those results must_complete: true repeat_prompt: true agents: - - GitHubSecurityLab/seclab-taskflow-agent/personalities/assistant + - personalities.assistant user_prompt: | What is the integer value of {{ RESULT }}? diff --git a/taskflows/examples/example_repeat_prompt_async.yaml b/taskflows/examples/example_repeat_prompt_async.yaml index c5f40fd..c289158 100644 --- a/taskflows/examples/example_repeat_prompt_async.yaml +++ b/taskflows/examples/example_repeat_prompt_async.yaml @@ -1,14 +1,13 @@ seclab-taskflow-agent: version: 1 filetype: taskflow - filekey: GitHubSecurityLab/seclab-taskflow-agent/taskflows/examples/example_repeat_prompt_async taskflow: - task: max_steps: 5 must_complete: true agents: - - GitHubSecurityLab/seclab-taskflow-agent/personalities/assistant + - personalities.assistant user_prompt: | Store the json array [1, 2, 3] in memory under the `test_repeat_prompt` key as a json object, then retrieve @@ -17,7 +16,7 @@ taskflow: MEMCACHE_STATE_DIR: "example_repeat_prompt_taskflow/" MEMCACHE_BACKEND: "dictionary_file" toolboxes: - - GitHubSecurityLab/seclab-taskflow-agent/toolboxes/memcache + - toolboxes.memcache - task: # if the last mcp tool result is iterable # repeat_prompt can iter those results @@ -28,6 +27,6 @@ taskflow: # you can also limit the max concurrent tasks (default 5) async_limit: 2 agents: - - GitHubSecurityLab/seclab-taskflow-agent/personalities/assistant + - personalities.assistant user_prompt: | What is the integer value of {{ RESULT }}? diff --git a/taskflows/examples/example_repeat_prompt_dictionary.yaml b/taskflows/examples/example_repeat_prompt_dictionary.yaml index f26d2e8..a9e23f6 100644 --- a/taskflows/examples/example_repeat_prompt_dictionary.yaml +++ b/taskflows/examples/example_repeat_prompt_dictionary.yaml @@ -1,14 +1,13 @@ seclab-taskflow-agent: version: 1 filetype: taskflow - filekey: GitHubSecurityLab/seclab-taskflow-agent/taskflows/examples/example_repeat_prompt_dictionary taskflow: - task: max_steps: 5 must_complete: true agents: - - GitHubSecurityLab/seclab-taskflow-agent/personalities/assistant + - personalities.assistant user_prompt: | Store the json array [{index : 1, value : 2}, {index : 2, value : 4}, {index : 3, value : 6}] in memory under the `test_repeat_prompt` key @@ -18,13 +17,13 @@ taskflow: MEMCACHE_STATE_DIR: "example_repeat_prompt_taskflow/" MEMCACHE_BACKEND: "dictionary_file" toolboxes: - - GitHubSecurityLab/seclab-taskflow-agent/toolboxes/memcache + - toolboxes.memcache - task: # if the last mcp tool result is iterable # repeat_prompt can iter those results must_complete: true repeat_prompt: true agents: - - GitHubSecurityLab/seclab-taskflow-agent/personalities/assistant + - personalities.assistant user_prompt: | What is the value of {{ RESULT_index }} + {{ RESULT_value }}? diff --git a/taskflows/examples/example_reusable_prompt.yaml b/taskflows/examples/example_reusable_prompt.yaml index a8f5262..845bcec 100644 --- a/taskflows/examples/example_reusable_prompt.yaml +++ b/taskflows/examples/example_reusable_prompt.yaml @@ -1,13 +1,12 @@ seclab-taskflow-agent: version: 1 filetype: taskflow - filekey: GitHubSecurityLab/seclab-taskflow-agent/taskflows/examples/example_reusable_prompt taskflow: - task: agents: - - GitHubSecurityLab/seclab-taskflow-agent/personalities/examples/fruit_expert + - personalities.examples.fruit_expert user_prompt: | Tell me more about apples. - {{ PROMPTS_GitHubSecurityLab/seclab-taskflow-agent/prompts/examples/example_prompt }} + {{ PROMPTS_prompts.examples.example_prompt }} diff --git a/taskflows/examples/example_reusable_taskflows.yaml b/taskflows/examples/example_reusable_taskflows.yaml index 33cc43c..6597e57 100644 --- a/taskflows/examples/example_reusable_taskflows.yaml +++ b/taskflows/examples/example_reusable_taskflows.yaml @@ -1,11 +1,10 @@ seclab-taskflow-agent: version: 1 filetype: taskflow - filekey: GitHubSecurityLab/seclab-taskflow-agent/taskflows/examples/example_reusable_taskflows taskflow: - task: # with the `uses` directive we can reuse single task taskflows - uses: GitHubSecurityLab/seclab-taskflow-agent/taskflows/examples/single_step_taskflow + uses: taskflows.examples.single_step_taskflow # and optionally override any of its configurations model: gpt-4o diff --git a/taskflows/examples/example_triage_taskflow.yaml b/taskflows/examples/example_triage_taskflow.yaml index 4ae689b..bd5672e 100644 --- a/taskflows/examples/example_triage_taskflow.yaml +++ b/taskflows/examples/example_triage_taskflow.yaml @@ -2,40 +2,39 @@ seclab-taskflow-agent: version: 1 filetype: taskflow - filekey: GitHubSecurityLab/seclab-taskflow-agent/taskflows/examples/example_triage_taskflow taskflow: - task: must_complete: true agents: - - GitHubSecurityLab/seclab-taskflow-agent/personalities/assistant + - personalities.assistant user_prompt: | Store the json array ["apples", "oranges", "bananas"] in the `fruits` memory key. env: MEMCACHE_STATE_DIR: "example_taskflow/" MEMCACHE_BACKEND: "dictionary_file" toolboxes: - - GitHubSecurityLab/seclab-taskflow-agent/toolboxes/memcache + - toolboxes.memcache - task: must_complete: true agents: - - GitHubSecurityLab/seclab-taskflow-agent/personalities/assistant + - personalities.assistant user_prompt: | Retrieve the contents of the `fruits` memory key. env: MEMCACHE_STATE_DIR: "example_taskflow/" MEMCACHE_BACKEND: "dictionary_file" toolboxes: - - GitHubSecurityLab/seclab-taskflow-agent/toolboxes/memcache + - toolboxes.memcache - task: model: gpt-4.1 repeat_prompt: true agents: # primary agent for this task - - GitHubSecurityLab/seclab-taskflow-agent/personalities/examples/example_triage_agent + - personalities.examples.example_triage_agent # handoff agents - - GitHubSecurityLab/seclab-taskflow-agent/personalities/examples/apple_expert - - GitHubSecurityLab/seclab-taskflow-agent/personalities/examples/orange_expert - - GitHubSecurityLab/seclab-taskflow-agent/personalities/examples/banana_expert + - personalities.examples.apple_expert + - personalities.examples.orange_expert + - personalities.examples.banana_expert user_prompt: | Tell me more about how {{ RESULT }} are grown. diff --git a/taskflows/examples/single_step_taskflow.yaml b/taskflows/examples/single_step_taskflow.yaml index 3914ac8..455cc91 100644 --- a/taskflows/examples/single_step_taskflow.yaml +++ b/taskflows/examples/single_step_taskflow.yaml @@ -1,12 +1,11 @@ seclab-taskflow-agent: version: 1 filetype: taskflow - filekey: GitHubSecurityLab/seclab-taskflow-agent/taskflows/examples/single_step_taskflow taskflow: - task: model: gpt-4.1 agents: - - GitHubSecurityLab/seclab-taskflow-agent/personalities/assistant + - personalities.assistant user_prompt: | Explain the plot of William Gibson's Neuromancer in a single paragraph. diff --git a/tests/data/test_yaml_parser_personality000.yaml b/tests/data/test_yaml_parser_personality000.yaml new file mode 100644 index 0000000..ec84f89 --- /dev/null +++ b/tests/data/test_yaml_parser_personality000.yaml @@ -0,0 +1,8 @@ +seclab-taskflow-agent: + version: 1 + filetype: personality + +personality: | + You are a helpful assistant. +task: | + Answer any question. diff --git a/tests/test_yaml_parser.py b/tests/test_yaml_parser.py index b9d1a42..8bcde74 100644 --- a/tests/test_yaml_parser.py +++ b/tests/test_yaml_parser.py @@ -8,85 +8,37 @@ import tempfile from pathlib import Path import yaml -from yaml_parser import YamlParser - +from available_tools import AvailableTools class TestYamlParser: """Test suite for YamlParser class.""" def test_yaml_parser_basic_functionality(self): """Test basic YAML parsing functionality.""" - # create a temporary directory with test yaml files - with tempfile.TemporaryDirectory() as temp_dir: - temp_path = Path(temp_dir) - test_yaml_content = { - 'seclab-taskflow-agent': { - 'type': 'taskflow', - 'version': 1 - }, - 'taskflow': [ - { - 'task': { - 'agents': ['assistant'], - 'user_prompt': 'Test prompt' - } - } - ] - } - - test_file = temp_path / 'test_taskflow.yaml' - with open(test_file, 'w') as f: - yaml.dump(test_yaml_content, f) - - parser = YamlParser(temp_path) - # get all yaml files in the directory - yaml_files = temp_path.glob('*.yaml') - result = parser.get_yaml_dict(yaml_files) - - assert 'test_taskflow.yaml' in result - assert result['test_taskflow.yaml']['seclab-taskflow-agent']['type'] == 'taskflow' - assert len(result['test_taskflow.yaml']['taskflow']) == 1 - assert result['test_taskflow.yaml']['taskflow'][0]['task']['agents'] == ['assistant'] - + available_tools = AvailableTools() + personality000 = available_tools.get_personality( + "tests.data.test_yaml_parser_personality000") + + assert personality000['seclab-taskflow-agent']['version'] == 1 + assert personality000['seclab-taskflow-agent']['filetype'] == 'personality' + assert personality000['personality'] == 'You are a helpful assistant.\n' + assert personality000['task'] == 'Answer any question.\n' class TestRealTaskflowFiles: """Test parsing of actual taskflow files in the project.""" def test_parse_example_taskflows(self): - """Test parsing the actual example taskflow files.""" + """Test parsing example taskflow files.""" # this test uses the actual taskflows in the project - examples_path = Path('taskflows/examples').absolute() - parser = YamlParser(examples_path) - - # Get all YAML files in the examples directory - yaml_files = examples_path.glob('*.yaml') - result = parser.get_yaml_dict(yaml_files) - - # should contain example files - assert len(result) > 0 + available_tools = AvailableTools() # check that example.yaml is parsed correctly - example_task_flow = result['example.yaml'] + example_task_flow = available_tools.get_taskflow( + "taskflows.examples.example") assert 'taskflow' in example_task_flow assert isinstance(example_task_flow['taskflow'], list) assert len(example_task_flow['taskflow']) == 4 # 4 tasks in taskflow assert example_task_flow['taskflow'][0]['task']['max_steps'] == 20 - def test_parse_all_taskflows(self): - """Test parsing all example taskflow files in the project.""" - taskflows_path = Path('taskflows').absolute() - parser = YamlParser(taskflows_path) - - yaml_files = taskflows_path.rglob('*.yaml') - result = parser.get_yaml_dict(yaml_files) - - # should contain all taskflow files (including subdirs) - assert len(result) == 13 - - # check access for files with directory structure in names - example_files = [key for key in result.keys() if 'examples/' in key] - assert len(example_files) > 0 - - if __name__ == '__main__': - pytest.main([__file__, '-v']) \ No newline at end of file + pytest.main([__file__, '-v']) diff --git a/toolboxes/codeql.yaml b/toolboxes/codeql.yaml index 865e9fa..702190b 100644 --- a/toolboxes/codeql.yaml +++ b/toolboxes/codeql.yaml @@ -1,7 +1,6 @@ seclab-taskflow-agent: version: 1 filetype: toolbox - filekey: GitHubSecurityLab/seclab-taskflow-agent/toolboxes/codeql server_params: kind: streamable diff --git a/toolboxes/echo.yaml b/toolboxes/echo.yaml index 2dc3c99..9032724 100644 --- a/toolboxes/echo.yaml +++ b/toolboxes/echo.yaml @@ -1,7 +1,6 @@ seclab-taskflow-agent: version: 1 filetype: toolbox - filekey: GitHubSecurityLab/seclab-taskflow-agent/toolboxes/echo server_params: kind: stdio diff --git a/toolboxes/github_official.yaml b/toolboxes/github_official.yaml index e5f718c..60f6a59 100644 --- a/toolboxes/github_official.yaml +++ b/toolboxes/github_official.yaml @@ -1,7 +1,6 @@ seclab-taskflow-agent: version: 1 filetype: toolbox - filekey: GitHubSecurityLab/seclab-taskflow-agent/toolboxes/github_official server_params: kind: streamable diff --git a/toolboxes/logbook.yaml b/toolboxes/logbook.yaml index c7142a4..f2c4dfe 100644 --- a/toolboxes/logbook.yaml +++ b/toolboxes/logbook.yaml @@ -1,7 +1,6 @@ seclab-taskflow-agent: version: 1 filetype: toolbox - filekey: GitHubSecurityLab/seclab-taskflow-agent/toolboxes/logbook server_params: kind: stdio diff --git a/toolboxes/memcache.yaml b/toolboxes/memcache.yaml index 9a643e7..6105461 100644 --- a/toolboxes/memcache.yaml +++ b/toolboxes/memcache.yaml @@ -1,7 +1,6 @@ seclab-taskflow-agent: version: 1 filetype: toolbox - filekey: GitHubSecurityLab/seclab-taskflow-agent/toolboxes/memcache server_params: kind: stdio diff --git a/yaml_parser.py b/yaml_parser.py deleted file mode 100644 index ea6b0a9..0000000 --- a/yaml_parser.py +++ /dev/null @@ -1,29 +0,0 @@ -import os -import yaml -from pathlib import Path -from collections.abc import Iterator -import logging - -class YamlParser: - def __init__(self, rootpath: Path): - self.rootpath = rootpath.absolute() - - def get_yaml_dict(self, files: Iterator) -> dict: - """Return a directory of yaml files as a dictionary of file name indexed yaml dicts""" - yaml_dict = {} - for f in files: - if f.is_file() and not f.name.startswith('.') and f.suffix in ['.yaml', '.yml']: - try: - with open(f) as stream: - try: - p = yaml.safe_load(stream) - # Use the relative path as the key in the dict. - namespaced_stem = str(f.relative_to(self.rootpath)) - yaml_dict[namespaced_stem] = p - except yaml.YAMLError as e: - logging.error(e) - except FileNotFoundError as e: - # deal with editor temp files etc. that might have disappeared - logging.error(f"File not found (editor tmp file?): {e}") - continue - return yaml_dict