AsyncPostgresStore prepare model inputs and pass it prebuilt create_react_agent

[aditya-ladawa]

imports:

from api.store_memory_ops import extract_paper_metadata_with_chain, DB_URI_CHECKPOINTER, DB_URI_STORE, store_pool_config, connection_kwargs, make_json_serializable, ensure_serializable_dict, prepare_system_message
from psycopg_pool import AsyncConnectionPool
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
from langgraph.store.postgres import AsyncPostgresStore
from langgraph.store.postgres.aio import PoolConfig

... necessary fast api imports...

my fastapi lifespan:

@asynccontextmanager
async def lifespan(app: FastAPI):
    await init_db()
    print('\nStarted SQL db')

    await initialize_redis()
    print('\nStarted Redis connection')

    async with AsyncConnectionPool(conninfo=DB_URI_CHECKPOINTER, max_size=20, kwargs=connection_kwargs) as pool:
        checkpointer = AsyncPostgresSaver(pool)
        await checkpointer.setup()
        app.state.checkpointer = checkpointer
        print('\nInitialized Postgres Checkpointer\n')

        async with AsyncPostgresStore.from_conn_string(conn_string=DB_URI_STORE, pool_config=store_pool_config) as store:
            await store.setup()
            app.state.store = store
            print('\nInitialized Postgres Store\n')

            yield  # This is the active lifespan context

        del app.state.store

    del app.state.checkpointer

    await close_redis_connection()
    print("\nClosed Redis connection")

What i want to do is, in my FastApi websocket route I want to use the below function to prepare model inputs by fetching the memories from store and prepare them accordingly for the create_react_agent.

async def prepare_system_message(state, config: RunnableConfig, store: BaseStore):
    user_id = config.get("configurable", {}).get("user_id")
    conversation_id = config.get("configurable", {}).get("thread_id")
    
    user_data_item = await store.aget(("users",), user_id)
    user_data = user_data_item.value if user_data_item else {}

    conversation_memories = await store.asearch(("conversation_metadata", user_id, conversation_id))
    memory_info = "\n".join(
        f"{item.key.replace('metadata_', '')} | "
        f"{item.value.get('title', 'Unknown')} | "
        f"{', '.join(item.value.get('authors', []))} | "
        f"{item.value.get('description', 'Unknown')}"
        for item in conversation_memories
    )

    conversation_info_item = await store.aget(("conversation_info", user_id), conversation_id)
    conversation_info = conversation_info_item.value if conversation_info_item else {}

    user_info = f"Name: {user_data.get('name', 'Unknown')}, Email: {user_data.get('email', 'Unknown')}."
    conv_info = f"Topic: {conversation_info.get('topic', 'Unknown')}, Conversation Name: {conversation_info.get('conversation_name', 'Unknown')}."

    
    prompt = "You are a helpful research assistant. Help user to the best of your abilities. Provide concise but accurate and up to point answers. As of now you have these tools in your arsenal: qdrant_retriever_tool (content retrieval from vector database), arxiv_search_tool (search research papers), tavily_search tool (internet search). If you do not know the answer, then simply say 'I don't know. If you need clarification on what exactly user wants, then ask the user again. If you know the answer to user's query then answer yourself, else you can also rely on tools you have.\n"

    msg = f"User Info: {user_info}\nConversation Info: {conv_info}\nFiles in vector DB and storage:\n{memory_info}"

    system_msg = prompt + msg

    return [{"role": "system", "content": system_msg}] + state['messages']

this is my Fastapi route which is a websocket to chat with agent:

