mirror of
https://github.com/hwchase17/langchain
synced 2024-10-29 17:07:25 +00:00
83 lines
3.1 KiB
Python
83 lines
3.1 KiB
Python
|
"""Test LLM Math functionality."""
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
from langchain.chains.llm_symbolic_math.base import LLMSymbolicMathChain
|
||
|
|
from langchain.chains.llm_symbolic_math.prompt import _PROMPT_TEMPLATE
|
||
|
|
from tests.unit_tests.llms.fake_llm import FakeLLM
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.fixture
|
||
|
|
@pytest.mark.requires("sympy")
|
||
|
|
def fake_llm_symbolic_math_chain() -> LLMSymbolicMathChain:
|
||
|
|
"""Fake LLM Math chain for testing."""
|
||
|
|
queries = {
|
||
|
|
_PROMPT_TEMPLATE.format(question="What is 1 plus 1?"): "Answer: 2",
|
||
|
|
_PROMPT_TEMPLATE.format(
|
||
|
|
question="What is the square root of 2?"
|
||
|
|
): "```text\nsqrt(2)\n```",
|
||
|
|
_PROMPT_TEMPLATE.format(
|
||
|
|
question="What is the limit of sin(x) / x as x goes to 0?"
|
||
|
|
): "```text\nlimit(sin(x)/x,x,0)\n```",
|
||
|
|
_PROMPT_TEMPLATE.format(
|
||
|
|
question="What is the integral of e^-x from 0 to infinity?"
|
||
|
|
): "```text\nintegrate(exp(-x), (x, 0, oo))\n```",
|
||
|
|
_PROMPT_TEMPLATE.format(
|
||
|
|
question="What are the solutions to this equation x**2 - x?"
|
||
|
|
): "```text\nsolveset(x**2 - x, x)\n```",
|
||
|
|
_PROMPT_TEMPLATE.format(question="foo"): "foo",
|
||
|
|
}
|
||
|
|
fake_llm = FakeLLM(queries=queries)
|
||
|
|
return LLMSymbolicMathChain.from_llm(fake_llm, input_key="q", output_key="a")
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.requires("sympy")
|
||
|
|
def test_simple_question(fake_llm_symbolic_math_chain: LLMSymbolicMathChain) -> None:
|
||
|
|
"""Test simple question that should not need python."""
|
||
|
|
question = "What is 1 plus 1?"
|
||
|
|
output = fake_llm_symbolic_math_chain.run(question)
|
||
|
|
assert output == "Answer: 2"
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.requires("sympy")
|
||
|
|
def test_root_question(fake_llm_symbolic_math_chain: LLMSymbolicMathChain) -> None:
|
||
|
|
"""Test irrational number that should need sympy."""
|
||
|
|
import sympy
|
||
|
|
|
||
|
|
question = "What is the square root of 2?"
|
||
|
|
output = fake_llm_symbolic_math_chain.run(question)
|
||
|
|
assert output == f"Answer: {sympy.sqrt(2)}"
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.requires("sympy")
|
||
|
|
def test_limit_question(fake_llm_symbolic_math_chain: LLMSymbolicMathChain) -> None:
|
||
|
|
"""Test question about limits that needs sympy"""
|
||
|
|
question = "What is the limit of sin(x) / x as x goes to 0?"
|
||
|
|
output = fake_llm_symbolic_math_chain.run(question)
|
||
|
|
assert output == "Answer: 1"
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.requires("sympy")
|
||
|
|
def test_integration_question(
|
||
|
|
fake_llm_symbolic_math_chain: LLMSymbolicMathChain,
|
||
|
|
) -> None:
|
||
|
|
"""Test question about integration that needs sympy"""
|
||
|
|
question = "What is the integral of e^-x from 0 to infinity?"
|
||
|
|
output = fake_llm_symbolic_math_chain.run(question)
|
||
|
|
assert output == "Answer: 1"
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.requires("sympy")
|
||
|
|
def test_solver_question(fake_llm_symbolic_math_chain: LLMSymbolicMathChain) -> None:
|
||
|
|
"""Test question about solving algebraic equations that needs sympy"""
|
||
|
|
question = "What are the solutions to this equation x**2 - x?"
|
||
|
|
output = fake_llm_symbolic_math_chain.run(question)
|
||
|
|
assert output == "Answer: {0, 1}"
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.requires("sympy")
|
||
|
|
def test_error(fake_llm_symbolic_math_chain: LLMSymbolicMathChain) -> None:
|
||
|
|
"""Test question that raises error."""
|
||
|
|
with pytest.raises(ValueError):
|
||
|
|
fake_llm_symbolic_math_chain.run("foo")
|