feat: Refactor GPT-Image client to support multiple models and enhance deployment tracking (#41)

Author: alisoliman

backend/core/__init__.py

@@ -25,16 +25,17 @@
     logger.error(f"Failed to initialize Sora 2 client: {str(e)}")
     sora_client = None
 
-# Initialize GPT-Image-1 client
+# Initialize GPT-Image client (using default model)
 try:
-    # Using OpenAI API directly for GPT-Image-1
+    # Using OpenAI API directly for GPT-Image
     dalle_client = GPTImageClient(
         api_key=settings.OPENAI_API_KEY,
-        organization_id=settings.OPENAI_ORG_ID if settings.OPENAI_ORG_ID else None
+        organization_id=settings.OPENAI_ORG_ID if settings.OPENAI_ORG_ID else None,
+        model=settings.DEFAULT_IMAGE_MODEL
     )
-    logger.info("Initialized GPT-Image-1 client using OpenAI API.")
+    logger.info(f"Initialized GPT-Image client using OpenAI API with default model: {settings.DEFAULT_IMAGE_MODEL}")
 except Exception as e:
-    logger.error(f"Failed to initialize GPT-Image-1 client: {str(e)}")
+    logger.error(f"Failed to initialize GPT-Image client: {str(e)}")
     dalle_client = None
 
 # Initialize LLM client (sync)

backend/core/config.py

@@ -25,10 +25,16 @@ class Settings(BaseSettings):
     # Azure OpenAI for Image Generation
     # The Azure OpenAI resource name for image generation
     IMAGEGEN_AOAI_RESOURCE: Optional[str] = None
-    # The image generation deployment name
+    # The image generation deployment name (gpt-image-1)
     IMAGEGEN_DEPLOYMENT: Optional[str] = None
+    # The gpt-image-1.5 deployment name
+    IMAGEGEN_15_DEPLOYMENT: Optional[str] = None
+    # The gpt-image-1-mini deployment name
+    IMAGEGEN_1_MINI_DEPLOYMENT: Optional[str] = None
     # The Azure OpenAI API key for image generation
     IMAGEGEN_AOAI_API_KEY: Optional[str] = None
+    # Default image generation model
+    DEFAULT_IMAGE_MODEL: str = "gpt-image-1"
 
     # OpenAI API for Image Generation with GPT-Image-1
     OPENAI_API_KEY: Optional[str] = None

backend/core/gpt_image.py

@@ -23,16 +23,19 @@ class GPTImageClient:
     Client for GPT-Image-1 generation and editing using the official OpenAI or Azure OpenAI Python client.
     """
 
-    def __init__(self, api_key: Optional[str] = None, organization_id: Optional[str] = None, provider: Optional[str] = None):
+    def __init__(self, api_key: Optional[str] = None, organization_id: Optional[str] = None, provider: Optional[str] = None, deployment_name: Optional[str] = None, model: Optional[str] = None):
         """
         Initialize the GPT Image client with either OpenAI or Azure OpenAI client
 
         Args:
             api_key: The API key to use (optional, will use from settings if not provided)
             organization_id: The organization ID for OpenAI (optional)
             provider: The provider to use ('openai' or 'azure', defaults to settings.MODEL_PROVIDER)
+            deployment_name: Specific deployment name to use (for Azure, optional)
+            model: Model name to use (gpt-image-1, gpt-image-1.5, gpt-image-1.5-mini)
         """
         provider = provider or settings.MODEL_PROVIDER
+        self.model = model or settings.DEFAULT_IMAGE_MODEL
 
         if provider.lower() == "azure":
             # Use Azure OpenAI
@@ -46,11 +49,11 @@ def __init__(self, api_key: Optional[str] = None, organization_id: Optional[str]
                 api_key=self.api_key,
                 api_version=settings.AOAI_API_VERSION
             )
-            # Set deployment name for later use
-            self.deployment_name = settings.IMAGEGEN_DEPLOYMENT
+            # Set deployment name: use provided or map from model
+            self.deployment_name = deployment_name or self._get_deployment_for_model(self.model)
             self.provider = "azure"
             logger.info(
-                "Initialized GPT-Image-1 client with Azure OpenAI Python SDK")
+                f"Initialized GPT Image client with Azure OpenAI Python SDK (model: {self.model}, deployment: {self.deployment_name})")
         else:
             # Use direct OpenAI
             self.api_key = api_key or settings.OPENAI_API_KEY
@@ -62,8 +65,32 @@ def __init__(self, api_key: Optional[str] = None, organization_id: Optional[str]
                 organization=organization_id or settings.OPENAI_ORG_ID
             )
             self.provider = "openai"
+            self.deployment_name = None  # Not used for OpenAI
             logger.info(
-                "Initialized GPT-Image-1 client with OpenAI Python SDK")
+                f"Initialized GPT Image client with OpenAI Python SDK (model: {self.model})")
+    
+    def _get_deployment_for_model(self, model: str) -> str:
+        """
+        Map model name to Azure deployment name
+        
+        Args:
+            model: Model identifier (gpt-image-1, gpt-image-1.5, gpt-image-1-mini)
+            
+        Returns:
+            Deployment name from settings
+        """
+        mapping = {
+            "gpt-image-1": settings.IMAGEGEN_DEPLOYMENT,
+            "gpt-image-1.5": settings.IMAGEGEN_15_DEPLOYMENT,
+            "gpt-image-1-mini": settings.IMAGEGEN_1_MINI_DEPLOYMENT,
+        }
+        deployment = mapping.get(model)
+        
+        if not deployment:
+            logger.warning(f"No deployment configured for model {model}, falling back to default")
+            deployment = settings.IMAGEGEN_DEPLOYMENT
+            
+        return deployment
 
     def generate_image(self, prompt: str, model: str = None, n: int = 1,
                        size: str = "auto", response_format: str = "b64_json",
@@ -168,6 +195,10 @@ def generate_image(self, prompt: str, model: str = None, n: int = 1,
 
                 formatted_response["data"].append(image_data)
 
+            # Add deployment metadata for tracking
+            formatted_response["_deployment_name"] = self.deployment_name
+            formatted_response["_model"] = self.model
+
             return formatted_response
 
         except Exception as e:
@@ -257,8 +288,11 @@ def edit_image(self, **kwargs):
                     url, headers=headers, files=files, data=data)
                 response.raise_for_status()
 
-                # Parse the response
-                return response.json()
+                # Parse the response and add deployment metadata
+                result = response.json()
+                result["_deployment_name"] = self.deployment_name
+                result["_model"] = self.model
+                return result
             else:
                 # Handle model parameter for OpenAI provider
                 if "model" not in kwargs:
@@ -277,7 +311,11 @@ def edit_image(self, **kwargs):
 
                 # Call the OpenAI API to edit the image
                 response = self.client.images.edit(**kwargs)
-                return response.model_dump()
+                result = response.model_dump()
+                # Add deployment metadata
+                result["_deployment_name"] = self.deployment_name
+                result["_model"] = self.model
+                return result
 
         except Exception as e:
             logger.error(f"Error editing image: {str(e)}")
@@ -455,7 +493,11 @@ async def process_edit_image_upload(self, prompt: str, model: str = None, n: int
                     def _sync_post():
                         resp = requests.post(url, headers=headers, files=files, data=data)
                         resp.raise_for_status()
-                        return resp.json()
+                        result = resp.json()
+                        # Add deployment metadata
+                        result["_deployment_name"] = self.deployment_name
+                        result["_model"] = self.model
+                        return result
 
                     return await asyncio.to_thread(_sync_post)
                 finally: