diff --git a/README.md b/README.md index f50746a..24ab58d 100644 --- a/README.md +++ b/README.md @@ -77,6 +77,18 @@ Example: deploying a Taskflow: hatch run main -t examples.taskflows.example ``` +Example: deploying a Taskflow with command line global variables: + +```sh +hatch run main -t examples.taskflows.example_globals -g fruit=apples +``` + +Multiple global variables can be set: + +```sh +hatch run main -t examples.taskflows.example_globals -g fruit=apples -g color=red +``` + ## Deploying from Docker You can deploy the Taskflow Agent via its Docker image using `docker/run.sh`. @@ -104,6 +116,13 @@ Example: deploying a Taskflow (example.yaml): ```sh docker/run.sh -t example ``` + +Example: deploying a Taskflow with global variables: + +```sh +docker/run.sh -t example_globals -g fruit=apples +``` + Example: deploying a custom taskflow (custom_taskflow.yaml): ```sh diff --git a/doc/GRAMMAR.md b/doc/GRAMMAR.md index 4bac01a..17ca0ec 100644 --- a/doc/GRAMMAR.md +++ b/doc/GRAMMAR.md @@ -340,6 +340,20 @@ taskflow: Tell me more about {{ GLOBALS_fruit }}. ``` +Global variables can also be set or overridden from the command line using the `-g` or `--global` flag: + +```sh +hatch run main -t examples.taskflows.example_globals -g fruit=apples +``` + +Multiple global variables can be set by repeating the flag: + +```sh +hatch run main -t examples.taskflows.example_globals -g fruit=apples -g color=red +``` + +Command line globals override any globals defined in the taskflow YAML file, allowing you to reuse taskflows with different parameter values without editing the files. + ### Reusable Tasks Tasks can reuse single step taskflows and optionally override any of its configurations. This is done by setting a `uses` field with a link to the single step taskflow YAML file as its value. diff --git a/src/seclab_taskflow_agent/__main__.py b/src/seclab_taskflow_agent/__main__.py index 5c59466..aadacf3 100644 --- a/src/seclab_taskflow_agent/__main__.py +++ b/src/seclab_taskflow_agent/__main__.py @@ -67,21 +67,35 @@ def parse_prompt_args(available_tools: AvailableTools, group.add_argument("-p", help="The personality to use (mutex with -t)", required=False) group.add_argument("-t", help="The taskflow to use (mutex with -p)", required=False) group.add_argument("-l", help="List available tool call models and exit", action='store_true', required=False) + parser.add_argument("-g", "--global", dest="globals", action='append', help="Set global variable (KEY=VALUE). Can be used multiple times.", required=False) parser.add_argument('prompt', nargs=argparse.REMAINDER) #parser.add_argument('remainder', nargs=argparse.REMAINDER, help="Remaining args") help_msg = parser.format_help() help_msg += "\nExamples:\n\n" help_msg += "`-p assistant explain modems to me please`\n" + help_msg += "`-t example -g fruit=apples`\n" + help_msg += "`-t example -g fruit=apples -g color=red`\n" try: args = parser.parse_known_args(user_prompt.split(' ') if user_prompt else None) except SystemExit as e: if e.code == 2: logging.error(f"User provided incomplete prompt: {user_prompt}") - return None, None, None, help_msg + return None, None, None, None, help_msg p = args[0].p.strip() if args[0].p else None t = args[0].t.strip() if args[0].t else None l = args[0].l - return p, t, l, ' '.join(args[0].prompt), help_msg + + # Parse global variables from command line + cli_globals = {} + if args[0].globals: + for g in args[0].globals: + if '=' not in g: + logging.error(f"Invalid global variable format: {g}. Expected KEY=VALUE") + return None, None, None, None, None, help_msg + key, value = g.split('=', 1) + cli_globals[key.strip()] = value.strip() + + return p, t, l, cli_globals, ' '.join(args[0].prompt), help_msg async def deploy_task_agents(available_tools: AvailableTools, agents: dict, @@ -378,7 +392,7 @@ async def _run_streamed(): async def main(available_tools: AvailableTools, - p: str | None, t: str | None, prompt: str | None): + p: str | None, t: str | None, cli_globals: dict, prompt: str | None): last_mcp_tool_results = [] # XXX: memleaky async def on_tool_end_hook( @@ -418,7 +432,10 @@ async def on_handoff_hook( await render_model_output(f"** 🤖💪 Running Task Flow: {t}\n") # optional global vars available for the taskflow tasks + # Start with globals from taskflow file, then override with CLI globals global_variables = taskflow.get('globals', {}) + if cli_globals: + global_variables.update(cli_globals) model_config = taskflow.get('model_config', {}) model_keys = [] if model_config: @@ -646,7 +663,7 @@ async def _deploy_task_agents(resolved_agents, prompt): cwd = pathlib.Path.cwd() available_tools = AvailableTools() - p, t, l, user_prompt, help_msg = parse_prompt_args(available_tools) + p, t, l, cli_globals, user_prompt, help_msg = parse_prompt_args(available_tools) if l: tool_models = list_tool_call_models(os.getenv('COPILOT_TOKEN')) @@ -658,4 +675,4 @@ async def _deploy_task_agents(resolved_agents, prompt): print(help_msg) sys.exit(1) - asyncio.run(main(available_tools, p, t, user_prompt), debug=True) + asyncio.run(main(available_tools, p, t, cli_globals, user_prompt), debug=True) diff --git a/tests/data/test_globals_taskflow.yaml b/tests/data/test_globals_taskflow.yaml new file mode 100644 index 0000000..537ee9e --- /dev/null +++ b/tests/data/test_globals_taskflow.yaml @@ -0,0 +1,13 @@ +# SPDX-FileCopyrightText: 2025 GitHub +# SPDX-License-Identifier: MIT + +seclab-taskflow-agent: + version: 1 + filetype: taskflow + +globals: + test_var: default_value +taskflow: + - task: + run: | + echo "test" diff --git a/tests/test_yaml_parser.py b/tests/test_yaml_parser.py index 743937b..9887a54 100644 --- a/tests/test_yaml_parser.py +++ b/tests/test_yaml_parser.py @@ -43,5 +43,66 @@ def test_parse_example_taskflows(self): assert len(example_task_flow['taskflow']) == 4 # 4 tasks in taskflow assert example_task_flow['taskflow'][0]['task']['max_steps'] == 20 +class TestCliGlobals: + """Test CLI global variable parsing.""" + + def test_parse_single_global(self): + """Test parsing a single global variable from command line.""" + from seclab_taskflow_agent.__main__ import parse_prompt_args + available_tools = AvailableTools() + + p, t, l, cli_globals, user_prompt, _ = parse_prompt_args( + available_tools, "-t example -g fruit=apples") + + assert t == "example" + assert cli_globals == {"fruit": "apples"} + assert p is None + assert l is False + + def test_parse_multiple_globals(self): + """Test parsing multiple global variables from command line.""" + from seclab_taskflow_agent.__main__ import parse_prompt_args + available_tools = AvailableTools() + + p, t, l, cli_globals, user_prompt, _ = parse_prompt_args( + available_tools, "-t example -g fruit=apples -g color=red") + + assert t == "example" + assert cli_globals == {"fruit": "apples", "color": "red"} + assert p is None + assert l is False + + def test_parse_global_with_spaces(self): + """Test parsing global variables with spaces in values.""" + from seclab_taskflow_agent.__main__ import parse_prompt_args + available_tools = AvailableTools() + + p, t, l, cli_globals, user_prompt, _ = parse_prompt_args( + available_tools, "-t example -g message=hello world") + + assert t == "example" + # "world" becomes part of the prompt, not the value + assert cli_globals == {"message": "hello"} + assert "world" in user_prompt + + def test_parse_global_with_equals_in_value(self): + """Test parsing global variables with equals sign in value.""" + from seclab_taskflow_agent.__main__ import parse_prompt_args + available_tools = AvailableTools() + + p, t, l, cli_globals, user_prompt, _ = parse_prompt_args( + available_tools, "-t example -g equation=x=5") + + assert t == "example" + assert cli_globals == {"equation": "x=5"} + + def test_globals_in_taskflow_file(self): + """Test that globals can be read from taskflow file.""" + available_tools = AvailableTools() + + taskflow = available_tools.get_taskflow("tests.data.test_globals_taskflow") + assert 'globals' in taskflow + assert taskflow['globals']['test_var'] == 'default_value' + if __name__ == '__main__': pytest.main([__file__, '-v'])