Merge pull request #30 from GitHubSecurityLab/add_model_config

add model config

Author: m-y-mo

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:

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

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)
 

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