Merge pull request #310 from angeligareta/add_condition_nodes_as_optional

Make condition nodes optional
This commit is contained in:
Nuno Campos
2024-04-17 09:21:49 -07:00
committed by GitHub
3 changed files with 142 additions and 7 deletions
+14 -5
View File
@@ -364,7 +364,11 @@ class CompiledGraph(Pregel):
self.nodes[end].channels.append(channel_name)
def get_graph(
self, config: Optional[RunnableConfig] = None, *, xray: bool = False
self,
config: Optional[RunnableConfig] = None,
*,
xray: bool = False,
add_condition_nodes: bool = True,
) -> RunnableGraph:
"""Returns a drawable representation of the computation graph."""
graph = RunnableGraph()
@@ -400,13 +404,18 @@ class CompiledGraph(Pregel):
graph.add_edge(start_nodes[start], end_nodes[end])
for start, branches in self.graph.branches.items():
for name, branch in branches.items():
cond = graph.add_node(branch.condition, name)
graph.add_edge(start_nodes[start], cond)
ends = branch.ends or {
**{k: k for k in self.graph.nodes},
END: END,
}
for label, end in ends.items():
graph.add_edge(cond, end_nodes[end], label)
if add_condition_nodes is True:
cond = graph.add_node(branch.condition, name)
graph.add_edge(start_nodes[start], cond)
for label, end in ends.items():
graph.add_edge(cond, end_nodes[end], label)
else:
for label, end in ends.items():
graph.add_edge(start_nodes[start], end_nodes[end], label)
return graph
+126 -2
View File
@@ -410,6 +410,25 @@
'''
# ---
# name: test_conditional_graph[end_of_run].2
'''
+-----------+
| __start__ |
+-----------+
*
*
*
+-------+
| agent |
+-------+
* *
** **
* *
+-------+ +---------+
| tools | | __end__ |
+-------+ +---------+
'''
# ---
# name: test_conditional_graph[end_of_run].3
'''
{
"nodes": [
@@ -563,7 +582,7 @@
}
'''
# ---
# name: test_conditional_graph[end_of_run].3
# name: test_conditional_graph[end_of_run].4
'''
+-----------+
| __start__ |
@@ -612,6 +631,49 @@
+---------+ +-------+
'''
# ---
# name: test_conditional_graph[end_of_run].5
'''
+-----------+
| __start__ |
+-----------+
*
*
*
+------------------------------+
| Parallel<agent_outcome>Input |
+------------------------------+
***** * *****
****** * ****
*** * *****
+----------------+ * ***
| PromptTemplate | * *
+----------------+ * *
* * *
* * *
* * *
+----------------------+ * *
| FakeStreamingListLLM | * *
+----------------------+ * *
* * *
* * *
* * *
+----------------------+ +-------------+ *
| Lambda(agent_parser) | | Passthrough | *
+----------------------+ +-------------+ *
*** *** *
** ** *
** ** *
+-------------------------------+ *
| Parallel<agent_outcome>Output | *
+-------------------------------+***** *
* ********* *
* ******** *
* ***** *
+---------+ +-------+
| __end__ | | tools |
+---------+ +-------+
'''
# ---
# name: test_conditional_graph[end_of_step]
'''
{
@@ -719,6 +781,25 @@
'''
# ---
# name: test_conditional_graph[end_of_step].2
'''
+-----------+
| __start__ |
+-----------+
*
*
*
+-------+
| agent |
+-------+
* *
** **
* *
+-------+ +---------+
| tools | | __end__ |
+-------+ +---------+
'''
# ---
# name: test_conditional_graph[end_of_step].3
'''
{
"nodes": [
@@ -872,7 +953,7 @@
}
'''
# ---
# name: test_conditional_graph[end_of_step].3
# name: test_conditional_graph[end_of_step].4
'''
+-----------+
| __start__ |
@@ -921,6 +1002,49 @@
+---------+ +-------+
'''
# ---
# name: test_conditional_graph[end_of_step].5
'''
+-----------+
| __start__ |
+-----------+
*
*
*
+------------------------------+
| Parallel<agent_outcome>Input |
+------------------------------+
***** * *****
****** * ****
*** * *****
+----------------+ * ***
| PromptTemplate | * *
+----------------+ * *
* * *
* * *
* * *
+----------------------+ * *
| FakeStreamingListLLM | * *
+----------------------+ * *
* * *
* * *
* * *
+----------------------+ +-------------+ *
| Lambda(agent_parser) | | Passthrough | *
+----------------------+ +-------------+ *
*** *** *
** ** *
** ** *
+-------------------------------+ *
| Parallel<agent_outcome>Output | *
+-------------------------------+***** *
* ********* *
* ******** *
* ***** *
+---------+ +-------+
| __end__ | | tools |
+---------+ +-------+
'''
# ---
# name: test_conditional_graph_state[end_of_run]
'{"title": "LangGraphInput", "$ref": "#/definitions/AgentState", "definitions": {"AgentAction": {"title": "AgentAction", "description": "A full description of an action for an ActionAgent to execute.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "The final return value of an ActionAgent.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}, "AgentState": {"title": "AgentState", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}}}}'
# ---
+2
View File
@@ -796,8 +796,10 @@ def test_conditional_graph(
assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
assert app.get_graph().draw_ascii() == snapshot
assert app.get_graph(add_condition_nodes=False).draw_ascii() == snapshot
assert json.dumps(app.get_graph(xray=True).to_json(), indent=2) == snapshot
assert app.get_graph(xray=True).draw_ascii() == snapshot
assert app.get_graph(xray=True, add_condition_nodes=False).draw_ascii() == snapshot
assert app.invoke({"input": "what is weather in sf"}) == {
"input": "what is weather in sf",