feat: filter models to only image generation models

Author: danielcopper

app/core/config.py

@@ -42,13 +42,15 @@ class Settings(BaseSettings):
 
     # Storage Configuration
     STORAGE_PATH: str = "./generated_images"
-    BASE_URL: str = "http://localhost:8000"
+    IMAGE_BASE_URL: str = "http://localhost:8000"  # URL where this API serves images
 
     # Security (Optional)
     API_BEARER_TOKEN: str | None = None
 
     # Model Registry
     MODEL_CACHE_TTL: int = 3600  # Cache models for 1 hour
+    FILTER_IMAGE_MODELS: bool = True  # Only return image generation models from LiteLLM
+    DEFAULT_MODEL: str | None = None  # Default model for image generation
 
     # Server Configuration
     HOST: str = "0.0.0.0"

app/services/model_registry.py

@@ -1,4 +1,5 @@
 import logging
+import re
 import time
 
 import httpx
@@ -15,6 +16,14 @@ class ModelRegistry:
     Caches results to avoid frequent API calls.
     """
 
+    # Patterns to identify image generation models
+    IMAGE_MODEL_PATTERNS = [
+        r"dall-e",
+        r"gpt-image",
+        r"gemini.*image",
+        r"imagen",
+    ]
+
     # Model capabilities database (fallback for known models)
     KNOWN_CAPABILITIES: dict[str, ModelCapabilities] = {
         "dall-e-2": ModelCapabilities(
@@ -124,6 +133,11 @@ async def _load_from_litellm(self) -> list[ModelInfo]:
             if not model_id:
                 continue
 
+            # Filter to only image generation models if enabled
+            if settings.FILTER_IMAGE_MODELS and not self._is_image_model(model_id):
+                logger.debug(f"Skipping non-image model: {model_id}")
+                continue
+
             # Determine provider from model ID
             provider = self._infer_provider(model_id)
 
@@ -167,6 +181,13 @@ def _get_static_models(self) -> list[ModelInfo]:
 
         return models
 
+    def _is_image_model(self, model_id: str) -> bool:
+        """
+        Check if a model is an image generation model.
+        """
+        model_lower = model_id.lower()
+        return any(re.search(pattern, model_lower) for pattern in self.IMAGE_MODEL_PATTERNS)
+
     def _infer_provider(self, model_id: str) -> str:
         """
         Infer provider from model ID.

tests/test_config.py

@@ -0,0 +1,74 @@
+import os
+from unittest.mock import patch
+
+
+def test_image_base_url_from_env():
+    """
+    Test IMAGE_BASE_URL is loaded from environment.
+    """
+    with patch.dict(
+        os.environ,
+        {
+            "IMAGE_BASE_URL": "http://image-api:8000",
+        },
+        clear=True,
+    ):
+        from importlib import reload
+
+        import app.core.config
+
+        reload(app.core.config)
+        settings = app.core.config.Settings()
+
+        assert settings.IMAGE_BASE_URL == "http://image-api:8000"
+
+
+def test_image_base_url_default():
+    """
+    Test default IMAGE_BASE_URL when not set.
+    """
+    with patch.dict(os.environ, {}, clear=True):
+        from importlib import reload
+
+        import app.core.config
+
+        reload(app.core.config)
+        settings = app.core.config.Settings()
+
+        assert settings.IMAGE_BASE_URL == "http://localhost:8000"
+
+
+def test_default_model_config():
+    """
+    Test DEFAULT_MODEL configuration.
+    """
+    with patch.dict(
+        os.environ,
+        {
+            "DEFAULT_MODEL": "gemini/gemini-2.0-flash-exp-image-generation",
+        },
+        clear=True,
+    ):
+        from importlib import reload
+
+        import app.core.config
+
+        reload(app.core.config)
+        settings = app.core.config.Settings()
+
+        assert settings.DEFAULT_MODEL == "gemini/gemini-2.0-flash-exp-image-generation"
+
+
+def test_filter_image_models_default():
+    """
+    Test FILTER_IMAGE_MODELS defaults to True.
+    """
+    with patch.dict(os.environ, {}, clear=True):
+        from importlib import reload
+
+        import app.core.config
+
+        reload(app.core.config)
+        settings = app.core.config.Settings()
+
+        assert settings.FILTER_IMAGE_MODELS is True

tests/test_model_registry.py

@@ -119,3 +119,68 @@ def test_infer_provider():
     assert registry._infer_provider("gemini-2.0-flash") == "gemini"
     assert registry._infer_provider("imagen-3.0") == "gemini"
     assert registry._infer_provider("unknown") == "unknown"
+
+
+def test_is_image_model():
+    """
+    Test image model detection.
+    """
+    registry = ModelRegistry()
+
+    # Should match image models
+    assert registry._is_image_model("dall-e-3") is True
+    assert registry._is_image_model("gpt-image-1") is True
+    assert registry._is_image_model("gemini-2.0-flash-preview-image-generation") is True
+    assert registry._is_image_model("imagen-3.0") is True
+
+    # Should NOT match non-image models
+    assert registry._is_image_model("gpt-4o") is False
+    assert registry._is_image_model("claude-3-opus") is False
+    assert registry._is_image_model("gemini-pro") is False
+
+
+@pytest.mark.asyncio
+async def test_filter_image_models():
+    """
+    Test that only image models are returned when filter is enabled.
+    """
+    registry = ModelRegistry()
+
+    mock_response = {
+        "data": [
+            {"id": "dall-e-3"},
+            {"id": "gpt-4o"},
+            {"id": "claude-3-opus"},
+            {"id": "gemini-2.0-flash-preview-image-generation"},
+        ]
+    }
+
+    with patch("app.services.model_registry.settings") as mock_settings:
+        mock_settings.LITELLM_BASE_URL = "http://litellm:4000"
+        mock_settings.LITELLM_API_KEY = None
+        mock_settings.litellm_available = True
+        mock_settings.FILTER_IMAGE_MODELS = True
+        mock_settings.MODEL_CACHE_TTL = 3600
+
+        # Mock HTTP response
+        mock_resp = MagicMock()
+        mock_resp.status_code = 200
+        mock_resp.json.return_value = mock_response
+        mock_resp.raise_for_status.return_value = None
+
+        mock_client_instance = MagicMock()
+        mock_client_instance.get = AsyncMock(return_value=mock_resp)
+        mock_client = MagicMock()
+        mock_client.__aenter__ = AsyncMock(return_value=mock_client_instance)
+        mock_client.__aexit__ = AsyncMock(return_value=None)
+
+        with patch("app.services.model_registry.httpx.AsyncClient", return_value=mock_client):
+            models = await registry.load_models()
+
+            # Should only have image models
+            assert len(models) == 2
+            model_ids = [m.id for m in models]
+            assert "dall-e-3" in model_ids
+            assert "gemini-2.0-flash-preview-image-generation" in model_ids
+            assert "gpt-4o" not in model_ids
+            assert "claude-3-opus" not in model_ids