Extend test

This commit is contained in:
Nuno Campos 2023-08-18 16:22:01 +01:00
parent 9777c2801d
commit fdcd50aab4

View File

@ -1113,10 +1113,16 @@ def test_each(snapshot: SnapshotAssertion) -> None:
+ "{question}"
)
first_llm = FakeStreamingListLLM(responses=["first item, second item, third item"])
parser = FakeSplitIntoListParser()
second_llm = FakeStreamingListLLM(responses=["this", "is", "a", "test"])
chain = prompt | first_llm | FakeSplitIntoListParser() | second_llm.map()
chain = prompt | first_llm | parser | second_llm.map()
assert dumps(chain, pretty=True) == snapshot
output = chain.invoke({"question": "What up"})
assert output == ["this", "is", "a"]
assert (parser | second_llm.map()).invoke("first item, second item") == [
"test",
"this",
]