Skip to content
1 change: 1 addition & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@ AI_API_TOKEN=<your_github_token>
# MCP configs
GH_TOKEN=<your_github_token>
CODEQL_DBS_BASE_PATH="/app/my_data/codeql_databases"
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
CODEQL_DBS_BASE_PATH="/app/my_data/codeql_databases"
CODEQL_DBS_BASE_PATH="/app/data/codeql_databases"

Should we change it to data since we removed my_?

AI_API_ENDPOINT="https://models.github.ai/inference"
```

## Deploying from Source
Expand Down
5 changes: 4 additions & 1 deletion src/seclab_taskflow_agent/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,10 @@
case AI_API_ENDPOINT_ENUM.AI_API_MODELS_GITHUB:
default_model = 'openai/gpt-4o'
case _:
raise ValueError(f"Unsupported Model Endpoint: {api_endpoint}")
raise ValueError(
f"Unsupported Model Endpoint: {api_endpoint}\n"
f"Supported endpoints: {[e.to_url() for e in AI_API_ENDPOINT_ENUM]}"
)

DEFAULT_MODEL = os.getenv('COPILOT_DEFAULT_MODEL', default=default_model)

Expand Down
28 changes: 21 additions & 7 deletions src/seclab_taskflow_agent/capi.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,8 +11,20 @@

# 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'
AI_API_MODELS_GITHUB = 'models.github.ai'
AI_API_GITHUBCOPILOT = 'api.githubcopilot.com'

def to_url(self):

Check notice

Code scanning / CodeQL

Explicit returns mixed with implicit (fall through) returns Note

Mixing implicit and explicit returns may indicate an error, as implicit returns always return None.
"""
Convert the endpoint to its full URL.
"""
match self:
case AI_API_ENDPOINT_ENUM.AI_API_GITHUBCOPILOT:
return f"https://{self}"
case AI_API_ENDPOINT_ENUM.AI_API_MODELS_GITHUB:
return f"https://{self}/inference"
case _:
raise ValueError(f"Unsupported endpoint: {self}")

COPILOT_INTEGRATION_ID = 'vscode-chat'

Expand All @@ -21,7 +33,7 @@
# 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')
return os.getenv('AI_API_ENDPOINT', default='https://models.github.ai/inference')

def get_AI_token():
"""
Expand Down Expand Up @@ -50,7 +62,8 @@
case AI_API_ENDPOINT_ENUM.AI_API_MODELS_GITHUB:
models_catalog = 'catalog/models'
case _:
raise ValueError(f"Unsupported Model Endpoint: {api_endpoint}")
raise ValueError(f"Unsupported Model Endpoint: {api_endpoint}\n"
f"Supported endpoints: {[e.to_url() for e in AI_API_ENDPOINT_ENUM]}")
r = httpx.get(httpx.URL(api_endpoint).join(models_catalog),
headers={
'Accept': 'application/json',
Expand All @@ -64,8 +77,6 @@
models_list = r.json().get('data', [])
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:
Expand All @@ -88,7 +99,10 @@
return 'tool-calling' in models.get(model, {}).\
get('capabilities', [])
case _:
raise ValueError(f"Unsupported Model Endpoint: {api_endpoint}")
raise ValueError(
f"Unsupported Model Endpoint: {api_endpoint}\n"
f"Supported endpoints: {[e.to_url() for e in AI_API_ENDPOINT_ENUM]}"
)

def list_tool_call_models(token: str) -> dict[str, dict]:
models = list_capi_models(token)
Expand Down
23 changes: 22 additions & 1 deletion tests/test_api_endpoint_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
import pytest
import os
from urllib.parse import urlparse
from seclab_taskflow_agent.capi import get_AI_endpoint, AI_API_ENDPOINT_ENUM
from seclab_taskflow_agent.capi import get_AI_endpoint, AI_API_ENDPOINT_ENUM, list_capi_models

class TestAPIEndpoint:
"""Test API endpoint configuration."""
Expand Down Expand Up @@ -43,5 +43,26 @@
if original_env:
os.environ['AI_API_ENDPOINT'] = original_env

def test_to_url_models_github(self):
"""Test to_url method for models.github.ai endpoint."""
endpoint = AI_API_ENDPOINT_ENUM.AI_API_MODELS_GITHUB
assert endpoint.to_url() == 'https://models.github.ai/inference'

def test_to_url_githubcopilot(self):
"""Test to_url method for GitHub Copilot endpoint."""
endpoint = AI_API_ENDPOINT_ENUM.AI_API_GITHUBCOPILOT
assert endpoint.to_url() == 'https://api.githubcopilot.com'

def test_unsupported_endpoint(self, monkeypatch):
"""Test that unsupported API endpoint raises ValueError."""
api_endpoint = 'https://unsupported.example.com'
monkeypatch.setenv('AI_API_ENDPOINT', api_endpoint)
with pytest.raises(ValueError) as excinfo:
list_capi_models("abc")
msg = str(excinfo.value)
assert 'Unsupported Model Endpoint' in msg
assert 'https://models.github.ai/inference' in msg
assert 'https://api.githubcopilot.com' in msg

if __name__ == '__main__':
pytest.main([__file__, '-v'])