Merge branch 'main' into remove-werkzeug

kevinbackhouse · web-flow · commit 77097720e287 · 2025-12-03T12:26:16.000Z
diff --git a/.devcontainer/post-attach.sh b/.devcontainer/post-attach.sh
@@ -4,8 +4,8 @@ set -e
 # If running in Codespaces, check for necessary secrets and print error if missing
 if [ -v CODESPACES ]; then
     echo "🔐 Running in Codespaces - injecting secrets from Codespaces settings..."
-    if [ ! -v COPILOT_TOKEN ]; then
-        echo "⚠️ Running in Codespaces - please add COPILOT_TOKEN to your Codespaces secrets"
+    if [ ! -v AI_API_TOKEN ]; then
+        echo "⚠️ Running in Codespaces - please add AI_API_TOKEN to your Codespaces secrets"
     fi
     if [ ! -v GITHUB_PERSONAL_ACCESS_TOKEN ]; then
         echo "⚠️ Running in Codespaces - please add GITHUB_PERSONAL_ACCESS_TOKEN to your Codespaces secrets"
diff --git a/.github/workflows/smoketest.yaml b/.github/workflows/smoketest.yaml
@@ -52,7 +52,7 @@ jobs:
 
       - name: Run tests
         env:
-          COPILOT_TOKEN: ${{ secrets.COPILOT_TOKEN }}
+          AI_API_ENDPOINT: ${{ secrets.AI_API_ENDPOINT }}
           GITHUB_AUTH_HEADER: "Bearer ${{ secrets.GITHUB_TOKEN }}"
 
         run: |
diff --git a/README.md b/README.md
@@ -36,15 +36,15 @@ Python >= 3.9 or Docker
 
 ## Configuration
 
