Async DB Hooks

Async database lifecycle hooks for non-blocking persistence in async agents

Async DB hooks enable non-blocking database operations in async agents through automatic async/sync detection and asyncio.to_thread fallback.

from praisonaiagents import Agent, MemoryConfig, PraisonAIDB

agent = Agent(
    name="Assistant",
    instructions="You are a helpful assistant.",
    memory=MemoryConfig(
        db=PraisonAIDB("sqlite:///agent_data.db"),
    ),
    async_mode=True)

# In async code: await agent.achat("Remember my preferences")

The user chats asynchronously; DB hooks persist messages without blocking the event loop.

Quick Start

1. Async Context Manager

   Use async DB operations with the async context manager for automatic cleanup.

   from praisonaiagents import Agent, MemoryConfig, PraisonAIDB
   
   async def main():
       async with PraisonAIDB("sqlite:///agent_data.db") as db:
           # Store user message
           await db.aon_user_message(
               session_id="chat_123",
               content="Hello, how can you help me?",
               metadata={"user_id": "user_456"}
           )
   
           # Store agent message  
           await db.aon_agent_message(
               session_id="chat_123", 
               content="I can help with various tasks.",
               metadata={"agent_name": "assistant"}
           )

2. Wire to Async Agent

   Connect async DB hooks to an async agent for seamless persistence.

   from praisonaiagents import Agent, MemoryConfig
   
   agent = Agent(
       name="Assistant",
       instructions="You are a helpful assistant.",
       memory=MemoryConfig(
           db=PraisonAIDB("postgresql://user:pass@host/db"),
       ),
       async_mode=True
   )
   
   async def chat():
       async with agent.db:
           response = await agent.run_async("What's the weather like?")
           print(response)

How It Works

Shared dispatch for both entry points

Both persistence entry points flow through the same async-store dispatch.

As of PR #4394, the async-store dispatch (isinstance(store, AsyncConversationStore) → await store.method(...) vs. asyncio.to_thread(store.method, ...)) lives on PraisonAIDB, in its internal _call_store / _dispatch_async helpers.

