How to Use PostgresStateStore for Persistent Agent State in aisuite

Use the PostgresStateStore.from_dsn() factory method to connect to a PostgreSQL database, set create_schema=True to initialize the required table automatically, and pass the resulting instance to the Agent constructor via the state_store parameter.

aisuite is an open-source Python library developed by Andrew Ng's team that unifies interactions with multiple LLM providers. The PostgresStateStore class, located in aisuite/agents/postgres_state_store.py, provides durable, PostgreSQL-backed persistence for agent state, ensuring that conversational context and working memory survive process restarts, crashes, and distributed deployments.

What is PostgresStateStore?

PostgresStateStore is one of three built-in state store implementations in aisuite, alongside InMemoryStateStore and FileStateStore. Unlike volatile memory or local file storage, this class persists every key-value revision to a PostgreSQL table named aisuite_state. The implementation supports thread-level isolation, automatic schema creation, and compaction for space management, making it suitable for production deployments where high availability and data durability are required.

Instantiating PostgresStateStore

Connecting via from_dsn()

The primary entry point is the from_dsn() class method, which accepts a PostgreSQL connection string and returns a configured store instance. The method signature is:

from_dsn(cls, dsn: str, *, create_schema: bool = False) -> PostgresStateStore

Provide a standard PostgreSQL DSN following the format postgresql://user:password@host:port/database. If create_schema=True, the method automatically creates the aisuite_state table if it does not exist.

from aisuite.agents import PostgresStateStore

DSN = "postgresql://aisuite_user:secret@localhost:5432/aisuite"

# Initialize store and create schema on first run

store = PostgresStateStore.from_dsn(DSN, create_schema=True)

Schema Validation and Error Handling

If create_schema=False (the default) and the aisuite_state table is missing, from_dsn() raises a RuntimeError with the message: "PostgresStateStore.from_dsn() requires the aisuite_state table – set create_schema=True". This guard prevents accidental connections to incorrect databases or uninitialized environments.

Attaching the Store to an Agent

Pass the PostgresStateStore instance to the Agent constructor using the state_store parameter. Once attached, all internal state operations—including conversation history and tool outputs—persist automatically to PostgreSQL.

from aisuite import Agent

agent = Agent(
    name="research-assistant",
    # ... other configuration (tools, LLM, etc.)

    state_store=store,  # Persistent PostgreSQL backing

)

agent.run()

Core API Operations

Storing Values with set()

Use set() to insert or update a value. Each call creates a new revision row in the database.

store.set("conversation_history", [{"role": "user", "content": "Hello"}])

The method signature includes an optional thread_id for isolation:

set(self, key: str, value: Any, thread_id: int = 0) -> None

Retrieving Values with get()

Fetch the latest revision of a key using get(). If the key does not exist, the method raises KeyError.

history = store.get("conversation_history")
print(history)  # -> [{'role': 'user', 'content': 'Hello'}]

Removing Data with delete()

Remove all revisions of a specific key permanently from the database.

store.delete("conversation_history")

Optimizing Storage with compact()

Over time, the append-only storage accumulates multiple revisions per key. Call compact() to rewrite the table, keeping only the newest revision of each key for the specified thread_id. The store also supports automatic compaction that triggers after a configurable number of writes.

store.compact()  # Reclaim space by removing old revisions

Closing Connections with close()

Explicitly close the underlying database connection when the store is no longer needed, such as during application shutdown.

store.close()

Thread Isolation for Multi-Agent Deployments

The thread_id parameter (defaulting to 0) enables multiple agents to share a single PostgreSQL database without namespace collisions. When calling get(), set(), or delete(), the store filters by both key and thread_id, effectively creating isolated storage partitions.

# Agent A operates in thread 1

store.set("session_id", "A-123", thread_id=1)

# Agent B operates in thread 2

store.set("session_id", "B-456", thread_id=2)

assert store.get("session_id", thread_id=1) == "A-123"
assert store.get("session_id", thread_id=2) == "B-456"

This mechanism allows horizontal scaling of agents across multiple containers or servers while using one centralized PostgreSQL cluster.

Summary

• Instantiate PostgresStateStore via from_dsn() using a valid PostgreSQL DSN string.
• Initialize schemas automatically by setting create_schema=True on first deployment.
• Attach persistence to any Agent by passing the store to the state_store constructor argument.
• Isolate agents using distinct thread_id values to prevent data leakage in shared databases.
• Maintain performance by calling compact() periodically to purge obsolete revisions and reclaim disk space.

Frequently Asked Questions

What is the difference between PostgresStateStore and FileStateStore?

PostgresStateStore persists data to a PostgreSQL table (aisuite_state), enabling durability across process restarts and support for distributed multi-host deployments. FileStateStore writes to local disk and cannot be shared between servers. Additionally, PostgresStateStore supports thread isolation via the thread_id parameter and offers compaction strategies for storage optimization.

Does PostgresStateStore require manual database migrations?

No manual migrations are required if you instantiate the store using create_schema=True in the from_dsn() method. This parameter automatically creates the necessary aisuite_state table with the correct schema. If the table is missing and create_schema=False, the method raises a RuntimeError directing you to enable schema creation.

How does thread_id isolation work in PostgresStateStore?

The thread_id parameter acts as a namespace qualifier within the database table. Every get(), set(), and delete() operation filters by the provided thread_id (default 0), ensuring that agents with different identifiers access distinct data partitions even when sharing the same database connection and table.

When should I call compact() on a PostgresStateStore instance?

Call compact() after batch write operations or during scheduled maintenance windows to remove historical revisions of keys and reclaim disk space. While the store supports automatic compaction after a configurable number of writes, manual compaction provides explicit control over when database resources are allocated to cleanup operations.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}