-Provide a GitHub token for an account that is entitled to use GitHub Copilot via the `COPILOT_TOKEN` environment variable. Further configuration is use case dependent, i.e. pending which MCP servers you'd like to use in your taskflows.
+Provide a GitHub token for an account that is entitled to use [GitHub Models](https://models.github.ai) via the `AI_API_ENDPOINT` environment variable. Further configuration is use case dependent, i.e. pending which MCP servers you'd like to use in your taskflows.
 
 You can set persisting environment variables via an `.env` file in the project root.
 
 Example:
 
 ```sh
 # Tokens
-COPILOT_TOKEN=<your_github_token>
+AI_API_ENDPOINT=<your_github_token>
 # MCP configs
 GITHUB_PERSONAL_ACCESS_TOKEN=<your_github_token>
 CODEQL_DBS_BASE_PATH="/app/my_data/codeql_databases"
diff --git a/docker/run.sh b/docker/run.sh
@@ -11,7 +11,7 @@
 #
 #   git clone https://github.com/GitHubSecurityLab/seclab-taskflow-agent.git
 #   cd seclab-taskflow-agent/src
-#   export COPILOT_TOKEN=<My GitHub PAT>
+#   export AI_API_TOKEN=<My GitHub PAT>
 #   export GITHUB_AUTH_HEADER=<My GitHub PAT>
 #   sudo -E ../docker/run.sh -p seclab_taskflow_agent.personalities.assistant 'explain modems to me please'
 
@@ -23,5 +23,5 @@ docker run -i \
        --mount type=bind,src="$PWD",dst=/app \
        -e DATA_DIR=/app/data \
        -e GITHUB_PERSONAL_ACCESS_TOKEN="$GITHUB_PERSONAL_ACCESS_TOKEN" \
-       -e COPILOT_TOKEN="$COPILOT_TOKEN" \
+       -e AI_API_TOKEN="$AI_API_TOKEN" \
        "ghcr.io/githubsecuritylab/seclab-taskflow-agent" "$@"
diff --git a/examples/taskflows/echo.yaml b/examples/taskflows/echo.yaml
@@ -7,7 +7,6 @@ seclab-taskflow-agent:
 
 taskflow:
   - task:
-      model: claude-3.5-sonnet
       max_steps: 5
       must_complete: true
       agents:
diff --git a/examples/taskflows/example.yaml b/examples/taskflows/example.yaml
@@ -5,10 +5,13 @@ seclab-taskflow-agent:
   version: 1
   filetype: taskflow
 
+# Import settings from a model_config file.
+model_config: examples.model_configs.model_config
+
 taskflow:
   - task:
       # taskflows can optionally choose any of the support CAPI models for a task
-      model: gpt-4.1
+      model: gpt_default
       # taskflows can optionally limit the max allowed number of Agent task loop
       # iterations to complete a task, this defaults to 50 when not provided
       max_steps: 20
@@ -41,7 +44,6 @@ taskflow:
         - seclab_taskflow_agent.toolboxes.codeql
   - task:
       must_complete: true
-      model: gpt-4.1
       agents:
         - seclab_taskflow_agent.personalities.c_auditer
       user_prompt: |
diff --git a/examples/taskflows/example_reusable_taskflows.yaml b/examples/taskflows/example_reusable_taskflows.yaml
@@ -5,9 +5,11 @@ seclab-taskflow-agent:
   version: 1
   filetype: taskflow
 
+model_config: examples.model_configs.model_config
+
 taskflow:
   - task:
       # 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-4o
+      model: gpt_latest
diff --git a/examples/taskflows/example_triage_taskflow.yaml b/examples/taskflows/example_triage_taskflow.yaml
@@ -30,7 +30,6 @@ taskflow:
       toolboxes:
         - seclab_taskflow_agent.toolboxes.memcache
   - task:
-      model: gpt-4.1
       repeat_prompt: true
       agents:
         # primary agent for this task
diff --git a/examples/taskflows/single_step_taskflow.yaml b/examples/taskflows/single_step_taskflow.yaml
@@ -7,7 +7,6 @@ seclab-taskflow-agent:
 
 taskflow:
   - task:
-      model: gpt-4.1
       agents:
         - seclab_taskflow_agent.personalities.assistant
       user_prompt: |
diff --git a/pyproject.toml b/pyproject.toml
@@ -97,6 +97,7 @@ dependencies = [
   "SQLAlchemy==2.0.41",
   "sse-starlette==2.4.1",
   "starlette==0.49.1",
+  "strenum==0.4.15",
   "tqdm==4.67.1",
   "typer==0.16.0",
   "types-requests==2.32.4.20250611",
diff --git a/src/seclab_taskflow_agent/__main__.py b/src/seclab_taskflow_agent/__main__.py
@@ -31,7 +31,7 @@
 from .render_utils import render_model_output, flush_async_output
 from .env_utils import TmpEnv
 from .agent import TaskAgent
-from .capi import list_tool_call_models
+from .capi import list_tool_call_models, get_AI_token
 from .available_tools import AvailableTools
 
 load_dotenv(find_dotenv(usecwd=True))
@@ -686,7 +686,7 @@ async def _deploy_task_agents(resolved_agents, prompt):
     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'))
+        tool_models = list_tool_call_models(get_AI_token())
         for model in tool_models:
             print(model)
         sys.exit(0)
diff --git a/src/seclab_taskflow_agent/agent.py b/src/seclab_taskflow_agent/agent.py
@@ -15,18 +15,19 @@
 from agents.run import RunHooks
 from agents import Agent, Runner, AgentHooks, RunHooks, result, function_tool, Tool, RunContextWrapper, TContext, OpenAIChatCompletionsModel, set_default_openai_client, set_default_openai_api, set_tracing_disabled
 
-from .capi import COPILOT_INTEGRATION_ID, COPILOT_API_ENDPOINT
+from .capi import COPILOT_INTEGRATION_ID, get_AI_endpoint, get_AI_token, AI_API_ENDPOINT_ENUM
 
 # grab our secrets from .env, this must be in .gitignore
 load_dotenv(find_dotenv(usecwd=True))
 
-match urlparse(COPILOT_API_ENDPOINT).netloc:
-    case 'api.githubcopilot.com':
+api_endpoint = get_AI_endpoint()
+match urlparse(api_endpoint).netloc:
+    case AI_API_ENDPOINT_ENUM.AI_API_GITHUBCOPILOT:
         default_model = 'gpt-4o'
-    case 'models.github.ai':
+    case AI_API_ENDPOINT_ENUM.AI_API_MODELS_GITHUB:
         default_model = 'openai/gpt-4o'
     case _:
-        raise ValueError(f"Unsupported Model Endpoint: {COPILOT_API_ENDPOINT}")
+        raise ValueError(f"Unsupported Model Endpoint: {api_endpoint}")
 
 DEFAULT_MODEL = os.getenv('COPILOT_DEFAULT_MODEL', default=default_model)
 
@@ -148,8 +149,8 @@ def __init__(self,
                  model_settings: ModelSettings | None = None,
                  run_hooks: TaskRunHooks | None = None,
                  agent_hooks: TaskAgentHooks | None = None):
-        client = AsyncOpenAI(base_url=COPILOT_API_ENDPOINT,
-                             api_key=os.getenv('COPILOT_TOKEN'),
+        client = AsyncOpenAI(base_url=api_endpoint,
+                             api_key=get_AI_token(),
                              default_headers={'Copilot-Integration-Id': COPILOT_INTEGRATION_ID})
         set_default_openai_client(client)
         # CAPI does not yet support the Responses API: https://github.com/github/copilot-api/issues/11185
diff --git a/src/seclab_taskflow_agent/capi.py b/src/seclab_taskflow_agent/capi.py
@@ -6,40 +6,66 @@
 import json
 import logging
 import os
+from strenum import StrEnum
 from urllib.parse import urlparse
 
-# you can also set https://models.github.ai/inference if you prefer
+# Enumeration of currently supported API endpoints.
+class AI_API_ENDPOINT_ENUM(StrEnum):
+  AI_API_MODELS_GITHUB = 'models.github.ai'
+  AI_API_GITHUBCOPILOT = 'api.githubcopilot.com'
+
+COPILOT_INTEGRATION_ID = 'vscode-chat'
+
+# you can also set https://api.githubcopilot.com if you prefer
 # but beware that your taskflows need to reference the correct model id
-# since the Modeld API uses it's own id schema, use -l with your desired
+# since different APIs use their own id schema, use -l with your desired
 # endpoint to retrieve the correct id names to use for your taskflow
-COPILOT_API_ENDPOINT = os.getenv('COPILOT_API_ENDPOINT', default='https://api.githubcopilot.com')
-COPILOT_INTEGRATION_ID = 'vscode-chat'
+def get_AI_endpoint():
+  return os.getenv('AI_API_ENDPOINT', default='https://models.github.ai/inference')
+
+def get_AI_token():
+    """
+    Get the token for the AI API from the environment.
+    The environment variable can be named either AI_API_TOKEN
+    or COPILOT_TOKEN.
+    """
+    token = os.getenv('AI_API_TOKEN')
+    if token:
+        return token
+    token = os.getenv('COPILOT_TOKEN')
+    if token:
+        return token
+    raise RuntimeError("AI_API_TOKEN environment variable is not set.")
 
 # assume we are >= python 3.9 for our type hints
 def list_capi_models(token: str) -> dict[str, dict]:
     """Retrieve a dictionary of available CAPI models"""
     models = {}
     try:
-        match urlparse(COPILOT_API_ENDPOINT).netloc:
-            case 'api.githubcopilot.com':
+        api_endpoint = get_AI_endpoint()
+        netloc = urlparse(api_endpoint).netloc
+        match netloc:
+            case AI_API_ENDPOINT_ENUM.AI_API_GITHUBCOPILOT:
                 models_catalog = 'models'
-            case 'models.github.ai':
+            case AI_API_ENDPOINT_ENUM.AI_API_MODELS_GITHUB:
                 models_catalog = 'catalog/models'
             case _:
-                raise ValueError(f"Unsupported Model Endpoint: {COPILOT_API_ENDPOINT}")
-        r = httpx.get(httpx.URL(COPILOT_API_ENDPOINT).join(models_catalog),
+                raise ValueError(f"Unsupported Model Endpoint: {api_endpoint}")
+        r = httpx.get(httpx.URL(api_endpoint).join(models_catalog),
                       headers={
                           'Accept': 'application/json',
                           'Authorization': f'Bearer {token}',
                           'Copilot-Integration-Id': COPILOT_INTEGRATION_ID
                       })
         r.raise_for_status()
         # CAPI vs Models API
-        match urlparse(COPILOT_API_ENDPOINT).netloc:
-            case 'api.githubcopilot.com':
+        match netloc:
+            case AI_API_ENDPOINT_ENUM.AI_API_GITHUBCOPILOT:
                 models_list = r.json().get('data', [])
-            case 'models.github.ai':
+            case AI_API_ENDPOINT_ENUM.AI_API_MODELS_GITHUB:
                 models_list = r.json()
+            case _:
+                raise ValueError(f"Unsupported Model Endpoint: {api_endpoint}")
         for model in models_list:
             models[model.get('id')] = dict(model)
     except httpx.RequestError as e:
@@ -51,17 +77,18 @@ def list_capi_models(token: str) -> dict[str, dict]:
     return models
 
 def supports_tool_calls(model: str, models: dict) -> bool:
-    match urlparse(COPILOT_API_ENDPOINT).netloc:
-        case 'api.githubcopilot.com':
+    api_endpoint = get_AI_endpoint()
+    match urlparse(api_endpoint).netloc:
+        case AI_API_ENDPOINT_ENUM.AI_API_GITHUBCOPILOT:
             return models.get(model, {}).\
                 get('capabilities', {}).\
                 get('supports', {}).\
                 get('tool_calls', False)
-        case 'models.github.ai':
+        case AI_API_ENDPOINT_ENUM.AI_API_MODELS_GITHUB:
             return 'tool-calling' in models.get(model, {}).\
                 get('capabilities', [])
         case _:
-            raise ValueError(f"Unsupported Model Endpoint: {COPILOT_API_ENDPOINT}")
+            raise ValueError(f"Unsupported Model Endpoint: {api_endpoint}")
 
 def list_tool_call_models(token: str) -> dict[str, dict]:
     models = list_capi_models(token)
diff --git a/tests/test_api_endpoint_config.py b/tests/test_api_endpoint_config.py
@@ -0,0 +1,47 @@
+# SPDX-FileCopyrightText: 2025 GitHub
+# SPDX-License-Identifier: MIT
+
+"""
+Test API endpoint configuration.
+"""
+
+import pytest
+import os
+from urllib.parse import urlparse
+from seclab_taskflow_agent.capi import get_AI_endpoint, AI_API_ENDPOINT_ENUM
+
+class TestAPIEndpoint:
+    """Test API endpoint configuration."""
+
+    def test_default_api_endpoint(self):
+        """Test that default API endpoint is set to models.github.ai/inference."""
+        # When no env var is set, it should default to models.github.ai/inference
+        try:
+            # Save original env
+            original_env = os.environ.pop('AI_API_ENDPOINT', None)
+            endpoint = get_AI_endpoint()
+            assert endpoint is not None
+            assert isinstance(endpoint, str)
+            assert urlparse(endpoint).netloc == AI_API_ENDPOINT_ENUM.AI_API_MODELS_GITHUB
+        finally:
+            # Restore original env
+            if original_env:
+                os.environ['AI_API_ENDPOINT'] = original_env
+
+    def test_api_endpoint_env_override(self):
+        """Test that AI_API_ENDPOINT can be overridden by environment variable."""
+        try:
+            # Save original env
+            original_env = os.environ.pop('AI_API_ENDPOINT', None)
+            # Set different endpoint
+            test_endpoint = 'https://api.githubcopilot.com'
+            os.environ['AI_API_ENDPOINT'] = test_endpoint
+
+            assert get_AI_endpoint() == test_endpoint
+        finally:
+            # Restore original env
+            if original_env:
+                os.environ['AI_API_ENDPOINT'] = original_env
+
+if __name__ == '__main__':
+    pytest.main([__file__, '-v'])
diff --git a/tests/test_cli_parser.py b/tests/test_cli_parser.py
@@ -0,0 +1,73 @@
+# SPDX-FileCopyrightText: 2025 GitHub
+# SPDX-License-Identifier: MIT
+
+"""
+Test CLI global variable parsing.
+"""
+
+import pytest
+from seclab_taskflow_agent.available_tools import AvailableTools
+
+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'])
diff --git a/tests/test_yaml_parser.py b/tests/test_yaml_parser.py
