diff --git a/components/chat.tsx b/components/chat.tsx
index a8d188f..c43d42a 100644
--- a/components/chat.tsx
+++ b/components/chat.tsx
@@ -7,9 +7,8 @@ import { EmptyScreen } from '@/components/empty-screen'
 import { useLocalStorage } from '@/lib/hooks/use-local-storage'
 import { useEffect, useState } from 'react'
 import { useUIState, useAIState } from 'ai/rsc'
-import { Session } from '@/lib/types'
+import { Message, Session } from '@/lib/types'
 import { usePathname, useRouter } from 'next/navigation'
-import { Message } from '@/lib/chat/actions'
 import { useScrollAnchor } from '@/lib/hooks/use-scroll-anchor'
 import { toast } from 'sonner'
 
@@ -71,7 +70,7 @@ export function Chat({ id, className, session, missingKeys }: ChatProps) {
         ) : (
           <EmptyScreen />
         )}
-        <div className="h-px w-full" ref={visibilityRef} />
+        <div className="w-full h-px" ref={visibilityRef} />
       </div>
       <ChatPanel
         id={id}
diff --git a/components/stocks/message.tsx b/components/stocks/message.tsx
index 5c7e7ca..8a75245 100644
--- a/components/stocks/message.tsx
+++ b/components/stocks/message.tsx
@@ -7,7 +7,7 @@ import { CodeBlock } from '../ui/codeblock'
 import { MemoizedReactMarkdown } from '../markdown'
 import remarkGfm from 'remark-gfm'
 import remarkMath from 'remark-math'
-import { StreamableValue } from 'ai/rsc'
+import { StreamableValue, useStreamableValue } from 'ai/rsc'
 import { useStreamableText } from '@/lib/hooks/use-streamable-text'
 
 // Different types of message bubbles.
