Patch LLMMathChain exec vulnerability

langchain/chains/llm_math/base.py

@@ -1,4 +1,5 @@
 """Chain that interprets a prompt and executes python code to do math."""
+import math  # noqa: F401
 from typing import Dict, List
 
 from pydantic import BaseModel, Extra
@@ -8,7 +9,6 @@
 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
 
 
 class LLMMathChain(Chain, BaseModel):
@@ -51,12 +51,11 @@ def output_keys(self) -> List[str]:
         return [self.output_key]
 
     def _process_llm_result(self, t: str) -> Dict[str, str]:
-        python_executor = PythonREPL()
         self.callback_manager.on_text(t, color="green", verbose=self.verbose)
         t = t.strip()
         if t.startswith("```python"):
             code = t[9:-4]
-            output = python_executor.run(code)
+            output = str(eval(code))
             self.callback_manager.on_text("\nAnswer: ", verbose=self.verbose)
             self.callback_manager.on_text(output, color="yellow", verbose=self.verbose)
             answer = "Answer: " + output

langchain/chains/llm_math/prompt.py

@@ -5,11 +5,13 @@
 
 You can do basic math, and your memorization abilities are impressive, but you can't do any complex calculations that a human could not do in their head. You also have an annoying tendency to just make up highly specific, but wrong, answers.
 
+Do not import any libraries, if you need advanced functions simply use `math.<function_name>`. Return the expression in one line.
+
 So we hooked you up to a Python 3 kernel, and now you can execute code. If anyone gives you a hard math problem, just use this format and we’ll take care of the rest:
 
 Question: ${{Question with hard calculation.}}
 ```python
-${{Code that prints what you need to know}}
+${{Code that evaluates the mathematical expression}}
 ```
 ```output
 ${{Output of your code}}
@@ -26,13 +28,23 @@
 Question: What is 37593 * 67?
 
 ```python
-print(37593 * 67)
+37593 * 67
 ```
 ```output
 2518731
 ```
 Answer: 2518731
 
+Question: What is the square root of 256 minus 10?
+
+```python
+math.sqrt(256) - 10
+```
+```output
+6
+```
+Answer: 6
+
 Question: {question}
 """
 

tests/unit_tests/chains/test_llm_math.py

@@ -13,7 +13,7 @@ def fake_llm_math_chain() -> LLMMathChain:
     complex_question = _PROMPT_TEMPLATE.format(question="What is the square root of 2?")
     queries = {
         _PROMPT_TEMPLATE.format(question="What is 1 plus 1?"): "Answer: 2",
-        complex_question: "```python\nprint(2**.5)\n```",
+        complex_question: "```python\n2**.5\n```",
         _PROMPT_TEMPLATE.format(question="foo"): "foo",
     }
     fake_llm = FakeLLM(queries=queries)
@@ -31,7 +31,7 @@ def test_complex_question(fake_llm_math_chain: LLMMathChain) -> None:
     """Test complex question that should need python."""
     question = "What is the square root of 2?"
     output = fake_llm_math_chain.run(question)
-    assert output == f"Answer: {2**.5}\n"
+    assert output == f"Answer: {2**.5}"
 
 
 def test_error(fake_llm_math_chain: LLMMathChain) -> None: