diff --git a/src/seclab_taskflow_agent/agent.py b/src/seclab_taskflow_agent/agent.py index 3e6fc75..4191cd5 100644 --- a/src/seclab_taskflow_agent/agent.py +++ b/src/seclab_taskflow_agent/agent.py @@ -29,7 +29,7 @@ set_tracing_disabled, ) -from .capi import COPILOT_INTEGRATION_ID, get_AI_endpoint, get_AI_token, AI_API_ENDPOINT_ENUM +from .capi import get_AI_endpoint, get_AI_token, get_custom_header, AI_API_ENDPOINT_ENUM # grab our secrets from .env, this must be in .gitignore load_dotenv(find_dotenv(usecwd=True)) @@ -150,7 +150,7 @@ def __init__( client = AsyncOpenAI( base_url=api_endpoint, api_key=get_AI_token(), - default_headers={"Copilot-Integration-Id": COPILOT_INTEGRATION_ID}, + default_headers=get_custom_header(), ) 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 index 2fc1d0d..ad0843f 100644 --- a/src/seclab_taskflow_agent/capi.py +++ b/src/seclab_taskflow_agent/capi.py @@ -56,6 +56,30 @@ def get_AI_token(): return token raise RuntimeError("AI_API_TOKEN environment variable is not set.") +def get_custom_header() -> dict[str, str]: + """ + Get custom header from environment variable AI_API_CUSTOM_HEADER. + Expected format: name:value + Returns a dictionary that can be merged into request headers. + """ + custom_header = os.getenv('AI_API_CUSTOM_HEADER') + if not custom_header: + return {} + + # Split on first colon to handle values that might contain colons + parts = custom_header.split(':', 1) + if len(parts) != 2: + logging.warning(f"Invalid AI_API_CUSTOM_HEADER format. Expected 'name:value', got: {custom_header}") + return {} + + name, value = parts + name = name.strip() + value = value.strip() + if not name or not value: + logging.warning(f"Invalid AI_API_CUSTOM_HEADER: header name and value must be non-empty after stripping. Got: '{custom_header}'") + return {} + return {name: value} + # assume we are >= python 3.9 for our type hints def list_capi_models(token: str) -> dict[str, dict]: @@ -76,13 +100,13 @@ def list_capi_models(token: str) -> dict[str, dict]: f"Unsupported Model Endpoint: {api_endpoint}\n" f"Supported endpoints: {[e.to_url() for e in AI_API_ENDPOINT_ENUM]}" ) + headers = { + 'Accept': 'application/json', + 'Authorization': f'Bearer {token}', + } | get_custom_header() r = httpx.get( httpx.URL(api_endpoint).join(models_catalog), - headers={ - "Accept": "application/json", - "Authorization": f"Bearer {token}", - "Copilot-Integration-Id": COPILOT_INTEGRATION_ID, - }, + headers=headers, ) r.raise_for_status() # CAPI vs Models API