@app.websocket("/api/llm_chat/{conversation_id}")
async def websocket_llm_chat(
    conversation_id: str,
    websocket: WebSocket,
    current_user: dict = Depends(get_authenticated_user_websocket),
    checkpointer=Depends(get_psql_checkpointer),
    store=Depends(get_psql_store)
):  
    config = {"configurable": {"user_id": current_user.get("user_id"), "thread_id": conversation_id}}

    # Create the agent and pass the prepared system message for initialization
    agent_graph = create_react_agent(
        model=llm, 
        tools=tools_for_agent, 
        state_modifier=prepare_system_message, 
        store=store, 
        checkpointer=checkpointer
    )

    await websocket.accept()
    try:
        await websocket.send_text("Connected to LLM WebSocket! Start sending your queries.")

        seen_tool_calls = set()

        while True: 
            user_query = await websocket.receive_text()

            query = {"messages": [HumanMessage(content=user_query)]}

            async for event in agent_graph.astream(query, stream_mode="values", config=config):
                if "messages" not in event:
                    continue

                latest_msg = event["messages"][-1]  # Get the latest message

                # Check for tool usage in the latest message
                if isinstance(latest_msg, AIMessage) and not latest_msg.content:
                    tool_calls = latest_msg.additional_kwargs.get("tool_calls", [])
                    for tool_call in tool_calls:
                        tool_name = tool_call["function"]["name"]
                        args = tool_call["function"]["arguments"]

                        tool_call_id = (tool_name, str(args))
                        if tool_call_id not in seen_tool_calls:
                            seen_tool_calls.add(tool_call_id)
                            await websocket.send_text(f"Calling tool: {tool_name}\nTool arguments: {args}")

                # Send the latest AI-generated content
                elif isinstance(latest_msg, AIMessage) and latest_msg.content:
                    await websocket.send_text(latest_msg.content)

    except WebSocketDisconnect:
        print("WebSocket connection closed.")

Now, prepare_system_message actually returns a coroutine, so I dont know how to handle this case for my route.

I also have doubts and i hope Icould get the answer for them as well:

1. When i upload new .pdf file it is embedded in vector db and the metadata of file (like title, authors, and description of abstract) is stored in AsyncPostgresStore (store) for agent memory, (which will help it to perform RAG if necessary else it will opt for web search)
   the question is: Given my websocket_llm_chat route will the agent get to know regarding the new info saved in memory without closing the socket and re opening it? If not then what is the alternative

