diff --git a/advanced_ai_agents/single_agent_apps/ai_streaming_chatbot/docs/images/streaming-ai-chatbot.gif b/advanced_ai_agents/single_agent_apps/ai_streaming_chatbot/docs/images/streaming-ai-chatbot.gif
index 28059dc..3b01147 100644
Binary files a/advanced_ai_agents/single_agent_apps/ai_streaming_chatbot/docs/images/streaming-ai-chatbot.gif and b/advanced_ai_agents/single_agent_apps/ai_streaming_chatbot/docs/images/streaming-ai-chatbot.gif differ
diff --git a/advanced_ai_agents/single_agent_apps/ai_streaming_chatbot/package.json b/advanced_ai_agents/single_agent_apps/ai_streaming_chatbot/package.json
index c906abd..a7d47d8 100644
--- a/advanced_ai_agents/single_agent_apps/ai_streaming_chatbot/package.json
+++ b/advanced_ai_agents/single_agent_apps/ai_streaming_chatbot/package.json
@@ -16,7 +16,7 @@
     "openai"
   ],
   "dependencies": {
-    "motia": "^0.2.1-beta.73",
+    "motia": "0.2.2",
     "openai": "^4.102.0",
     "zod": "^3.25.20"
   },
diff --git a/advanced_ai_agents/single_agent_apps/ai_streaming_chatbot/steps/ai-response.step.ts b/advanced_ai_agents/single_agent_apps/ai_streaming_chatbot/steps/ai-response.step.ts
index e6816bb..d8b7686 100644
--- a/advanced_ai_agents/single_agent_apps/ai_streaming_chatbot/steps/ai-response.step.ts
+++ b/advanced_ai_agents/single_agent_apps/ai_streaming_chatbot/steps/ai-response.step.ts
@@ -1,6 +1,7 @@
 import { EventConfig, Handlers } from 'motia'
 import { OpenAI } from 'openai'
 import { z } from 'zod'
+// import { AzureOpenAI } from 'openai'
 
 export const config: EventConfig = {
   type: 'event',
@@ -22,6 +23,14 @@ export const handler: Handlers['AiResponse'] = async (input, context) => {
 
   logger.info('Generating AI response', { conversationId })
 
+  // For Azure OpenAI
+  // const openai = new AzureOpenAI({
+  //   endpoint: process.env.AZURE_OPENAI_ENDPOINT || 'demo-key',
+  //   apiKey: process.env.AZURE_OPENAI_API_KEY || 'demo-key',
+  //   deployment: 'gpt-4o-mini',
+  //   apiVersion: '2024-12-01-preview'
+  // })
+
   const openai = new OpenAI({ 
     apiKey: process.env.OPENAI_API_KEY,
     baseURL: process.env.OPENAI_BASE_URL || 'https://api.openai.com/v1'
diff --git a/advanced_ai_agents/single_agent_apps/ai_streaming_chatbot/steps/chat-api.step.ts b/advanced_ai_agents/single_agent_apps/ai_streaming_chatbot/steps/chat-api.step.ts
index b764aa5..00257fd 100644
--- a/advanced_ai_agents/single_agent_apps/ai_streaming_chatbot/steps/chat-api.step.ts
+++ b/advanced_ai_agents/single_agent_apps/ai_streaming_chatbot/steps/chat-api.step.ts
@@ -1,5 +1,6 @@
 import { ApiRouteConfig, Handlers } from 'motia'
 import { z } from 'zod'
+import { conversationSchema } from './conversation.stream'
 
 const inputSchema = z.object({
   message: z.string().min(1, 'Message is required'),
@@ -15,11 +16,7 @@ export const config: ApiRouteConfig = {
   emits: ['chat-message'],
   bodySchema: inputSchema,
   responseSchema: {
-    200: z.object({
-      conversationId: z.string(),
-      message: z.string(),
-      status: z.enum(['created', 'streaming', 'completed']).optional(),
-    })
+    200: conversationSchema
   },
   flows: ['chat'],
 }
@@ -41,7 +38,7 @@ export const handler: Handlers['ChatApi'] = async (req, { logger, emit, streams
     timestamp: new Date().toISOString(),
   })
 
-  await streams.conversation.set(conversationId, assistantMessageId, {
+  const aiResponse = await streams.conversation.set(conversationId, assistantMessageId, {
     message: '',
     from: 'assistant',
     status: 'created',
@@ -57,46 +54,13 @@ export const handler: Handlers['ChatApi'] = async (req, { logger, emit, streams
     },
   })
 
-  const maxWaitTime = 3000 // 3 seconds
-  const startTime = Date.now()
-  let aiResponse = null
-
-  // Initial delay to allow AI to start processing
-  await new Promise(resolve => setTimeout(resolve, 100))
-
-  while (Date.now() - startTime < maxWaitTime) {
-    const response = await streams.conversation.get(conversationId, assistantMessageId)
-    if (response) {
-      aiResponse = response
-      // Break only if we have a completed status
-      if (response.status === 'completed') {
-        break
-      }
-    }
-    await new Promise(resolve => setTimeout(resolve, 100))
-  }
-
-  // Get final response state
-  if (!aiResponse?.status || aiResponse.status !== 'completed') {
-    const finalResponse = await streams.conversation.get(conversationId, assistantMessageId)
-    if (finalResponse?.status === 'completed') {
-      aiResponse = finalResponse
-    }
-  }
-
   logger.info('Returning chat response', { 
     conversationId,
     messageId: assistantMessageId,
-    status: aiResponse?.status,
-    hasMessage: !!aiResponse?.message
   })
 
   return {
     status: 200,
-    body: {
-      conversationId,
-      message: aiResponse?.status === 'completed' ? aiResponse.message : 'Message received, AI is responding...',
-      status: aiResponse?.status || 'streaming',
-    },
+    body: aiResponse,
   }
 }
diff --git a/advanced_ai_agents/single_agent_apps/ai_streaming_chatbot/steps/conversation.stream.ts b/advanced_ai_agents/single_agent_apps/ai_streaming_chatbot/steps/conversation.stream.ts
index c0e8755..632d263 100644
--- a/advanced_ai_agents/single_agent_apps/ai_streaming_chatbot/steps/conversation.stream.ts
+++ b/advanced_ai_agents/single_agent_apps/ai_streaming_chatbot/steps/conversation.stream.ts
@@ -1,13 +1,15 @@
-import { StateStreamConfig } from 'motia'
+import { StreamConfig } from 'motia'
 import { z } from 'zod'
 
-export const config: StateStreamConfig = {
+export const conversationSchema = z.object({
+  message: z.string(),
+  from: z.enum(['user', 'assistant']),
+  status: z.enum(['created', 'streaming', 'completed']),
+  timestamp: z.string(),
+})
+
+export const config: StreamConfig = {
   name: 'conversation',
-  schema: z.object({
-    message: z.string(),
-    from: z.enum(['user', 'assistant']),
-    status: z.enum(['created', 'streaming', 'completed']),
-    timestamp: z.string(),
-  }),
-  baseConfig: { storageType: 'state' },
+  schema: conversationSchema,
+  baseConfig: { storageType: 'default' },
 }