diff --git a/python/samples/getting_started/agents/azure_ai/azure_ai_with_image_generation.py b/python/samples/getting_started/agents/azure_ai/azure_ai_with_image_generation.py
index 849e44761a..2fcc0d09c1 100644
--- a/python/samples/getting_started/agents/azure_ai/azure_ai_with_image_generation.py
+++ b/python/samples/getting_started/agents/azure_ai/azure_ai_with_image_generation.py
@@ -1,16 +1,14 @@
 # Copyright (c) Microsoft. All rights reserved.
-
 import asyncio
 from pathlib import Path
 
 import aiofiles
 from agent_framework import DataContent
 from agent_framework.azure import AzureAIClient
-from azure.ai.projects.models import ImageGenTool
 from azure.identity.aio import AzureCliCredential
 
 """
-Azure AI Agent With Image Generation
+Azure AI Agent with Image Generation Example
 
 This sample demonstrates basic usage of AzureAIClient to create an agent
 that can generate images based on user requirements.
@@ -30,7 +28,14 @@ async def main() -> None:
         AzureAIClient(credential=credential).create_agent(
             name="ImageGenAgent",
             instructions="Generate images based on user requirements.",
-            tools=[ImageGenTool(quality="low", size="1024x1024")],
+            tools=[
+                {
+                    "type": "image_generation",
+                    "model": "gpt-image-1-mini",
+                    "quality": "low",
+                    "size": "1024x1024",
+                }
+            ],
         ) as agent,
     ):
         query = "Generate an image of Microsoft logo."
@@ -39,7 +44,7 @@ async def main() -> None:
             query,
             # These additional options are required for image generation
             additional_chat_options={
-                "extra_headers": {"x-ms-oai-image-generation-deployment": "gpt-image-1"},
+                "extra_headers": {"x-ms-oai-image-generation-deployment": "gpt-image-1-mini"},
             },
         )
         print(f"Agent: {result}\n")