Code
from typing import Annotated, TypedDict
from langchain_openai import AzureChatOpenAI
from langgraph.graph import StateGraph
from langgraph.graph.message import add_messages
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
from psycopg.rows import dict_row
class State(TypedDict):
messages: Annotated[list, add_messages]
data_to_save: dict
async def setup_graph(postgres_config: dict):
llm = AzureChatOpenAI(model="gpt-4o-mini")
def ask(state: State):
response = llm.invoke(state["messages"])
state["messages"].append(response)
state["data_to_save"] = {
"messages": state["messages"],
"latest_response": response
}
return state
conn_string = " ".join([f"{key}={value}" for key, value in postgres_config.items()])
async_conn_pool = AsyncConnectionPool(
conninfo=conn_string,
kwargs={
"autocommit": True,
"prepare_threshold": 0,
"row_factory": dict_row,
"keepalives": 1,
"keepalives_idle": 30,
"keepalives_interval": 10,
"keepalives_count": 5
},
min_size=5,
max_size=15,
timeout=30
)
asyncpostgress_memory = AsyncPostgresSaver(async_connection=async_conn_pool)
await asyncpostgress_memory.setup()
builder = StateGraph(State)
builder.add_node("ask", ask)
builder.set_entry_point("ask")
builder.set_finish_point("ask")
return builder.compile(checkpointer=asyncpostgress_memory)
async def main():
postgres_config = {
"host": "some_value",
"port": "some_value",
"user": "some_value",
"password": "some_value",
"dbname": "some_value"
}
graph = await setup_graph(postgres_config)
input_state = {
"messages": [
{"role": "user", "content": "What is the weather today?"}
],
"data_to_save": {}
}
config = {
"retry_policy": {
"max_retries": 3,
"backoff_factor": 2
}
}
await graph.ainvoke(
input=input_state,
config=config,
)
Description
Each time I execute the main() function, RAM usage increases and the memory is successfully released upon the completion of main(). However, when graph.ainvoke results in error, the memory is not released.
Installed Dependencies
langchain==0.3.26
langchain-core==0.3.69
langchain-openai==0.3.28
langchain-experimental==0.3.4
langgraph==0.5.3
langgraph-checkpoint==2.1.1
langgraph-checkpoint-postgres==2.0.23