diff --git a/main.py b/main.py
index a426a7d..82e3f70 100644
--- a/main.py
+++ b/main.py
@@ -292,7 +292,7 @@ def create_simple_agent(system_prompt: Optional[str] = None):
     Create a simple chat agent for providers without full tool-calling support (e.g., Ollama).
     This agent just uses direct chat without advanced tool calling features.
     """
-    from langchain.schema import HumanMessage, AIMessage, SystemMessage
+    from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
     from langchain.chat_models import init_chat_model
     
     model = load_model_from_config()
@@ -302,7 +302,7 @@ def create_simple_agent(system_prompt: Optional[str] = None):
         
         def __init__(self, llm, prompt):
             self.llm = llm
-            self.messages = []
+            self.messages: List = []
             if prompt:
                 self.messages.append(SystemMessage(content=prompt))