feat(litellm): add direct provider fallback for unsupported features

LiteLLM doesn't support all provider features (e.g., Gemini aspect
ratios other than 1:1). New DIRECT_PROVIDER_FALLBACK env var enables
automatic fallback to native provider APIs when needed.

- Add DIRECT_PROVIDER_FALLBACK config option (default: false)
- Implement fallback logic in LiteLLMService for Gemini aspect ratios
- Add helper methods for Gemini model detection and normalization
- Add 14 tests for fallback logic

Author: danielcopper

.env.example

@@ -53,6 +53,17 @@ MARKDOWN_EMBED_IMAGES=false
 # Bearer token for API authentication (leave empty to disable auth)
 API_BEARER_TOKEN=
 
+# -----------------------------------------------------------------------------
+# Provider Fallback
+# -----------------------------------------------------------------------------
+# Use direct provider API when LiteLLM doesn't support a feature (default: false)
+# When enabled, requests using unsupported features will bypass LiteLLM and
+# use the provider's native API directly. Requires the respective API key.
+#
+# Currently applies to:
+# - Gemini aspect ratios other than 1:1 (requires GEMINI_API_KEY)
+DIRECT_PROVIDER_FALLBACK=false
+
 # -----------------------------------------------------------------------------
 # Model Configuration
 # -----------------------------------------------------------------------------

README.md

