WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-13 13:17:52 +02:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'main' into wfh/gc_test_all

Browse Source
This commit is contained in:
William Fu-Hinthorn
2025-03-19 15:32:00 -07:00
parent ee66e9165c 03fc695d60
commit f42112b223
5 changed files with 1212 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/poetry.lock
Generated
+164 -6
View File
File diff suppressed because it is too large Load Diff
libs/langgraph/tests/__snapshots__/test_pregel.ambr
+968
View File
@@ -1370,6 +1370,426 @@
'type': 'object',
})
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_input[memory]
'''
graph TD;
__start__ --> rewrite_query;
analyzer_one --> retriever_one;
qa --> __end__;
retriever_one --> qa;
retriever_two --> qa;
rewrite_query --> analyzer_one;
rewrite_query -.-> retriever_two;
'''
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_input[memory].1
dict({
'$defs': dict({
'InnerObject': dict({
'properties': dict({
'yo': dict({
'title': 'Yo',
'type': 'integer',
}),
}),
'required': list([
'yo',
]),
'title': 'InnerObject',
'type': 'object',
}),
}),
'properties': dict({
'inner': dict({
'$ref': '#/$defs/InnerObject',
}),
'query': dict({
'title': 'Query',
'type': 'string',
}),
}),
'required': list([
'query',
'inner',
]),
'title': 'Input',
'type': 'object',
})
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_input[memory].2
dict({
'properties': dict({
'answer': dict({
'title': 'Answer',
'type': 'string',
}),
'docs': dict({
'items': dict({
'type': 'string',
}),
'title': 'Docs',
'type': 'array',
}),
}),
'required': list([
'answer',
'docs',
]),
'title': 'Output',
'type': 'object',
})
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_input[postgres]
'''
graph TD;
__start__ --> rewrite_query;
analyzer_one --> retriever_one;
qa --> __end__;
retriever_one --> qa;
retriever_two --> qa;
rewrite_query --> analyzer_one;
rewrite_query -.-> retriever_two;
'''
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_input[postgres].1
dict({
'$defs': dict({
'InnerObject': dict({
'properties': dict({
'yo': dict({
'title': 'Yo',
'type': 'integer',
}),
}),
'required': list([
'yo',
]),
'title': 'InnerObject',
'type': 'object',
}),
}),
'properties': dict({
'inner': dict({
'$ref': '#/$defs/InnerObject',
}),
'query': dict({
'title': 'Query',
'type': 'string',
}),
}),
'required': list([
'query',
'inner',
]),
'title': 'Input',
'type': 'object',
})
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_input[postgres].2
dict({
'properties': dict({
'answer': dict({
'title': 'Answer',
'type': 'string',
}),
'docs': dict({
'items': dict({
'type': 'string',
}),
'title': 'Docs',
'type': 'array',
}),
}),
'required': list([
'answer',
'docs',
]),
'title': 'Output',
'type': 'object',
})
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_input[postgres_pipe]
'''
graph TD;
__start__ --> rewrite_query;
analyzer_one --> retriever_one;
qa --> __end__;
retriever_one --> qa;
retriever_two --> qa;
rewrite_query --> analyzer_one;
rewrite_query -.-> retriever_two;
'''
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_input[postgres_pipe].1
dict({
'$defs': dict({
'InnerObject': dict({
'properties': dict({
'yo': dict({
'title': 'Yo',
'type': 'integer',
}),
}),
'required': list([
'yo',
]),
'title': 'InnerObject',
'type': 'object',
}),
}),
'properties': dict({
'inner': dict({
'$ref': '#/$defs/InnerObject',
}),
'query': dict({
'title': 'Query',
'type': 'string',
}),
}),
'required': list([
'query',
'inner',
]),
'title': 'Input',
'type': 'object',
})
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_input[postgres_pipe].2
dict({
'properties': dict({
'answer': dict({
'title': 'Answer',
'type': 'string',
}),
'docs': dict({
'items': dict({
'type': 'string',
}),
'title': 'Docs',
'type': 'array',
}),
}),
'required': list([
'answer',
'docs',
]),
'title': 'Output',
'type': 'object',
})
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_input[postgres_pool]
'''
graph TD;
__start__ --> rewrite_query;
analyzer_one --> retriever_one;
qa --> __end__;
retriever_one --> qa;
retriever_two --> qa;
rewrite_query --> analyzer_one;
rewrite_query -.-> retriever_two;
'''
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_input[postgres_pool].1
dict({
'$defs': dict({
'InnerObject': dict({
'properties': dict({
'yo': dict({
'title': 'Yo',
'type': 'integer',
}),
}),
'required': list([
'yo',
]),
'title': 'InnerObject',
'type': 'object',
}),
}),
'properties': dict({
'inner': dict({
'$ref': '#/$defs/InnerObject',
}),
'query': dict({
'title': 'Query',
'type': 'string',
}),
}),
'required': list([
'query',
'inner',
]),
'title': 'Input',
'type': 'object',
})
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_input[postgres_pool].2
dict({
'properties': dict({
'answer': dict({
'title': 'Answer',
'type': 'string',
}),
'docs': dict({
'items': dict({
'type': 'string',
}),
'title': 'Docs',
'type': 'array',
}),
}),
'required': list([
'answer',
'docs',
]),
'title': 'Output',
'type': 'object',
})
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_input[postgres_shallow]
'''
graph TD;
__start__ --> rewrite_query;
analyzer_one --> retriever_one;
qa --> __end__;
retriever_one --> qa;
retriever_two --> qa;
rewrite_query --> analyzer_one;
rewrite_query -.-> retriever_two;
'''
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_input[postgres_shallow].1
dict({
'$defs': dict({
'InnerObject': dict({
'properties': dict({
'yo': dict({
'title': 'Yo',
'type': 'integer',
}),
}),
'required': list([
'yo',
]),
'title': 'InnerObject',
'type': 'object',
}),
}),
'properties': dict({
'inner': dict({
'$ref': '#/$defs/InnerObject',
}),
'query': dict({
'title': 'Query',
'type': 'string',
}),
}),
'required': list([
'query',
'inner',
]),
'title': 'Input',
'type': 'object',
})
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_input[postgres_shallow].2
dict({
'properties': dict({
'answer': dict({
'title': 'Answer',
'type': 'string',
}),
'docs': dict({
'items': dict({
'type': 'string',
}),
'title': 'Docs',
'type': 'array',
}),
}),
'required': list([
'answer',
'docs',
]),
'title': 'Output',
'type': 'object',
})
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_input[sqlite]
'''
graph TD;
__start__ --> rewrite_query;
analyzer_one --> retriever_one;
qa --> __end__;
retriever_one --> qa;
retriever_two --> qa;
rewrite_query --> analyzer_one;
rewrite_query -.-> retriever_two;
'''
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_input[sqlite].1
dict({
'$defs': dict({
'InnerObject': dict({
'properties': dict({
'yo': dict({
'title': 'Yo',
'type': 'integer',
}),
}),
'required': list([
'yo',
]),
'title': 'InnerObject',
'type': 'object',
}),
}),
'properties': dict({
'inner': dict({
'$ref': '#/$defs/InnerObject',
}),
'query': dict({
'title': 'Query',
'type': 'string',
}),
}),
'required': list([
'query',
'inner',
]),
'title': 'Input',
'type': 'object',
})
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_input[sqlite].2
dict({
'properties': dict({
'answer': dict({
'title': 'Answer',
'type': 'string',
}),
'docs': dict({
'items': dict({
'type': 'string',
}),
'title': 'Docs',
'type': 'array',
}),
}),
'required': list([
'answer',
'docs',
]),
'title': 'Output',
'type': 'object',
})
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[memory]
'''
graph TD;
@@ -1461,6 +1881,240 @@
'''
# ---
# name: test_multiple_sinks_subgraphs
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([<p>__start__</p>]):::first
uno(uno)
dos(dos)
__start__ --> uno;
uno -.-> dos;
uno -.-> subgraph_one;
subgraph subgraph
subgraph_one(one)
subgraph_two(two)
subgraph_three(three)
subgraph_one -.-> subgraph_two;
subgraph_one -.-> subgraph_three;
end
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_nested_graph
'''
graph TD;
__start__ --> inner;
inner --> side;
side --> __end__;
'''
# ---
# name: test_nested_graph.1
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([<p>__start__</p>]):::first
side(side)
__end__([<p>__end__</p>]):::last
__start__ --> inner_up;
inner_up --> side;
side --> __end__;
subgraph inner
inner_up(up)
end
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_nested_graph_xray
dict({
'edges': list([
dict({
'conditional': True,
'source': 'tool_two:__start__',
'target': 'tool_two:tool_two_slow',
}),
dict({
'source': 'tool_two:tool_two_slow',
'target': 'tool_two:__end__',
}),
dict({
'conditional': True,
'source': 'tool_two:__start__',
'target': 'tool_two:tool_two_fast',
}),
dict({
'source': 'tool_two:tool_two_fast',
'target': 'tool_two:__end__',
}),
dict({
'conditional': True,
'source': '__start__',
'target': 'tool_one',
}),
dict({
'source': 'tool_one',
'target': '__end__',
}),
dict({
'conditional': True,
'source': '__start__',
'target': 'tool_two:__start__',
}),
dict({
'source': 'tool_two:__end__',
'target': '__end__',
}),
dict({
'conditional': True,
'source': '__start__',
'target': 'tool_three',
}),
dict({
'source': 'tool_three',
'target': '__end__',
}),
]),
'nodes': list([
dict({
'data': '__start__',
'id': '__start__',
'type': 'schema',
}),
dict({
'data': dict({
'id': list([
'langgraph',
'utils',
'runnable',
'RunnableCallable',
]),
'name': 'tool_one',
}),
'id': 'tool_one',
'type': 'runnable',
}),
dict({
'data': 'tool_two:__start__',
'id': 'tool_two:__start__',
'type': 'schema',
}),
dict({
'data': dict({
'id': list([
'langgraph',
'utils',
'runnable',
'RunnableCallable',
]),
'name': 'tool_two:tool_two_slow',
}),
'id': 'tool_two:tool_two_slow',
'type': 'runnable',
}),
dict({
'data': dict({
'id': list([
'langgraph',
'utils',
'runnable',
'RunnableCallable',
]),
'name': 'tool_two:tool_two_fast',
}),
'id': 'tool_two:tool_two_fast',
'type': 'runnable',
}),
dict({
'data': 'tool_two:__end__',
'id': 'tool_two:__end__',
'type': 'schema',
}),
dict({
'data': dict({
'id': list([
'langgraph',
'utils',
'runnable',
'RunnableCallable',
]),
'name': 'tool_three',
}),
'id': 'tool_three',
'type': 'runnable',
}),
dict({
'data': '__end__',
'id': '__end__',
'type': 'schema',
}),
]),
})
# ---
# name: test_nested_graph_xray.1
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([<p>__start__</p>]):::first
tool_one(tool_one)
tool_three(tool_three)
__end__([<p>__end__</p>]):::last
__start__ -.-> tool_one;
tool_one --> __end__;
__start__ -.-> tool_two___start__;
tool_two___end__ --> __end__;
__start__ -.-> tool_three;
tool_three --> __end__;
subgraph tool_two
tool_two___start__(<p>__start__</p>)
tool_two_tool_two_slow(tool_two_slow)
tool_two_tool_two_fast(tool_two_fast)
tool_two___end__(<p>__end__</p>)
tool_two___start__ -.-> tool_two_tool_two_slow;
tool_two_tool_two_slow --> tool_two___end__;
tool_two___start__ -.-> tool_two_tool_two_fast;
tool_two_tool_two_fast --> tool_two___end__;
end
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_repeat_condition
'''
graph TD;
__start__ --> Researcher;
Researcher -. &nbsp;continue&nbsp; .-> Chart_Generator;
Researcher -. &nbsp;call_tool&nbsp; .-> Call_Tool;
Researcher -. &nbsp;end&nbsp; .-> __end__;
Chart_Generator -. &nbsp;continue&nbsp; .-> Researcher;
Chart_Generator -. &nbsp;call_tool&nbsp; .-> Call_Tool;
Chart_Generator -. &nbsp;end&nbsp; .-> __end__;
Call_Tool -.-> Researcher;
Call_Tool -.-> Chart_Generator;
Researcher -. &nbsp;redo&nbsp; .-> Researcher;
'''
# ---
# name: test_simple_multi_edge
'''
graph TD;
__start__ --> up;
down --> __end__;
side --> down;
up --> down;
up --> other;
up --> side;
'''
# ---
# name: test_state_graph_w_config_inherited_state_keys
'{"$defs": {"Config": {"properties": {"tools": {"items": {"type": "string"}, "title": "Tools", "type": "array"}}, "title": "Config", "type": "object"}}, "properties": {"configurable": {"$ref": "#/$defs/Config", "default": null}}, "title": "LangGraphConfig", "type": "object"}'
# ---
@@ -1470,3 +2124,317 @@
# name: test_state_graph_w_config_inherited_state_keys.2
'{"$defs": {"AgentAction": {"description": "Represents a request to execute an action by an agent.\\n\\nThe action consists of the name of the tool to execute and the input to pass\\nto the tool. The log is used to pass along extra information about the action.", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"anyOf": [{"type": "string"}, {"type": "object"}], "title": "Tool Input"}, "log": {"title": "Log", "type": "string"}, "type": {"const": "AgentAction", "default": "AgentAction", "enum": ["AgentAction"], "title": "Type", "type": "string"}}, "required": ["tool", "tool_input", "log"], "title": "AgentAction", "type": "object"}, "AgentFinish": {"description": "Final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"const": "AgentFinish", "default": "AgentFinish", "enum": ["AgentFinish"], "title": "Type", "type": "string"}}, "required": ["return_values", "log"], "title": "AgentFinish", "type": "object"}}, "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"anyOf": [{"$ref": "#/$defs/AgentAction"}, {"$ref": "#/$defs/AgentFinish"}, {"type": "null"}], "default": null, "title": "Agent Outcome"}, "intermediate_steps": {"default": null, "items": {"maxItems": 2, "minItems": 2, "prefixItems": [{"$ref": "#/$defs/AgentAction"}, {"type": "string"}], "type": "array"}, "title": "Intermediate Steps", "type": "array"}}, "required": ["input"], "title": "LangGraphOutput", "type": "object"}'
# ---
# name: test_xray_bool
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([<p>__start__</p>]):::first
gp_one(gp_one)
__end__([<p>__end__</p>]):::last
__start__ --> gp_one;
gp_two___end__ --> gp_one;
gp_one -. &nbsp;0&nbsp; .-> gp_two___start__;
gp_one -. &nbsp;1&nbsp; .-> __end__;
subgraph gp_two
gp_two___start__(<p>__start__</p>)
gp_two_p_one(p_one)
gp_two___end__(<p>__end__</p>)
gp_two___start__ --> gp_two_p_one;
gp_two_p_two___end__ --> gp_two_p_one;
gp_two_p_one -. &nbsp;0&nbsp; .-> gp_two_p_two___start__;
gp_two_p_one -. &nbsp;1&nbsp; .-> gp_two___end__;
subgraph p_two
gp_two_p_two___start__(<p>__start__</p>)
gp_two_p_two_c_one(c_one)
gp_two_p_two_c_two(c_two)
gp_two_p_two___end__(<p>__end__</p>)
gp_two_p_two___start__ --> gp_two_p_two_c_one;
gp_two_p_two_c_two --> gp_two_p_two_c_one;
gp_two_p_two_c_one -. &nbsp;0&nbsp; .-> gp_two_p_two_c_two;
gp_two_p_two_c_one -. &nbsp;1&nbsp; .-> gp_two_p_two___end__;
end
end
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_xray_issue
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([<p>__start__</p>]):::first
p_one(p_one)
__end__([<p>__end__</p>]):::last
__start__ --> p_one;
p_two___end__ --> p_one;
p_one -. &nbsp;0&nbsp; .-> p_two___start__;
p_one -. &nbsp;1&nbsp; .-> __end__;
subgraph p_two
p_two___start__(<p>__start__</p>)
p_two_c_one(c_one)
p_two_c_two(c_two)
p_two___end__(<p>__end__</p>)
p_two___start__ --> p_two_c_one;
p_two_c_two --> p_two_c_one;
p_two_c_one -. &nbsp;0&nbsp; .-> p_two_c_two;
p_two_c_one -. &nbsp;1&nbsp; .-> p_two___end__;
end
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_xray_lance
dict({
'edges': list([
dict({
'source': '__start__',
'target': 'ask_question',
}),
dict({
'source': 'ask_question',
'target': 'answer_question',
}),
dict({
'conditional': True,
'source': 'answer_question',
'target': 'ask_question',
}),
dict({
'conditional': True,
'source': 'answer_question',
'target': '__end__',
}),
]),
'nodes': list([
dict({
'data': '__start__',
'id': '__start__',
'type': 'schema',
}),
dict({
'data': dict({
'id': list([
'langgraph',
'utils',
'runnable',
'RunnableCallable',
]),
'name': 'ask_question',
}),
'id': 'ask_question',
'type': 'runnable',
}),
dict({
'data': dict({
'id': list([
'langgraph',
'utils',
'runnable',
'RunnableCallable',
]),
'name': 'answer_question',
}),
'id': 'answer_question',
'type': 'runnable',
}),
dict({
'data': '__end__',
'id': '__end__',
'type': 'schema',
}),
]),
})
# ---
# name: test_xray_lance.1
dict({
'edges': list([
dict({
'source': '__start__',
'target': 'generate_analysts',
}),
dict({
'source': 'conduct_interview',
'target': 'generate_sections',
}),
dict({
'source': 'generate_sections',
'target': '__end__',
}),
dict({
'conditional': True,
'source': 'generate_analysts',
'target': 'conduct_interview',
}),
]),
'nodes': list([
dict({
'data': '__start__',
'id': '__start__',
'type': 'schema',
}),
dict({
'data': dict({
'id': list([
'langgraph',
'utils',
'runnable',
'RunnableCallable',
]),
'name': 'generate_analysts',
}),
'id': 'generate_analysts',
'type': 'runnable',
}),
dict({
'data': dict({
'id': list([
'langgraph',
'graph',
'state',
'CompiledStateGraph',
]),
'name': 'conduct_interview',
}),
'id': 'conduct_interview',
'type': 'runnable',
}),
dict({
'data': dict({
'id': list([
'langgraph',
'utils',
'runnable',
'RunnableCallable',
]),
'name': 'generate_sections',
}),
'id': 'generate_sections',
'type': 'runnable',
}),
dict({
'data': '__end__',
'id': '__end__',
'type': 'schema',
}),
]),
})
# ---
# name: test_xray_lance.2
dict({
'edges': list([
dict({
'source': 'conduct_interview:__start__',
'target': 'conduct_interview:ask_question',
}),
dict({
'source': 'conduct_interview:ask_question',
'target': 'conduct_interview:answer_question',
}),
dict({
'conditional': True,
'source': 'conduct_interview:answer_question',
'target': 'conduct_interview:ask_question',
}),
dict({
'conditional': True,
'source': 'conduct_interview:answer_question',
'target': 'conduct_interview:__end__',
}),
dict({
'source': '__start__',
'target': 'generate_analysts',
}),
dict({
'source': 'conduct_interview:__end__',
'target': 'generate_sections',
}),
dict({
'source': 'generate_sections',
'target': '__end__',
}),
dict({
'conditional': True,
'source': 'generate_analysts',
'target': 'conduct_interview:__start__',
}),
]),
'nodes': list([
dict({
'data': '__start__',
'id': '__start__',
'type': 'schema',
}),
dict({
'data': dict({
'id': list([
'langgraph',
'utils',
'runnable',
'RunnableCallable',
]),
'name': 'generate_analysts',
}),
'id': 'generate_analysts',
'type': 'runnable',
}),
dict({
'data': 'conduct_interview:__start__',
'id': 'conduct_interview:__start__',
'type': 'schema',
}),
dict({
'data': dict({
'id': list([
'langgraph',
'utils',
'runnable',
'RunnableCallable',
]),
'name': 'conduct_interview:ask_question',
}),
'id': 'conduct_interview:ask_question',
'type': 'runnable',
}),
dict({
'data': dict({
'id': list([
'langgraph',
'utils',
'runnable',
'RunnableCallable',
]),
'name': 'conduct_interview:answer_question',
}),
'id': 'conduct_interview:answer_question',
'type': 'runnable',
}),
dict({
'data': 'conduct_interview:__end__',
'id': 'conduct_interview:__end__',
'type': 'schema',
}),
dict({
'data': dict({
'id': list([
'langgraph',
'utils',
'runnable',
'RunnableCallable',
]),
'name': 'generate_sections',
}),
'id': 'generate_sections',
'type': 'runnable',
}),
dict({
'data': '__end__',
'id': '__end__',
'type': 'schema',
}),
]),
})
# ---
libs/langgraph/tests/conftest.py
+22 -22
View File
@@ -20,7 +20,8 @@ from langgraph.checkpoint.postgres.aio import (
from langgraph.checkpoint.serde.encrypted import EncryptedSerializer
from langgraph.checkpoint.sqlite import SqliteSaver
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
from langgraph.pregel.loop import PregelTaskWrites
from langgraph.pregel.loop import AsyncPregelLoop, PregelTaskWrites, SyncPregelLoop
from langgraph.pregel.runner import PregelRunner
from langgraph.store.base import BaseStore
from langgraph.store.memory import InMemoryStore
from langgraph.store.postgres import AsyncPostgresStore, PostgresStore
@@ -451,27 +452,26 @@ async def awith_store(store_name: Optional[str]) -> AsyncIterator[BaseStore]:
@pytest.fixture(autouse=True)
def check_live_objects() -> None:
"""Check for live objects after each test."""
try:
gc.disable()
# Should we be yielding?
gc.collect()
leaked_objs = [
o
for o in gc.get_objects()
if isinstance(
o,
(
PregelExecutableTask,
PregelTask,
PregelScratchpad,
PregelTaskWrites,
StateSnapshot,
),
)
]
assert not leaked_objs, f"{len(leaked_objs)} leaked objects at end of test."
finally:
gc.enable()
# TODO: Ideally we should be yielding
gc.collect()
leaked_objs = [
o
for o in gc.get_objects()
if isinstance(
o,
(
PregelExecutableTask,
PregelTask,
PregelScratchpad,
PregelTaskWrites,
StateSnapshot,
SyncPregelLoop,
AsyncPregelLoop,
PregelRunner,
),
)
]
assert not leaked_objs, f"{len(leaked_objs)} leaked objects at end of test."
SHALLOW_CHECKPOINTERS_SYNC = ["postgres_shallow"]
libs/langgraph/tests/test_pregel.py
+27 -18
View File
@@ -7623,28 +7623,37 @@ def test_parallel_interrupts_double(
def test_pregel_loop_refcount():
class State(TypedDict):
messages: Annotated[list, add_messages]
gc.collect()
try:
gc.disable()
graph_builder = StateGraph(State)
class State(TypedDict):
messages: Annotated[list, add_messages]
def chatbot(state: State):
return {"messages": [("ai", "HIYA")]}
graph_builder = StateGraph(State)
graph_builder.add_node("chatbot", chatbot)
graph_builder.set_entry_point("chatbot")
graph_builder.set_finish_point("chatbot")
graph = graph_builder.compile()
def chatbot(state: State):
return {"messages": [("ai", "HIYA")]}
for _ in range(5):
graph.invoke({"messages": [{"role": "user", "content": "hi"}]})
assert (
len([obj for obj in gc.get_objects() if isinstance(obj, SyncPregelLoop)])
== 0
)
assert (
len([obj for obj in gc.get_objects() if isinstance(obj, PregelRunner)]) == 0
)
graph_builder.add_node("chatbot", chatbot)
graph_builder.set_entry_point("chatbot")
graph_builder.set_finish_point("chatbot")
graph = graph_builder.compile()
for _ in range(5):
graph.invoke({"messages": [{"role": "user", "content": "hi"}]})
assert (
len(
[obj for obj in gc.get_objects() if isinstance(obj, SyncPregelLoop)]
)
== 0
)
assert (
len([obj for obj in gc.get_objects() if isinstance(obj, PregelRunner)])
== 0
)
finally:
gc.enable()
@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC)
libs/langgraph/tests/test_pregel_async.py
+31 -18
View File
@@ -7848,28 +7848,41 @@ async def test_handles_multiple_interrupts_from_tasks() -> None:
async def test_pregel_loop_refcount():
class State(TypedDict):
messages: Annotated[list, add_messages]
gc.collect()
try:
gc.disable()
graph_builder = StateGraph(State)
class State(TypedDict):
messages: Annotated[list, add_messages]
async def chatbot(state: State):
return {"messages": [("ai", "HIYA")]}
graph_builder = StateGraph(State)
graph_builder.add_node("chatbot", chatbot)
graph_builder.set_entry_point("chatbot")
graph_builder.set_finish_point("chatbot")
graph = graph_builder.compile()
async def chatbot(state: State):
return {"messages": [("ai", "HIYA")]}
for _ in range(5):
await graph.ainvoke({"messages": [{"role": "user", "content": "hi"}]})
assert (
len([obj for obj in gc.get_objects() if isinstance(obj, AsyncPregelLoop)])
== 0
)
assert (
len([obj for obj in gc.get_objects() if isinstance(obj, PregelRunner)]) == 0
)
graph_builder.add_node("chatbot", chatbot)
graph_builder.set_entry_point("chatbot")
graph_builder.set_finish_point("chatbot")
graph = graph_builder.compile()
for _ in range(5):
await graph.ainvoke({"messages": [{"role": "user", "content": "hi"}]})
assert (
len(
[
obj
for obj in gc.get_objects()
if isinstance(obj, AsyncPregelLoop)
]
)
== 0
)
assert (
len([obj for obj in gc.get_objects() if isinstance(obj, PregelRunner)])
== 0
)
finally:
gc.enable()
@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC)