Toggle Menu
Prev Up Next

# Source code for langchain.chains.llm_math.base

    """Chain that interprets a prompt and executes python code to do math."""
    from __future__ import annotations

    import math
    import re
    import warnings
    from typing import Any, Dict, List, Optional

    from langchain_core.callbacks import (
        AsyncCallbackManagerForChainRun,
        CallbackManagerForChainRun,
    )
    from langchain_core.language_models import BaseLanguageModel
    from langchain_core.prompts import BasePromptTemplate
    from langchain_core.pydantic_v1 import Extra, root_validator

    from langchain.chains.base import Chain
    from langchain.chains.llm import LLMChain
    from langchain.chains.llm_math.prompt import PROMPT

    class LLMMathChain(Chain):
        """Chain that interprets a prompt and executes python code to do math.

        Example:
            .. code-block:: python

                from langchain.chains import LLMMathChain
                from langchain_community.llms import OpenAI
                llm_math = LLMMathChain.from_llm(OpenAI())
        """

        llm_chain: LLMChain
        llm: Optional[BaseLanguageModel] = None
        """[Deprecated] LLM wrapper to use."""
        prompt: BasePromptTemplate = PROMPT
        """[Deprecated] Prompt to use to translate to python if necessary."""
        input_key: str = "question"  #: :meta private:
        output_key: str = "answer"  #: :meta private:

        class Config:
            """Configuration for this pydantic object."""

            extra = Extra.forbid
            arbitrary_types_allowed = True

        @root_validator(pre=True)
        def raise_deprecation(cls, values: Dict) -> Dict:
            try:
                import numexpr  # noqa: F401
            except ImportError:
                raise ImportError(
                    "LLMMathChain requires the numexpr package. "
                    "Please install it with `pip install numexpr`."
                )
            if "llm" in values:
                warnings.warn(
                    "Directly instantiating an LLMMathChain with an llm is deprecated. "
                    "Please instantiate with llm_chain argument or using the from_llm "
                    "class method."
                )
                if "llm_chain" not in values and values["llm"] is not None:
                    prompt = values.get("prompt", PROMPT)
                    values["llm_chain"] = LLMChain(llm=values["llm"], prompt=prompt)
            return values

        @property
        def input_keys(self) -> List[str]:
            """Expect input key.

            :meta private:
            """
            return [self.input_key]

        @property
        def output_keys(self) -> List[str]:
            """Expect output key.

            :meta private:
            """
            return [self.output_key]

        def _evaluate_expression(self, expression: str) -> str:
            import numexpr

            try:
                local_dict = {"pi": math.pi, "e": math.e}
                output = str(
                    numexpr.evaluate(
                        expression.strip(),
                        global_dict={},  # restrict access to globals
                        local_dict=local_dict,  # add common mathematical functions
                    )
                )
            except Exception as e:
                raise ValueError(
                    f'LLMMathChain._evaluate("{expression}") raised error: {e}.'
                    " Please try again with a valid numerical expression"
                )

            # Remove any leading and trailing brackets from the output
            return re.sub(r"^\[|\]$", "", output)

        def _process_llm_result(
            self, llm_output: str, run_manager: CallbackManagerForChainRun
        ) -> Dict[str, str]:
            run_manager.on_text(llm_output, color="green", verbose=self.verbose)
            llm_output = llm_output.strip()
            text_match = re.search(r"^```text(.*?)```", llm_output, re.DOTALL)
            if text_match:
                expression = text_match.group(1)
                output = self._evaluate_expression(expression)
                run_manager.on_text("\nAnswer: ", verbose=self.verbose)
                run_manager.on_text(output, color="yellow", verbose=self.verbose)
                answer = "Answer: " + output
            elif llm_output.startswith("Answer:"):
                answer = llm_output
            elif "Answer:" in llm_output:
                answer = "Answer: " + llm_output.split("Answer:")[-1]
            else:
                raise ValueError(f"unknown format from LLM: {llm_output}")
            return {self.output_key: answer}

        async def _aprocess_llm_result(
            self,
            llm_output: str,
            run_manager: AsyncCallbackManagerForChainRun,
        ) -> Dict[str, str]:
            await run_manager.on_text(llm_output, color="green", verbose=self.verbose)
            llm_output = llm_output.strip()
            text_match = re.search(r"^```text(.*?)```", llm_output, re.DOTALL)
            if text_match:
                expression = text_match.group(1)
                output = self._evaluate_expression(expression)
                await run_manager.on_text("\nAnswer: ", verbose=self.verbose)
                await run_manager.on_text(output, color="yellow", verbose=self.verbose)
                answer = "Answer: " + output
            elif llm_output.startswith("Answer:"):
                answer = llm_output
            elif "Answer:" in llm_output:
                answer = "Answer: " + llm_output.split("Answer:")[-1]
            else:
                raise ValueError(f"unknown format from LLM: {llm_output}")
            return {self.output_key: answer}

        def _call(
            self,
            inputs: Dict[str, str],
            run_manager: Optional[CallbackManagerForChainRun] = None,
        ) -> Dict[str, str]:
            _run_manager = run_manager or CallbackManagerForChainRun.get_noop_manager()
            _run_manager.on_text(inputs[self.input_key])
            llm_output = self.llm_chain.predict(
                question=inputs[self.input_key],
                stop=["```output"],
                callbacks=_run_manager.get_child(),
            )
            return self._process_llm_result(llm_output, _run_manager)

        async def _acall(
            self,
            inputs: Dict[str, str],
            run_manager: Optional[AsyncCallbackManagerForChainRun] = None,
        ) -> Dict[str, str]:
            _run_manager = run_manager or AsyncCallbackManagerForChainRun.get_noop_manager()
            await _run_manager.on_text(inputs[self.input_key])
            llm_output = await self.llm_chain.apredict(
                question=inputs[self.input_key],
                stop=["```output"],
                callbacks=_run_manager.get_child(),
            )
            return await self._aprocess_llm_result(llm_output, _run_manager)

        @property
        def _chain_type(self) -> str:
            return "llm_math_chain"

        @classmethod
        def from_llm(
            cls,
            llm: BaseLanguageModel,
            prompt: BasePromptTemplate = PROMPT,
            **kwargs: Any,
        ) -> LLMMathChain:
            llm_chain = LLMChain(llm=llm, prompt=prompt)
            return cls(llm_chain=llm_chain, **kwargs)