Improvements to the code written by Copilot.

Author: kevinbackhouse

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",

src/seclab_taskflow_agent/capi.py

@@ -6,24 +6,31 @@
 import json
 import logging
 import os
+from strenum import StrEnum
 from urllib.parse import urlparse
 
 # you can also set https://api.githubcopilot.com if you prefer
 # but beware that your taskflows need to reference the correct model id
 # 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
 AI_API_ENDPOINT = os.getenv('AI_API_ENDPOINT', default='https://models.github.ai/inference')
+
+class AI_API_ENDPOINT_ENUM(StrEnum):
+  AI_API_MODELS_GITHUB = 'models.github.ai'
+  AI_API_GITHUBCOPILOT = 'api.githubcopilot.com'
+
 COPILOT_INTEGRATION_ID = 'vscode-chat'
 
 # 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(AI_API_ENDPOINT).netloc:
-            case 'api.githubcopilot.com':
+        netloc = urlparse(AI_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: {AI_API_ENDPOINT}")
@@ -35,11 +42,13 @@ def list_capi_models(token: str) -> dict[str, dict]:
                       })
         r.raise_for_status()
         # CAPI vs Models API
-        match urlparse(AI_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: {AI_API_ENDPOINT}")
         for model in models_list:
             models[model.get('id')] = dict(model)
     except httpx.RequestError as e:
@@ -52,12 +61,12 @@ def list_capi_models(token: str) -> dict[str, dict]:
 
 def supports_tool_calls(model: str, models: dict) -> bool:
     match urlparse(AI_API_ENDPOINT).netloc:
-        case 'api.githubcopilot.com':
+        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 _:

tests/test_api_endpoint_config.py

@@ -0,0 +1,61 @@
+# SPDX-FileCopyrightText: 2025 GitHub
+# SPDX-License-Identifier: MIT
+
+"""
+Test API endpoint configuration.
+"""
+
+import pytest
+import tempfile
+from pathlib import Path
+import yaml
+import os
+from urllib.parse import urlparse
+from seclab_taskflow_agent.available_tools import AvailableTools
+
+class TestAPIEndpoint:
+    """Test API endpoint configuration."""
+    
+    @staticmethod
+    def _reload_capi_module():
+        """Helper method to reload the capi module."""
+        import importlib
+        import seclab_taskflow_agent.capi
+        importlib.reload(seclab_taskflow_agent.capi)
+    
+    def test_default_api_endpoint(self):
+        """Test that default API endpoint is set to models.github.ai/inference."""
+        from seclab_taskflow_agent.capi import AI_API_ENDPOINT, AI_API_ENDPOINT_ENUM
+        # When no env var is set, it should default to models.github.ai/inference
+        # Note: We can't easily test this without manipulating the environment
+        # so we'll just import and verify the constant exists
+        assert AI_API_ENDPOINT is not None
+        assert isinstance(AI_API_ENDPOINT, str)
+        assert urlparse(AI_API_ENDPOINT).netloc == AI_API_ENDPOINT_ENUM.AI_API_MODELS_GITHUB
+    
+    def test_api_endpoint_env_override(self):
+        """Test that AI_API_ENDPOINT can be overridden by environment variable."""
+        # Save original env
+        original_env = os.environ.get('AI_API_ENDPOINT')
+        
+        try:
+            # Set custom endpoint
+            test_endpoint = 'https://test.example.com'
+            os.environ['AI_API_ENDPOINT'] = test_endpoint
+            
+            # Reload the module to pick up the new env var
+            self._reload_capi_module()
+            
+            from seclab_taskflow_agent.capi import AI_API_ENDPOINT
+            assert AI_API_ENDPOINT == test_endpoint
+        finally:
+            # Restore original env
+            if original_env is None:
+                os.environ.pop('AI_API_ENDPOINT', None)
+            else:
+                os.environ['AI_API_ENDPOINT'] = original_env
+            # Reload again to restore original state
+            self._reload_capi_module()
+
+if __name__ == '__main__':
+    pytest.main([__file__, '-v'])

tests/test_cli_parser.py

@@ -0,0 +1,78 @@
+# SPDX-FileCopyrightText: 2025 GitHub
+# SPDX-License-Identifier: MIT
+
+"""
+Test CLI global variable parsing.
+"""
+
+import pytest
+import tempfile
+from pathlib import Path
+import yaml
+import os
+from urllib.parse import urlparse
+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'])

tests/test_yaml_parser.py

@@ -12,6 +12,7 @@
 from pathlib import Path
 import yaml
 import os
+from urllib.parse import urlparse
 from seclab_taskflow_agent.available_tools import AvailableTools
 
 class TestYamlParser:
@@ -44,109 +45,5 @@ def test_parse_example_taskflows(self):
         assert len(example_task_flow['taskflow']) == 4  # 4 tasks in taskflow
         assert example_task_flow['taskflow'][0]['task']['max_steps'] == 20
 
-class TestCliGlobals:
-    """Test CLI global variable parsing."""
-    
-    def test_parse_single_global(self):
-        """Test parsing a single global variable from command line."""
-        from seclab_taskflow_agent.__main__ import parse_prompt_args
-        available_tools = AvailableTools()
-        
-        p, t, l, cli_globals, user_prompt, _ = parse_prompt_args(
-            available_tools, "-t example -g fruit=apples")
-        
-        assert t == "example"
-        assert cli_globals == {"fruit": "apples"}
-        assert p is None
-        assert l is False
-    
-    def test_parse_multiple_globals(self):
-        """Test parsing multiple global variables from command line."""
-        from seclab_taskflow_agent.__main__ import parse_prompt_args
-        available_tools = AvailableTools()
-        
-        p, t, l, cli_globals, user_prompt, _ = parse_prompt_args(
-            available_tools, "-t example -g fruit=apples -g color=red")
-        
-        assert t == "example"
-        assert cli_globals == {"fruit": "apples", "color": "red"}
-        assert p is None
-        assert l is False
-    
-    def test_parse_global_with_spaces(self):
-        """Test parsing global variables with spaces in values."""
-        from seclab_taskflow_agent.__main__ import parse_prompt_args
-        available_tools = AvailableTools()
-        
-        p, t, l, cli_globals, user_prompt, _ = parse_prompt_args(
-            available_tools, "-t example -g message=hello world")
-        
-        assert t == "example"
-        # "world" becomes part of the prompt, not the value
-        assert cli_globals == {"message": "hello"}
-        assert "world" in user_prompt
-    
-    def test_parse_global_with_equals_in_value(self):
-        """Test parsing global variables with equals sign in value."""
-        from seclab_taskflow_agent.__main__ import parse_prompt_args
-        available_tools = AvailableTools()
-        
-        p, t, l, cli_globals, user_prompt, _ = parse_prompt_args(
-            available_tools, "-t example -g equation=x=5")
-        
-        assert t == "example"
-        assert cli_globals == {"equation": "x=5"}
-    
-    def test_globals_in_taskflow_file(self):
-        """Test that globals can be read from taskflow file."""
-        available_tools = AvailableTools()
-        
-        taskflow = available_tools.get_taskflow("tests.data.test_globals_taskflow")
-        assert 'globals' in taskflow
-        assert taskflow['globals']['test_var'] == 'default_value'
-
-class TestAPIEndpoint:
-    """Test API endpoint configuration."""
-    
-    @staticmethod
-    def _reload_capi_module():
-        """Helper method to reload the capi module."""
-        import importlib
-        import seclab_taskflow_agent.capi
-        importlib.reload(seclab_taskflow_agent.capi)
-    
-    def test_default_api_endpoint(self):
-        """Test that default API endpoint is set to models.github.ai/inference."""
-        from seclab_taskflow_agent.capi import AI_API_ENDPOINT
-        # When no env var is set, it should default to models.github.ai/inference
-        # Note: We can't easily test this without manipulating the environment
-        # so we'll just import and verify the constant exists
-        assert AI_API_ENDPOINT is not None
-        assert isinstance(AI_API_ENDPOINT, str)
-    
-    def test_api_endpoint_env_override(self):
-        """Test that AI_API_ENDPOINT can be overridden by environment variable."""
-        # Save original env
-        original_env = os.environ.get('AI_API_ENDPOINT')
-        
-        try:
-            # Set custom endpoint
-            test_endpoint = 'https://test.example.com'
-            os.environ['AI_API_ENDPOINT'] = test_endpoint
-            
-            # Reload the module to pick up the new env var
-            self._reload_capi_module()
-            
-            from seclab_taskflow_agent.capi import AI_API_ENDPOINT
-            assert AI_API_ENDPOINT == test_endpoint
-        finally:
-            # Restore original env
-            if original_env is None:
-                os.environ.pop('AI_API_ENDPOINT', None)
-            else:
-                os.environ['AI_API_ENDPOINT'] = original_env
-            # Reload again to restore original state
-            self._reload_capi_module()
-
 if __name__ == '__main__':
     pytest.main([__file__, '-v'])