mirror of
https://github.com/hwchase17/langchain
synced 2024-11-06 03:20:49 +00:00
d3ec00b566
Co-authored-by: Nuno Campos <nuno@boringbits.io> Co-authored-by: Davis Chase <130488702+dev2049@users.noreply.github.com> Co-authored-by: Zander Chase <130414180+vowelparrot@users.noreply.github.com> Co-authored-by: Harrison Chase <hw.chase.17@gmail.com>
41 lines
1.4 KiB
Python
41 lines
1.4 KiB
Python
"""Test LLM Math functionality."""
|
|
|
|
import pytest
|
|
|
|
from langchain.chains.llm_math.base import LLMMathChain
|
|
from langchain.chains.llm_math.prompt import _PROMPT_TEMPLATE
|
|
from tests.unit_tests.llms.fake_llm import FakeLLM
|
|
|
|
|
|
@pytest.fixture
|
|
def fake_llm_math_chain() -> LLMMathChain:
|
|
"""Fake LLM Math chain for testing."""
|
|
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: "```text\n2**.5\n```",
|
|
_PROMPT_TEMPLATE.format(question="foo"): "foo",
|
|
}
|
|
fake_llm = FakeLLM(queries=queries)
|
|
return LLMMathChain.from_llm(fake_llm, input_key="q", output_key="a")
|
|
|
|
|
|
def test_simple_question(fake_llm_math_chain: LLMMathChain) -> None:
|
|
"""Test simple question that should not need python."""
|
|
question = "What is 1 plus 1?"
|
|
output = fake_llm_math_chain.run(question)
|
|
assert output == "Answer: 2"
|
|
|
|
|
|
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}"
|
|
|
|
|
|
def test_error(fake_llm_math_chain: LLMMathChain) -> None:
|
|
"""Test question that raises error."""
|
|
with pytest.raises(ValueError):
|
|
fake_llm_math_chain.run("foo")
|