From fd1fcb5a7d48cbe18b480b1493b66540e4709745 Mon Sep 17 00:00:00 2001
From: Ben Heckmann <79015931+benheckmann@users.noreply.github.com>
Date: Thu, 30 Mar 2023 16:52:58 +0200
Subject: [PATCH] fix typing for LLMMathChain (#2183)

Fix typing in LLMMathChain to allow chat models (#1834). Might have been
forgotten in related PR #1807.
---
 langchain/chains/llm_math/base.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

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."""