diff --git a/src/core/ai_agent.py b/src/core/ai_agent.py
index b6e954d..f24e852 100644
--- a/src/core/ai_agent.py
+++ b/src/core/ai_agent.py
@@ -3,6 +3,7 @@
 
 from llama_index.llms import OpenAI, ChatMessage
 from llama_index.tools import BaseTool, FunctionTool
+from lightning import Data, Channel
 
 import nest_asyncio