diff --git a/app/api/chat/route.ts b/app/api/chat/route.ts
index 54ca364..a0eacca 100644
--- a/app/api/chat/route.ts
+++ b/app/api/chat/route.ts
@@ -237,6 +237,7 @@ When asked a "What is" question, maintain the same format as the question and an
         parameters: z.object({
           title: z.string().optional().describe("The title of the code snippet."),
           code: z.string().describe("The Python code to execute."),
+          icon: z.enum(["stock", "date", "calculation", "default"]).describe("The icon to display for the code snippet."),
         }),
         execute: async ({ code }: { code: string }) => {
           const sandbox = await CodeInterpreter.create();
diff --git a/app/page.tsx b/app/page.tsx
index bdc1d88..a41951f 100644
--- a/app/page.tsx
+++ b/app/page.tsx
@@ -50,7 +50,10 @@ import {
   Terminal,
   Pause,
   Play,
-  RotateCw
+  RotateCw,
+  TrendingUpIcon,
+  Calendar,
+  Calculator
 } from 'lucide-react';
 import {
   HoverCard,
@@ -801,7 +804,10 @@ export default function Home() {
             <AccordionContent>
               <div className="w-full my-2 border border-gray-200 overflow-hidden rounded-md">
                 <div className="bg-gray-100 p-2 flex items-center">
-                  <Code className="h-5 w-5 text-gray-500 mr-2" />
+                  {args.icon === 'stock' && <TrendingUpIcon className="h-5 w-5 text-primary mr-2" />}
+                  {args.icon === 'default' && <Code className="h-5 w-5 text-primary mr-2" />}
+                  {args.icon === 'date' && <Calendar className="h-5 w-5 text-primary mr-2" />}
+                  {args.icon === 'calculation' && <Calculator className="h-5 w-5 text-primary mr-2" />}
                   <span className="text-sm font-medium">{args.title}</span>
                 </div>
                 <Tabs defaultValue="code" className="w-full">