The same isinstance() dispatch pattern now has formal AsyncStateStore / AsyncKnowledgeStore siblings (PR #5403), even though the built-in state and knowledge stores are not yet rewired onto them. See Async State Store and Async Knowledge Store.

PersistenceOrchestrator.on_message / aon_message / on_agent_end / aon_agent_end delegate to those helpers instead of keeping a second copy, so the async/sync-store behaviour is identical no matter which entry point you pick — MemoryConfig(db=PraisonAIDB(...)) or the orchestrator-based wrap_agent_with_persistence / PersistentAgent / create_persistent_session.

Sync reads from inside a running loop

Bare sync reads reached from inside a running event loop fail loudly instead of silently dropping the read.

Async completion write path

The four sync completion hooks persist correctly from inside a running loop by routing the whole read-modify-write through one bridge submission.

on_run_end, on_agent_end, on_trace_end, and on_span_end each call _merge_and_set, which builds a single coroutine that does get → merge → set on the state store. Since PR #5038, _merge_and_set hands that coroutine to dispatch_maybe_awaitable(..., kind=DispatchKind.WRITE, tracker=self._bg_writes, tracker_lock=self._bg_writes_lock, op_name=f"merge_and_set:{key}") instead of hand-rolling the submit/track/callback — same behaviour, one path. It becomes a tracked fire-and-forget write (recorded in self._bg_writes). Because the read never crosses the loop boundary, the write half always executes — the completion actually persists instead of leaving the record stuck at status="running".

The merge is uuid-keyed and idempotent, so deferring the write is safe by construction. Pre-existing fields (run_id, started_at, input_content, spans, events) are preserved; completion fields (output_content, status, ended_at) are merged in.

Ordering guarantee. Callers who must see the persisted record before returning (tests, request handlers reading the state store synchronously) should flush the pending write first.

import asyncio
from praisonai.db.adapter import PraisonAIDB

async def main():
    db = PraisonAIDB("sqlite:///agent_data.db")

    # Sync completion hook from inside a running loop — persists, never raises.
    db.on_run_end("chat_123", "run_1", output_content="done", status="completed")

    # The write is fire-and-forget on the bridge; flush it before reading back.
    await asyncio.to_thread(db.flush_pending_writes)

asyncio.run(main())

Only completion hooks that read internally changed. Bare reads (get, get_session, get_messages, list_sessions) still fail loudly from inside a running loop — PR #4821 semantics are unchanged.

Which hook to call from where

Pick sync or async by whether you sit inside a running loop and whether you read or write.

State-store lifecycle key

The wrapper writes both start and end under agent:{session_id}, so the end transition updates the same record the start wrote instead of a disconnected one.

Start previously used agent:{session_id}:{agent_id or name} and end wrote a separate record. Consumers who query the state store directly must update their key format to agent:{session_id}.

Configuration Options

All async hooks support these signatures from the DB adapter:

async def aon_agent_start(
    self,
    agent_name: str,
    session_id: str,
    user_id: Optional[str] = None,
    metadata: Optional[Dict[str, Any]] = None,
) -> List   # previous messages for resume; empty list on new session

async def aon_user_message(
    self, 
    session_id: str, 
    content: str, 
    metadata: Optional[Dict[str, Any]] = None
) -> None

async def aon_agent_message(
    self, 
    session_id: str, 
    content: str, 
    metadata: Optional[Dict[str, Any]] = None
) -> None

async def aon_tool_call(
    self, 
    session_id: str, 
    tool_name: str, 
    arguments: Dict[str, Any], 
    result: Any = None, 
    metadata: Optional[Dict[str, Any]] = None
) -> None

async def aon_agent_end(
    self,
    session_id: str,
    metadata: Optional[Dict[str, Any]] = None,
) -> None

async def aclose(self) -> None

Common Patterns

Complete Async Lifecycle

async def agent_lifecycle():
    async with PraisonAIDB("sqlite:///lifecycle.db") as db:
        session_id = "session_789"

        # Start agent session — agent_name is first, user_id is persisted
        messages = await db.aon_agent_start(
            "DataAgent",           # agent_name (now first)
            session_id,
            user_id="user_456",    # new; previously silently dropped
            metadata={"env": "prod"},
        )

        # User interaction
        await db.aon_user_message(
            session_id=session_id,
            content="Analyze this data",
            metadata={"timestamp": "2024-01-01T10:00:00Z"}
        )

        # Tool execution
        await db.aon_tool_call(
            session_id=session_id,
            tool_name="data_analyzer",
            arguments={"file": "data.csv"},
            result={"rows": 1000, "columns": 5}
        )

        # Agent response
        await db.aon_agent_message(
            session_id=session_id,
            content="Analysis complete: 1000 rows, 5 columns found."
        )

        # End session — no longer takes name / agent_id
        await db.aon_agent_end(session_id, metadata={"outcome": "ok"})

Sync Store Compatibility

# Works with existing sync stores
class MySQLStore:
    def add_message(self, content, metadata):
        # Sync implementation
        pass

    # No async_add_message - will use asyncio.to_thread

async with PraisonAIDB(MySQLStore()) as db:
    # Still works - automatically wrapped
    await db.aon_user_message("session_1", "Hello")

Native Async Store

# Implement async methods for true async performance
class AsyncPostgreStore:
    async def async_add_message(self, content, metadata):
        async with self.pool.acquire() as conn:
            await conn.execute("INSERT INTO messages...", content)

    def add_message(self, content, metadata):
        # Sync fallback if needed
        pass

# Will use native async methods
async with PraisonAIDB(AsyncPostgreStore()) as db:
    await db.aon_user_message("session_1", "Hello")  # True async

Best Practices

# ✅ Good: cleanup runs even on error
async with PraisonAIDB(store) as db:
    await db.aon_user_message("session", "Hello")

# ✅ Also fine: aclose() is idempotent and safe to call
db = PraisonAIDB(store)
await db.aon_user_message("session", "Hello")
await db.aclose()

db = PraisonAIDB("sqlite:///agent.db")
with db:
    db.on_user_message("session", "Hello")   # first use
# ...later in the same process...
with db:
    db.on_user_message("session", "Later")   # cleanly re-initializes

class OptimizedStore:
    # ✅ Native async - no thread overhead
    async def async_add_message(self, content, metadata):
        await self.async_conn.execute(query, content)

    # ✅ Sync fallback for compatibility
    def add_message(self, content, metadata):
        self.sync_conn.execute(query, content)

metadata = {
    "user_id": "user_123",
    "session_type": "chat",
    "timestamp": datetime.now(timezone.utc).isoformat(),
    "environment": "production"
}

await db.aon_user_message(
    session_id="session_456", 
    content="User input",
    metadata=metadata
)

# Sync context (for mixed sync/async code)
with PraisonAIDB(store) as db:
    # Use sync methods
    db.on_user_message("session", "Hello")

# Async context (for async agents)  
async with PraisonAIDB(store) as db:
    # Use async methods
    await db.aon_user_message("session", "Hello")

Related