diff --git a/packages/forge/blocks/difyAi/actions/createChatMessage.ts b/packages/forge/blocks/difyAi/actions/createChatMessage.ts
index d51f877de..7ffaea7fe 100644
--- a/packages/forge/blocks/difyAi/actions/createChatMessage.ts
+++ b/packages/forge/blocks/difyAi/actions/createChatMessage.ts
@@ -5,6 +5,7 @@ import { defaultBaseUrl } from '../constants'
 import { Chunk } from '../types'
 import ky, { HTTPError } from 'ky'
 import { deprecatedCreateChatMessageOptions } from '../deprecated'
+import { formatStreamPart } from 'ai'
 
 export const createChatMessage = createAction({
   auth,
@@ -120,7 +121,7 @@ export const createChatMessage = createAction({
                     onMessage: (message) => {
                       controller.enqueue(
                         new TextEncoder().encode(
-                          '0:"' + message.replace(/"/g, '\\"') + '"\n'
+                          formatStreamPart('text', message)
                         )
                       )
                     },
@@ -299,7 +300,7 @@ const processDifyStream = async (
       totalTokens,
       conversationId,
     }: {
-      totalTokens: number
+      totalTokens?: number
       conversationId: string
     }) => void
   }
@@ -333,7 +334,7 @@ const processDifyStream = async (
         }
         if (data.event === 'message_end') {
           callbacks.onMessageEnd?.({
-            totalTokens: data.metadata.usage.total_tokens,
+            totalTokens: data.metadata.usage?.total_tokens,
             conversationId: data.conversation_id,
           })
         }
diff --git a/packages/forge/blocks/difyAi/package.json b/packages/forge/blocks/difyAi/package.json
index 8635e2e48..69dd0ea3d 100644
--- a/packages/forge/blocks/difyAi/package.json
+++ b/packages/forge/blocks/difyAi/package.json
@@ -12,5 +12,8 @@
     "@types/react": "18.2.15",
     "ky": "1.2.3",
     "typescript": "5.4.5"
+  },
+  "dependencies": {
+    "ai": "3.1.12"
   }
 }
diff --git a/packages/forge/blocks/difyAi/types.ts b/packages/forge/blocks/difyAi/types.ts
index 27139c55d..a5465049b 100644
--- a/packages/forge/blocks/difyAi/types.ts
+++ b/packages/forge/blocks/difyAi/types.ts
@@ -11,7 +11,7 @@ type AgentMessageChunk = {
 type MessageEndChunk = {
   event: 'message_end'
   metadata: {
-    usage: {
+    usage?: {
       total_tokens: number
     }
   }
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
index c0409fb4e..8ef349532 100644
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@@ -1347,6 +1347,10 @@ importers:
         version: 5.4.5
 
   packages/forge/blocks/difyAi:
+    dependencies:
+      ai:
+        specifier: 3.1.12
+        version: 3.1.12(openai@4.47.1)(react@18.2.0)(solid-js@1.7.8)(svelte@4.2.12)(vue@3.4.21)(zod@3.22.4)
     devDependencies:
       '@typebot.io/forge':
         specifier: workspace:*