diff --git a/app/api/chat/route.ts b/app/api/chat/route.ts
index 0856577..59ff037 100644
--- a/app/api/chat/route.ts
+++ b/app/api/chat/route.ts
@@ -1,7 +1,8 @@
 import { bedrock } from '@ai-sdk/amazon-bedrock';
 import { openai } from '@ai-sdk/openai';
 import { google } from '@ai-sdk/google';
-import { streamText } from "ai";
+import { smoothStream, streamText } from 'ai';
+
 import { z } from "zod";
 
 export const maxDuration = 60
@@ -59,7 +60,6 @@ ${lastMessage.content}
   const result = streamText({
     // model: google("gemini-2.5-pro-exp-03-25"),
     // model: google("gemini-2.0-flash-001"),
-    // model: openrouter("google/gemini-2.0-flash-exp:free"),
     model: bedrock('anthropic.claude-3-5-sonnet-20241022-v2:0'),
     // model: openai("gpt-4o"),
     toolCallStreaming: true,
@@ -84,5 +84,7 @@ ${lastMessage.content}
     },
     temperature: 0,
   });
-  return result.toDataStreamResponse();
+  return result.toDataStreamResponse({
+
+  });
 }
diff --git a/components/chat-message-display.tsx b/components/chat-message-display.tsx
index a246b8b..608edd9 100644
--- a/components/chat-message-display.tsx
+++ b/components/chat-message-display.tsx
@@ -54,6 +54,13 @@ export function ChatMessageDisplay({
             >
                 <div className="flex flex-col gap-2">
                     <div className="text-xs">Tool: display_diagram</div>
+                    {args && (
+                        <div className="mt-1 font-mono text-xs overflow-hidden">
+                            {typeof args === "object" &&
+                                Object.keys(args).length > 0 &&
+                                `Args: ${JSON.stringify(args, null, 2)}`}
+                        </div>
+                    )}
                     <div className="mt-2 text-sm">
                         {state === "partial-call" ? (
                             <div className="h-4 w-4 border-2 border-primary border-t-transparent rounded-full animate-spin" />