open-notebook/open_notebook/utils/graph_utils.py
2026-01-29 23:00:22 -03:00

18 lines
691 B
Python

from langchain_core.runnables import RunnableConfig
from loguru import logger
async def get_session_message_count(graph, session_id: str) -> int:
"""Get message count from LangGraph state, returns 0 on error."""
try:
thread_state = await graph.aget_state( # async version
config=RunnableConfig(configurable={"thread_id": session_id})
)
if (
thread_state
and thread_state.values
and "messages" in thread_state.values
):
return len(thread_state.values["messages"])
except Exception as e:
logger.warning(f"Could not fetch message count for session {session_id}: {e}")
return 0