diff --git a/app/api/chat/route.ts b/app/api/chat/route.ts
index c66a71d..69538f2 100644
--- a/app/api/chat/route.ts
+++ b/app/api/chat/route.ts
@@ -1,13 +1,40 @@
 import { google } from "@ai-sdk/google";
+import { Message } from "ai/react";
 import { streamText } from "ai";
-
+import { z } from "zod";
 export const maxDuration = 30;
 
+// Define tool interfaces
+interface DisplayFlowChartArgs {
+  xml: string;
+}
+
+interface ToolContext {
+  getCurrentXML: () => string;
+  displayChart: (xml: string) => void;
+}
+
 export async function POST(req: Request) {
   const { messages } = await req.json();
-  const result = streamText({
+  const response = streamText({
     model: google("gemini-2.0-flash"),
     messages,
+    tools: {
+      display_flow_chart: {
+        description: "Display a flowchart on draw.io",
+        parameters: z.object(
+          {
+            xml: z.string().describe("XML string to be displayed on draw.io"),
+          },
+        )
+      },
+      fetch_flow_chart: {
+        description: "Get the current flowchart XML from draw.io",
+        parameters: z.object({}),
+      }
+    },
+    temperature: 0,
   });
-  return result.toDataStreamResponse();
-}
\ No newline at end of file
+
+  return response.toDataStreamResponse();
+}
diff --git a/app/page.tsx b/app/page.tsx
index 65cb4dc..b2fbb09 100644
--- a/app/page.tsx
+++ b/app/page.tsx
@@ -9,24 +9,23 @@ import ChatPanel from "@/components/chatPanel";
 export default function Home() {
     const drawioRef = useRef<DrawIoEmbedRef>(null);
     const [chartXML, setChartXML] = useState<string>("");
-    const [diagram, setDiagram] = useState<string>("");
-    // const handleExport = () => {};
+
     const handleExport = () => {
-        // use this function to export the diagramxml from the drawio editor
         if (drawioRef.current) {
             drawioRef.current.exportDiagram({
                 format: "xmlsvg",
             });
         }
     };
+
     const loadDiagram = (chart: string) => {
-        // use this function to display the diagramxml in the drawio editor
         if (drawioRef.current) {
             drawioRef.current.load({
                 xml: chart,
             });
         }
     };
+
     return (
         <div className="flex h-screen bg-gray-100">
             <div className="w-2/3 p-1">
@@ -34,7 +33,6 @@ export default function Home() {
                     ref={drawioRef}
                     onExport={(data) => setChartXML(extractDiagramXML(data.data))}
                     urlParameters={{
-                        // ui: "kennedy",
                         spin: true,
                         libraries: false,
                         saveAndExit: false,
@@ -42,8 +40,14 @@ export default function Home() {
                     }}
                 />
             </div>
-            <div className="w-1/3 p-1  border-gray-300 ">
-                <ChatPanel />
+            <div className="w-1/3 p-1 border-gray-300">
+                <ChatPanel
+                    onDisplayChart={(xml) => loadDiagram(xml)}
+                    onFetchChart={() => {
+                        handleExport();
+                        return chartXML;
+                    }}
+                />
             </div>
         </div>
     );
diff --git a/components/chatPanel.tsx b/components/chatPanel.tsx
index 2b29abf..d49faa7 100644
--- a/components/chatPanel.tsx
+++ b/components/chatPanel.tsx
@@ -1,20 +1,41 @@
 "use client"
 
 import type React from "react"
-
 import { useRef, useEffect } from "react"
-import { useChat } from "@ai-sdk/react"
+
 import { Button } from "@/components/ui/button"
 import { Input } from "@/components/ui/input"
 import { Card, CardContent, CardFooter, CardHeader, CardTitle } from "@/components/ui/card"
 import { ScrollArea } from "@/components/ui/scroll-area"
 import { Loader2, Send } from "lucide-react"
+import { useChat } from '@ai-sdk/react';
+import { ToolInvocation } from 'ai';
 
 
-export default function ChatPanel() {
-    const { messages, input, handleInputChange, handleSubmit, isLoading } = useChat({
-        api: "/api/chat",
-        maxSteps: 5, // Allow multiple steps for complex diagram generation
+interface ChatPanelProps {
+    onDisplayChart: (xml: string) => void;
+    onFetchChart: () => string;
+}
+
+export default function ChatPanel({ onDisplayChart, onFetchChart }: ChatPanelProps) {
+    const { messages, input, handleInputChange, handleSubmit, isLoading, error, addToolResult } = useChat({
+        async onToolCall({ toolCall }) {
+            console.log("Tool call:", toolCall);
+            if (toolCall.toolName === "display_flow_chart") {
+                const { xml } = toolCall.args as { xml: string };
+                onDisplayChart(xml);
+            } else if (toolCall.toolName === "fetch_flow_chart") {
+                const currentXML = onFetchChart();
+                console.log("Current XML:", currentXML);
+                addToolResult({
+                    toolCallId: toolCall.toolCallId,
+                    result: currentXML
+                });
+            }
+        },
+        onError: (error) => {
+            console.error("Chat error:", error);
+        }
     })
     const messagesEndRef = useRef<HTMLDivElement>(null)
 
@@ -25,10 +46,9 @@ export default function ChatPanel() {
         }
     }, [messages])
 
-
     const onFormSubmit = (e: React.FormEvent<HTMLFormElement>) => {
         e.preventDefault()
-        if (input.trim()) {
+        if (input.trim() && !isLoading) {
             handleSubmit(e)
         }
     }
@@ -42,7 +62,7 @@ export default function ChatPanel() {
                 <ScrollArea className="h-full pr-4">
                     {messages.length === 0 ? (
                         <div className="text-center text-gray-500 mt-8">
-                            <p>Start a conversation to generate a diagram.</p>
+                            <p>Start a conversation to generate or modify diagrams.</p>
                             <p className="text-sm mt-2">Try: "Create a flowchart for user authentication"</p>
                         </div>
                     ) : (
@@ -54,16 +74,21 @@ export default function ChatPanel() {
                                 >
                                     {message.content}
                                 </div>
-                                {message.toolInvocations?.map((tool, index) => (
-                                    <div key={index} className="mt-2 text-left">
+                                {(message as any).function_call && (
+                                    <div className="mt-2 text-left">
                                         <div className="text-xs text-gray-500">
-                                            {tool.state === "call" ? "Generating diagram..." : "Diagram generated"}
+                                            Using tool: {(message as any).function_call.name}...
                                         </div>
                                     </div>
-                                ))}
+                                )}
                             </div>
                         ))
                     )}
+                    {error && (
+                        <div className="text-red-500 text-sm mt-2">
+                            Error: {error.message}
+                        </div>
+                    )}
                     <div ref={messagesEndRef} />
                 </ScrollArea>
             </CardContent>
@@ -72,7 +97,7 @@ export default function ChatPanel() {
                     <Input
                         value={input}
                         onChange={handleInputChange}
-                        placeholder="Describe the diagram you want to create..."
+                        placeholder="Describe what changes you want to make to the diagram..."
                         disabled={isLoading}
                         className="flex-grow"
                     />
@@ -84,4 +109,3 @@ export default function ChatPanel() {
         </Card>
     )
 }
-