Add universal gc check

This commit is contained in:
William Fu-Hinthorn
2025-03-19 14:30:15 -07:00
parent 53f8558914
commit 7c7bb9c627
4 changed files with 71 additions and 1129 deletions
+35 -161
View File
File diff suppressed because it is too large Load Diff
+1
View File
@@ -38,6 +38,7 @@ pyperf = "^2.7.0"
py-spy = "^0.3.14"
types-requests = "^2.32.0.20240914"
pycryptodome = "^3.21.0"
referrers = "^0.12.0"
[tool.ruff]
lint.select = [ "E", "F", "I", "TID251" ]
@@ -1370,426 +1370,6 @@
'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;
@@ -1881,240 +1461,6 @@
'''
# ---
# 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"}'
# ---
@@ -2124,317 +1470,3 @@
# 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',
}),
]),
})
# ---
+35
View File
@@ -1,9 +1,11 @@
import gc
import sys
from contextlib import asynccontextmanager
from typing import AsyncIterator, Optional
from uuid import UUID, uuid4
import pytest
import referrers
from langchain_core import __version__ as core_version
from packaging import version
from psycopg import AsyncConnection, Connection
@@ -19,9 +21,16 @@ 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.store.base import BaseStore
from langgraph.store.memory import InMemoryStore
from langgraph.store.postgres import AsyncPostgresStore, PostgresStore
from langgraph.types import (
PregelExecutableTask,
PregelScratchpad,
PregelTask,
StateSnapshot,
)
pytest.register_assert_rewrite("tests.memory_assert")
@@ -440,6 +449,32 @@ async def awith_store(store_name: Optional[str]) -> AsyncIterator[BaseStore]:
raise NotImplementedError(f"Unknown store {store_name}")
@pytest.fixture(autouse=True)
def check_live_objects() -> None:
"""Check for live objects after each test."""
gc.collect()
gc.disable()
try:
yield
leaked_objs = [
o
for o in gc.get_objects()
if isinstance(
o,
(
PregelExecutableTask,
PregelTask,
PregelScratchpad,
PregelTaskWrites,
StateSnapshot,
),
)
]
assert not leaked_objs, f"""Leaked objects at end of test: {referrers.get_referrer_graph_for_list(leaked_objs, exclude_object_ids=(id(leaked_objs),))}"""
finally:
gc.enable()
SHALLOW_CHECKPOINTERS_SYNC = ["postgres_shallow"]
REGULAR_CHECKPOINTERS_SYNC = [
"memory",