Skip to main content
`Ctrl`+`K`

#

# Source code for langchain.chains.llm_math.base

    """Chain that interprets a prompt and executes python code to do math."""
    import math
    import re
    from typing import Dict, List

    import numexpr
    from pydantic import Extra

    from langchain.chains.base import Chain
    from langchain.chains.llm import LLMChain
    from langchain.chains.llm_math.prompt import PROMPT
    from langchain.prompts.base import BasePromptTemplate
    from langchain.schema import BaseLanguageModel

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

        Example:
            .. code-block:: python

                from langchain import LLMMathChain, OpenAI
                llm_math = LLMMathChain(llm=OpenAI())
        """

        llm: BaseLanguageModel
        """LLM wrapper to use."""
        prompt: BasePromptTemplate = PROMPT
        """Prompt to use to translate to python if neccessary."""
        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

        @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:
            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"{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) -> Dict[str, str]:
            self.callback_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)
                self.callback_manager.on_text("\nAnswer: ", verbose=self.verbose)
                self.callback_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) -> Dict[str, str]:
            if self.callback_manager.is_async:
                await self.callback_manager.on_text(
                    llm_output, color="green", verbose=self.verbose
                )
            else:
                self.callback_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)
                if self.callback_manager.is_async:
                    await self.callback_manager.on_text("\nAnswer: ", verbose=self.verbose)
                    await self.callback_manager.on_text(
                        output, color="yellow", verbose=self.verbose
                    )
                else:
                    self.callback_manager.on_text("\nAnswer: ", verbose=self.verbose)
                    self.callback_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]) -> Dict[str, str]:
            llm_executor = LLMChain(
                prompt=self.prompt, llm=self.llm, callback_manager=self.callback_manager
            )
            self.callback_manager.on_text(inputs[self.input_key], verbose=self.verbose)
            llm_output = llm_executor.predict(
                question=inputs[self.input_key], stop=["```output"]
            )
            return self._process_llm_result(llm_output)

        async def _acall(self, inputs: Dict[str, str]) -> Dict[str, str]:
            llm_executor = LLMChain(
                prompt=self.prompt, llm=self.llm, callback_manager=self.callback_manager
            )
            if self.callback_manager.is_async:
                await self.callback_manager.on_text(
                    inputs[self.input_key], verbose=self.verbose
                )
            else:
                self.callback_manager.on_text(inputs[self.input_key], verbose=self.verbose)
            llm_output = await llm_executor.apredict(
                question=inputs[self.input_key], stop=["```output"]
            )
            return await self._aprocess_llm_result(llm_output)

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