Remove None default for missing keys when generating input for each node

This commit is contained in:
Nuno Campos
2024-08-30 13:59:36 -07:00
parent 35807068ef
commit c5c10a41e5
3 changed files with 20 additions and 52 deletions
+12 -12
View File
@@ -384,7 +384,7 @@ def prepare_next_tasks(
try:
val = next(
_proc_input(
step, name, proc, managed, channels, for_execution=for_execution
step, proc, managed, channels, for_execution=for_execution
)
)
except StopIteration:
@@ -474,7 +474,6 @@ def prepare_next_tasks(
def _proc_input(
step: int,
name: str,
proc: PregelNode,
managed: ManagedValueMapping,
channels: Mapping[str, BaseChannel],
@@ -485,16 +484,17 @@ def _proc_input(
# then invoke the process with the values of all non-empty channels
if isinstance(proc.channels, dict):
try:
val: dict = {
k: read_channel(
channels,
chan,
catch=chan not in proc.triggers,
)
if chan in channels
else managed[k](step)
for k, chan in proc.channels.items()
}
val: dict[str, Any] = {}
for k, chan in proc.channels.items():
if chan in proc.triggers:
val[k] = read_channel(channels, chan, catch=False)
elif chan in channels:
try:
val[k] = read_channel(channels, chan, catch=False)
except EmptyChannelError:
continue
else:
val[k] = managed[k](step)
except EmptyChannelError:
return
elif isinstance(proc.channels, list):
+4 -20
View File
@@ -294,7 +294,6 @@ def test_node_schemas_custom_output() -> None:
def node_b(state: StateForB):
assert state == {
"bye": "world",
"now": None,
}
return {
"now": 123,
@@ -6364,11 +6363,7 @@ def test_in_one_fan_out_out_one_graph_state() -> None:
"payload": {
"id": "592f3430-c17c-5d1c-831f-fecebb2c05bf",
"name": "rewrite_query",
"input": {
"query": "what is weather in sf",
"answer": None,
"docs": [],
},
"input": {"query": "what is weather in sf", "docs": []},
"triggers": ["start:rewrite_query"],
},
},
@@ -6399,11 +6394,7 @@ def test_in_one_fan_out_out_one_graph_state() -> None:
"payload": {
"id": "7db5e9d8-e132-5079-ab99-ced15e67d48b",
"name": "retriever_one",
"input": {
"query": "query: what is weather in sf",
"answer": None,
"docs": [],
},
"input": {"query": "query: what is weather in sf", "docs": []},
"triggers": ["rewrite_query"],
},
},
@@ -6417,11 +6408,7 @@ def test_in_one_fan_out_out_one_graph_state() -> None:
"payload": {
"id": "96965ed0-2c10-52a1-86eb-081ba6de73b2",
"name": "retriever_two",
"input": {
"query": "query: what is weather in sf",
"answer": None,
"docs": [],
},
"input": {"query": "query: what is weather in sf", "docs": []},
"triggers": ["rewrite_query"],
},
},
@@ -6482,7 +6469,6 @@ def test_in_one_fan_out_out_one_graph_state() -> None:
"name": "qa",
"input": {
"query": "query: what is weather in sf",
"answer": None,
"docs": ["doc1", "doc2", "doc3", "doc4"],
},
"triggers": ["retriever_one", "retriever_two"],
@@ -8906,9 +8892,7 @@ def test_nested_graph_state(
},
metadata={
"source": "input",
"writes": {
"__start__": {"my_key": "hi my value", "other_parent_key": None}
},
"writes": {"__start__": {"my_key": "hi my value"}},
"step": -1,
"parents": {"": AnyStr()},
},
+4 -20
View File
@@ -570,7 +570,6 @@ async def test_node_schemas_custom_output() -> None:
async def node_b(state: StateForB):
assert state == {
"bye": "world",
"now": None,
}
return {
"now": 123,
@@ -4987,11 +4986,7 @@ async def test_in_one_fan_out_out_one_graph_state() -> None:
"payload": {
"id": "592f3430-c17c-5d1c-831f-fecebb2c05bf",
"name": "rewrite_query",
"input": {
"query": "what is weather in sf",
"answer": None,
"docs": [],
},
"input": {"query": "what is weather in sf", "docs": []},
"triggers": ["start:rewrite_query"],
},
},
@@ -5022,11 +5017,7 @@ async def test_in_one_fan_out_out_one_graph_state() -> None:
"payload": {
"id": "7db5e9d8-e132-5079-ab99-ced15e67d48b",
"name": "retriever_one",
"input": {
"query": "query: what is weather in sf",
"answer": None,
"docs": [],
},
"input": {"query": "query: what is weather in sf", "docs": []},
"triggers": ["rewrite_query"],
},
},
@@ -5040,11 +5031,7 @@ async def test_in_one_fan_out_out_one_graph_state() -> None:
"payload": {
"id": "96965ed0-2c10-52a1-86eb-081ba6de73b2",
"name": "retriever_two",
"input": {
"query": "query: what is weather in sf",
"answer": None,
"docs": [],
},
"input": {"query": "query: what is weather in sf", "docs": []},
"triggers": ["rewrite_query"],
},
},
@@ -5105,7 +5092,6 @@ async def test_in_one_fan_out_out_one_graph_state() -> None:
"name": "qa",
"input": {
"query": "query: what is weather in sf",
"answer": None,
"docs": ["doc1", "doc2", "doc3", "doc4"],
},
"triggers": ["retriever_one", "retriever_two"],
@@ -7403,9 +7389,7 @@ async def test_nested_graph_state(
},
metadata={
"source": "input",
"writes": {
"__start__": {"my_key": "hi my value", "other_parent_key": None}
},
"writes": {"__start__": {"my_key": "hi my value"}},
"step": -1,
"parents": {"": AnyStr()},
},