diff --git a/langchain/chains/llm_math/base.py b/langchain/chains/llm_math/base.py index 3faebe3af2..3f26254b04 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."""