diff --git a/src/praisonai/praisonai/db/adapter.py b/src/praisonai/praisonai/db/adapter.py index a78311b18..c7530bd34 100644 --- a/src/praisonai/praisonai/db/adapter.py +++ b/src/praisonai/praisonai/db/adapter.py @@ -299,24 +299,25 @@ def on_agent_start( return [] from ..persistence.conversation.base import ConversationSession + from ..persistence.conversation._ops import resume_or_create_session - # Check if session exists - session = self._conversation_store.get_session(session_id) - - if session is None: - # Create new session - session = ConversationSession( + store = self._conversation_store + messages = resume_or_create_session( + store, + store.get_session(session_id), + session_id, + build_session=lambda: ConversationSession( session_id=session_id, user_id=user_id or "default", agent_id=agent_name, name=f"Session {session_id}", - metadata=metadata or {} - ) - self._conversation_store.create_session(session) - return [] + metadata=metadata or {}, + ), + get_messages=lambda: store.get_messages(session_id), + ) - # Resume existing session - get messages - messages = self._conversation_store.get_messages(session_id) + if messages is None: + return [] # Convert to DbMessage format from praisonaiagents.db.protocol import DbMessage @@ -737,31 +738,32 @@ async def aon_agent_start( if self._conversation_store: from ..persistence.conversation.base import ConversationSession + from ..persistence.conversation._ops import aresume_or_create_session + store = self._conversation_store session = await self._dispatch_async( - self._conversation_store, "get_session", "async_get_session", session_id + store, "get_session", "async_get_session", session_id ) - if session is None: - new_session = ConversationSession( + raw = await aresume_or_create_session( + store, + session, + session_id, + build_session=lambda: ConversationSession( session_id=session_id, agent_id=agent_id or name, name=f"Session {session_id}", metadata=metadata or {}, - ) - await self._dispatch_async( - self._conversation_store, - "create_session", - "async_create_session", - new_session, - ) - else: - raw = await self._dispatch_async( - self._conversation_store, - "get_messages", - "async_get_messages", - session_id, - ) + ), + create_session=lambda s: self._dispatch_async( + store, "create_session", "async_create_session", s + ), + get_messages=lambda: self._dispatch_async( + store, "get_messages", "async_get_messages", session_id + ), + ) + + if raw is not None: from praisonaiagents.db.protocol import DbMessage messages = [ diff --git a/src/praisonai/praisonai/persistence/conversation/_ops.py b/src/praisonai/praisonai/persistence/conversation/_ops.py new file mode 100644 index 000000000..05ea84cad --- /dev/null +++ b/src/praisonai/praisonai/persistence/conversation/_ops.py @@ -0,0 +1,67 @@ +""" +Shared conversation-store operations. + +Single owner for the *create-or-resume session* flow over a +``ConversationStore``. Both ``PraisonAIDB`` (``db/adapter.py``) and +``PersistenceOrchestrator`` (``persistence/orchestrator.py``) call these helpers +so the store-level session semantics live in one place instead of being copied +across the sync/async surfaces of both classes. + +The helpers only touch the store (``get_session`` / ``create_session`` / +``get_messages``); each caller keeps its own return-type contract, caching, and +lock/cooldown machinery by wrapping the result. +""" + +from typing import Any, Awaitable, Callable, List, Optional + +from .base import ConversationSession, ConversationMessage + + +def resume_or_create_session( + store: Any, + session: Optional[ConversationSession], + session_id: str, + build_session: Callable[[], ConversationSession], + get_messages: Callable[[], List[ConversationMessage]], +) -> Optional[List[ConversationMessage]]: + """Create the session if missing, else return its messages (sync). + + Args: + store: The conversation store. + session: Result of the caller's ``get_session`` lookup (``None`` when + the session does not yet exist). + session_id: The session identifier (unused directly; kept for parity + and readability at call sites). + build_session: Factory returning the ``ConversationSession`` to create + when ``session`` is ``None`` — lets each caller keep its own name/ + metadata conventions. + get_messages: Callable returning the existing messages when resuming. + + Returns: + ``None`` when a new session was created (no history), otherwise the list + of previously persisted messages. + """ + if session is None: + store.create_session(build_session()) + return None + return get_messages() + + +async def aresume_or_create_session( + store: Any, + session: Optional[ConversationSession], + session_id: str, + build_session: Callable[[], ConversationSession], + create_session: Callable[[ConversationSession], Awaitable[Any]], + get_messages: Callable[[], Awaitable[List[ConversationMessage]]], +) -> Optional[List[ConversationMessage]]: + """Async variant of :func:`resume_or_create_session`. + + ``create_session`` and ``get_messages`` are awaitables supplied by the + caller so each keeps its own async-dispatch discipline (dedicated async + method vs. ``asyncio.to_thread`` off-loading). + """ + if session is None: + await create_session(build_session()) + return None + return await get_messages() diff --git a/src/praisonai/praisonai/persistence/orchestrator.py b/src/praisonai/praisonai/persistence/orchestrator.py index f258641e0..b90fe5d8d 100644 --- a/src/praisonai/praisonai/persistence/orchestrator.py +++ b/src/praisonai/praisonai/persistence/orchestrator.py @@ -165,34 +165,53 @@ def on_agent_start( logger.debug("No conversation store configured, skipping session load") return [] + from .conversation._ops import resume_or_create_session + # Try to load existing session session = None if resume: session = self._sync(self.conversation.get_session(session_id)) - + if session: logger.info(f"Resuming session: {session_id}") self._current_session = session self._cache_put(session) - - # Load previous messages - messages = self._sync(self.conversation.get_messages(session_id)) - return messages - else: - # Create new session + + # Build the new session lazily inside the factory so the resume path + # never touches the agent's identity or constructs a discarded object. + # The factory captures the exact instance the helper persists so the + # same object (identical timestamps/identity) is cached below. + created: List[ConversationSession] = [] + + def _build_session() -> ConversationSession: agent_id = getattr(agent, "name", None) or getattr(agent, "agent_id", None) - session = ConversationSession( + new_session = ConversationSession( session_id=session_id, user_id=user_id, agent_id=agent_id, name=f"Session {session_id[:8]}", metadata={"agent_type": type(agent).__name__}, ) - self._sync(self.conversation.create_session(session)) + created.append(new_session) + return new_session + + messages = resume_or_create_session( + self.conversation, + session, + session_id, + build_session=_build_session, + get_messages=lambda: self._sync(self.conversation.get_messages(session_id)), + ) + + if messages is None: + # New session was created inside the helper; capture and cache it. + new_session = created[0] logger.info(f"Created new session: {session_id}") - self._current_session = session - self._cache_put(session) + self._current_session = new_session + self._cache_put(new_session) return [] + + return messages def on_message( self, @@ -291,46 +310,72 @@ async def aon_agent_start( logger.debug("No conversation store configured, skipping session load") return [] + from .conversation._ops import aresume_or_create_session + + is_async = isinstance(self.conversation, AsyncConversationStore) + + async def _get_session(): + if is_async: + return await self.conversation.get_session(session_id) + # Run blocking store off the loop so we don't block multi-agent execution + return await asyncio.to_thread(self.conversation.get_session, session_id) + + async def _create_session(s): + if is_async: + return await self.conversation.create_session(s) + return await asyncio.to_thread(self.conversation.create_session, s) + + async def _get_messages(): + if is_async: + return await self.conversation.get_messages(session_id) + return await asyncio.to_thread(self.conversation.get_messages, session_id) + # Try to load existing session session = None if resume: - if isinstance(self.conversation, AsyncConversationStore): - session = await self.conversation.get_session(session_id) - else: - # Run blocking store off the loop so we don't block multi-agent execution - session = await asyncio.to_thread(self.conversation.get_session, session_id) - + session = await _get_session() + if session: logger.info(f"Resuming session: {session_id}") self._current_session = session self._cache_put(session) - - # Load previous messages - if isinstance(self.conversation, AsyncConversationStore): - messages = await self.conversation.get_messages(session_id) - else: - messages = await asyncio.to_thread(self.conversation.get_messages, session_id) - return messages - else: - # Create new session + + # Build the new session lazily inside the factory so the resume path + # never touches the agent's identity or constructs a discarded object. + # The factory captures the exact instance the helper persists so the + # same object (identical timestamps/identity) is cached below. + created: List[ConversationSession] = [] + + def _build_session() -> ConversationSession: agent_id = getattr(agent, "name", None) or getattr(agent, "agent_id", None) - session = ConversationSession( + new_session = ConversationSession( session_id=session_id, user_id=user_id, agent_id=agent_id, name=f"Session {session_id[:8]}", metadata={"agent_type": type(agent).__name__}, ) - - if isinstance(self.conversation, AsyncConversationStore): - await self.conversation.create_session(session) - else: - await asyncio.to_thread(self.conversation.create_session, session) - + created.append(new_session) + return new_session + + messages = await aresume_or_create_session( + self.conversation, + session, + session_id, + build_session=_build_session, + create_session=_create_session, + get_messages=_get_messages, + ) + + if messages is None: + # New session was created inside the helper; capture and cache it. + new_session = created[0] logger.info(f"Created new session: {session_id}") - self._current_session = session - self._cache_put(session) + self._current_session = new_session + self._cache_put(new_session) return [] + + return messages async def aon_message( self,