diff --git a/available_tools.py b/available_tools.py index 273ed34..2df1437 100644 --- a/available_tools.py +++ b/available_tools.py @@ -25,6 +25,7 @@ def __init__(self, yamls: dict): self.taskflows = {} self.prompts = {} self.toolboxes = {} + self.model_config = {} # Iterate through all the yaml files and divide them into categories. # Each file should contain a header like this: @@ -49,6 +50,8 @@ def __init__(self, yamls: dict): 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: diff --git a/configs/model_config.yaml b/configs/model_config.yaml new file mode 100644 index 0000000..a037365 --- /dev/null +++ b/configs/model_config.yaml @@ -0,0 +1,9 @@ +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 + gpt_default: gpt-4.1 + gpt_latest: gpt-5 \ No newline at end of file diff --git a/main.py b/main.py index 479a060..5986c50 100644 --- a/main.py +++ b/main.py @@ -429,6 +429,16 @@ async def on_handoff_hook( # optional global vars available for the taskflow tasks global_variables = taskflow.get('globals', {}) + model_config = taskflow.get('model_config', {}) + 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 = model_dict.get('models', {}) + if model_dict: + if not isinstance(model_dict, dict): + raise ValueError(f"Models section of the model_config file {model_config} must be a dictionary") + model_keys = model_dict.keys() for task in taskflow['taskflow']: @@ -448,7 +458,9 @@ async def on_handoff_hook( for k,v in reusable_taskflow['taskflow'][0]['task'].items(): if k not in task_body: task_body[k] = v - + model = task_body.get('model', DEFAULT_MODEL) + if model in model_keys: + model = model_dict[model] # parse our taskflow grammar name = task_body.get('name', 'taskflow') # placeholder, not used yet description = task_body.get('description', 'taskflow') # placeholder not used yet @@ -465,7 +477,6 @@ async def on_handoff_hook( toolboxes_override = task_body.get('toolboxes', []) env = task_body.get('env', {}) repeat_prompt = task_body.get('repeat_prompt', False) - model = task_body.get('model', DEFAULT_MODEL) # this will set Agent 'stop_on_first_tool' tool use behavior, which prevents output back to llm exclude_from_context = task_body.get('exclude_from_context', False) # this allows you to run repeated prompts concurrently with a limit @@ -600,6 +611,7 @@ async def _deploy_task_agents(resolved_agents, prompt): run_hooks=TaskRunHooks( on_tool_end=on_tool_end_hook, on_tool_start=on_tool_start_hook), + model = model, agent_hooks=TaskAgentHooks( on_handoff=on_handoff_hook)) return result @@ -643,7 +655,8 @@ async def _deploy_task_agents(resolved_agents, prompt): 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/'toolboxes').rglob('*')) | + YamlParser(cwd).get_yaml_dict((cwd/'configs').rglob('*'))) p, t, l, user_prompt, help_msg = parse_prompt_args(available_tools) diff --git a/taskflows/CVE-2023-2283/CVE-2023-2283.yaml b/taskflows/CVE-2023-2283/CVE-2023-2283.yaml index 9f30f29..c6617ff 100644 --- a/taskflows/CVE-2023-2283/CVE-2023-2283.yaml +++ b/taskflows/CVE-2023-2283/CVE-2023-2283.yaml @@ -3,6 +3,8 @@ seclab-taskflow-agent: filetype: taskflow filekey: GitHubSecurityLab/seclab-taskflow-agent/taskflows/CVE-2023-2283/CVE-2023-2283 +model_config: GitHubSecurityLab/seclab-taskflow-agent/configs/model_config + taskflow: - task: must_complete: true @@ -14,7 +16,7 @@ taskflow: toolboxes: - GitHubSecurityLab/seclab-taskflow-agent/toolboxes/memcache - task: - model: gpt-4.1 + model: gpt_latest must_complete: false agents: - GitHubSecurityLab/seclab-taskflow-agent/personalities/c_auditer