diff --git a/README.md b/README.md index 2730a71..1878e77 100644 --- a/README.md +++ b/README.md @@ -6,6 +6,24 @@ The Taskflow Agent is built on top of the [OpenAI Agents SDK](https://openai.git While the Taskflow Agent does not integrate into the GitHub Dotcom Copilot UX, it does operate using the Copilot API (CAPI) as its backend, similar to Copilot IDE extensions. +## Template Syntax Migration (v2) + +**Breaking Change:** Taskflow YAML files now use Jinja2 templating (version 2). Version 1 files are no longer supported and will be rejected at load time. + +**New Jinja2 syntax:** +- `{{ globals.key }}` instead of `{{ GLOBALS_key }}` +- `{{ inputs.key }}` instead of `{{ INPUTS_key }}` +- `{{ result }}` / `{{ result.key }}` instead of `{{ RESULT }}` / `{{ RESULT_key }}` +- `{{ env('VAR') }}` instead of `{{ env VAR }}` +- `{% include 'path' %}` instead of `{{ PROMPTS_path }}` + +**To migrate existing taskflows:** +```bash +python scripts/migrate_to_jinja2.py /path/to/your/taskflows +``` + +See [doc/MIGRATION.md](doc/MIGRATION.md) for detailed migration instructions and new Jinja2 features. + ## Core Concepts The Taskflow Agent leverages a GitHub Workflow-esque YAML based grammar to perform a series of tasks using a set of Agents. @@ -296,10 +314,10 @@ server_params: url: https://api.githubcopilot.com/mcp/ #See https://github.com/github/github-mcp-server/blob/main/docs/remote-server.md headers: - Authorization: "{{ env GITHUB_AUTH_HEADER }}" + Authorization: "{{ env('GITHUB_AUTH_HEADER') }}" optional_headers: - X-MCP-Toolsets: "{{ env GITHUB_MCP_TOOLSETS }}" - X-MCP-Readonly: "{{ env GITHUB_MCP_READONLY }}" + X-MCP-Toolsets: "{{ env('GITHUB_MCP_TOOLSETS') }}" + X-MCP-Readonly: "{{ env('GITHUB_MCP_READONLY') }}" ``` You can force certain tools within a `toolbox` to require user confirmation to run. This can be helpful if a tool may perform irreversible actions and should require user approval prior to its use. This is done by including the name of the tool (function) in the MCP server in the `confirm` section: @@ -347,7 +365,7 @@ taskflow: Finally, why are apples and oranges healthy to eat? # taskflows can set temporary environment variables, these support the general - # "{{ env FROM_EXISTING_ENVIRONMENT }" pattern we use elsewhere as well + # "{{ env('FROM_EXISTING_ENVIRONMENT') }}" pattern we use elsewhere as well # these environment variables can then be made available to any stdio mcp server # through its respective yaml configuration, see memcache.yaml for an example # you can use these to override top-level environment variables on a per-task basis @@ -494,12 +512,12 @@ Files of types `taskflow` and `toolbox` allow environment variables to be passed server_params: ... env: - CODEQL_DBS_BASE_PATH: "{{ env CODEQL_DBS_BASE_PATH }}" + CODEQL_DBS_BASE_PATH: "{{ env('CODEQL_DBS_BASE_PATH') }}" # prevent git repo operations on gh codeql executions GH_NO_UPDATE_NOTIFIER: "disable" ``` -For `toolbox`, `env` can be used inside `server_params`. A template of the form `{{ env ENV_VARIABLE_NAME }}` can be used to pass values of the environment variable from the current process to the MCP server. So in the above, the MCP server is run with `GH_NO_UPDATE_NOTIFIER=disable` and passes the value of `CODEQL_DBS_BASE_PATH` from the current process to the MCP server. The templated paramater `{{ env CODEQL_DBS_BASE_PATH }}` is replaced by the value of the environment variable `CODEQL_DBS_BASE_PATH` in the current process. +For `toolbox`, `env` can be used inside `server_params`. A template of the form `{{ env('ENV_VARIABLE_NAME') }}` can be used to pass values of the environment variable from the current process to the MCP server. So in the above, the MCP server is run with `GH_NO_UPDATE_NOTIFIER=disable` and passes the value of `CODEQL_DBS_BASE_PATH` from the current process to the MCP server. The templated parameter `{{ env('CODEQL_DBS_BASE_PATH') }}` is replaced by the value of the environment variable `CODEQL_DBS_BASE_PATH` in the current process. Similarly, environment variables can be passed to a `task` in a `taskflow`: @@ -516,9 +534,9 @@ taskflow: MEMCACHE_BACKEND: "dictionary_file" ``` -This overwrites the environment variables `MEMCACHE_STATE_DIR` and `MEMCACHE_BACKEND` for the task only. A template `{{ env ENV_VARIABLE_NAME }}` can also be used. +This overwrites the environment variables `MEMCACHE_STATE_DIR` and `MEMCACHE_BACKEND` for the task only. A template `{{ env('ENV_VARIABLE_NAME') }}` can also be used. -Note that when using the template `{{ env ENV_VARIABLE_NAME }}`, `ENV_VARIABLE_NAME` must be the name of an environment variable in the current process. +Note that when using the template `{{ env('ENV_VARIABLE_NAME') }}`, `ENV_VARIABLE_NAME` must be the name of an environment variable in the current process. ## Import paths diff --git a/doc/GRAMMAR.md b/doc/GRAMMAR.md index c0581f4..1ebfaf5 100644 --- a/doc/GRAMMAR.md +++ b/doc/GRAMMAR.md @@ -133,10 +133,10 @@ Often we may want to iterate through the same tasks with different inputs. For e agents: - seclab_taskflow_agent.personalities.c_auditer user_prompt: | - The function has name {{ RESULT_name }} and body {{ RESULT_body }} analyze the function. + The function has name {{ result.name }} and body {{ result.body }} analyze the function. ``` -In the above, the first task fetches functions in the code base and creates a json list object, with each entry having a `name` and `body` field. In the next task, `repeat_prompt` is set to true, meaning that a task is created for each individual object in the list and the object fields are referenced in the templated prompt using `{{ RESULT_ }}`. In other words, `{{ RESULT_name }}` in the prompt is replaced with the value of the `name` field of the object etc. For example, if the list of functions fetched from the first task is: +In the above, the first task fetches functions in the code base and creates a json list object, with each entry having a `name` and `body` field. In the next task, `repeat_prompt` is set to true, meaning that a task is created for each individual object in the list and the object fields are referenced in the templated prompt using `{{ result.fieldname }}`. In other words, `{{ result.name }}` in the prompt is replaced with the value of the `name` field of the object etc. For example, if the list of functions fetched from the first task is: ```javascript [{'name' : foo, 'body' : foo(){return 1;}}, {'name' : bar, 'body' : bar(a) {return a + 1;}}] @@ -152,7 +152,7 @@ etc. Note that when using `repeat_prompt`, the last tool call result of the previous task is used as the iterable. It is recommended to keep the task that creates the iterable short and simple (e.g. just make one tool call to fetch a list of results) to avoid wrong results being passed to the repeat prompt. -The iterable can also contain a list of primitives like string or number, in which case, the template `{{ RESULT }}` can be used in the `repeat_prompt` prompt to parse the results instead: +The iterable can also contain a list of primitives like string or number, in which case, the template `{{ result }}` can be used in the `repeat_prompt` prompt to parse the results instead: ```yaml - task: @@ -173,7 +173,7 @@ The iterable can also contain a list of primitives like string or number, in whi agents: - seclab_taskflow_agent.personalities.assistant user_prompt: | - What is the integer value of {{ RESULT }}? + What is the integer value of {{ result }}? ``` Repeat prompt can be run in parallel by setting the `async` field to `true`: @@ -185,7 +185,7 @@ Repeat prompt can be run in parallel by setting the `async` field to `true`: agents: - seclab_taskflow_agent.personalities.c_auditer user_prompt: | - The function has name {{ RESULT_name }} and body {{ RESULT_body }} analyze the function. + The function has name {{ result.name }} and body {{ result.body }} analyze the function. ``` An optional limit can be set to limit the number of asynchronous tasks via `async_limit`. If not set, the default value (5) is used. @@ -198,7 +198,7 @@ An optional limit can be set to limit the number of asynchronous tasks via `asyn agents: - seclab_taskflow_agent.personalities.c_auditer user_prompt: | - The function has name {{ RESULT_name }} and body {{ RESULT_body }} analyze the function. + The function has name {{ result.name }} and body {{ result.body }} analyze the function. ``` Both `async` and `async_limit` have no effect when used outside of a `repeat_prompt`. @@ -211,7 +211,7 @@ At the moment, we do not support nested `repeat_prompt`. So the following is not agents: - seclab_taskflow_agent.personalities.c_auditer user_prompt: | - The function has name {{ RESULT_name }} and body {{ RESULT_body }} analyze the function. + The function has name {{ result.name }} and body {{ result.body }} analyze the function. - task: repeat_prompt: true ... @@ -233,7 +233,7 @@ For example: agents: - seclab_taskflow_agent.personalities.assistant user_prompt: | - What kind of fruit is {{ RESULT }}? + What kind of fruit is {{ result }}? ``` The string `["apple", "banana", "orange"]` is then passed directly to the next task. @@ -349,7 +349,7 @@ taskflow: agents: - examples.personalities.fruit_expert user_prompt: | - Tell me more about {{ GLOBALS_fruit }}. + Tell me more about {{ globals.fruit }}. ``` Global variables can also be set or overridden from the command line using the `-g` or `--global` flag: @@ -422,10 +422,10 @@ A reusable taskflow can also have a templated prompt that takes inputs from its agents: - examples.personalities.fruit_expert user_prompt: | - Tell me more about {{ INPUTS_fruit }}. + Tell me more about {{ inputs.fruit }}. ``` -In this case, the template parameter `{{ INPUTS_fruit }}` is replaced by the value of `fruit` from the `inputs` of the user, which is apples in this case: +In this case, the template parameter `{{ inputs.fruit }}` is replaced by the value of `fruit` from the `inputs` of the user, which is apples in this case: ```yaml - task: @@ -437,9 +437,9 @@ In this case, the template parameter `{{ INPUTS_fruit }}` is replaced by the val ### Reusable Prompts -Reusable prompts are defined in files of `filetype` `prompts`. These are like macros that get replaced when a templated parameter of the form `{{ PROMPTS_ }}` is encountered. +Reusable prompts are defined in files of `filetype` `prompts`. These are like macros that get included using Jinja2's `{% include %}` directive. -Tasks can incorporate templated prompts which are then replaced by the actual prompt. For example: +Tasks can incorporate reusable prompts using the include directive. For example: Example: @@ -449,8 +449,8 @@ Example: - examples.personalities.fruit_expert user_prompt: | Tell me more about apples. - - {{ PROMPTS_examples.prompts.example_prompt }} + + {% include 'examples.prompts.example_prompt' %} ``` and `examples.prompts.example_prompt` is the following: diff --git a/doc/MIGRATION.md b/doc/MIGRATION.md new file mode 100644 index 0000000..4310cb4 --- /dev/null +++ b/doc/MIGRATION.md @@ -0,0 +1,259 @@ +# Jinja2 Templating Migration Guide + +This guide explains how to migrate taskflow YAML files from version 1 (custom template syntax) to version 2 (Jinja2 templating). + +## Overview + +Version 2 replaces the custom regex-based template processing with Jinja2, providing: +- More powerful templating features (filters, conditionals, loops) +- Better error messages with clear variable undefined errors +- Industry-standard syntax familiar to many developers +- Extensibility for future template features + +## Syntax Changes + +### 1. Global Variables + +**Version 1:** +```yaml +globals: + fruit: apples +taskflow: + - task: + user_prompt: | + Tell me about {{ GLOBALS_fruit }}. +``` + +**Version 2:** +```yaml +globals: + fruit: apples +taskflow: + - task: + user_prompt: | + Tell me about {{ globals.fruit }}. +``` + +**Nested structures:** +```yaml +globals: + config: + model: gpt-4 + temperature: 0.7 +taskflow: + - task: + user_prompt: | + Using {{ globals.config.model }} with temp {{ globals.config.temperature }} +``` + +### 2. Input Variables + +**Version 1:** +```yaml +user_prompt: | + Color: {{ INPUTS_color }} +``` + +**Version 2:** +```yaml +user_prompt: | + Color: {{ inputs.color }} +``` + +### 3. Result Variables + +**Version 1 (primitives):** +```yaml +repeat_prompt: true +user_prompt: | + Process {{ RESULT }} +``` + +**Version 2:** +```yaml +repeat_prompt: true +user_prompt: | + Process {{ result }} +``` + +**Version 1 (dictionary keys):** +```yaml +user_prompt: | + Function {{ RESULT_name }} has body {{ RESULT_body }} +``` + +**Version 2:** +```yaml +user_prompt: | + Function {{ result.name }} has body {{ result.body }} +``` + +### 4. Environment Variables + +**Version 1:** +```yaml +env: + DATABASE: "{{ env DATABASE_URL }}" +``` + +**Version 2:** +```yaml +env: + DATABASE: "{{ env('DATABASE_URL') }}" +``` + +**With defaults (new feature):** +```yaml +env: + DATABASE: "{{ env('DATABASE_URL', 'localhost:5432') }}" +``` + +### 5. Reusable Prompts + +**Version 1:** +```yaml +user_prompt: | + Main task. + {{ PROMPTS_examples.prompts.shared }} +``` + +**Version 2:** +```yaml +user_prompt: | + Main task. + {% include 'examples.prompts.shared' %} +``` + +## New Jinja2 Features + +### Filters + +Transform values with filters: + +```yaml +user_prompt: | + Uppercase: {{ globals.name | upper }} + Lowercase: {{ globals.name | lower }} + Default: {{ globals.optional | default('N/A') }} + List length: {{ globals.items | length }} +``` + +### Conditionals + +Add conditional logic: + +```yaml +user_prompt: | + {% if globals.debug_mode %} + Running in debug mode + {% else %} + Running in production mode + {% endif %} + + {% if result.score > 0.8 %} + High confidence result + {% endif %} +``` + +### Loops + +Iterate over collections: + +```yaml +user_prompt: | + Analyze these functions: + {% for func in result.functions %} + - {{ func.name }}: {{ func.complexity }} + {% endfor %} +``` + +### Math Operations + +Perform calculations: + +```yaml +user_prompt: | + Sum: {{ result.a + result.b }} + Product: {{ result.count * 2 }} + Comparison: {% if result.score > 0.5 %}Pass{% else %}Fail{% endif %} +``` + +## Automated Migration + +Use the provided migration script: + +```bash +# Migrate all YAML files in directory +python scripts/migrate_to_jinja2.py /path/to/taskflows + +# Preview changes without writing +python scripts/migrate_to_jinja2.py --dry-run /path/to/taskflows + +# Migrate specific file +python scripts/migrate_to_jinja2.py myflow.yaml +``` + +## Manual Migration Checklist + +1. Update YAML version from `1` to `2` +2. Replace `{{ GLOBALS_` with `{{ globals.` +3. Replace `{{ INPUTS_` with `{{ inputs.` +4. Replace `{{ RESULT_` with `{{ result.` +5. Replace `{{ RESULT }}` with `{{ result }}` +6. Replace `{{ env VAR }}` with `{{ env('VAR') }}` +7. Replace `{{ PROMPTS_` with `{% include '` and add closing `' %}` +8. Test taskflow execution + +## Testing Your Migration + +```bash +# Run specific taskflow +python -m seclab_taskflow_agent -t your.taskflow.name + +# Run with globals +python -m seclab_taskflow_agent -t your.taskflow.name -g key=value +``` + +## Common Issues + +### Issue: `UndefinedError: 'globals' is undefined` + +**Cause:** Using `{{ globals.key }}` when no globals are defined + +**Fix:** Either define globals in taskflow or use Jinja2's get method: +```yaml +{{ globals.get('key', 'default') }} +``` + +### Issue: `TemplateNotFound: examples.prompts.mypromp` + +**Cause:** Typo in include path + +**Fix:** Verify path matches file location exactly + +### Issue: Environment variable errors + +**Cause:** Required env var not set + +**Fix:** Set env var or make it optional: +```yaml +{{ env('VAR', 'default') }} +``` + +## Backwards Compatibility + +Version 1 syntax is no longer supported. Attempting to load a v1 file will fail with: + +``` +VersionException: YAML file uses unsupported version 1 template syntax. +Version 2 (Jinja2) is required. +Migrate using: python scripts/migrate_to_jinja2.py +``` + +All v1 files must be migrated to v2 before use. + +## Additional Resources + +- [Jinja2 Documentation](https://jinja.palletsprojects.com/) +- [Jinja2 Template Designer Documentation](https://jinja.palletsprojects.com/en/3.1.x/templates/) +- Example taskflows in `examples/taskflows/` diff --git a/examples/model_configs/model_config.yaml b/examples/model_configs/model_config.yaml index c04dfd4..aa149d0 100644 --- a/examples/model_configs/model_config.yaml +++ b/examples/model_configs/model_config.yaml @@ -2,7 +2,7 @@ # SPDX-License-Identifier: MIT seclab-taskflow-agent: - version: 1 + version: "1.0" filetype: model_config models: sonnet_default: claude-sonnet-4 diff --git a/examples/personalities/apple_expert.yaml b/examples/personalities/apple_expert.yaml index 2c0b4cb..371a317 100644 --- a/examples/personalities/apple_expert.yaml +++ b/examples/personalities/apple_expert.yaml @@ -2,7 +2,7 @@ # SPDX-License-Identifier: MIT seclab-taskflow-agent: - version: 1 + version: "1.0" filetype: personality personality: | diff --git a/examples/personalities/banana_expert.yaml b/examples/personalities/banana_expert.yaml index 7e18c44..e98ccf7 100644 --- a/examples/personalities/banana_expert.yaml +++ b/examples/personalities/banana_expert.yaml @@ -2,7 +2,7 @@ # SPDX-License-Identifier: MIT seclab-taskflow-agent: - version: 1 + version: "1.0" filetype: personality personality: | diff --git a/examples/personalities/echo.yaml b/examples/personalities/echo.yaml index a5006e2..c9491df 100644 --- a/examples/personalities/echo.yaml +++ b/examples/personalities/echo.yaml @@ -3,7 +3,7 @@ # personalities define the system prompt level directives for this Agent seclab-taskflow-agent: - version: 1 + version: "1.0" filetype: personality personality: | diff --git a/examples/personalities/example_triage_agent.yaml b/examples/personalities/example_triage_agent.yaml index 8fe8b14..e6d8f4f 100644 --- a/examples/personalities/example_triage_agent.yaml +++ b/examples/personalities/example_triage_agent.yaml @@ -2,7 +2,7 @@ # SPDX-License-Identifier: MIT seclab-taskflow-agent: - version: 1 + version: "1.0" filetype: personality personality: | diff --git a/examples/personalities/fruit_expert.yaml b/examples/personalities/fruit_expert.yaml index 243340a..cd976a8 100644 --- a/examples/personalities/fruit_expert.yaml +++ b/examples/personalities/fruit_expert.yaml @@ -2,7 +2,7 @@ # SPDX-License-Identifier: MIT seclab-taskflow-agent: - version: 1 + version: "1.0" filetype: personality personality: | diff --git a/examples/personalities/orange_expert.yaml b/examples/personalities/orange_expert.yaml index a651f1e..5f3d20c 100644 --- a/examples/personalities/orange_expert.yaml +++ b/examples/personalities/orange_expert.yaml @@ -2,7 +2,7 @@ # SPDX-License-Identifier: MIT seclab-taskflow-agent: - version: 1 + version: "1.0" filetype: personality personality: | diff --git a/examples/prompts/example_prompt.yaml b/examples/prompts/example_prompt.yaml index c3da6b6..9f4734b 100644 --- a/examples/prompts/example_prompt.yaml +++ b/examples/prompts/example_prompt.yaml @@ -2,7 +2,7 @@ # SPDX-License-Identifier: MIT seclab-taskflow-agent: - version: 1 + version: "1.0" filetype: prompt prompt: | diff --git a/examples/taskflows/CVE-2023-2283.yaml b/examples/taskflows/CVE-2023-2283.yaml index 984a3f1..f408c1d 100644 --- a/examples/taskflows/CVE-2023-2283.yaml +++ b/examples/taskflows/CVE-2023-2283.yaml @@ -2,7 +2,7 @@ # SPDX-License-Identifier: MIT seclab-taskflow-agent: - version: 1 + version: "1.0" filetype: taskflow model_config: examples.model_configs.model_config diff --git a/examples/taskflows/echo.yaml b/examples/taskflows/echo.yaml index 3056704..2afcf0f 100644 --- a/examples/taskflows/echo.yaml +++ b/examples/taskflows/echo.yaml @@ -2,7 +2,7 @@ # SPDX-License-Identifier: MIT seclab-taskflow-agent: - version: 1 + version: "1.0" filetype: taskflow taskflow: diff --git a/examples/taskflows/example.yaml b/examples/taskflows/example.yaml index 65c2c62..e8abca6 100644 --- a/examples/taskflows/example.yaml +++ b/examples/taskflows/example.yaml @@ -2,7 +2,7 @@ # SPDX-License-Identifier: MIT seclab-taskflow-agent: - version: 1 + version: "1.0" filetype: taskflow # Import settings from a model_config file. @@ -29,7 +29,7 @@ taskflow: Finally, why are apples and oranges healthy to eat? # taskflows can set temporary environment variables, these support the general - # "{{ env FROM_EXISTING_ENVIRONMENT }" pattern we use elsewhere as well + # "{{ env('FROM_EXISTING_ENVIRONMENT') }}" pattern we use elsewhere as well # these environment variables can then be made available to any stdio mcp server # through its respective yaml configuration, see memcache.yaml for an example # you can use these to override top-level environment variables on a per-task basis @@ -70,4 +70,4 @@ taskflow: agents: - seclab_taskflow_agent.personalities.assistant user_prompt: | - What kind of fruit is {{ RESULT }}? + What kind of fruit is {{ result }}? diff --git a/examples/taskflows/example_globals.yaml b/examples/taskflows/example_globals.yaml index a68b4a7..9ab73fc 100644 --- a/examples/taskflows/example_globals.yaml +++ b/examples/taskflows/example_globals.yaml @@ -2,7 +2,7 @@ # SPDX-License-Identifier: MIT seclab-taskflow-agent: - version: 1 + version: "1.0" filetype: taskflow globals: @@ -12,4 +12,4 @@ taskflow: agents: - examples.personalities.fruit_expert user_prompt: | - Tell me more about {{ GLOBALS_fruit }}. + Tell me more about {{ globals.fruit }}. diff --git a/examples/taskflows/example_inputs.yaml b/examples/taskflows/example_inputs.yaml index a8f9997..3c0ed8c 100644 --- a/examples/taskflows/example_inputs.yaml +++ b/examples/taskflows/example_inputs.yaml @@ -2,7 +2,7 @@ # SPDX-License-Identifier: MIT seclab-taskflow-agent: - version: 1 + version: "1.0" filetype: taskflow taskflow: @@ -12,5 +12,5 @@ taskflow: inputs: fruit: apples user_prompt: | - Tell me more about {{ INPUTS_fruit }}. + Tell me more about {{ inputs.fruit }}. diff --git a/examples/taskflows/example_large_list_result_iter.yaml b/examples/taskflows/example_large_list_result_iter.yaml index 7d1ce6c..9fe4bc4 100644 --- a/examples/taskflows/example_large_list_result_iter.yaml +++ b/examples/taskflows/example_large_list_result_iter.yaml @@ -2,7 +2,7 @@ # SPDX-License-Identifier: MIT seclab-taskflow-agent: - version: 1 + version: "1.0" filetype: taskflow taskflow: @@ -24,4 +24,4 @@ taskflow: agents: - seclab_taskflow_agent.personalities.assistant user_prompt: | - Echo this: The title is {{ RESULT_title }} and the url is {{ RESULT_url }}. + Echo this: The title is {{ result.title }} and the url is {{ result.url }}. diff --git a/examples/taskflows/example_repeat_prompt.yaml b/examples/taskflows/example_repeat_prompt.yaml index 2336a93..68b6154 100644 --- a/examples/taskflows/example_repeat_prompt.yaml +++ b/examples/taskflows/example_repeat_prompt.yaml @@ -2,7 +2,7 @@ # SPDX-License-Identifier: MIT seclab-taskflow-agent: - version: 1 + version: "1.0" filetype: taskflow taskflow: @@ -28,4 +28,4 @@ taskflow: agents: - seclab_taskflow_agent.personalities.assistant user_prompt: | - What is the integer value of {{ RESULT }}? + What is the integer value of {{ result }}? diff --git a/examples/taskflows/example_repeat_prompt_async.yaml b/examples/taskflows/example_repeat_prompt_async.yaml index faaad56..16d9326 100644 --- a/examples/taskflows/example_repeat_prompt_async.yaml +++ b/examples/taskflows/example_repeat_prompt_async.yaml @@ -2,7 +2,7 @@ # SPDX-License-Identifier: MIT seclab-taskflow-agent: - version: 1 + version: "1.0" filetype: taskflow taskflow: @@ -32,4 +32,4 @@ taskflow: agents: - seclab_taskflow_agent.personalities.assistant user_prompt: | - What is the integer value of {{ RESULT }}? + What is the integer value of {{ result }}? diff --git a/examples/taskflows/example_repeat_prompt_dictionary.yaml b/examples/taskflows/example_repeat_prompt_dictionary.yaml index 1ceeecf..8c8dd05 100644 --- a/examples/taskflows/example_repeat_prompt_dictionary.yaml +++ b/examples/taskflows/example_repeat_prompt_dictionary.yaml @@ -2,7 +2,7 @@ # SPDX-License-Identifier: MIT seclab-taskflow-agent: - version: 1 + version: "1.0" filetype: taskflow taskflow: @@ -29,4 +29,4 @@ taskflow: agents: - seclab_taskflow_agent.personalities.assistant user_prompt: | - What is the value of {{ RESULT_index }} + {{ RESULT_value }}? + What is the value of {{ result.index }} + {{ result.value }}? diff --git a/examples/taskflows/example_reusable_prompt.yaml b/examples/taskflows/example_reusable_prompt.yaml index eebb71e..5a8e46f 100644 --- a/examples/taskflows/example_reusable_prompt.yaml +++ b/examples/taskflows/example_reusable_prompt.yaml @@ -2,7 +2,7 @@ # SPDX-License-Identifier: MIT seclab-taskflow-agent: - version: 1 + version: "1.0" filetype: taskflow taskflow: @@ -12,4 +12,4 @@ taskflow: user_prompt: | Tell me more about apples. - {{ PROMPTS_examples.prompts.example_prompt }} + {% include 'examples.prompts.example_prompt' %} diff --git a/examples/taskflows/example_reusable_taskflows.yaml b/examples/taskflows/example_reusable_taskflows.yaml index 06748b2..95722ab 100644 --- a/examples/taskflows/example_reusable_taskflows.yaml +++ b/examples/taskflows/example_reusable_taskflows.yaml @@ -2,7 +2,7 @@ # SPDX-License-Identifier: MIT seclab-taskflow-agent: - version: 1 + version: "1.0" filetype: taskflow model_config: examples.model_configs.model_config @@ -12,4 +12,4 @@ taskflow: # with the `uses` directive we can reuse single task taskflows uses: examples.taskflows.single_step_taskflow # and optionally override any of its configurations - model: gpt_latest + model: gpt_default diff --git a/examples/taskflows/example_triage_taskflow.yaml b/examples/taskflows/example_triage_taskflow.yaml index 17adbe2..ec4cdf7 100644 --- a/examples/taskflows/example_triage_taskflow.yaml +++ b/examples/taskflows/example_triage_taskflow.yaml @@ -3,7 +3,7 @@ # a simple example of the triage Agent pattern seclab-taskflow-agent: - version: 1 + version: "1.0" filetype: taskflow taskflow: @@ -39,4 +39,4 @@ taskflow: - examples.personalities.orange_expert - examples.personalities.banana_expert user_prompt: | - Tell me more about how {{ RESULT }} are grown. + Tell me more about how {{ result }} are grown. diff --git a/examples/taskflows/single_step_taskflow.yaml b/examples/taskflows/single_step_taskflow.yaml index a85483a..d434f59 100644 --- a/examples/taskflows/single_step_taskflow.yaml +++ b/examples/taskflows/single_step_taskflow.yaml @@ -2,7 +2,7 @@ # SPDX-License-Identifier: MIT seclab-taskflow-agent: - version: 1 + version: "1.0" filetype: taskflow taskflow: diff --git a/pyproject.toml b/pyproject.toml index 3e254e4..acd5943 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -53,6 +53,7 @@ dependencies = [ "importlib_metadata==8.7.0", "isodate==0.7.2", "jedi==0.19.2", + "Jinja2>=3.1.0", "jiter==0.10.0", "jsonschema==4.24.0", "jsonschema-path==0.3.4", diff --git a/scripts/migrate_to_jinja2.py b/scripts/migrate_to_jinja2.py new file mode 100755 index 0000000..1671175 --- /dev/null +++ b/scripts/migrate_to_jinja2.py @@ -0,0 +1,211 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: 2025 GitHub +# SPDX-License-Identifier: MIT + +""" +Automated migration script for converting taskflow YAML files +from custom template syntax to Jinja2 syntax. + +Usage: + python scripts/migrate_to_jinja2.py /path/to/taskflows + python scripts/migrate_to_jinja2.py --dry-run taskflow.yaml +""" + +import re +import sys +import argparse +from pathlib import Path +from typing import List, Tuple + + +class TemplateMigrator: + """Migrates custom template syntax to Jinja2.""" + + def __init__(self, dry_run: bool = False): + self.dry_run = dry_run + self.transformations: List[Tuple[str, str]] = [] + + def migrate_content(self, content: str) -> str: + """Apply all template transformations to content.""" + original = content + + # 0. Version number: version: 1 or version: 2 -> version: "1.0" + content = re.sub( + r'^(\s*version:\s*)(?:1|2)\s*$', + r'\1"1.0"', + content, + flags=re.MULTILINE + ) + + # 1. Environment variables: {{ env VAR }} -> {{ env('VAR') }} + content = re.sub( + r'\{\{\s*env\s+([A-Z0-9_]+)\s*\}\}', + r"{{ env('\1') }}", + content + ) + + # 2. Global variables: {{ GLOBALS_key }} -> {{ globals.key }} + content = re.sub( + r'\{\{\s*GLOBALS_([a-zA-Z0-9_\.]+)\s*\}\}', + r'{{ globals.\1 }}', + content + ) + + # 3. Input variables: {{ INPUTS_key }} -> {{ inputs.key }} + content = re.sub( + r'\{\{\s*INPUTS_([a-zA-Z0-9_\.]+)\s*\}\}', + r'{{ inputs.\1 }}', + content + ) + + # 4. Result dict keys: {{ RESULT_key }} -> {{ result.key }} + content = re.sub( + r'\{\{\s*RESULT_([a-zA-Z0-9_\.]+)\s*\}\}', + r'{{ result.\1 }}', + content + ) + + # 5. Result primitive: {{ RESULT }} -> {{ result }} + content = re.sub( + r'\{\{\s*RESULT\s*\}\}', + r'{{ result }}', + content + ) + + # 6. Reusable prompts: {{ PROMPTS_path }} -> {% include 'path' %} + content = re.sub( + r'\{\{\s*PROMPTS_([a-zA-Z0-9_\.]+)\s*\}\}', + r"{% include '\1' %}", + content + ) + + if content != original: + self.transformations.append((original, content)) + + return content + + def migrate_file(self, file_path: Path) -> bool: + """Migrate a single YAML file. + + Returns: + True if file was modified, False otherwise + """ + if not file_path.suffix == '.yaml': + print(f"Skipping non-YAML file: {file_path}") + return False + + try: + with open(file_path, 'r') as f: + original_content = f.read() + + migrated_content = self.migrate_content(original_content) + + if migrated_content == original_content: + print(f"No changes needed: {file_path}") + return False + + if self.dry_run: + print(f"\n{'='*60}") + print(f"Would modify: {file_path}") + print(f"{'='*60}") + self._show_diff(original_content, migrated_content) + return True + + # Write migrated content + with open(file_path, 'w') as f: + f.write(migrated_content) + + print(f"Migrated: {file_path}") + return True + + except Exception as e: + print(f"Error migrating {file_path}: {e}", file=sys.stderr) + return False + + def _show_diff(self, original: str, migrated: str): + """Show simplified diff between original and migrated.""" + orig_lines = original.splitlines() + mig_lines = migrated.splitlines() + + for i, (orig, mig) in enumerate(zip(orig_lines, mig_lines), 1): + if orig != mig: + print(f"Line {i}:") + print(f" - {orig}") + print(f" + {mig}") + + def migrate_directory(self, directory: Path, recursive: bool = True) -> int: + """Migrate all YAML files in directory. + + Returns: + Number of files modified + """ + pattern = '**/*.yaml' if recursive else '*.yaml' + yaml_files = list(directory.glob(pattern)) + + if not yaml_files: + print(f"No YAML files found in {directory}") + return 0 + + print(f"Found {len(yaml_files)} YAML files") + + modified_count = 0 + for yaml_file in yaml_files: + if self.migrate_file(yaml_file): + modified_count += 1 + + return modified_count + + +def main(): + parser = argparse.ArgumentParser( + description='Migrate taskflow YAML files to Jinja2 syntax' + ) + parser.add_argument( + 'paths', + nargs='+', + type=Path, + help='YAML files or directories to migrate' + ) + parser.add_argument( + '--dry-run', + action='store_true', + help='Show changes without modifying files' + ) + parser.add_argument( + '--no-recursive', + action='store_true', + help='Do not recurse into subdirectories' + ) + + args = parser.parse_args() + + migrator = TemplateMigrator(dry_run=args.dry_run) + + total_modified = 0 + for path in args.paths: + if not path.exists(): + print(f"Path not found: {path}", file=sys.stderr) + continue + + if path.is_file(): + if migrator.migrate_file(path): + total_modified += 1 + elif path.is_dir(): + modified = migrator.migrate_directory( + path, + recursive=not args.no_recursive + ) + total_modified += modified + else: + print(f"Invalid path: {path}", file=sys.stderr) + + print(f"\n{'='*60}") + if args.dry_run: + print(f"Dry run complete. {total_modified} files would be modified.") + else: + print(f"Migration complete. {total_modified} files modified.") + print(f"{'='*60}") + + +if __name__ == '__main__': + main() diff --git a/scripts/test_examples.sh b/scripts/test_examples.sh new file mode 100755 index 0000000..adc47d7 --- /dev/null +++ b/scripts/test_examples.sh @@ -0,0 +1,144 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: 2025 GitHub +# SPDX-License-Identifier: MIT + +# Colors for output +RED='\033[0;31m' +GREEN='\033[0;32m' +YELLOW='\033[1;33m' +NC='\033[0m' # No Color + +# Check if gh cli is available +if ! command -v gh &> /dev/null; then + echo -e "${RED}Error: gh cli not found. Install from https://cli.github.com${NC}" + exit 1 +fi + +# Get API token +echo "Getting GitHub API token..." +export AI_API_TOKEN="$(gh auth token)" +if [ -z "$AI_API_TOKEN" ]; then + echo -e "${RED}Error: Failed to get GitHub API token${NC}" + exit 1 +fi + +# Activate venv if exists +if [ -d ".venv" ]; then + echo "Activating virtual environment..." + source .venv/bin/activate +else + echo -e "${YELLOW}Warning: No .venv found, using system Python${NC}" +fi + +# Track test results +PASSED=0 +FAILED=0 +FAILED_TESTS=() + +# Test function +run_test() { + local name="$1" + local taskflow="$2" + local args="${3:-}" + local timeout="${4:-30}" + + echo -e "\n${YELLOW}Testing: $name${NC}" + echo -e "${YELLOW}========================================${NC}" + + # Run command with output shown in real-time, capture to temp file for checking + local tmpfile=$(mktemp) + timeout "$timeout" python -m seclab_taskflow_agent -t "$taskflow" $args 2>&1 | tee "$tmpfile" || true + + echo -e "${YELLOW}========================================${NC}" + + # Check for error conditions first + if grep -qi "rate limit" "$tmpfile" || \ + grep -q "ERROR:" "$tmpfile" || \ + grep -q "Max rate limit backoff reached" "$tmpfile" || \ + grep -q "APITimeoutError" "$tmpfile" || \ + grep -q "Exception:" "$tmpfile"; then + echo -e "${RED}✗ $name failed (error detected)${NC}" + ((FAILED++)) + FAILED_TESTS+=("$name") + rm "$tmpfile" + return 1 + fi + + # Check for successful start + if grep -q "Running Task Flow" "$tmpfile"; then + echo -e "${GREEN}✓ $name passed${NC}" + ((PASSED++)) + rm "$tmpfile" + return 0 + else + echo -e "${RED}✗ $name failed (did not start)${NC}" + ((FAILED++)) + FAILED_TESTS+=("$name") + rm "$tmpfile" + return 1 + fi +} + +echo -e "${GREEN}Starting example taskflow tests...${NC}\n" + +# Test 1: Simple single-step taskflow +run_test "single_step_taskflow" "examples.taskflows.single_step_taskflow" + +# Test 2: Echo taskflow +run_test "echo" "examples.taskflows.echo" + +# Test 3: Globals example +run_test "example_globals" "examples.taskflows.example_globals" "-g fruit=apples" + +# Test 4: Inputs example +run_test "example_inputs" "examples.taskflows.example_inputs" + +# Test 5: Repeat prompt example +run_test "example_repeat_prompt" "examples.taskflows.example_repeat_prompt" "" "45" + +# Test 6: Reusable prompt example +run_test "example_reusable_prompt" "examples.taskflows.example_reusable_prompt" + +# Test 7: Full example taskflow +run_test "example" "examples.taskflows.example" + +# Test 8: Reusable taskflows (may fail on temperature setting, but should load YAML) +echo -e "\n${YELLOW}Testing: example_reusable_taskflows (YAML load test)${NC}" +echo -e "${YELLOW}========================================${NC}" +tmpfile=$(mktemp) +timeout 30 python -m seclab_taskflow_agent -t examples.taskflows.example_reusable_taskflows 2>&1 | tee "$tmpfile" || true +echo -e "${YELLOW}========================================${NC}" + +# Check for errors first +if grep -qi "rate limit" "$tmpfile" || \ + grep -q "Max rate limit backoff reached" "$tmpfile" || \ + grep -q "APITimeoutError" "$tmpfile"; then + echo -e "${RED}✗ example_reusable_taskflows failed (error detected)${NC}" + ((FAILED++)) + FAILED_TESTS+=("example_reusable_taskflows") +elif grep -q "Running Task Flow" "$tmpfile"; then + echo -e "${GREEN}✓ example_reusable_taskflows YAML loaded correctly${NC}" + ((PASSED++)) +else + echo -e "${YELLOW}⚠ example_reusable_taskflows - may have API issues but YAML loaded${NC}" + ((PASSED++)) +fi +rm "$tmpfile" + +# Print summary +echo -e "\n${GREEN}========================================${NC}" +echo -e "${GREEN}Test Summary${NC}" +echo -e "${GREEN}========================================${NC}" +echo -e "Passed: ${GREEN}$PASSED${NC}" +echo -e "Failed: ${RED}$FAILED${NC}" + +if [ $FAILED -gt 0 ]; then + echo -e "\n${RED}Failed tests:${NC}" + for test in "${FAILED_TESTS[@]}"; do + echo -e " - $test" + done + exit 1 +else + echo -e "\n${GREEN}All tests passed!${NC}" + exit 0 +fi diff --git a/src/seclab_taskflow_agent/__about__.py b/src/seclab_taskflow_agent/__about__.py index 12596b5..a2ce705 100644 --- a/src/seclab_taskflow_agent/__about__.py +++ b/src/seclab_taskflow_agent/__about__.py @@ -1,3 +1,3 @@ # SPDX-FileCopyrightText: 2025 GitHub # SPDX-License-Identifier: MIT -__version__ = "0.0.9" +__version__ = "0.1.0" diff --git a/src/seclab_taskflow_agent/__main__.py b/src/seclab_taskflow_agent/__main__.py index 1072907..0a73515 100644 --- a/src/seclab_taskflow_agent/__main__.py +++ b/src/seclab_taskflow_agent/__main__.py @@ -9,7 +9,6 @@ from dotenv import load_dotenv, find_dotenv import logging from logging.handlers import RotatingFileHandler -from pprint import pprint, pformat import re import json import uuid @@ -34,6 +33,8 @@ from .capi import list_tool_call_models, get_AI_token from .available_tools import AvailableTools from .path_utils import log_file_name +from .template_utils import render_template +import jinja2 load_dotenv(find_dotenv(usecwd=True)) @@ -507,86 +508,65 @@ 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: 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) - 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', - lambda key: available_tools.get_prompt(key), - 'prompt') - - # pre-process the prompt for any inputs - if prompt and 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', - lambda key: global_variables.get(key)) + # Render prompt template with Jinja2 (skip if repeat_prompt since result is not yet available) + if prompt and not repeat_prompt: + try: + prompt = render_template( + template_str=prompt, + available_tools=available_tools, + globals_dict=global_variables, + inputs_dict=inputs, + ) + except jinja2.TemplateError as e: + logging.error(f"Template rendering error: {e}") + raise ValueError(f"Failed to render prompt template: {e}") with TmpEnv(env): prompts_to_run = [] if repeat_prompt: - pattern = r"\{\{\s+RESULT_*(.*?|)\s+\}\}" - m = re.search(pattern, prompt) - # if last mcp tool result is an iterable it becomes available for repeat prompts - if not m: - logging.critical("Expected templated prompt, aborting!") - break + # Check if prompt contains result template variable + if 'result' not in prompt.lower(): + logging.warning("repeat_prompt enabled but no {{ result }} in prompt") + try: - # if this is json loadable, then it might be an iter, so check for that + # Get last MCP tool result last_result = json.loads(last_mcp_tool_results.pop()) text = last_result.get('text', '') try: iterable_result = json.loads(text) except json.decoder.JSONDecodeError as exc: - e = f"Could not json.loads result text: {text}" - logging.critical(e) - raise ValueError(e) from exc - iter(iterable_result) + logging.critical(f"Could not parse result text: {text}") + raise ValueError(f"Result text is not valid JSON") from exc + + # Verify iterable + try: + iter(iterable_result) + except TypeError: + logging.critical("Last MCP tool result is not iterable") + raise except IndexError: - logging.critical("No last mcp tool result available, aborting!") - raise - except ValueError: - logging.critical("Could not json.loads last mcp tool results, aborting!") - raise - except TypeError: - logging.critical("Last mcp tool results are not iterable, aborting!") + logging.critical("No last MCP tool result available") raise + if not iterable_result: await render_model_output("** 🤖❗MCP tool result iterable is empty!\n") else: - # we use our own template marker here so prompts are not limited to use {} - logging.debug(f"Entering templated prompt loop for results: {iterable_result}") + logging.debug(f"Rendering templated prompts for results: {iterable_result}") + + # Render template for each result value for value in iterable_result: - # support RESULT_key -> value swap format as well - if isinstance(value, dict) and m.group(1): - _prompt = prompt - for full_match in re.findall(r"\{\{\s+RESULT_(?:.*?)\s+\}\}", prompt): - _m = re.search(r"\{\{\s+RESULT_(.*?)\s+\}\}", full_match) - if _m and _m.group(1) in value: - _prompt = _prompt.replace( - full_match, - pformat(value.get(_m.group(1)))) - prompts_to_run.append(_prompt) - else: - prompts_to_run.append( - prompt.replace( - m.group(0), - pformat(value))) + try: + rendered_prompt = render_template( + template_str=prompt, + available_tools=available_tools, + globals_dict=global_variables, + inputs_dict=inputs, + result_value=value, + ) + prompts_to_run.append(rendered_prompt) + except jinja2.TemplateError as e: + logging.error(f"Error rendering template for result {value}: {e}") + raise ValueError(f"Template rendering failed: {e}") else: prompts_to_run.append(prompt) diff --git a/src/seclab_taskflow_agent/available_tools.py b/src/seclab_taskflow_agent/available_tools.py index 3750b0d..0904dab 100644 --- a/src/seclab_taskflow_agent/available_tools.py +++ b/src/seclab_taskflow_agent/available_tools.py @@ -70,8 +70,21 @@ def get_tool(self, tooltype: AvailableToolType, toolname: str): y = yaml.safe_load(s) header = y['seclab-taskflow-agent'] version = header['version'] - if version != 1: - raise VersionException(str(version)) + # Support both string and int for backwards compatibility during migration + version_str = str(version) + + # Reject old integer version 1 (pre-Jinja2) + if version == 1: + raise VersionException( + f"YAML file {f} uses unsupported version 1 template syntax. " + f"Version 1.0+ (Jinja2) is required. " + f"Migrate using: python scripts/migrate_to_jinja2.py {f}" + ) + # Accept version "1.0" or newer (string format following semver) + elif not version_str.startswith("1."): + raise VersionException( + f"Unsupported version: {version}. Only version 1.x is supported." + ) filetype = header['filetype'] if filetype != tooltype.value: raise FileTypeException( diff --git a/src/seclab_taskflow_agent/env_utils.py b/src/seclab_taskflow_agent/env_utils.py index 39a28b6..64df406 100644 --- a/src/seclab_taskflow_agent/env_utils.py +++ b/src/seclab_taskflow_agent/env_utils.py @@ -1,14 +1,41 @@ # SPDX-FileCopyrightText: 2025 GitHub # SPDX-License-Identifier: MIT -import re import os +import jinja2 -def swap_env(s): - match = re.search(r"{{\s*(env)\s+([A-Z0-9_]+)\s*}}", s) - if match and not os.getenv(match.group(2)): - raise LookupError(f"Requested {match.group(2)} from env but it does not exist!") - return os.getenv(match.group(2)) if match else s + +def swap_env(s: str) -> str: + """Replace {{ env('VAR') }} patterns in string with environment values. + + Args: + s: String potentially containing env templates + + Returns: + String with env templates replaced + + Raises: + LookupError: If required env var not found + """ + # Quick check if templating needed + if '{{' not in s: + return s + + try: + # Import here to avoid circular dependency + from .template_utils import create_jinja_environment + from .available_tools import AvailableTools + + available_tools = AvailableTools() + jinja_env = create_jinja_environment(available_tools) + template = jinja_env.from_string(s) + return template.render() + except jinja2.UndefinedError as e: + # Convert Jinja undefined to LookupError for compatibility + raise LookupError(str(e)) + except jinja2.TemplateError: + # Not a template or failed to render, return as-is + return s class TmpEnv: def __init__(self, env): diff --git a/src/seclab_taskflow_agent/personalities/assistant.yaml b/src/seclab_taskflow_agent/personalities/assistant.yaml index 88e51fc..55c9544 100644 --- a/src/seclab_taskflow_agent/personalities/assistant.yaml +++ b/src/seclab_taskflow_agent/personalities/assistant.yaml @@ -2,7 +2,7 @@ # SPDX-License-Identifier: MIT seclab-taskflow-agent: - version: 1 + version: "1.0" filetype: personality personality: | diff --git a/src/seclab_taskflow_agent/personalities/c_auditer.yaml b/src/seclab_taskflow_agent/personalities/c_auditer.yaml index 64d1994..6855586 100644 --- a/src/seclab_taskflow_agent/personalities/c_auditer.yaml +++ b/src/seclab_taskflow_agent/personalities/c_auditer.yaml @@ -2,7 +2,7 @@ # SPDX-License-Identifier: MIT seclab-taskflow-agent: - version: 1 + version: "1.0" filetype: personality personality: | diff --git a/src/seclab_taskflow_agent/template_utils.py b/src/seclab_taskflow_agent/template_utils.py new file mode 100644 index 0000000..2b7e0c9 --- /dev/null +++ b/src/seclab_taskflow_agent/template_utils.py @@ -0,0 +1,154 @@ +# SPDX-FileCopyrightText: 2025 GitHub +# SPDX-License-Identifier: MIT + +"""Jinja2 template utilities for taskflow template rendering.""" + +import os +import jinja2 +from typing import Any, Dict, Optional + + +class PromptLoader(jinja2.BaseLoader): + """Custom Jinja2 loader for reusable prompts.""" + + def __init__(self, available_tools): + """Initialize the prompt loader. + + Args: + available_tools: AvailableTools instance for prompt loading + """ + self.available_tools = available_tools + + def get_source(self, environment, template): + """Load prompt from available_tools by path. + + Args: + environment: Jinja2 environment + template: Template path (e.g., 'examples.prompts.example_prompt') + + Returns: + Tuple of (source, filename, uptodate_func) + + Raises: + jinja2.TemplateNotFound: If prompt not found + """ + try: + prompt_data = self.available_tools.get_prompt(template) + if not prompt_data: + raise jinja2.TemplateNotFound(template) + source = prompt_data.get('prompt', '') + # Return: (source, filename, uptodate_func) + return source, None, lambda: True + except Exception: + raise jinja2.TemplateNotFound(template) + + +def env_function(var_name: str, default: Optional[str] = None, required: bool = True) -> str: + """Jinja2 function to access environment variables. + + Args: + var_name: Name of environment variable + default: Default value if not found + required: If True, raises error when not found and no default + + Returns: + Environment variable value or default + + Raises: + LookupError: If required var not found + + Examples: + {{ env('LOG_DIR') }} + {{ env('OPTIONAL_VAR', 'default_value') }} + {{ env('OPTIONAL_VAR', required=False) }} + """ + value = os.getenv(var_name, default) + if value is None and required: + raise LookupError(f"Required environment variable {var_name} not found!") + return value or "" + + +def create_jinja_environment(available_tools) -> jinja2.Environment: + """Create configured Jinja2 environment for taskflow templates. + + Args: + available_tools: AvailableTools instance for prompt loading + + Returns: + Configured Jinja2 Environment + """ + env = jinja2.Environment( + loader=PromptLoader(available_tools), + # Use same delimiters as custom system + variable_start_string='{{', + variable_end_string='}}', + block_start_string='{%', + block_end_string='%}', + # Disable auto-escaping (YAML context doesn't need HTML escaping) + autoescape=False, + # Keep whitespace for prompt formatting + trim_blocks=True, + lstrip_blocks=True, + # Raise errors for undefined variables + undefined=jinja2.StrictUndefined, + ) + + # Register custom functions + env.globals['env'] = env_function + + return env + + +def render_template( + template_str: str, + available_tools, + globals_dict: Optional[Dict[str, Any]] = None, + inputs_dict: Optional[Dict[str, Any]] = None, + result_value: Optional[Any] = None, +) -> str: + """Render a template string with provided context. + + Args: + template_str: Template string to render + available_tools: AvailableTools instance + globals_dict: Global variables dict + inputs_dict: Input variables dict + result_value: Result value for repeat_prompt + + Returns: + Rendered template string + + Raises: + jinja2.TemplateError: On template rendering errors + + Examples: + # Render with globals + render_template("{{ globals.fruit }}", tools, globals_dict={'fruit': 'apple'}) + + # Render with result + render_template("{{ result.name }}", tools, result_value={'name': 'test'}) + + # Render with all context types + render_template( + "{{ globals.x }} {{ inputs.y }} {{ result.z }}", + tools, + globals_dict={'x': 1}, + inputs_dict={'y': 2}, + result_value={'z': 3} + ) + """ + jinja_env = create_jinja_environment(available_tools) + + # Build template context + context = { + 'globals': globals_dict or {}, + 'inputs': inputs_dict or {}, + } + + # Add result if provided + if result_value is not None: + context['result'] = result_value + + # Render template + template = jinja_env.from_string(template_str) + return template.render(**context) diff --git a/src/seclab_taskflow_agent/toolboxes/codeql.yaml b/src/seclab_taskflow_agent/toolboxes/codeql.yaml index 98185f9..b90fa45 100644 --- a/src/seclab_taskflow_agent/toolboxes/codeql.yaml +++ b/src/seclab_taskflow_agent/toolboxes/codeql.yaml @@ -2,7 +2,7 @@ # SPDX-License-Identifier: MIT seclab-taskflow-agent: - version: 1 + version: "1.0" filetype: toolbox server_params: @@ -12,12 +12,12 @@ server_params: command: python args: ["-m", "seclab_taskflow_agent.mcp_servers.codeql.mcp_server"] env: - CODEQL_DBS_BASE_PATH: "{{ env CODEQL_DBS_BASE_PATH }}" + CODEQL_DBS_BASE_PATH: "{{ env('CODEQL_DBS_BASE_PATH') }}" # prevent git repo operations on gh codeql executions GH_NO_UPDATE_NOTIFIER: "Disable" GH_NO_EXTENSION_UPDATE_NOTIFIER: "Disable" - CODEQL_CLI: "{{ env CODEQL_CLI }}" - LOG_DIR: "{{ env LOG_DIR }}" + CODEQL_CLI: "{{ env('CODEQL_CLI') }}" + LOG_DIR: "{{ env('LOG_DIR') }}" server_prompt: | ## CodeQL Supported Programming Languages diff --git a/src/seclab_taskflow_agent/toolboxes/echo.yaml b/src/seclab_taskflow_agent/toolboxes/echo.yaml index cacb591..9328728 100644 --- a/src/seclab_taskflow_agent/toolboxes/echo.yaml +++ b/src/seclab_taskflow_agent/toolboxes/echo.yaml @@ -3,7 +3,7 @@ # stdio mcp server configuration seclab-taskflow-agent: - version: 1 + version: "1.0" filetype: toolbox server_params: @@ -12,4 +12,4 @@ server_params: args: ["-m", "seclab_taskflow_agent.mcp_servers.echo.echo"] env: TEST: value - LOG_DIR: "{{ env LOG_DIR }}" + LOG_DIR: "{{ env('LOG_DIR') }}" diff --git a/src/seclab_taskflow_agent/toolboxes/github_official.yaml b/src/seclab_taskflow_agent/toolboxes/github_official.yaml index 971c16e..73ee9c5 100644 --- a/src/seclab_taskflow_agent/toolboxes/github_official.yaml +++ b/src/seclab_taskflow_agent/toolboxes/github_official.yaml @@ -2,7 +2,7 @@ # SPDX-License-Identifier: MIT seclab-taskflow-agent: - version: 1 + version: "1.0" filetype: toolbox server_params: @@ -10,7 +10,7 @@ server_params: url: https://api.githubcopilot.com/mcp/ #See https://github.com/github/github-mcp-server/blob/main/docs/remote-server.md headers: - Authorization: "{{ env GITHUB_AUTH_HEADER }}" + Authorization: "{{ env('GITHUB_AUTH_HEADER') }}" optional_headers: - X-MCP-Toolsets: "{{ env GITHUB_MCP_TOOLSETS }}" - X-MCP-Readonly: "{{ env GITHUB_MCP_READONLY }}" \ No newline at end of file + X-MCP-Toolsets: "{{ env('GITHUB_MCP_TOOLSETS') }}" + X-MCP-Readonly: "{{ env('GITHUB_MCP_READONLY') }}" \ No newline at end of file diff --git a/src/seclab_taskflow_agent/toolboxes/logbook.yaml b/src/seclab_taskflow_agent/toolboxes/logbook.yaml index ceaa109..50b1ef6 100644 --- a/src/seclab_taskflow_agent/toolboxes/logbook.yaml +++ b/src/seclab_taskflow_agent/toolboxes/logbook.yaml @@ -2,7 +2,7 @@ # SPDX-License-Identifier: MIT seclab-taskflow-agent: - version: 1 + version: "1.0" filetype: toolbox server_params: @@ -10,8 +10,8 @@ server_params: command: python args: ["-m", "seclab_taskflow_agent.mcp_servers.logbook.logbook"] env: - LOGBOOK_STATE_DIR: "{{ env LOGBOOK_STATE_DIR }}" - LOG_DIR: "{{ env LOG_DIR }}" + LOGBOOK_STATE_DIR: "{{ env('LOGBOOK_STATE_DIR') }}" + LOG_DIR: "{{ env('LOG_DIR') }}" # the list of tools that you want the framework to confirm with the user before executing # use this to guard rail any potentially dangerous functions from MCP servers confirm: diff --git a/src/seclab_taskflow_agent/toolboxes/memcache.yaml b/src/seclab_taskflow_agent/toolboxes/memcache.yaml index 0fcc3cb..37f70ea 100644 --- a/src/seclab_taskflow_agent/toolboxes/memcache.yaml +++ b/src/seclab_taskflow_agent/toolboxes/memcache.yaml @@ -2,7 +2,7 @@ # SPDX-License-Identifier: MIT seclab-taskflow-agent: - version: 1 + version: "1.0" filetype: toolbox server_params: @@ -10,9 +10,9 @@ server_params: command: python args: ["-m", "seclab_taskflow_agent.mcp_servers.memcache.memcache"] env: - MEMCACHE_STATE_DIR: "{{ env MEMCACHE_STATE_DIR }}" - MEMCACHE_BACKEND: "{{ env MEMCACHE_BACKEND }}" - LOG_DIR: "{{ env LOG_DIR }}" + MEMCACHE_STATE_DIR: "{{ env('MEMCACHE_STATE_DIR') }}" + MEMCACHE_BACKEND: "{{ env('MEMCACHE_BACKEND') }}" + LOG_DIR: "{{ env('LOG_DIR') }}" # the list of tools that you want the framework to confirm with the user before executing # use this to guard rail any potentially dangerous functions from MCP servers confirm: diff --git a/tests/data/test_globals_taskflow.yaml b/tests/data/test_globals_taskflow.yaml index 537ee9e..61e0a7f 100644 --- a/tests/data/test_globals_taskflow.yaml +++ b/tests/data/test_globals_taskflow.yaml @@ -2,7 +2,7 @@ # SPDX-License-Identifier: MIT seclab-taskflow-agent: - version: 1 + version: "1.0" filetype: taskflow globals: diff --git a/tests/data/test_prompt_with_variables.yaml b/tests/data/test_prompt_with_variables.yaml new file mode 100644 index 0000000..f2c4612 --- /dev/null +++ b/tests/data/test_prompt_with_variables.yaml @@ -0,0 +1,11 @@ +# SPDX-FileCopyrightText: 2025 GitHub +# SPDX-License-Identifier: MIT + +seclab-taskflow-agent: + version: "1.0" + filetype: prompt + +prompt: | + This is a reusable prompt. + Global variable: {{ globals.test_global }} + Input variable: {{ inputs.test_input }} diff --git a/tests/data/test_reusable_taskflow_with_variables.yaml b/tests/data/test_reusable_taskflow_with_variables.yaml new file mode 100644 index 0000000..8a5eb0a --- /dev/null +++ b/tests/data/test_reusable_taskflow_with_variables.yaml @@ -0,0 +1,15 @@ +# SPDX-FileCopyrightText: 2025 GitHub +# SPDX-License-Identifier: MIT + +seclab-taskflow-agent: + version: "1.0" + filetype: taskflow + +taskflow: + - task: + agents: + - tests.data.test_yaml_parser_personality000 + user_prompt: | + This is a reusable taskflow. + Global: {{ globals.reusable_global }} + Input: {{ inputs.reusable_input }} diff --git a/tests/data/test_taskflow_using_reusable.yaml b/tests/data/test_taskflow_using_reusable.yaml new file mode 100644 index 0000000..0056b08 --- /dev/null +++ b/tests/data/test_taskflow_using_reusable.yaml @@ -0,0 +1,15 @@ +# SPDX-FileCopyrightText: 2025 GitHub +# SPDX-License-Identifier: MIT + +seclab-taskflow-agent: + version: "1.0" + filetype: taskflow + +globals: + reusable_global: "global_from_parent" + +taskflow: + - task: + uses: tests.data.test_reusable_taskflow_with_variables + inputs: + reusable_input: "input_from_parent" diff --git a/tests/data/test_yaml_parser_personality000.yaml b/tests/data/test_yaml_parser_personality000.yaml index 0dfbbfb..4009683 100644 --- a/tests/data/test_yaml_parser_personality000.yaml +++ b/tests/data/test_yaml_parser_personality000.yaml @@ -2,7 +2,7 @@ # SPDX-License-Identifier: MIT seclab-taskflow-agent: - version: 1 + version: "1.0" filetype: personality personality: | diff --git a/tests/test_template_utils.py b/tests/test_template_utils.py new file mode 100644 index 0000000..e8de5d0 --- /dev/null +++ b/tests/test_template_utils.py @@ -0,0 +1,301 @@ +# SPDX-FileCopyrightText: 2025 GitHub +# SPDX-License-Identifier: MIT + +"""Tests for Jinja2 template utilities.""" + +import pytest +import os +import jinja2 +from seclab_taskflow_agent.template_utils import ( + env_function, + create_jinja_environment, + render_template, + PromptLoader, +) +from seclab_taskflow_agent.available_tools import AvailableTools + + +class TestEnvFunction: + """Test environment variable function.""" + + def test_env_existing_var(self): + """Test accessing existing environment variable.""" + os.environ['TEST_VAR_JINJA'] = 'test_value' + try: + assert env_function('TEST_VAR_JINJA') == 'test_value' + finally: + del os.environ['TEST_VAR_JINJA'] + + def test_env_missing_required(self): + """Test error on missing required variable.""" + with pytest.raises(LookupError, match="Required environment variable"): + env_function('NONEXISTENT_VAR_JINJA') + + def test_env_with_default(self): + """Test default value for missing variable.""" + result = env_function('NONEXISTENT_VAR_JINJA', default='default_value', required=False) + assert result == 'default_value' + + def test_env_optional_missing(self): + """Test optional variable returns empty string.""" + result = env_function('NONEXISTENT_VAR_JINJA', required=False) + assert result == '' + + def test_env_with_default_exists(self): + """Test that existing var takes precedence over default.""" + os.environ['TEST_VAR_DEFAULT'] = 'actual_value' + try: + result = env_function('TEST_VAR_DEFAULT', default='default_value') + assert result == 'actual_value' + finally: + del os.environ['TEST_VAR_DEFAULT'] + + +class TestJinjaEnvironment: + """Test Jinja2 environment setup.""" + + def test_create_environment(self): + """Test environment creation.""" + available_tools = AvailableTools() + env = create_jinja_environment(available_tools) + assert isinstance(env, jinja2.Environment) + assert 'env' in env.globals + + def test_strict_undefined(self): + """Test undefined variables raise errors.""" + available_tools = AvailableTools() + env = create_jinja_environment(available_tools) + template = env.from_string("{{ undefined_var }}") + with pytest.raises(jinja2.UndefinedError): + template.render() + + def test_env_function_in_template(self): + """Test env function works in template.""" + os.environ['TEST_TEMPLATE_VAR'] = 'template_value' + try: + available_tools = AvailableTools() + env = create_jinja_environment(available_tools) + template = env.from_string("{{ env('TEST_TEMPLATE_VAR') }}") + result = template.render() + assert result == 'template_value' + finally: + del os.environ['TEST_TEMPLATE_VAR'] + + +class TestRenderTemplate: + """Test template rendering.""" + + def test_render_globals(self): + """Test rendering with global variables.""" + available_tools = AvailableTools() + template_str = "Tell me about {{ globals.fruit }}" + result = render_template( + template_str, + available_tools, + globals_dict={'fruit': 'apples'} + ) + assert result == "Tell me about apples" + + def test_render_inputs(self): + """Test rendering with input variables.""" + available_tools = AvailableTools() + template_str = "Color: {{ inputs.color }}" + result = render_template( + template_str, + available_tools, + inputs_dict={'color': 'red'} + ) + assert result == "Color: red" + + def test_render_result_primitive(self): + """Test rendering with primitive result value.""" + available_tools = AvailableTools() + template_str = "Value: {{ result }}" + result = render_template( + template_str, + available_tools, + result_value=42 + ) + assert result == "Value: 42" + + def test_render_result_dict(self): + """Test rendering with dictionary result value.""" + available_tools = AvailableTools() + template_str = "Name: {{ result.name }}, Age: {{ result.age }}" + result = render_template( + template_str, + available_tools, + result_value={'name': 'Alice', 'age': 30} + ) + assert result == "Name: Alice, Age: 30" + + def test_render_with_env(self): + """Test rendering with env function.""" + os.environ['TEST_ENV_VAR'] = 'env_value' + try: + available_tools = AvailableTools() + template_str = "Env: {{ env('TEST_ENV_VAR') }}" + result = render_template(template_str, available_tools) + assert result == "Env: env_value" + finally: + del os.environ['TEST_ENV_VAR'] + + def test_render_complex(self): + """Test rendering with multiple variable types.""" + os.environ['TEST_MODEL'] = 'gpt-4' + try: + available_tools = AvailableTools() + template_str = """Model: {{ env('TEST_MODEL') }} +Fruit: {{ globals.fruit }} +Color: {{ inputs.color }} +Result: {{ result.value }}""" + result = render_template( + template_str, + available_tools, + globals_dict={'fruit': 'banana'}, + inputs_dict={'color': 'yellow'}, + result_value={'value': 123} + ) + assert 'gpt-4' in result + assert 'banana' in result + assert 'yellow' in result + assert '123' in result + finally: + del os.environ['TEST_MODEL'] + + def test_render_undefined_error(self): + """Test error on undefined variable.""" + available_tools = AvailableTools() + template_str = "{{ globals.undefined }}" + with pytest.raises(jinja2.UndefinedError): + render_template(template_str, available_tools) + + def test_render_nested_dict(self): + """Test rendering with nested dictionary access.""" + available_tools = AvailableTools() + template_str = "Config: {{ globals.config.model }}" + result = render_template( + template_str, + available_tools, + globals_dict={'config': {'model': 'claude-3'}} + ) + assert result == "Config: claude-3" + + def test_render_with_filter(self): + """Test Jinja2 filters work.""" + available_tools = AvailableTools() + template_str = "{{ globals.name | upper }}" + result = render_template( + template_str, + available_tools, + globals_dict={'name': 'alice'} + ) + assert result == "ALICE" + + def test_render_empty_context(self): + """Test rendering with no context variables.""" + available_tools = AvailableTools() + template_str = "Static text" + result = render_template(template_str, available_tools) + assert result == "Static text" + + +class TestPromptLoader: + """Test custom prompt loader.""" + + def test_load_existing_prompt(self): + """Test loading existing prompt.""" + available_tools = AvailableTools() + loader = PromptLoader(available_tools) + env = jinja2.Environment(loader=loader) + + # Test with actual example prompt + template = env.get_template('examples.prompts.example_prompt') + result = template.render() + assert 'bananas' in result.lower() + + def test_load_nonexistent_prompt(self): + """Test error on nonexistent prompt.""" + available_tools = AvailableTools() + loader = PromptLoader(available_tools) + env = jinja2.Environment(loader=loader) + + with pytest.raises(jinja2.TemplateNotFound): + env.get_template('nonexistent.prompt') + + def test_include_prompt(self): + """Test {% include %} directive.""" + available_tools = AvailableTools() + env = create_jinja_environment(available_tools) + + template_str = """Main content. +{% include 'examples.prompts.example_prompt' %}""" + template = env.from_string(template_str) + result = template.render() + assert 'Main content' in result + assert 'bananas' in result.lower() + + def test_include_with_context(self): + """Test that included templates have access to context.""" + available_tools = AvailableTools() + env = create_jinja_environment(available_tools) + + # Create a template that uses context + template_str = """Value: {{ globals.test }} +{% include 'examples.prompts.example_prompt' %}""" + template = env.from_string(template_str) + result = template.render(globals={'test': 'context_value'}) + assert 'context_value' in result + + def test_include_prompt_with_globals_and_inputs(self): + """Test that included prompts render globals and inputs correctly.""" + available_tools = AvailableTools() + + # Use render_template to ensure context is properly set + template_str = """Main task prompt. +{% include 'tests.data.test_prompt_with_variables' %} +End of prompt.""" + + result = render_template( + template_str, + available_tools, + globals_dict={'test_global': 'global_value'}, + inputs_dict={'test_input': 'input_value'} + ) + + assert 'Main task prompt' in result + assert 'global_value' in result + assert 'input_value' in result + assert 'End of prompt' in result + + def test_reusable_taskflow_prompt_renders_variables(self): + """Test that reusable taskflow prompts render globals and inputs correctly. + + This simulates what happens when a taskflow uses another taskflow: + 1. Parent taskflow defines globals and task inputs + 2. Reusable taskflow's user_prompt uses those variables + 3. The prompt should render with the parent's context + """ + available_tools = AvailableTools() + + # Load the reusable taskflow + reusable_taskflow = available_tools.get_taskflow('tests.data.test_reusable_taskflow_with_variables') + + # Get the user_prompt from the reusable taskflow's task + user_prompt = reusable_taskflow['taskflow'][0]['task']['user_prompt'] + + # Render it with parent's globals and inputs (simulating what __main__.py does) + result = render_template( + user_prompt, + available_tools, + globals_dict={'reusable_global': 'parent_global_value'}, + inputs_dict={'reusable_input': 'parent_input_value'} + ) + + assert 'This is a reusable taskflow' in result + assert 'parent_global_value' in result + assert 'parent_input_value' in result + + +if __name__ == '__main__': + pytest.main([__file__, '-v']) diff --git a/tests/test_yaml_parser.py b/tests/test_yaml_parser.py index c035da6..ea3332a 100644 --- a/tests/test_yaml_parser.py +++ b/tests/test_yaml_parser.py @@ -19,7 +19,7 @@ def test_yaml_parser_basic_functionality(self): personality000 = available_tools.get_personality( "tests.data.test_yaml_parser_personality000") - assert personality000['seclab-taskflow-agent']['version'] == 1 + assert personality000['seclab-taskflow-agent']['version'] == "1.0" assert personality000['seclab-taskflow-agent']['filetype'] == 'personality' assert personality000['personality'] == 'You are a helpful assistant.\n' assert personality000['task'] == 'Answer any question.\n'