forked from Archives/langchain
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
41 lines
1.2 KiB
Python
41 lines
1.2 KiB
Python
"""Test transform chain."""
|
|
from typing import Dict
|
|
|
|
import pytest
|
|
|
|
from langchain.chains.transform import TransformChain
|
|
|
|
|
|
def dummy_transform(inputs: Dict[str, str]) -> Dict[str, str]:
|
|
"""Transform a dummy input for tests."""
|
|
outputs = inputs
|
|
outputs["greeting"] = f"{inputs['first_name']} {inputs['last_name']} says hello"
|
|
del outputs["first_name"]
|
|
del outputs["last_name"]
|
|
return outputs
|
|
|
|
|
|
def test_tranform_chain() -> None:
|
|
"""Test basic transform chain."""
|
|
transform_chain = TransformChain(
|
|
input_variables=["first_name", "last_name"],
|
|
output_variables=["greeting"],
|
|
transform=dummy_transform,
|
|
)
|
|
input_dict = {"first_name": "Leroy", "last_name": "Jenkins"}
|
|
response = transform_chain(input_dict)
|
|
expected_response = {"greeting": "Leroy Jenkins says hello"}
|
|
assert response == expected_response
|
|
|
|
|
|
def test_transform_chain_bad_inputs() -> None:
|
|
"""Test basic transform chain."""
|
|
transform_chain = TransformChain(
|
|
input_variables=["first_name", "last_name"],
|
|
output_variables=["greeting"],
|
|
transform=dummy_transform,
|
|
)
|
|
input_dict = {"name": "Leroy", "last_name": "Jenkins"}
|
|
with pytest.raises(ValueError):
|
|
_ = transform_chain(input_dict)
|