diff --git a/src/mindpilot/app/callback_handler/agent_callback_handler.py b/src/mindpilot/app/callback_handler/agent_callback_handler.py
index 1bb13df..f1d69d2 100644
--- a/src/mindpilot/app/callback_handler/agent_callback_handler.py
+++ b/src/mindpilot/app/callback_handler/agent_callback_handler.py
@@ -81,12 +81,11 @@ class AgentExecutorAsyncIteratorCallbackHandler(AsyncIteratorCallbackHandler):
         self.queue.put_nowait(dumps(data))
 
     async def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
-        pass
-        # data = {
-        #     "status": AgentStatus.llm_end,
-        #     "text": response.generations[0][0].message.content,
-        # }
-        # self.queue.put_nowait(dumps(data))
+        data = {
+            "status": AgentStatus.llm_end,
+            "text": response.generations[0][0].message.content,
+        }
+        self.queue.put_nowait(dumps(data))
 
     async def on_llm_error(
             self, error: Exception | KeyboardInterrupt, **kwargs: Any
@@ -153,42 +152,42 @@ class AgentExecutorAsyncIteratorCallbackHandler(AsyncIteratorCallbackHandler):
         # self.done.clear()
         self.queue.put_nowait(dumps(data))
 
-    # async def on_agent_action(
-    #         self,
-    #         action: AgentAction,
-    #         *,
-    #         run_id: UUID,
-    #         parent_run_id: Optional[UUID] = None,
-    #         tags: Optional[List[str]] = None,
-    #         **kwargs: Any,
-    # ) -> None:
-    #     data = {
-    #         "status": AgentStatus.agent_action,
-    #         "tool_name": action.tool,
-    #         "tool_input": action.tool_input,
-    #         "text": action.log,
-    #     }
-    #     self.queue.put_nowait(dumps(data))
-
-    # async def on_agent_finish(
-    #         self,
-    #         finish: AgentFinish,
-    #         *,
-    #         run_id: UUID,
-    #         parent_run_id: Optional[UUID] = None,
-    #         tags: Optional[List[str]] = None,
-    #         **kwargs: Any,
-    # ) -> None:
-    #     if "Thought:" in finish.return_values["output"]:
-    #         finish.return_values["output"] = finish.return_values["output"].replace(
-    #             "Thought:", ""
-    #         )
-    #
-    #     data = {
-    #         "status": AgentStatus.agent_finish,
-    #         "text": finish.return_values["output"],
-    #     }
-    #     self.queue.put_nowait(dumps(data))
+    async def on_agent_action(
+            self,
+            action: AgentAction,
+            *,
+            run_id: UUID,
+            parent_run_id: Optional[UUID] = None,
+            tags: Optional[List[str]] = None,
+            **kwargs: Any,
+    ) -> None:
+        data = {
+            "status": AgentStatus.agent_action,
+            "tool_name": action.tool,
+            "tool_input": action.tool_input,
+            "text": action.log,
+        }
+        self.queue.put_nowait(dumps(data))
+
+    async def on_agent_finish(
+            self,
+            finish: AgentFinish,
+            *,
+            run_id: UUID,
+            parent_run_id: Optional[UUID] = None,
+            tags: Optional[List[str]] = None,
+            **kwargs: Any,
+    ) -> None:
+        if "Thought:" in finish.return_values["output"]:
+            finish.return_values["output"] = finish.return_values["output"].replace(
+                "Thought:", ""
+            )
+
+        data = {
+            "status": AgentStatus.agent_finish,
+            "text": finish.return_values["output"],
+        }
+        self.queue.put_nowait(dumps(data))
 
     async def on_chain_end(
             self,
diff --git a/src/mindpilot/app/chat/chat.py b/src/mindpilot/app/chat/chat.py
index be086f9..e339567 100644
--- a/src/mindpilot/app/chat/chat.py
+++ b/src/mindpilot/app/chat/chat.py
@@ -164,13 +164,13 @@ async def chat(
                         data["message_type"] = message_type
                 except:
                     ...
-            # elif data["status"] == AgentStatus.agent_finish:
-            #     try:
-            #         tool_output = json.loads(data["text"])
-            #         if message_type := tool_output.get("message_type"):
-            #             data["message_type"] = message_type
-            #     except:
-            #         ...
+            elif data["status"] == AgentStatus.agent_finish:
+                try:
+                    tool_output = json.loads(data["text"])
+                    if message_type := tool_output.get("message_type"):
+                        data["message_type"] = message_type
+                except:
+                    ...
 
             ret = OpenAIChatOutput(
                 id=f"chat{uuid.uuid4()}",