Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 19 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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`.
Expand Down Expand Up @@ -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
Expand Down
14 changes: 14 additions & 0 deletions doc/GRAMMAR.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
27 changes: 22 additions & 5 deletions src/seclab_taskflow_agent/__main__.py
Original file line number Diff line number Diff line change
Expand Up @@ -59,29 +59,43 @@
MAX_API_RETRY = 5
MCP_CLEANUP_TIMEOUT = 5

def parse_prompt_args(available_tools: AvailableTools,
user_prompt: str | None = None):

Check notice

Code scanning / CodeQL

Returning tuples with varying lengths Note

parse_prompt_args returns
tuple of size 5
and
tuple of size 6
.
parse_prompt_args returns
tuple of size 5
and
tuple of size 6
.
parser = argparse.ArgumentParser(add_help=False, description="SecLab Taskflow Agent")
parser.prog = ''
group = parser.add_mutually_exclusive_group()
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,
Expand Down Expand Up @@ -378,7 +392,7 @@


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(
Expand Down Expand Up @@ -418,7 +432,10 @@
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:
Expand Down Expand Up @@ -646,7 +663,7 @@
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'))
Expand All @@ -658,4 +675,4 @@
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)
13 changes: 13 additions & 0 deletions tests/data/test_globals_taskflow.yaml
Original file line number Diff line number Diff line change
@@ -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"
61 changes: 61 additions & 0 deletions tests/test_yaml_parser.py
Original file line number Diff line number Diff line change
Expand Up @@ -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'])