test_run_llm_or_chain_with_input_mapper() — langchain Function Reference

Source Code

libs/langchain/tests/unit_tests/smith/evaluation/test_runner_utils.py lines 161–211

def test_run_llm_or_chain_with_input_mapper() -> None:
    example = Example(
        id=uuid.uuid4(),
        created_at=_CREATED_AT,
        inputs={"the wrong input": "1", "another key": "2"},
        outputs={"output": "2"},
        dataset_id=str(uuid.uuid4()),
    )

    def run_val(inputs: dict) -> dict:
        assert "the right input" in inputs
        return {"output": "2"}

    mock_chain = TransformChain(
        input_variables=["the right input"],
        output_variables=["output"],
        transform=run_val,
    )

    def input_mapper(inputs: dict) -> dict:
        assert "the wrong input" in inputs
        return {"the right input": inputs["the wrong input"]}

    result = _run_llm_or_chain(
        example,
        {"callbacks": [], "tags": []},
        llm_or_chain_factory=lambda: mock_chain,
        input_mapper=input_mapper,
    )
    assert result == {"output": "2", "the right input": "1"}
    bad_result = _run_llm_or_chain(
        example,
        {"callbacks": [], "tags": []},
        llm_or_chain_factory=lambda: mock_chain,
    )
    assert "Error" in bad_result

    # Try with LLM
    def llm_input_mapper(inputs: dict) -> str:
        assert "the wrong input" in inputs
        return "the right input"

    mock_llm = FakeLLM(queries={"the right input": "somenumber"})
    llm_result = _run_llm_or_chain(
        example,
        {"callbacks": [], "tags": []},
        llm_or_chain_factory=mock_llm,
        input_mapper=llm_input_mapper,
    )
    assert isinstance(llm_result, str)
    assert llm_result == "somenumber"