diff --git a/src/backend/base/langflow/components/models/OllamaModel.py b/src/backend/base/langflow/components/models/OllamaModel.py
index 3b2872951..4173b9e13 100644
--- a/src/backend/base/langflow/components/models/OllamaModel.py
+++ b/src/backend/base/langflow/components/models/OllamaModel.py
@@ -1,5 +1,5 @@
 from langchain_community.chat_models import ChatOllama
-
+from langchain_core.language_models.chat_models import BaseChatModel
 from langflow.base.constants import STREAM_INFO_TEXT
 from langflow.base.models.model import LCModelComponent
 from langflow.field_typing import BaseLanguageModel, Text
@@ -176,7 +176,7 @@ class ChatOllamaComponent(LCModelComponent):
         self.status = result
         return result
 
-    def build_model(self) -> BaseLanguageModel:
+    def build_model(self) -> BaseLanguageModel | BaseChatModel:
         # Mapping mirostat settings to their corresponding values
         mirostat_options = {"Mirostat": 1, "Mirostat 2.0": 2}