Add optional direction parameter to PipelineTask.queue_frame() and queue_frames()

Allow pushing frames upstream through the pipeline by passing
FrameDirection.UPSTREAM. Downstream frames use the existing push queue,
while upstream frames are pushed directly from the pipeline sink.
This commit is contained in:
Aleix Conchillo Flaqué
2026-02-28 17:11:07 -08:00
parent 17205c1647
commit f37fd39cdb
3 changed files with 84 additions and 10 deletions

View File

@@ -292,6 +292,63 @@ class TestPipelineTask(unittest.IsolatedAsyncioTestCase):
assert upstream_received
assert downstream_received
async def test_task_queue_frame_upstream(self):
upstream_received = False
pipeline = Pipeline([IdentityFilter()])
task = PipelineTask(pipeline, cancel_on_idle_timeout=False)
task.set_reached_upstream_filter((TextFrame,))
@task.event_handler("on_frame_reached_upstream")
async def on_frame_reached_upstream(task, frame):
nonlocal upstream_received
if isinstance(frame, TextFrame) and frame.text == "Hello Upstream!":
upstream_received = True
@task.event_handler("on_pipeline_started")
async def on_pipeline_started(task, frame):
await task.queue_frame(TextFrame(text="Hello Upstream!"), FrameDirection.UPSTREAM)
try:
await asyncio.wait_for(
task.run(PipelineTaskParams(loop=asyncio.get_event_loop())),
timeout=1.0,
)
except asyncio.TimeoutError:
pass
assert upstream_received
async def test_task_queue_frames_upstream(self):
upstream_texts = []
pipeline = Pipeline([IdentityFilter()])
task = PipelineTask(pipeline, cancel_on_idle_timeout=False)
task.set_reached_upstream_filter((TextFrame,))
@task.event_handler("on_frame_reached_upstream")
async def on_frame_reached_upstream(task, frame):
if isinstance(frame, TextFrame):
upstream_texts.append(frame.text)
@task.event_handler("on_pipeline_started")
async def on_pipeline_started(task, frame):
await task.queue_frames(
[TextFrame(text="First"), TextFrame(text="Second")],
FrameDirection.UPSTREAM,
)
try:
await asyncio.wait_for(
task.run(PipelineTaskParams(loop=asyncio.get_event_loop())),
timeout=1.0,
)
except asyncio.TimeoutError:
pass
assert "First" in upstream_texts
assert "Second" in upstream_texts
async def test_task_heartbeats(self):
heartbeats_counter = 0