Add new environment variable named AI_API_TOKEN.

Author: kevinbackhouse

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)

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, AI_API_ENDPOINT, AI_API_ENDPOINT_ENUM
+from .capi import COPILOT_INTEGRATION_ID, get_AI_endpoint, AI_API_ENDPOINT_ENUM
 
 # grab our secrets from .env, this must be in .gitignore
 load_dotenv(find_dotenv(usecwd=True))
 
-match urlparse(AI_API_ENDPOINT).netloc:
+api_endpoint = get_AI_endpoint()
+match urlparse(api_endpoint).netloc:
     case AI_API_ENDPOINT_ENUM.AI_API_GITHUBCOPILOT:
         default_model = 'gpt-4o'
     case AI_API_ENDPOINT_ENUM.AI_API_MODELS_GITHUB:
         default_model = 'openai/gpt-4o'
     case _:
-        raise ValueError(f"Unsupported Model Endpoint: {AI_API_ENDPOINT}")
+        raise ValueError(f"Unsupported Model Endpoint: {api_endpoint}")
 
 DEFAULT_MODEL = os.getenv('COPILOT_DEFAULT_MODEL', default=default_model)
 
@@ -148,7 +149,7 @@ def __init__(self,
                  model_settings: ModelSettings | None = None,
                  run_hooks: TaskRunHooks | None = None,
                  agent_hooks: TaskAgentHooks | None = None):
-        client = AsyncOpenAI(base_url=AI_API_ENDPOINT,
+        client = AsyncOpenAI(base_url=api_endpoint,
                              api_key=os.getenv('COPILOT_TOKEN'),
                              default_headers={'Copilot-Integration-Id': COPILOT_INTEGRATION_ID})
         set_default_openai_client(client)

src/seclab_taskflow_agent/capi.py

@@ -9,33 +9,49 @@
 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')
-
 # 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 different APIs use their own id schema, use -l with your desired
+# endpoint to retrieve the correct id names to use for your taskflow
+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:
-        netloc = urlparse(AI_API_ENDPOINT).netloc
+        api_endpoint = get_AI_endpoint()
+        netloc = urlparse(endpoint).netloc
         match netloc:
             case AI_API_ENDPOINT_ENUM.AI_API_GITHUBCOPILOT:
                 models_catalog = 'models'
             case AI_API_ENDPOINT_ENUM.AI_API_MODELS_GITHUB:
                 models_catalog = 'catalog/models'
             case _:
-                raise ValueError(f"Unsupported Model Endpoint: {AI_API_ENDPOINT}")
-        r = httpx.get(httpx.URL(AI_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}',
@@ -49,7 +65,7 @@ def list_capi_models(token: str) -> dict[str, dict]:
             case AI_API_ENDPOINT_ENUM.AI_API_MODELS_GITHUB:
                 models_list = r.json()
             case _:
-                raise ValueError(f"Unsupported Model Endpoint: {AI_API_ENDPOINT}")
+                raise ValueError(f"Unsupported Model Endpoint: {api_endpoint}")
         for model in models_list:
             models[model.get('id')] = dict(model)
     except httpx.RequestError as e:
@@ -61,7 +77,8 @@ def list_capi_models(token: str) -> dict[str, dict]:
     return models
 
 def supports_tool_calls(model: str, models: dict) -> bool:
-    match urlparse(AI_API_ENDPOINT).netloc:
+    api_endpoint = get_AI_endpoint()
+    match urlparse(api_endpoint).netloc:
         case AI_API_ENDPOINT_ENUM.AI_API_GITHUBCOPILOT:
             return models.get(model, {}).\
                 get('capabilities', {}).\
@@ -71,7 +88,7 @@ def supports_tool_calls(model: str, models: dict) -> bool:
             return 'tool-calling' in models.get(model, {}).\
                 get('capabilities', [])
         case _:
-            raise ValueError(f"Unsupported Model Endpoint: {AI_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)

tests/test_api_endpoint_config.py

@@ -8,51 +8,39 @@
 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."""
-    
-    @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."""
-        import seclab_taskflow_agent.capi
         # 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
-        endpoint = seclab_taskflow_agent.capi.AI_API_ENDPOINT
-        assert endpoint is not None
-        assert isinstance(endpoint, str)
-        assert urlparse(endpoint).netloc == seclab_taskflow_agent.capi.AI_API_ENDPOINT_ENUM.AI_API_MODELS_GITHUB
-    
+        try:
+            # Save original env
+            original_env = os.environ
+            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
+            os.environ = original_env
+
     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'
+            # Save original env
+            original_env = os.environ
+            # Set different endpoint
+            test_endpoint = 'https://api.githubcopilot.com'
             os.environ['AI_API_ENDPOINT'] = test_endpoint
-            
-            # Reload the module to pick up the new env var
-            self._reload_capi_module()
-            
-            import seclab_taskflow_agent.capi
-            assert seclab_taskflow_agent.capi.AI_API_ENDPOINT == test_endpoint
+
+            assert get_AI_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()
+            os.environ = original_env
 
 if __name__ == '__main__':
     pytest.main([__file__, '-v'])