core[patch]: Add unit test for RunnableGenerator for eventstream v2 (#21990)

No unit tests with runnable generator
This commit is contained in:
Eugene Yurtsev 2024-05-21 14:29:15 -04:00 committed by GitHub
parent fb6108c8f5
commit ded53297e0
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -35,6 +35,7 @@ from langchain_core.runnables import (
ConfigurableField, ConfigurableField,
Runnable, Runnable,
RunnableConfig, RunnableConfig,
RunnableGenerator,
RunnableLambda, RunnableLambda,
ensure_config, ensure_config,
) )
@ -1830,3 +1831,48 @@ async def test_astream_events_from_custom_runnable() -> None:
"tags": [], "tags": [],
}, },
] ]
async def test_runnable_generator() -> None:
"""Test async events from sync lambda."""
async def generator(inputs: AsyncIterator[str]) -> AsyncIterator[str]:
yield "1"
yield "2"
runnable: Runnable[str, str] = RunnableGenerator(transform=generator)
events = await _collect_events(runnable.astream_events("hello", version="v2"))
assert events == [
{
"data": {"input": "hello"},
"event": "on_chain_start",
"metadata": {},
"name": "generator",
"run_id": "",
"tags": [],
},
{
"data": {"chunk": "1"},
"event": "on_chain_stream",
"metadata": {},
"name": "generator",
"run_id": "",
"tags": [],
},
{
"data": {"chunk": "2"},
"event": "on_chain_stream",
"metadata": {},
"name": "generator",
"run_id": "",
"tags": [],
},
{
"data": {"output": "12"},
"event": "on_chain_end",
"metadata": {},
"name": "generator",
"run_id": "",
"tags": [],
},
]