AsyncPostgresStore cpu hight

[tianshangwuyun]

Checked other resources

• I added a very descriptive title to this question.
• I searched the LangChain documentation with the integrated search.
• I used the GitHub search to find a similar question and didn't find it.

Commit to Help

• I commit to help with one of those options 👆

Example Code

@router.get("/test-res")
async def test_exception(cache_service: CacheService = Depends(get_cache_service)):
from sse_starlette.sse import EventSourceResponse
from langchain.prompts import ChatPromptTemplate
from langchain.tools import tool
from langgraph.prebuilt import create_react_agent
from typing import Literal
from typing_extensions import TypedDict

from langgraph.graph import END, StateGraph, START
from langgraph.graph.message import MessagesState
from langchain_core.messages import BaseMessage, SystemMessage, HumanMessage
from contextlib import asynccontextmanager
from langgraph.store.postgres import AsyncPostgresStore
from psycopg.rows import dict_row
from agents.memory.custom_memory import CustomPgMemorySaver, delete_memory
from psycopg import AsyncConnection
from psycopg_pool import AsyncConnectionPool

from configs.config import DB_URI

model = llm_openrouter_free
final_model = llm_openrouter_free
final_model = final_model.with_config(tags=["final_node"])

def should_continue(state: MessagesState) -> Literal["final"]:
    messages = state["messages"]
    last_message = messages[-1]
    # If the LLM makes a tool call, then we route to the "tools" node
    # if last_message.tool_calls:
    #     return "tools"
    # Otherwise, we stop (reply to the user)
    return "final"


async def call_model(state: MessagesState):
    messages = state["messages"]
    response = await model.ainvoke(messages)
    # We return a list, because this will get added to the existing list
    return {"messages": [response]}


async def call_final_model(state: MessagesState):
    messages = state["messages"]
    last_ai_message = messages[-1]
    response = await final_model.ainvoke(
        [
            SystemMessage("Rewrite this in the voice of Al Roker"),
            HumanMessage(last_ai_message.content),
        ]
    )
    # overwrite the last AI message from the agent
    response.id = last_ai_message.id
    return {"messages": [response]}


@asynccontextmanager
async def compile_graph():
    connection_kwargs = {
        "autocommit": True,
        "prepare_threshold": 0,
        "row_factory": dict_row,
    }
    # yield create_agent(None,None)
    async with AsyncConnectionPool(
                        conninfo=DB_URI,
                        max_size=20,
                        kwargs=connection_kwargs,
                    ) as pool:
            store = AsyncPostgresStore(pool)
            await store.setup()
            memory = CustomPgMemorySaver(pool)
            await memory.setup()
            yield create_agent(memory,store)

def create_agent(checkpointer,store):
    builder = StateGraph(MessagesState)

    builder.add_node("agent", call_model)
    builder.add_node("final", call_final_model)

    builder.add_edge(START, "agent")
    builder.add_conditional_edges(
        "agent",
        should_continue,
    )

    builder.add_edge("final", END)

    graph = builder.compile(checkpointer=checkpointer,store=store)
    return graph

    
   
# chat = prompt | llm_openrouter_free
async def gen():
    async with compile_graph() as agent:
        config= {"configurable": {"thread_id": "123", "user_id": "123"}}
        async for event in agent.astream_events(
                {"messages": [{"role": "user", "content": "what's in the 1*10."}]}, version="v2",config=config
            ):
                kind = event["event"]
                tags = event.get("tags", [])
                # filter on the custom tag
                if kind == "on_chat_model_stream" and "final_node" in event.get("tags", []):
                    data = event["data"]
                    content = data["chunk"].content
                    if content:
                        # Empty content in the context of OpenAI or Anthropic usually means
                        # that the model is asking for a tool to be invoked.
                        # So we only print non-empty content
                        logger.info(f"test event: {content}")
                        yield content

async def _gen(res):
    async for chunk in res:
        logger.info(f"test chunk: {chunk}")
        yield chunk

msg_chunk = gen()


return EventSourceResponse(_gen(msg_chunk))

llm stream=true

### Description



I am using fastpapi+ Langgraph to build an agent and using Astream_ events for SSE output. 

when used 

async with AsyncConnectionPool(
conninfo=DB_URI,
max_size=20,
kwargs=connection_kwargs,
) as pool:
store = AsyncPostgresStore(pool)
await store.setup()
memory = CustomPgMemorySaver(pool)
await memory.setup()
yield create_agent(memory,store)

cpu 100% 
When using the store, the task is completed and the CPU will not decrease
maybe The _run method in Asynchronous BatchedBaseStore
The store object will not be reclaimed, resulting in the inability to call the _del_ method_ While True in the run method:
Awaiting asyncio. sleep (0) causes asyncio to constantly switch, resulting in CPU 100%
### System Info

langchain==0.3.3
langchain-anthropic==0.2.4
langchain-community==0.3.1
langchain-core==0.3.15
langchain-openai==0.2.2
langchain-text-splitters==0.3.0
langgraph==0.2.34