diff --git a/python/packages/autogen-ext/src/autogen_ext/tools/graphrag/_global_search.py b/python/packages/autogen-ext/src/autogen_ext/tools/graphrag/_global_search.py
index ccdedb44b21c..d99a845fad47 100644
--- a/python/packages/autogen-ext/src/autogen_ext/tools/graphrag/_global_search.py
+++ b/python/packages/autogen-ext/src/autogen_ext/tools/graphrag/_global_search.py
@@ -64,6 +64,7 @@ class GlobalSearchTool(BaseTool[GlobalSearchToolArgs, GlobalSearchToolReturn]):
 
         import asyncio
         from autogen_ext.models.openai import OpenAIChatCompletionClient
+        from autogen_agentchat.messages import TextMessage
         from autogen_ext.tools.graphrag import GlobalSearchTool
         from autogen_agentchat.agents import AssistantAgent
 
diff --git a/python/packages/autogen-ext/src/autogen_ext/tools/graphrag/_local_search.py b/python/packages/autogen-ext/src/autogen_ext/tools/graphrag/_local_search.py
index 6ca3b2ae15a2..d1b7985be48b 100644
--- a/python/packages/autogen-ext/src/autogen_ext/tools/graphrag/_local_search.py
+++ b/python/packages/autogen-ext/src/autogen_ext/tools/graphrag/_local_search.py
@@ -64,6 +64,7 @@ class LocalSearchTool(BaseTool[LocalSearchToolArgs, LocalSearchToolReturn]):
 
         import asyncio
         from autogen_ext.models.openai import OpenAIChatCompletionClient
+        from autogen_agentchat.messages import TextMessage
         from autogen_ext.tools.graphrag import LocalSearchTool
         from autogen_agentchat.agents import AssistantAgent
 
diff --git a/python/samples/agentchat_graphrag/app.py b/python/samples/agentchat_graphrag/app.py
index 6a6d1fd16fe4..6572ffc887cd 100644
--- a/python/samples/agentchat_graphrag/app.py
+++ b/python/samples/agentchat_graphrag/app.py
@@ -61,6 +61,7 @@ async def main(model_config: Dict[str, Any]) -> None:
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Run a GraphRAG search with an agent.")
     parser.add_argument("--verbose", action="store_true", help="Enable verbose logging.")
+    # See agentchat_chainlit/model_config_template.json for model_config.json
     parser.add_argument(
         "--model-config", type=str, help="Path to the model configuration file.", default="model_config.json"
     )