@@ -127,6 +127,18 @@ OPENAI_API_KEY=sk-...
 GEMINI_API_KEY=...
 ```
 
+### Direct Provider Fallback
+
+LiteLLM doesn't support all provider features. Enable `DIRECT_PROVIDER_FALLBACK` to automatically use the native provider API when needed:
+
+```env
+DIRECT_PROVIDER_FALLBACK=true
+GEMINI_API_KEY=...  # Required for Gemini fallback
+```
+
+Currently applies to:
+- **Gemini aspect ratios**: 16:9, 9:16, 4:3, 3:4 (LiteLLM only supports 1:1)
+
 ## Supported Models
 
 ### OpenAI

app/core/config.py

@@ -57,6 +57,9 @@ class Settings(BaseSettings):
     FILTER_IMAGE_MODELS: bool = True  # Only return image generation models from LiteLLM
     DEFAULT_MODEL: str | None = None  # Default model for image generation
 
+    # Provider Fallback
+    DIRECT_PROVIDER_FALLBACK: bool = False  # Use direct provider API for unsupported LiteLLM features
+
     # Server Configuration
     HOST: str = "0.0.0.0"
     PORT: int = 8000

app/services/litellm_service.py

@@ -58,6 +58,16 @@ async def generate_image(
         """
         logger.info(f"Generating {n} image(s) with {model} via LiteLLM")
 
+        # Fallback to direct provider API for unsupported features
+        if self._should_use_direct_provider(model, aspect_ratio):
+            return await self._generate_via_direct_provider(
+                prompt=prompt,
+                model=model,
+                aspect_ratio=aspect_ratio,
+                quality=quality,
+                n=n,
+            )
+
         # Get model capabilities to adjust parameters
         model_info = model_registry.get_model(model)
         size = self._get_size(aspect_ratio)
@@ -168,6 +178,71 @@ def _get_size(self, aspect_ratio: str) -> str:
         """
         return self.ASPECT_RATIO_SIZES.get(aspect_ratio, "1024x1024")
 
+    def _should_use_direct_provider(self, model: str, aspect_ratio: str) -> bool:
+        """
+        Check if we should use direct provider API instead of LiteLLM.
+
+        Currently applies to:
+        - Gemini models with non-square aspect ratios (LiteLLM doesn't support this)
+        """
+        if not settings.DIRECT_PROVIDER_FALLBACK:
+            return False
+
+        # Gemini: aspect ratios other than 1:1 not supported via LiteLLM
+        if self._is_gemini_model(model) and aspect_ratio != "1:1":
+            if settings.gemini_available:
+                return True
+            logger.warning(
+                f"DIRECT_PROVIDER_FALLBACK enabled but GEMINI_API_KEY not set. "
+                f"Falling back to LiteLLM (aspect_ratio={aspect_ratio} may not work)."
+            )
+
+        return False
+
+    def _is_gemini_model(self, model: str) -> bool:
+        """Check if model is a Gemini/Imagen model."""
+        model_lower = model.lower()
+        return "gemini" in model_lower or "imagen" in model_lower
+
+    def _normalize_gemini_model(self, model: str) -> str:
+        """
+        Normalize Gemini model name for direct API.
+
+        LiteLLM uses "gemini/model-name", direct API uses "model-name".
+        """
+        if model.startswith("gemini/"):
+            return model[7:]
+        return model
+
+    async def _generate_via_direct_provider(
+        self,
+        prompt: str,
+        model: str,
+        aspect_ratio: str,
+        quality: str,
+        n: int,
+    ) -> list[str]:
+        """
+        Generate images via direct provider API.
+
+        Used when LiteLLM doesn't support a required feature.
+        """
+        if self._is_gemini_model(model):
+            from app.services.gemini_service import get_gemini_service
+
+            logger.info(f"Using direct Gemini API for {model} (aspect_ratio={aspect_ratio})")
+            gemini_service = get_gemini_service()
+            return await gemini_service.generate_image(
+                prompt=prompt,
+                model=self._normalize_gemini_model(model),
+                aspect_ratio=aspect_ratio,
+                quality=quality,
+                n=n,
+            )
+
+        # Add other providers here as needed
+        raise ValueError(f"No direct provider fallback available for model: {model}")
+
 
 def get_litellm_service() -> LiteLLMService:
     """

tests/test_litellm_service.py

@@ -0,0 +1,179 @@
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+from app.services.litellm_service import LiteLLMService
+
+
+class TestIsGeminiModel:
+    """Tests for _is_gemini_model helper."""
+
+    @pytest.fixture
+    def service(self):
+        with patch("app.services.litellm_service.settings") as mock_settings:
+            mock_settings.LITELLM_BASE_URL = "http://localhost:4000"
+            mock_settings.LITELLM_API_KEY = "test-key"
+            return LiteLLMService()
+
+    def test_gemini_model(self, service):
+        assert service._is_gemini_model("gemini-2.0-flash-preview-image-generation") is True
+        assert service._is_gemini_model("gemini/gemini-2.0-flash-exp") is True
+        assert service._is_gemini_model("GEMINI-MODEL") is True
+
+    def test_imagen_model(self, service):
+        assert service._is_gemini_model("imagen-3.0-generate-002") is True
+        assert service._is_gemini_model("vertex/imagen-3") is True
+
+    def test_non_gemini_model(self, service):
+        assert service._is_gemini_model("dall-e-3") is False
+        assert service._is_gemini_model("gpt-image-1") is False
+        assert service._is_gemini_model("stable-diffusion") is False
+
+
+class TestNormalizeGeminiModel:
+    """Tests for _normalize_gemini_model helper."""
+
+    @pytest.fixture
+    def service(self):
+        with patch("app.services.litellm_service.settings") as mock_settings:
+            mock_settings.LITELLM_BASE_URL = "http://localhost:4000"
+            mock_settings.LITELLM_API_KEY = "test-key"
+            return LiteLLMService()
+
+    def test_removes_gemini_prefix(self, service):
+        assert service._normalize_gemini_model("gemini/gemini-2.0-flash-exp") == "gemini-2.0-flash-exp"
+
+    def test_keeps_model_without_prefix(self, service):
+        assert service._normalize_gemini_model("gemini-2.0-flash-exp") == "gemini-2.0-flash-exp"
+
+    def test_keeps_other_prefixes(self, service):
+        assert service._normalize_gemini_model("vertex/imagen-3") == "vertex/imagen-3"
+
+
+class TestShouldUseDirectProvider:
+    """Tests for _should_use_direct_provider logic."""
+
+    @pytest.fixture
+    def service(self):
+        with patch("app.services.litellm_service.settings") as mock_settings:
+            mock_settings.LITELLM_BASE_URL = "http://localhost:4000"
+            mock_settings.LITELLM_API_KEY = "test-key"
+            return LiteLLMService()
+
+    def test_disabled_when_fallback_false(self, service):
+        with patch("app.services.litellm_service.settings") as mock_settings:
+            mock_settings.DIRECT_PROVIDER_FALLBACK = False
+            assert service._should_use_direct_provider("gemini/model", "16:9") is False
+
+    def test_disabled_for_square_aspect_ratio(self, service):
+        with patch("app.services.litellm_service.settings") as mock_settings:
+            mock_settings.DIRECT_PROVIDER_FALLBACK = True
+            mock_settings.gemini_available = True
+            assert service._should_use_direct_provider("gemini/model", "1:1") is False
+
+    def test_enabled_for_gemini_non_square(self, service):
+        with patch("app.services.litellm_service.settings") as mock_settings:
+            mock_settings.DIRECT_PROVIDER_FALLBACK = True
+            mock_settings.gemini_available = True
+            assert service._should_use_direct_provider("gemini/model", "16:9") is True
+            assert service._should_use_direct_provider("gemini/model", "9:16") is True
+            assert service._should_use_direct_provider("imagen-3", "4:3") is True
+
+    def test_disabled_when_gemini_key_missing(self, service, caplog):
+        with patch("app.services.litellm_service.settings") as mock_settings:
+            mock_settings.DIRECT_PROVIDER_FALLBACK = True
+            mock_settings.gemini_available = False
+            assert service._should_use_direct_provider("gemini/model", "16:9") is False
+            assert "GEMINI_API_KEY not set" in caplog.text
+
+    def test_disabled_for_non_gemini_models(self, service):
+        with patch("app.services.litellm_service.settings") as mock_settings:
+            mock_settings.DIRECT_PROVIDER_FALLBACK = True
+            mock_settings.gemini_available = True
+            assert service._should_use_direct_provider("dall-e-3", "16:9") is False
+            assert service._should_use_direct_provider("gpt-image-1", "9:16") is False
+
+
+class TestGenerateImageWithFallback:
+    """Tests for generate_image with direct provider fallback."""
+
+    @pytest.fixture
+    def service(self):
+        with patch("app.services.litellm_service.settings") as mock_settings:
+            mock_settings.LITELLM_BASE_URL = "http://localhost:4000"
+            mock_settings.LITELLM_API_KEY = "test-key"
+            return LiteLLMService()
+
+    @pytest.mark.asyncio
+    async def test_uses_direct_gemini_for_non_square(self, service):
+        with patch("app.services.litellm_service.settings") as mock_settings:
+            mock_settings.DIRECT_PROVIDER_FALLBACK = True
+            mock_settings.gemini_available = True
+
+            mock_gemini_service = MagicMock()
+            mock_gemini_service.generate_image = AsyncMock(return_value=["http://example.com/image.png"])
+
+            with patch("app.services.gemini_service.get_gemini_service", return_value=mock_gemini_service):
+                result = await service.generate_image(
+                    prompt="test prompt",
+                    model="gemini/gemini-2.0-flash-exp",
+                    aspect_ratio="16:9",
+                    quality="standard",
+                    n=1,
+                )
+
+                assert result == ["http://example.com/image.png"]
+                mock_gemini_service.generate_image.assert_called_once_with(
+                    prompt="test prompt",
+                    model="gemini-2.0-flash-exp",  # Prefix removed
+                    aspect_ratio="16:9",
+                    quality="standard",
+                    n=1,
+                )
+
+    @pytest.mark.asyncio
+    async def test_uses_litellm_for_square_aspect_ratio(self, service, mock_openai_response):
+        with patch("app.services.litellm_service.settings") as mock_settings:
+            mock_settings.DIRECT_PROVIDER_FALLBACK = True
+            mock_settings.gemini_available = True
+
+            with patch.object(service.client.images, "generate", return_value=mock_openai_response):
+                with patch("app.services.litellm_service.model_registry") as mock_registry:
+                    mock_registry.get_model.return_value = None
+
+                    with patch("app.services.litellm_service.storage_service") as mock_storage:
+                        mock_storage.save_image = AsyncMock(return_value="http://example.com/image.png")
+
+                        result = await service.generate_image(
+                            prompt="test prompt",
+                            model="gemini/gemini-2.0-flash-exp",
+                            aspect_ratio="1:1",
+                            quality="standard",
+                            n=1,
+                        )
+
+                        assert result == ["http://example.com/image.png"]
+                        service.client.images.generate.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_uses_litellm_when_fallback_disabled(self, service, mock_openai_response):
+        with patch("app.services.litellm_service.settings") as mock_settings:
+            mock_settings.DIRECT_PROVIDER_FALLBACK = False
+
+            with patch.object(service.client.images, "generate", return_value=mock_openai_response):
+                with patch("app.services.litellm_service.model_registry") as mock_registry:
+                    mock_registry.get_model.return_value = None
+
+                    with patch("app.services.litellm_service.storage_service") as mock_storage:
+                        mock_storage.save_image = AsyncMock(return_value="http://example.com/image.png")
+
+                        result = await service.generate_image(
+                            prompt="test prompt",
+                            model="gemini/gemini-2.0-flash-exp",
+                            aspect_ratio="16:9",
+                            quality="standard",
+                            n=1,
+                        )
+
+                        assert result == ["http://example.com/image.png"]
+                        service.client.images.generate.assert_called_once()