diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/custom-agents.ipynb b/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/custom-agents.ipynb
index 4ca98a4336d5..3835d129ef8d 100644
--- a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/custom-agents.ipynb
+++ b/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/custom-agents.ipynb
@@ -319,6 +319,7 @@
                 "\n",
                 "from autogen_agentchat.agents import BaseChatAgent\n",
                 "from autogen_agentchat.base import Response\n",
+                "from autogen_agentchat.messages import TextMessage\n"
                 "from autogen_agentchat.messages import BaseAgentEvent, BaseChatMessage\n",
                 "from autogen_core import CancellationToken\n",
                 "from autogen_core.model_context import UnboundedChatCompletionContext\n",