SQLAlchemy for AI Agent Applications: ORM Patterns and Async Database Access

Why SQLAlchemy 2.0 for Agent Systems

SQLAlchemy 2.0 introduced native async support, a unified query interface, and Mapped type annotations that bring full IDE autocompletion to database models. For AI agent backends built with FastAPI or other async frameworks, this means you can write non-blocking database code without sacrificing the safety and expressiveness of an ORM.

The key advantage over raw SQL for agent systems is that SQLAlchemy handles connection lifecycle, relationship loading, and transaction boundaries — all areas where hand-written code tends to accumulate bugs over time.

Defining Agent Models

Start with a base class and define your models using SQLAlchemy 2.0 Mapped annotations:

from datetime import datetime
from uuid import uuid4
from sqlalchemy import ForeignKey, Text, CheckConstraint
from sqlalchemy.dialects.postgresql import JSONB, UUID
from sqlalchemy.orm import (
    DeclarativeBase,
    Mapped,
    mapped_column,
    relationship,
)


class Base(DeclarativeBase):
    pass


class Conversation(Base):
    __tablename__ = "conversations"

    id: Mapped[str] = mapped_column(
        UUID(as_uuid=False), primary_key=True, default=uuid4
    )
    user_id: Mapped[str] = mapped_column(UUID(as_uuid=False), nullable=False)
    agent_id: Mapped[str] = mapped_column(UUID(as_uuid=False), nullable=False)
    title: Mapped[str | None] = mapped_column(Text)
    status: Mapped[str] = mapped_column(
        Text, default="active"
    )
    metadata_: Mapped[dict] = mapped_column(
        "metadata", JSONB, default=dict
    )
    created_at: Mapped[datetime] = mapped_column(default=datetime.utcnow)

    messages: Mapped[list["Message"]] = relationship(
        back_populates="conversation",
        cascade="all, delete-orphan",
        order_by="Message.created_at",
    )

    __table_args__ = (
        CheckConstraint(
            "status IN ('active', 'archived', 'deleted')",
            name="ck_conversation_status",
        ),
    )


class Message(Base):
    __tablename__ = "messages"

    id: Mapped[str] = mapped_column(
        UUID(as_uuid=False), primary_key=True, default=uuid4
    )
    conversation_id: Mapped[str] = mapped_column(
        ForeignKey("conversations.id", ondelete="CASCADE")
    )
    role: Mapped[str] = mapped_column(Text, nullable=False)
    content: Mapped[str | None] = mapped_column(Text)
    token_count: Mapped[int | None]
    model: Mapped[str | None] = mapped_column(Text)
    tool_calls: Mapped[dict | None] = mapped_column(JSONB)
    created_at: Mapped[datetime] = mapped_column(default=datetime.utcnow)

    conversation: Mapped["Conversation"] = relationship(
        back_populates="messages"
    )

Note the metadata_ Python attribute mapped to the metadata column name. This avoids shadowing SQLAlchemy's internal metadata attribute on the base class.

Async Session Setup

Configure an async engine and session factory for use with FastAPI:

from sqlalchemy.ext.asyncio import (
    AsyncSession,
    async_sessionmaker,
    create_async_engine,
)

DATABASE_URL = "postgresql+asyncpg://user:pass@localhost:5432/agents"

engine = create_async_engine(
    DATABASE_URL,
    pool_size=20,
    max_overflow=10,
    pool_pre_ping=True,
)

async_session = async_sessionmaker(
    engine, class_=AsyncSession, expire_on_commit=False
)


async def get_db() -> AsyncSession:
    async with async_session() as session:
        yield session

The expire_on_commit=False setting prevents SQLAlchemy from lazily reloading attributes after commit. In async code, lazy loading raises errors because it triggers implicit IO. Always set this for async sessions.

Relationship Loading Strategies

The most common performance mistake in ORM-based agent code is N+1 queries. When you load a conversation and then iterate over its messages, SQLAlchemy issues one query per message by default. Fix this with eager loading:

from sqlalchemy import select
from sqlalchemy.orm import selectinload


async def get_conversation_with_messages(
    session: AsyncSession, conversation_id: str
) -> Conversation | None:
    stmt = (
        select(Conversation)
        .where(Conversation.id == conversation_id)
        .options(selectinload(Conversation.messages))
    )
    result = await session.execute(stmt)
    return result.scalar_one_or_none()

selectinload issues a second query (SELECT ... WHERE conversation_id IN (...)) to batch-load all messages. For one-to-one relationships, joinedload is more efficient as it uses a single JOIN.

Transaction Patterns

Agent operations often need atomicity — creating a message and its tool calls should succeed or fail together:

async def create_message_with_tool_calls(
    session: AsyncSession,
    conversation_id: str,
    content: str,
    tool_calls_data: list[dict],
) -> Message:
    async with session.begin():
        message = Message(
            conversation_id=conversation_id,
            role="assistant",
            content=content,
        )
        session.add(message)
        await session.flush()  # Get the message ID

        for tc in tool_calls_data:
            tool_call = ToolCall(
                message_id=message.id,
                tool_name=tc["name"],
                arguments=tc["args"],
            )
            session.add(tool_call)

    return message

The session.begin() context manager automatically commits on success and rolls back on exception. The flush() call writes to the database within the transaction so the message ID is available for the tool call foreign keys.

Alembic Migration Setup

Generate migrations automatically from model changes:

alembic init alembic

Configure alembic/env.py to import your models and use the async engine. Then generate and apply migrations:

alembic revision --autogenerate -m "add_conversations_and_messages"
alembic upgrade head

Always review autogenerated migrations before applying them. Alembic cannot detect renamed columns (it sees a drop and add), and it may miss index changes on JSONB columns.

FAQ

Should I use SQLAlchemy Core or ORM for agent systems?

Use the ORM for application code that creates and queries agent data — it handles relationships, transactions, and type safety well. Use Core (raw text() or select() without models) for analytics queries, bulk inserts, or performance-critical paths where ORM overhead matters.

How do I handle connection pool exhaustion during high agent concurrency?

Set pool_size and max_overflow based on your database's max connections divided by the number of application instances. Monitor with pool.status() in health checks. If agents make long-running LLM calls while holding a database session, restructure the code to release the session before the LLM call and re-acquire it afterward.

Can SQLAlchemy async work with SQLite for testing?

Yes, using aiosqlite as the driver: sqlite+aiosqlite:///test.db. However, SQLite does not support JSONB, arrays, or concurrent writes. Use a PostgreSQL container via testcontainers-python for integration tests that match production behavior.

---

#SQLAlchemy #Python #ORM #Async #AIAgents #Database #AgenticAI #LearnAI #AIEngineering