diff --git a/components/stocks/stock-purchase.tsx b/components/stocks/stock-purchase.tsx
index 8acb8e3..0bf922e 100644
--- a/components/stocks/stock-purchase.tsx
+++ b/components/stocks/stock-purchase.tsx
@@ -14,7 +14,7 @@ interface Purchase {
 }
 
 export function Purchase({
-  props: { numberOfShares, symbol, price, status = 'requires_action' }
+  props: { numberOfShares, symbol, price, status = 'expired' }
 }: {
   props: Purchase
 }) {
@@ -60,8 +60,8 @@ export function Purchase({
   }
 
   return (
-    <div className="rounded-xl border bg-zinc-950 p-4 text-green-400">
-      <div className="float-right inline-block rounded-full bg-white/10 px-2 py-1 text-xs">
+    <div className="p-4 text-green-400 border rounded-xl bg-zinc-950">
+      <div className="inline-block float-right px-2 py-1 text-xs rounded-full bg-white/10">
         +1.23% ↑
       </div>
       <div className="text-lg text-zinc-300">{symbol}</div>
@@ -70,7 +70,7 @@ export function Purchase({
         <div className="mt-4 text-zinc-200">{purchasingUI}</div>
       ) : status === 'requires_action' ? (
         <>
-          <div className="relative mt-6 pb-6">
+          <div className="relative pb-6 mt-6">
             <p>Shares to purchase</p>
             <input
               id="labels-range-input"
@@ -79,18 +79,18 @@ export function Purchase({
               onChange={onSliderChange}
               min="10"
               max="1000"
-              className="h-1 w-full cursor-pointer appearance-none rounded-lg bg-zinc-600 accent-green-500 dark:bg-zinc-700"
+              className="w-full h-1 rounded-lg appearance-none cursor-pointer bg-zinc-600 accent-green-500 dark:bg-zinc-700"
             />
-            <span className="absolute bottom-1 start-0 text-xs text-zinc-400">
+            <span className="absolute text-xs bottom-1 start-0 text-zinc-400">
               10
             </span>
-            <span className="absolute bottom-1 start-1/3 -translate-x-1/2 text-xs text-zinc-400 rtl:translate-x-1/2">
+            <span className="absolute text-xs -translate-x-1/2 bottom-1 start-1/3 text-zinc-400 rtl:translate-x-1/2">
               100
             </span>
-            <span className="absolute bottom-1 start-2/3 -translate-x-1/2 text-xs text-zinc-400 rtl:translate-x-1/2">
+            <span className="absolute text-xs -translate-x-1/2 bottom-1 start-2/3 text-zinc-400 rtl:translate-x-1/2">
               500
             </span>
-            <span className="absolute bottom-1 end-0 text-xs text-zinc-400">
+            <span className="absolute text-xs bottom-1 end-0 text-zinc-400">
               1000
             </span>
           </div>
@@ -98,27 +98,27 @@ export function Purchase({
           <div className="mt-6">
             <p>Total cost</p>
             <div className="flex flex-wrap items-center text-xl font-bold sm:items-end sm:gap-2 sm:text-3xl">
-              <div className="flex basis-1/3 flex-col tabular-nums sm:basis-auto sm:flex-row sm:items-center sm:gap-2">
+              <div className="flex flex-col basis-1/3 tabular-nums sm:basis-auto sm:flex-row sm:items-center sm:gap-2">
                 {value}
                 <span className="mb-1 text-sm font-normal text-zinc-600 sm:mb-0 dark:text-zinc-400">
                   shares
                 </span>
               </div>
-              <div className="basis-1/3 text-center sm:basis-auto">×</div>
-              <span className="flex basis-1/3 flex-col tabular-nums sm:basis-auto sm:flex-row sm:items-center sm:gap-2">
+              <div className="text-center basis-1/3 sm:basis-auto">×</div>
+              <span className="flex flex-col basis-1/3 tabular-nums sm:basis-auto sm:flex-row sm:items-center sm:gap-2">
                 ${price}
                 <span className="mb-1 ml-1 text-sm font-normal text-zinc-600 sm:mb-0 dark:text-zinc-400">
                   per share
                 </span>
               </span>
-              <div className="mt-2 basis-full border-t border-t-zinc-700 pt-2 text-center sm:mt-0 sm:basis-auto sm:border-0 sm:pt-0 sm:text-left">
+              <div className="pt-2 mt-2 text-center border-t basis-full border-t-zinc-700 sm:mt-0 sm:basis-auto sm:border-0 sm:pt-0 sm:text-left">
                 = <span>{formatNumber(value * price)}</span>
               </div>
             </div>
           </div>
 
           <button
-            className="mt-6 w-full rounded-lg bg-green-400 px-4 py-2 font-bold text-zinc-900 hover:bg-green-500"
+            className="w-full px-4 py-2 mt-6 font-bold bg-green-400 rounded-lg text-zinc-900 hover:bg-green-500"
             onClick={async () => {
               const response = await confirmPurchase(symbol, price, value)
               setPurchasingUI(response.purchasingUI)
diff --git a/lib/chat/actions.tsx b/lib/chat/actions.tsx
index 6c89e31..330b2d6 100644
--- a/lib/chat/actions.tsx
+++ b/lib/chat/actions.tsx
@@ -33,7 +33,7 @@ import {
 } from '@/lib/utils'
 import { saveChat } from '@/app/actions'
 import { SpinnerMessage, UserMessage } from '@/components/stocks/message'
-import { Chat } from '@/lib/types'
+import { Chat, Message } from '@/lib/types'
 import { auth } from '@/auth'
 
 async function confirmPurchase(symbol: string, price: number, amount: number) {
@@ -85,18 +85,7 @@ async function confirmPurchase(symbol: string, price: number, amount: number) {
     aiState.done({
       ...aiState.get(),
       messages: [
-        ...aiState.get().messages.slice(0, -1),
-        {
-          id: nanoid(),
-          role: 'function',
-          name: 'showStockPurchase',
-          content: JSON.stringify({
-            symbol,
-            price,
-            defaultAmount: amount,
-            status: 'completed'
-          })
-        },
+        ...aiState.get().messages,
         {
           id: nanoid(),
           role: 'system',
@@ -208,15 +197,35 @@ async function submitUserMessage(content: string) {
 
           await sleep(1000)
 
+          const toolCallId = nanoid()
+
           aiState.done({
             ...aiState.get(),
             messages: [
               ...aiState.get().messages,
               {
                 id: nanoid(),
-                role: 'function',
-                name: 'listStocks',
-                content: JSON.stringify(stocks)
+                role: 'assistant',
+                content: [
+                  {
+                    type: 'tool-call',
+                    toolName: 'listStocks',
+                    toolCallId,
+                    args: { stocks }
+                  }
+                ]
+              },
+              {
+                id: nanoid(),
+                role: 'tool',
+                content: [
+                  {
+                    type: 'tool-result',
+                    toolName: 'listStocks',
+                    toolCallId,
+                    result: stocks
+                  }
+                ]
               }
             ]
           })
@@ -249,15 +258,35 @@ async function submitUserMessage(content: string) {
 
           await sleep(1000)
 
+          const toolCallId = nanoid()
+
           aiState.done({
             ...aiState.get(),
             messages: [
               ...aiState.get().messages,
               {
                 id: nanoid(),
-                role: 'function',
-                name: 'showStockPrice',
-                content: JSON.stringify({ symbol, price, delta })
+                role: 'assistant',
+                content: [
+                  {
+                    type: 'tool-call',
+                    toolName: 'showStockPrice',
+                    toolCallId,
+                    args: { symbol, price, delta }
+                  }
+                ]
+              },
+              {
+                id: nanoid(),
+                role: 'tool',
+                content: [
+                  {
+                    type: 'tool-result',
+                    toolName: 'showStockPrice',
+                    toolCallId,
+                    result: { symbol, price, delta }
+                  }
+                ]
               }
             ]
           })
@@ -286,11 +315,42 @@ async function submitUserMessage(content: string) {
             )
         }),
         generate: async function* ({ symbol, price, numberOfShares = 100 }) {
+          const toolCallId = nanoid()
+
           if (numberOfShares <= 0 || numberOfShares > 1000) {
             aiState.done({
               ...aiState.get(),
               messages: [
                 ...aiState.get().messages,
+                {
+                  id: nanoid(),
+                  role: 'assistant',
+                  content: [
+                    {
+                      type: 'tool-call',
+                      toolName: 'showStockPurchase',
+                      toolCallId,
+                      args: { symbol, price, numberOfShares }
+                    }
+                  ]
+                },
+                {
+                  id: nanoid(),
+                  role: 'tool',
+                  content: [
+                    {
+                      type: 'tool-result',
+                      toolName: 'showStockPurchase',
+                      toolCallId,
+                      result: {
+                        symbol,
+                        price,
+                        numberOfShares,
+                        status: 'expired'
+                      }
+                    }
+                  ]
+                },
                 {
                   id: nanoid(),
                   role: 'system',
@@ -300,37 +360,55 @@ async function submitUserMessage(content: string) {
             })
 
             return <BotMessage content={'Invalid amount'} />
+          } else {
+            aiState.done({
+              ...aiState.get(),
+              messages: [
+                ...aiState.get().messages,
+                {
+                  id: nanoid(),
+                  role: 'assistant',
+                  content: [
+                    {
+                      type: 'tool-call',
+                      toolName: 'showStockPurchase',
+                      toolCallId,
+                      args: { symbol, price, numberOfShares }
+                    }
+                  ]
+                },
+                {
+                  id: nanoid(),
+                  role: 'tool',
+                  content: [
+                    {
+                      type: 'tool-result',
+                      toolName: 'showStockPurchase',
+                      toolCallId,
+                      result: {
+                        symbol,
+                        price,
+                        numberOfShares
+                      }
+                    }
+                  ]
+                }
+              ]
+            })
+
+            return (
+              <BotCard>
+                <Purchase
+                  props={{
+                    numberOfShares,
+                    symbol,
+                    price: +price,
+                    status: 'requires_action'
+                  }}
+                />
+              </BotCard>
+            )
           }
-
-          aiState.done({
-            ...aiState.get(),
-            messages: [
-              ...aiState.get().messages,
-              {
-                id: nanoid(),
-                role: 'function',
-                name: 'showStockPurchase',
-                content: JSON.stringify({
-                  symbol,
-                  price,
-                  numberOfShares
-                })
-              }
-            ]
-          })
-
-          return (
-            <BotCard>
-              <Purchase
-                props={{
-                  numberOfShares,
-                  symbol,
-                  price: +price,
-                  status: 'requires_action'
-                }}
-              />
-            </BotCard>
-          )
         }
       },
       getEvents: {
@@ -356,15 +434,35 @@ async function submitUserMessage(content: string) {
 
           await sleep(1000)
 
+          const toolCallId = nanoid()
+
           aiState.done({
             ...aiState.get(),
             messages: [
               ...aiState.get().messages,
               {
                 id: nanoid(),
-                role: 'function',
-                name: 'getEvents',
-                content: JSON.stringify(events)
+                role: 'assistant',
+                content: [
+                  {
+                    type: 'tool-call',
+                    toolName: 'getEvents',
+                    toolCallId,
+                    args: { events }
+                  }
+                ]
+              },
+              {
+                id: nanoid(),
+                role: 'tool',
+                content: [
+                  {
+                    type: 'tool-result',
+                    toolName: 'getEvents',
+                    toolCallId,
+                    result: events
+                  }
+                ]
               }
             ]
           })
@@ -385,13 +483,6 @@ async function submitUserMessage(content: string) {
   }
 }
 
-export type Message = {
-  role: 'user' | 'assistant' | 'system' | 'function' | 'data' | 'tool'
-  content: string
-  id: string
-  name?: string
-}
-
 export type AIState = {
   chatId: string
   messages: Message[]
@@ -425,7 +516,7 @@ export const AI = createAI<AIState, UIState>({
       return
     }
   },
-  onSetAIState: async ({ state, done }) => {
+  onSetAIState: async ({ state }) => {
     'use server'
 
     const session = await auth()
@@ -436,7 +527,9 @@ export const AI = createAI<AIState, UIState>({
       const createdAt = new Date()
       const userId = session.user.id as string
       const path = `/chat/${chatId}`
-      const title = messages[0].content.substring(0, 100)
+
+      const firstMessageContent = messages[0].content as string
+      const title = firstMessageContent.substring(0, 100)
 
       const chat: Chat = {
         id: chatId,
@@ -460,28 +553,36 @@ export const getUIStateFromAIState = (aiState: Chat) => {
     .map((message, index) => ({
       id: `${aiState.chatId}-${index}`,
       display:
-        message.role === 'function' ? (
-          message.name === 'listStocks' ? (
-            <BotCard>
-              <Stocks props={JSON.parse(message.content)} />
-            </BotCard>
-          ) : message.name === 'showStockPrice' ? (
-            <BotCard>
-              <Stock props={JSON.parse(message.content)} />
-            </BotCard>
-          ) : message.name === 'showStockPurchase' ? (
-            <BotCard>
-              <Purchase props={JSON.parse(message.content)} />
-            </BotCard>
-          ) : message.name === 'getEvents' ? (
-            <BotCard>
-              <Events props={JSON.parse(message.content)} />
-            </BotCard>
-          ) : null
+        message.role === 'tool' ? (
+          message.content.map(tool => {
+            return tool.toolName === 'listStocks' ? (
+              <BotCard>
+                {/* TODO: Infer types based on the tool result*/}
+                {/* @ts-expect-error */}
+                <Stocks props={tool.result} />
+              </BotCard>
+            ) : tool.toolName === 'showStockPrice' ? (
+              <BotCard>
+                {/* @ts-expect-error */}
+                <Stock props={tool.result} />
+              </BotCard>
+            ) : tool.toolName === 'showStockPurchase' ? (
+              <BotCard>
+                {/* @ts-expect-error */}
+                <Purchase props={tool.result} />
+              </BotCard>
+            ) : tool.toolName === 'getEvents' ? (
+              <BotCard>
+                {/* @ts-expect-error */}
+                <Events props={tool.result} />
+              </BotCard>
+            ) : null
+          })
         ) : message.role === 'user' ? (
-          <UserMessage>{message.content}</UserMessage>
-        ) : (
+          <UserMessage>{message.content as string}</UserMessage>
+        ) : message.role === 'assistant' &&
+          typeof message.content === 'string' ? (
           <BotMessage content={message.content} />
-        )
+        ) : null
     }))
 }
diff --git a/lib/types.ts b/lib/types.ts
index b65c941..d892a0c 100644
--- a/lib/types.ts
+++ b/lib/types.ts
@@ -1,4 +1,8 @@
-import { Message } from 'ai'
+import { CoreMessage } from 'ai'
+
+export type Message = CoreMessage & {
+  id: string
+}
 
 export interface Chat extends Record<string, any> {
   id: string
diff --git a/package.json b/package.json
index ce25b7a..15d2754 100644
--- a/package.json
+++ b/package.json
@@ -27,7 +27,7 @@
     "@vercel/analytics": "^1.1.2",
     "@vercel/kv": "^1.0.1",
     "@vercel/og": "^0.6.2",
-    "ai": "^3.1.1",
+    "ai": "^3.1.5",
     "class-variance-authority": "^0.7.0",
     "clsx": "^2.1.0",
     "d3-scale": "^4.0.2",
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
index 1d7b590..c0d545f 100644
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@@ -48,8 +48,8 @@ dependencies:
     specifier: ^0.6.2
     version: 0.6.2
   ai:
-    specifier: ^3.1.1
-    version: 3.1.1(react@18.2.0)(solid-js@1.8.11)(svelte@4.2.9)(vue@3.4.15)(zod@3.22.4)
+    specifier: ^3.1.5
+    version: 3.1.5(openai@4.24.7)(react@18.2.0)(solid-js@1.8.11)(svelte@4.2.9)(vue@3.4.15)(zod@3.22.4)
   class-variance-authority:
     specifier: ^0.7.0
     version: 0.7.0
@@ -216,6 +216,22 @@ packages:
       zod: 3.22.4
     dev: false
 
+  /@ai-sdk/provider-utils@0.0.6(zod@3.22.4):
+    resolution: {integrity: sha512-SxOZgSxnaVlW04/SjfMoAD45kWOWTWx0QcZrHaQnePooLhyk5AqQpgauPijL803uoJPCKfzd0UBv1gSKvWiU0A==}
+    engines: {node: '>=18'}
+    peerDependencies:
+      zod: ^3.0.0
+    peerDependenciesMeta:
+      zod:
+        optional: true
+    dependencies:
+      '@ai-sdk/provider': 0.0.3
+      eventsource-parser: 1.1.2
+      nanoid: 3.3.6
+      secure-json-parse: 2.7.0
+      zod: 3.22.4
+    dev: false
+
   /@ai-sdk/provider@0.0.3:
     resolution: {integrity: sha512-0B8P6VZpJ6F9yS9BpmJBYSqIaIfeRtL5tD5SP+qgR8y0pPwalIbRMUFiLz9YUT6g70MJsCLpm/2/fX3cfAYCJw==}
     engines: {node: '>=18'}
@@ -1593,16 +1609,19 @@ packages:
       humanize-ms: 1.2.1
     dev: false
 
-  /ai@3.1.1(react@18.2.0)(solid-js@1.8.11)(svelte@4.2.9)(vue@3.4.15)(zod@3.22.4):
-    resolution: {integrity: sha512-pJZc6q7SLd2/NenxN62iagMw9HHQ4Q8FyKqkrZUJntupRTHHgN3fL7exzJU/ICHDAbtn/EcJXOau6P61QgUtKQ==}
+  /ai@3.1.5(openai@4.24.7)(react@18.2.0)(solid-js@1.8.11)(svelte@4.2.9)(vue@3.4.15)(zod@3.22.4):
+    resolution: {integrity: sha512-uE8EfcigIoep4NCqIOWsNUP/alyZ93XtsDsLjiop51y4lEtgJg0GzpbBRfYKaEK1O74HiQ2CmaeOOBlNVBBiIQ==}
     engines: {node: '>=18'}
     peerDependencies:
+      openai: ^4.42.0
       react: ^18.2.0
       solid-js: ^1.7.7
       svelte: ^3.0.0 || ^4.0.0
       vue: ^3.3.4
       zod: ^3.0.0
     peerDependenciesMeta:
+      openai:
+        optional: true
       react:
         optional: true
       solid-js:
@@ -1615,11 +1634,12 @@ packages:
         optional: true
     dependencies:
       '@ai-sdk/provider': 0.0.3
-      '@ai-sdk/provider-utils': 0.0.5(zod@3.22.4)
+      '@ai-sdk/provider-utils': 0.0.6(zod@3.22.4)
       eventsource-parser: 1.1.2
       json-schema: 0.4.0
       jsondiffpatch: 0.6.0
       nanoid: 3.3.6
+      openai: 4.24.7
       react: 18.2.0
       secure-json-parse: 2.7.0
       solid-js: 1.8.11