Merge pull request #2354 from langchain-ai/nc/6nov/cached-tasks-output-timing

lib: For cached tasks, emit output events after task events
This commit is contained in:
Nuno Campos
2024-11-06 08:29:38 -08:00
committed by GitHub
2 changed files with 8 additions and 7 deletions
+5 -4
View File
@@ -395,10 +395,6 @@ class PregelLoop(LoopProtocol):
self.tasks[tid] = task._replace(scheduled=True)
else:
task.writes.append((k, v))
# print output for any tasks we applied previous writes to
for task in self.tasks.values():
if task.writes:
self._output_writes(task.id, task.writes, cached=True)
elif not self.skip_done_tasks:
# "not skip_done_tasks" only applies to first tick after resuming
self.skip_done_tasks = True
@@ -424,6 +420,11 @@ class PregelLoop(LoopProtocol):
if self.debug:
print_step_tasks(self.step, list(self.tasks.values()))
# print output for any tasks we applied previous writes to
for task in self.tasks.values():
if task.writes:
self._output_writes(task.id, task.writes, cached=True)
return True
# private
+3 -3
View File
@@ -8660,14 +8660,14 @@ async def test_stream_subgraphs_during_execution(checkpointer_name: str) -> None
),
(FloatBetween(0.2, 0.3), ((), {"outer_1": {"my_key": " and parallel"}})),
(
FloatBetween(0.5, 0.6),
FloatBetween(0.5, 0.7),
(
(AnyStr("inner:"),),
{"inner_2": {"my_key": " and there", "my_other_key": "got here"}},
),
),
(FloatBetween(0.5, 0.6), ((), {"inner": {"my_key": "got here and there"}})),
(FloatBetween(0.5, 0.6), ((), {"outer_2": {"my_key": " and back again"}})),
(FloatBetween(0.5, 0.7), ((), {"inner": {"my_key": "got here and there"}})),
(FloatBetween(0.5, 0.7), ((), {"outer_2": {"my_key": " and back again"}})),
]