feat: added image gen support

surfsense_backend/app/config/__init__.py

@@ -81,6 +81,56 @@ def load_router_settings():
         return default_settings
 
 
+def load_global_image_gen_configs():
+    """
+    Load global image generation configurations from YAML file.
+
+    Returns:
+        list: List of global image generation config dictionaries, or empty list
+    """
+    global_config_file = BASE_DIR / "app" / "config" / "global_llm_config.yaml"
+
+    if not global_config_file.exists():
+        return []
+
+    try:
+        with open(global_config_file, encoding="utf-8") as f:
+            data = yaml.safe_load(f)
+            return data.get("global_image_generation_configs", [])
+    except Exception as e:
+        print(f"Warning: Failed to load global image generation configs: {e}")
+        return []
+
+
+def load_image_gen_router_settings():
+    """
+    Load router settings for image generation Auto mode from YAML file.
+
+    Returns:
+        dict: Router settings dictionary
+    """
+    default_settings = {
+        "routing_strategy": "usage-based-routing",
+        "num_retries": 3,
+        "allowed_fails": 3,
+        "cooldown_time": 60,
+    }
+
+    global_config_file = BASE_DIR / "app" / "config" / "global_llm_config.yaml"
+
+    if not global_config_file.exists():
+        return default_settings
+
+    try:
+        with open(global_config_file, encoding="utf-8") as f:
+            data = yaml.safe_load(f)
+            settings = data.get("image_generation_router_settings", {})
+            return {**default_settings, **settings}
+    except Exception as e:
+        print(f"Warning: Failed to load image generation router settings: {e}")
+        return default_settings
+
+
 def initialize_llm_router():
     """
     Initialize the LLM Router service for Auto mode.
@@ -105,6 +155,33 @@ def initialize_llm_router():
         print(f"Warning: Failed to initialize LLM Router: {e}")
 
 
+def initialize_image_gen_router():
+    """
+    Initialize the Image Generation Router service for Auto mode.
+    This should be called during application startup.
+    """
+    image_gen_configs = load_global_image_gen_configs()
+    router_settings = load_image_gen_router_settings()
+
+    if not image_gen_configs:
+        print(
+            "Info: No global image generation configs found, "
+            "Image Generation Auto mode will not be available"
+        )
+        return
+
+    try:
+        from app.services.image_gen_router_service import ImageGenRouterService
+
+        ImageGenRouterService.initialize(image_gen_configs, router_settings)
+        print(
+            f"Info: Image Generation Router initialized with {len(image_gen_configs)} models "
+            f"(strategy: {router_settings.get('routing_strategy', 'usage-based-routing')})"
+        )
+    except Exception as e:
+        print(f"Warning: Failed to initialize Image Generation Router: {e}")
+
+
 class Config:
     # Check if ffmpeg is installed
     if not is_ffmpeg_installed():
@@ -216,6 +293,12 @@ class Config:
     # Router settings for Auto mode (LiteLLM Router load balancing)
     ROUTER_SETTINGS = load_router_settings()
 
+    # Global Image Generation Configurations (optional)
+    GLOBAL_IMAGE_GEN_CONFIGS = load_global_image_gen_configs()
+
+    # Router settings for Image Generation Auto mode
+    IMAGE_GEN_ROUTER_SETTINGS = load_image_gen_router_settings()
+
     # Chonkie Configuration | Edit this to your needs
     EMBEDDING_MODEL = os.getenv("EMBEDDING_MODEL")
     # Azure OpenAI credentials from environment variables

surfsense_backend/app/config/global_llm_config.example.yaml

@@ -183,6 +183,73 @@ global_llm_configs:
     use_default_system_instructions: true
     citations_enabled: true
 
+# =============================================================================
+# Image Generation Configuration
+# =============================================================================
+# These configurations power the image generation feature using litellm.aimage_generation().
+# Supported providers: OpenAI, Azure, Google AI Studio, Vertex AI, AWS Bedrock,
+# Recraft, OpenRouter, Xinference, Nscale
+#
+# Auto mode (ID 0) uses LiteLLM Router for load balancing across all image gen configs.
+
+# Router Settings for Image Generation Auto Mode
+image_generation_router_settings:
+  routing_strategy: "usage-based-routing"
+  num_retries: 3
+  allowed_fails: 3
+  cooldown_time: 60
+
+global_image_generation_configs:
+  # Example: OpenAI DALL-E 3
+  - id: -1
+    name: "Global DALL-E 3"
+    description: "OpenAI's DALL-E 3 for high-quality image generation"
+    provider: "OPENAI"
+    model_name: "dall-e-3"
+    api_key: "sk-your-openai-api-key-here"
+    api_base: ""
+    rpm: 50
+    tpm: 100000
+    litellm_params: {}
+
+  # Example: OpenAI GPT Image 1
+  - id: -2
+    name: "Global GPT Image 1"
+    description: "OpenAI's GPT Image 1 model"
+    provider: "OPENAI"
+    model_name: "gpt-image-1"
+    api_key: "sk-your-openai-api-key-here"
+    api_base: ""
+    rpm: 50
+    tpm: 100000
+    litellm_params: {}
+
+  # Example: Azure OpenAI DALL-E 3
+  - id: -3
+    name: "Global Azure DALL-E 3"
+    description: "Azure-hosted DALL-E 3 deployment"
+    provider: "AZURE_OPENAI"
+    model_name: "azure/dall-e-3-deployment"
+    api_key: "your-azure-api-key-here"
+    api_base: "https://your-resource.openai.azure.com"
+    api_version: "2024-02-15-preview"
+    rpm: 50
+    tpm: 100000
+    litellm_params:
+      base_model: "dall-e-3"
+
+  # Example: OpenRouter Gemini Image Generation
+  # - id: -4
+  #   name: "Global Gemini Image Gen"
+  #   description: "Google Gemini image generation via OpenRouter"
+  #   provider: "OPENROUTER"
+  #   model_name: "google/gemini-2.5-flash-image"
+  #   api_key: "your-openrouter-api-key-here"
+  #   api_base: ""
+  #   rpm: 30
+  #   tpm: 50000
+  #   litellm_params: {}
+
 # Notes:
 # - ID 0 is reserved for "Auto" mode - uses LiteLLM Router for load balancing
 # - Use negative IDs to distinguish global configs from user configs (NewLLMConfig in DB)
@@ -202,3 +269,10 @@ global_llm_configs:
 # - model_name format: "azure/<your-deployment-name>"
 # - api_version: Use a recent Azure API version (e.g., "2024-02-15-preview")
 # - See: https://docs.litellm.ai/docs/proxy/cost_tracking#spend-tracking-for-azure-openai-models
+#
+# IMAGE GENERATION NOTES:
+# - Image generation configs use the same ID scheme as LLM configs (negative for global)
+# - Supported models: dall-e-2, dall-e-3, gpt-image-1 (OpenAI), azure/* (Azure), 
+#   bedrock/* (AWS), vertex_ai/* (Google), recraft/* (Recraft), openrouter/* (OpenRouter)
+# - The router uses litellm.aimage_generation() for async image generation
+# - api_version is required for Azure image generation deployments