mirror of
https://github.com/prowler-cloud/prowler.git
synced 2025-12-19 05:17:47 +00:00
feat(lighthouse): filter out non-compatible OpenAI models (#9523)
Co-authored-by: Chandrapal Badshah <12944530+Chan9390@users.noreply.github.com> Co-authored-by: Adrián Jesús Peña Rodríguez <adrianjpr@gmail.com>
This commit is contained in:
committed by
GitHub
parent
4398b00801
commit
2b4b23c719
@@ -10,6 +10,7 @@ All notable changes to the **Prowler API** are documented in this file.
|
|||||||
|
|
||||||
### Changed
|
### Changed
|
||||||
- Endpoint `GET /overviews/attack-surfaces` no longer returns the related check IDs [(#9529)](https://github.com/prowler-cloud/prowler/pull/9529)
|
- Endpoint `GET /overviews/attack-surfaces` no longer returns the related check IDs [(#9529)](https://github.com/prowler-cloud/prowler/pull/9529)
|
||||||
|
- OpenAI provider to only load chat-compatible models with tool calling support [(#9523)](https://github.com/prowler-cloud/prowler/pull/9523)
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
|
|||||||
@@ -11,6 +11,41 @@ from api.models import LighthouseProviderConfiguration, LighthouseProviderModels
|
|||||||
|
|
||||||
logger = get_task_logger(__name__)
|
logger = get_task_logger(__name__)
|
||||||
|
|
||||||
|
# OpenAI model prefixes to exclude from Lighthouse model selection.
|
||||||
|
# These models don't support text chat completions and tool calling.
|
||||||
|
EXCLUDED_OPENAI_MODEL_PREFIXES = (
|
||||||
|
"dall-e", # Image generation
|
||||||
|
"whisper", # Audio transcription
|
||||||
|
"tts-", # Text-to-speech (tts-1, tts-1-hd, etc.)
|
||||||
|
"sora", # Text-to-video (sora-2, sora-2-pro, etc.)
|
||||||
|
"text-embedding", # Embeddings
|
||||||
|
"embedding", # Embeddings (alternative naming)
|
||||||
|
"text-moderation", # Content moderation
|
||||||
|
"omni-moderation", # Content moderation
|
||||||
|
"text-davinci", # Legacy completion models
|
||||||
|
"text-curie", # Legacy completion models
|
||||||
|
"text-babbage", # Legacy completion models
|
||||||
|
"text-ada", # Legacy completion models
|
||||||
|
"davinci", # Legacy completion models
|
||||||
|
"curie", # Legacy completion models
|
||||||
|
"babbage", # Legacy completion models
|
||||||
|
"ada", # Legacy completion models
|
||||||
|
"computer-use", # Computer control agent
|
||||||
|
"gpt-image", # Image generation
|
||||||
|
"gpt-audio", # Audio models
|
||||||
|
"gpt-realtime", # Realtime voice API
|
||||||
|
)
|
||||||
|
|
||||||
|
# OpenAI model substrings to exclude (patterns that can appear anywhere in model ID).
|
||||||
|
# These patterns identify non-chat model variants.
|
||||||
|
EXCLUDED_OPENAI_MODEL_SUBSTRINGS = (
|
||||||
|
"-audio-", # Audio preview models (gpt-4o-audio-preview, etc.)
|
||||||
|
"-realtime-", # Realtime preview models (gpt-4o-realtime-preview, etc.)
|
||||||
|
"-transcribe", # Transcription models (gpt-4o-transcribe, etc.)
|
||||||
|
"-tts", # TTS models (gpt-4o-mini-tts)
|
||||||
|
"-instruct", # Legacy instruct models (gpt-3.5-turbo-instruct, etc.)
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def _extract_error_message(e: Exception) -> str:
|
def _extract_error_message(e: Exception) -> str:
|
||||||
"""
|
"""
|
||||||
@@ -283,20 +318,41 @@ def _fetch_openai_models(api_key: str) -> Dict[str, str]:
|
|||||||
"""
|
"""
|
||||||
Fetch available models from OpenAI API.
|
Fetch available models from OpenAI API.
|
||||||
|
|
||||||
|
Filters out models that don't support text input/output and tool calling,
|
||||||
|
such as image generation (DALL-E), audio transcription (Whisper),
|
||||||
|
text-to-speech (TTS), embeddings, and moderation models.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
api_key: OpenAI API key for authentication.
|
api_key: OpenAI API key for authentication.
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
Dict mapping model_id to model_name. For OpenAI, both are the same
|
Dict mapping model_id to model_name. For OpenAI, both are the same
|
||||||
as the API doesn't provide separate display names.
|
as the API doesn't provide separate display names. Only includes
|
||||||
|
models that support text input, text output or tool calling.
|
||||||
|
|
||||||
Raises:
|
Raises:
|
||||||
Exception: If the API call fails.
|
Exception: If the API call fails.
|
||||||
"""
|
"""
|
||||||
client = openai.OpenAI(api_key=api_key)
|
client = openai.OpenAI(api_key=api_key)
|
||||||
models = client.models.list()
|
models = client.models.list()
|
||||||
# OpenAI uses model.id for both ID and display name
|
|
||||||
return {m.id: m.id for m in getattr(models, "data", [])}
|
# Filter models to only include those supporting chat completions + tool calling
|
||||||
|
filtered_models = {}
|
||||||
|
for model in getattr(models, "data", []):
|
||||||
|
model_id = model.id
|
||||||
|
|
||||||
|
# Skip if model ID starts with excluded prefixes
|
||||||
|
if model_id.startswith(EXCLUDED_OPENAI_MODEL_PREFIXES):
|
||||||
|
continue
|
||||||
|
|
||||||
|
# Skip if model ID contains excluded substrings
|
||||||
|
if any(substring in model_id for substring in EXCLUDED_OPENAI_MODEL_SUBSTRINGS):
|
||||||
|
continue
|
||||||
|
|
||||||
|
# Include model (supports chat completions + tool calling)
|
||||||
|
filtered_models[model_id] = model_id
|
||||||
|
|
||||||
|
return filtered_models
|
||||||
|
|
||||||
|
|
||||||
def _fetch_openai_compatible_models(base_url: str, api_key: str) -> Dict[str, str]:
|
def _fetch_openai_compatible_models(base_url: str, api_key: str) -> Dict[str, str]:
|
||||||
|
|||||||
Reference in New Issue
Block a user