diff --git a/langchain/chains/llm_math/base.py b/langchain/chains/llm_math/base.py
index 3faebe3a..3f26254b 100644
--- a/langchain/chains/llm_math/base.py
+++ b/langchain/chains/llm_math/base.py
@@ -6,9 +6,9 @@ from pydantic import BaseModel, Extra
 from langchain.chains.base import Chain
 from langchain.chains.llm import LLMChain
 from langchain.chains.llm_math.prompt import PROMPT
-from langchain.llms.base import BaseLLM
 from langchain.prompts.base import BasePromptTemplate
 from langchain.python import PythonREPL
+from langchain.schema import BaseLanguageModel
 
 
 class LLMMathChain(Chain, BaseModel):
@@ -21,7 +21,7 @@ class LLMMathChain(Chain, BaseModel):
             llm_math = LLMMathChain(llm=OpenAI())
     """
 
-    llm: BaseLLM
+    llm: BaseLanguageModel
     """LLM wrapper to use."""
     prompt: BasePromptTemplate = PROMPT
     """Prompt to use to translate to python if neccessary."""