here is the traceback i get:
[1] INFO: connection open
[1] ERROR: Exception in ASGI application
[1] Traceback (most recent call last):
[1] File "/home/aditya-ladawa/Aditya/z_projects/aireas/myenv/lib/python3.12/site-packages/uvicorn/protocols/websockets/websockets_impl.py", line 242, in run_asgi
[1] result = await self.app(self.scope, self.asgi_receive, self.asgi_send) # type: ignore[func-returns-value]
[1] ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[1] File "/home/aditya-ladawa/Aditya/z_projects/aireas/myenv/lib/python3.12/site-packages/uvicorn/middleware/proxy_headers.py", line 60, in call
[1] return await self.app(scope, receive, send)
[1] ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[1] File "/home/aditya-ladawa/Aditya/z_projects/aireas/myenv/lib/python3.12/site-packages/fastapi/applications.py", line 1054, in call
[1] await super().call(scope, receive, send)
[1] File "/home/aditya-ladawa/Aditya/z_projects/aireas/myenv/lib/python3.12/site-packages/starlette/applications.py", line 113, in call
[1] await self.middleware_stack(scope, receive, send)
[1] File "/home/aditya-ladawa/Aditya/z_projects/aireas/myenv/lib/python3.12/site-packages/starlette/middleware/errors.py", line 152, in call
[1] await self.app(scope, receive, send)
[1] File "/home/aditya-ladawa/Aditya/z_projects/aireas/myenv/lib/python3.12/site-packages/starlette/middleware/cors.py", line 77, in call
[1] await self.app(scope, receive, send)
[1] File "/home/aditya-ladawa/Aditya/z_projects/aireas/myenv/lib/python3.12/site-packages/starlette/middleware/exceptions.py", line 62, in call
[1] await wrap_app_handling_exceptions(self.app, conn)(scope, receive, send)
[1] File "/home/aditya-ladawa/Aditya/z_projects/aireas/myenv/lib/python3.12/site-packages/starlette/_exception_handler.py", line 53, in wrapped_app
[1] raise exc
[1] File "/home/aditya-ladawa/Aditya/z_projects/aireas/myenv/lib/python3.12/site-packages/starlette/_exception_handler.py", line 42, in wrapped_app
[1] await app(scope, receive, sender)
[1] File "/home/aditya-ladawa/Aditya/z_projects/aireas/myenv/lib/python3.12/site-packages/starlette/routing.py", line 715, in call
[1] await self.middleware_stack(scope, receive, send)
[1] File "/home/aditya-ladawa/Aditya/z_projects/aireas/myenv/lib/python3.12/site-packages/starlette/routing.py", line 735, in app
[1] await route.handle(scope, receive, send)
[1] File "/home/aditya-ladawa/Aditya/z_projects/aireas/myenv/lib/python3.12/site-packages/starlette/routing.py", line 362, in handle
[1] await self.app(scope, receive, send)
[1] File "/home/aditya-ladawa/Aditya/z_projects/aireas/myenv/lib/python3.12/site-packages/starlette/routing.py", line 95, in app
[1] await wrap_app_handling_exceptions(app, session)(scope, receive, send)
[1] File "/home/aditya-ladawa/Aditya/z_projects/aireas/myenv/lib/python3.12/site-packages/starlette/_exception_handler.py", line 53, in wrapped_app
[1] raise exc
[1] File "/home/aditya-ladawa/Aditya/z_projects/aireas/myenv/lib/python3.12/site-packages/starlette/_exception_handler.py", line 42, in wrapped_app
[1] await app(scope, receive, sender)
[1] File "/home/aditya-ladawa/Aditya/z_projects/aireas/myenv/lib/python3.12/site-packages/starlette/routing.py", line 93, in app
[1] await func(session)
[1] File "/home/aditya-ladawa/Aditya/z_projects/aireas/myenv/lib/python3.12/site-packages/fastapi/routing.py", line 383, in app
[1] await dependant.call(**solved_result.values)
[1] File "/home/aditya-ladawa/Aditya/z_projects/aireas/api/index.py", line 611, in websocket_llm_chat
[1] async for event in agent_graph.astream(query, stream_mode="values", config=config):
[1] File "/home/aditya-ladawa/Aditya/z_projects/aireas/myenv/lib/python3.12/site-packages/langgraph/pregel/init.py", line 1874, in astream
[1] async for _ in runner.atick(
[1] File "/home/aditya-ladawa/Aditya/z_projects/aireas/myenv/lib/python3.12/site-packages/langgraph/pregel/runner.py", line 362, in atick
[1] await arun_with_retry(
[1] File "/home/aditya-ladawa/Aditya/z_projects/aireas/myenv/lib/python3.12/site-packages/langgraph/pregel/retry.py", line 132, in arun_with_retry
[1] return await task.proc.ainvoke(task.input, config)
[1] ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[1] File "/home/aditya-ladawa/Aditya/z_projects/aireas/myenv/lib/python3.12/site-packages/langgraph/utils/runnable.py", line 445, in ainvoke
[1] input = await step.ainvoke(input, config, **kwargs)
[1] ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[1] File "/home/aditya-ladawa/Aditya/z_projects/aireas/myenv/lib/python3.12/site-packages/langgraph/utils/runnable.py", line 224, in ainvoke
[1] ret = await asyncio.create_task(coro, context=context)
[1] ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[1] File "/home/aditya-ladawa/Aditya/z_projects/aireas/myenv/lib/python3.12/site-packages/langgraph/prebuilt/chat_agent_executor.py", line 595, in acall_model
[1] response = await model_runnable.ainvoke(state, config)
[1] ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[1] File "/home/aditya-ladawa/Aditya/z_projects/aireas/myenv/lib/python3.12/site-packages/langchain_core/runnables/base.py", line 3066, in ainvoke
[1] input = await asyncio.create_task(part(), context=context) # type: ignore
[1] ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[1] File "/home/aditya-ladawa/Aditya/z_projects/aireas/myenv/lib/python3.12/site-packages/langgraph/utils/runnable.py", line 193, in ainvoke
[1] return self.invoke(input, config)
[1] ^^^^^^^^^^^^^^^^^^^^^^^^^^
[1] File "/home/aditya-ladawa/Aditya/z_projects/aireas/myenv/lib/python3.12/site-packages/langgraph/utils/runnable.py", line 176, in invoke
[1] ret = context.run(self.func, input, **kwargs)
[1] ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[1] File "/home/aditya-ladawa/Aditya/z_projects/aireas/api/store_memory_ops.py", line 124, in prepare_system_message
[1] user_data = user_data_item.value if user_data_item else {}
[1] ^^^^^^^^^^^^^^^^^^^^
[1] AttributeError: 'coroutine' object has no attribute 'value'
[1] During task with name 'agent' and id '9d994a88-6a25-8211-e1c4-4671591fc31e'
[1] INFO: connection closed