From bbc795b7528de8a27cbc3c99dd523c637af08811 Mon Sep 17 00:00:00 2001 From: Erick Friis Date: Thu, 7 Dec 2023 11:01:59 -0800 Subject: [PATCH] override llm config --- .../tests/unit_tests/chains/test_llm.py | 75 +++++++++++++++++++ 1 file changed, 75 insertions(+) create mode 100644 .scripts/community_split/libs/langchain/tests/unit_tests/chains/test_llm.py diff --git a/.scripts/community_split/libs/langchain/tests/unit_tests/chains/test_llm.py b/.scripts/community_split/libs/langchain/tests/unit_tests/chains/test_llm.py new file mode 100644 index 00000000000..0179cd135f2 --- /dev/null +++ b/.scripts/community_split/libs/langchain/tests/unit_tests/chains/test_llm.py @@ -0,0 +1,75 @@ +"""Test LLM chain.""" +from tempfile import TemporaryDirectory +from typing import Dict, List, Union +from unittest.mock import patch + +import pytest +from langchain_core.output_parsers import BaseOutputParser +from langchain_core.prompts import PromptTemplate + +from langchain.chains.llm import LLMChain +from tests.unit_tests.llms.fake_llm import FakeLLM + + +class FakeOutputParser(BaseOutputParser): + """Fake output parser class for testing.""" + + def parse(self, text: str) -> Union[str, List[str], Dict[str, str]]: + """Parse by splitting.""" + return text.split() + + +@pytest.fixture +def fake_llm_chain() -> LLMChain: + """Fake LLM chain for testing purposes.""" + prompt = PromptTemplate(input_variables=["bar"], template="This is a {bar}:") + return LLMChain(prompt=prompt, llm=FakeLLM(), output_key="text1") + + +@patch( + "langchain_community.llms.loading.get_type_to_cls_dict", + lambda: {"fake": lambda: FakeLLM}, +) +def test_serialization(fake_llm_chain: LLMChain) -> None: + """Test serialization.""" + from langchain.chains.loading import load_chain + + with TemporaryDirectory() as temp_dir: + file = temp_dir + "/llm.json" + fake_llm_chain.save(file) + loaded_chain = load_chain(file) + assert loaded_chain == fake_llm_chain + + +def test_missing_inputs(fake_llm_chain: LLMChain) -> None: + """Test error is raised if inputs are missing.""" + with pytest.raises(ValueError): + fake_llm_chain({"foo": "bar"}) + + +def test_valid_call(fake_llm_chain: LLMChain) -> None: + """Test valid call of LLM chain.""" + output = fake_llm_chain({"bar": "baz"}) + assert output == {"bar": "baz", "text1": "foo"} + + # Test with stop words. + output = fake_llm_chain({"bar": "baz", "stop": ["foo"]}) + # Response should be `bar` now. + assert output == {"bar": "baz", "stop": ["foo"], "text1": "bar"} + + +def test_predict_method(fake_llm_chain: LLMChain) -> None: + """Test predict method works.""" + output = fake_llm_chain.predict(bar="baz") + assert output == "foo" + + +def test_predict_and_parse() -> None: + """Test parsing ability.""" + prompt = PromptTemplate( + input_variables=["foo"], template="{foo}", output_parser=FakeOutputParser() + ) + llm = FakeLLM(queries={"foo": "foo bar"}) + chain = LLMChain(prompt=prompt, llm=llm) + output = chain.predict_and_parse(foo="foo") + assert output == ["foo", "bar"]