diff --git a/dist/index.js b/dist/index.js
index 44796dd..5fdf914 100644
--- a/dist/index.js
+++ b/dist/index.js
@@ -4,6 +4,12 @@ export class TetherClient {
     subscribedQueries = new Map();
     connect = (url) => {
         this.websocketHandler.startConnection(url);
+        this.websocketHandler.onQuery = (location, data) => {
+            if (location) {
+                const callback = this.subscribedQueries.get(location);
+                callback?.(data);
+            }
+        };
     };
     disconnect = () => {
         this.websocketHandler.close();
diff --git a/dist/utils/websocket.d.ts b/dist/utils/websocket.d.ts
index ae4be2f..c439029 100644
--- a/dist/utils/websocket.d.ts
+++ b/dist/utils/websocket.d.ts
@@ -1,9 +1,9 @@
 export declare class WebSocketHandler {
     private ws;
     private url;
-    private subscribedQueries;
-    private onOpen;
-    private onClose;
+    onOpen: () => void;
+    onQuery: (location: string | undefined, data: any) => void;
+    onClose: () => void;
     private reconnectAttempts;
     private maxReconnectAttempts;
     private reconnectInterval;
diff --git a/dist/utils/websocket.js b/dist/utils/websocket.js
index b15a41a..d463454 100644
--- a/dist/utils/websocket.js
+++ b/dist/utils/websocket.js
@@ -1,8 +1,8 @@
 export class WebSocketHandler {
     ws = null;
     url = '';
-    subscribedQueries = new Map();
     onOpen = () => { };
+    onQuery = () => { };
     onClose = () => { };
     reconnectAttempts = 0;
     maxReconnectAttempts = 5;
@@ -24,11 +24,7 @@ export class WebSocketHandler {
         ws.onmessage = (event) => {
             const data = JSON.parse(String(event.data));
             if (data.type === 'query') {
-                this.subscribedQueries.forEach((callback, query) => {
-                    if (data.query === query) {
-                        callback(data.data);
-                    }
-                });
+                this.onQuery(data.location, data.data);
             }
             else if (data.type === 'error') {
                 console.error(data.error);
diff --git a/src/index.ts b/src/index.ts
index a33c823..d495b3d 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -5,6 +5,12 @@ export class TetherClient {
 
     connect = (url: string) => {
         this.websocketHandler.startConnection(url);
+        this.websocketHandler.onQuery = (location, data) => {
+            if (location) {
+                const callback = this.subscribedQueries.get(location);
+                callback?.(data);
+            }
+        };
     };
     
     disconnect = () => {
diff --git a/src/utils/websocket.ts b/src/utils/websocket.ts
index 8bb8dec..20e1def 100644
--- a/src/utils/websocket.ts
+++ b/src/utils/websocket.ts
@@ -1,9 +1,9 @@
 export class WebSocketHandler {
     private ws: WebSocket | null = null;
     private url: string = '';
-    private subscribedQueries = new Map<string, (data: any) => void>();
-    private onOpen: () => void = () => {};
-    private onClose: () => void = () => {};
+    public onOpen: () => void = () => {};
+    public onQuery: (location: string | undefined, data: any) => void = () => {};
+    public onClose: () => void = () => {};
     private reconnectAttempts: number = 0;
     private maxReconnectAttempts: number = 5;
     private reconnectInterval: number = 1000;
@@ -27,16 +27,12 @@ export class WebSocketHandler {
         ws.onmessage = (event: MessageEvent) => {
             const data = JSON.parse(String(event.data)) as {
                 type: string;
-                query?: string;
+                location?: string;
                 data?: unknown;
                 error?: string;
             };
             if (data.type === 'query') {
-                this.subscribedQueries.forEach((callback, query) => {
-                    if (data.query === query) {
-                        callback(data.data);
-                    }
-                });
+                this.onQuery(data.location, data.data);
             } else if (data.type === 'error') {
                 console.error(data.error);
             }