diff --git a/custom-frontend/frontend/package-lock.json b/custom-frontend/frontend/package-lock.json
index 24de125e..51f37156 100644
--- a/custom-frontend/frontend/package-lock.json
+++ b/custom-frontend/frontend/package-lock.json
@@ -8,7 +8,7 @@
       "name": "vite-project",
       "version": "0.0.0",
       "dependencies": {
-        "@chainlit/react-client": "^0.0.1",
+        "@chainlit/react-client": "0.0.6",
         "@radix-ui/react-slot": "^1.0.2",
         "class-variance-authority": "^0.7.0",
         "clsx": "^2.0.0",
@@ -70,9 +70,9 @@
       }
     },
     "node_modules/@chainlit/react-client": {
-      "version": "0.0.1",
-      "resolved": "https://registry.npmjs.org/@chainlit/react-client/-/react-client-0.0.1.tgz",
-      "integrity": "sha512-iKle82o4jXnxc/MPsVdTq0kGYFm46Z5q0r9R5lIDVOuGXJes9IgzsNVaBUv4Ur//VIJvJ6dVWoufvDUWsLreaw==",
+      "version": "0.0.6",
+      "resolved": "https://registry.npmjs.org/@chainlit/react-client/-/react-client-0.0.6.tgz",
+      "integrity": "sha512-vyJOsmJuZfzS1wMrpeXGF3lN+zhfKV/8T+3o8YhCzMmBtlA3tBHEedQCTvfY5Ft0e22ZYY8YsKATNKwcXuYkEA==",
       "dependencies": {
         "jwt-decode": "^3.1.2",
         "lodash": "^4.17.21",
diff --git a/custom-frontend/frontend/package.json b/custom-frontend/frontend/package.json
index 507d9ef9..6fce9088 100644
--- a/custom-frontend/frontend/package.json
+++ b/custom-frontend/frontend/package.json
@@ -10,7 +10,7 @@
     "preview": "vite preview"
   },
   "dependencies": {
-    "@chainlit/react-client": "0.0.2",
+    "@chainlit/react-client": "0.0.6",
     "@radix-ui/react-slot": "^1.0.2",
     "class-variance-authority": "^0.7.0",
     "clsx": "^2.0.0",
diff --git a/custom-frontend/frontend/src/App.tsx b/custom-frontend/frontend/src/App.tsx
index 5ba36f14..8fa7eb02 100644
--- a/custom-frontend/frontend/src/App.tsx
+++ b/custom-frontend/frontend/src/App.tsx
@@ -1,32 +1,37 @@
 import { useEffect } from "react";
 
-import { ChainlitAPI, sessionState, useChatSession } from "@chainlit/react-client";
+import {
+  ChainlitAPI,
+  sessionState,
+  useChatSession,
+} from "@chainlit/react-client";
 import { Playground } from "./components/playground";
 import { useRecoilValue } from "recoil";
 
 const CHAINLIT_SERVER = "http://localhost:8000";
 const userEnv = {};
 
-const apiClient = new ChainlitAPI(CHAINLIT_SERVER);
+const apiClient = new ChainlitAPI(CHAINLIT_SERVER, "app");
 
 function App() {
   const { connect } = useChatSession();
   const session = useRecoilValue(sessionState);
-
   useEffect(() => {
     if (session?.socket.connected) {
-      return
+      return;
     }
-    fetch(apiClient
-      .buildEndpoint("/custom-auth")).then((res) => {
+    fetch(apiClient.buildEndpoint("/custom-auth"))
+      .then((res) => {
         return res.json();
-      }
-    ).then((data) => {
-      connect({ client: apiClient, userEnv, accessToken: `Bearer: ${data.token}` });
-    }
-    )
-  }, [session, connect]);
-
+      })
+      .then((data) => {
+        connect({
+          client: apiClient,
+          userEnv,
+          accessToken: `Bearer: ${data.token}`,
+        });
+      });
+  }, [connect]);
 
   return (
     <>
diff --git a/llama-index/app.py b/llama-index/app.py
index 7223a43e..e82addfe 100644
--- a/llama-index/app.py
+++ b/llama-index/app.py
@@ -11,6 +11,9 @@
 )
 from llama_index.llms.openai import OpenAI
 from llama_index.embeddings.openai import OpenAIEmbedding
+from llama_index.core.query_engine.retriever_query_engine import RetrieverQueryEngine
+from llama_index.core.callbacks import CallbackManager
+from llama_index.core.service_context import ServiceContext
 
 openai.api_key = os.environ.get("OPENAI_API_KEY")
 
@@ -33,7 +36,8 @@ async def start():
     Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-small")
     Settings.context_window = 4096
 
-    query_engine = index.as_query_engine(streaming=True, similarity_top_k=2)
+    service_context = ServiceContext.from_defaults(callback_manager=CallbackManager([cl.LlamaIndexCallbackHandler()]))
+    query_engine = index.as_query_engine(streaming=True, similarity_top_k=2, service_context=service_context)
     cl.user_session.set("query_engine", query_engine)
 
     await cl.Message(