Documentation
¶
Overview ¶
Package exp provides experimental PostgreSQL integrations for Genkit's agent runtime (see github.com/firebase/genkit/go/ai/exp).
The PostgresSessionStore persists agent session snapshots in PostgreSQL. It resolves its connection pool from the PostgreSQL plugin registered with the Genkit instance, then wires into an agent:
engine, err := postgresql.NewPostgresEngine(ctx,
postgresql.WithCloudSQLInstance("my-project", "us-central1", "my-instance"),
postgresql.WithDatabase("my-database"))
// handle err; postgresql.WithPool(pool) connects to any other PostgreSQL
g := genkit.Init(ctx,
genkit.WithExperimental(),
genkit.WithPlugins(&postgresql.Postgres{Engine: engine}))
store, err := exp.NewPostgresSessionStore[MyState](ctx, g)
// handle err
agent := genkitx.DefineAgent(g, "assistant",
aix.InlinePrompt{ai.WithModelName("googleai/gemini-flash-latest")},
aix.WithSessionStore(store))
APIs in this package are under active development and may change in any minor version release. Use with caution in production environments.
Index ¶
- type PostgresSessionStore
- func (s *PostgresSessionStore[State]) GetLatestSnapshot(ctx context.Context, sessionID string) (*aix.SessionSnapshot[State], error)
- func (s *PostgresSessionStore[State]) GetLatestSnapshotMetadata(ctx context.Context, sessionID string) (*aix.SessionSnapshot[State], error)
- func (s *PostgresSessionStore[State]) GetSnapshot(ctx context.Context, snapshotID string) (*aix.SessionSnapshot[State], error)
- func (s *PostgresSessionStore[State]) GetSnapshotMetadata(ctx context.Context, snapshotID string) (*aix.SessionSnapshot[State], error)
- func (s *PostgresSessionStore[State]) OnSnapshotStatusChange(ctx context.Context, snapshotID string) <-chan aix.SnapshotStatus
- func (s *PostgresSessionStore[State]) SaveSnapshot(ctx context.Context, id string, ...) (*aix.SessionSnapshot[State], error)
- type SessionStoreOption
- func WithCheckpointInterval(turns int) SessionStoreOption
- func WithPollInterval(d time.Duration) SessionStoreOption
- func WithSchemaName(name string) SessionStoreOption
- func WithSnapshotPathPrefix(fn func(ctx context.Context) string) SessionStoreOption
- func WithTableName(name string) SessionStoreOption
Constants ¶
This section is empty.
Variables ¶
This section is empty.
Functions ¶
This section is empty.
Types ¶
type PostgresSessionStore ¶
type PostgresSessionStore[State any] struct { // contains filtered or unexported fields }
PostgresSessionStore is a PostgreSQL-backed aix.SessionStore that persists session snapshots as JSON Patch diffs anchored to periodic full-state checkpoints, like the Firestore store, in one table:
- A checkpoint row holds the full state in its state column.
- A diff row holds the JSON Patch from its parent's state in its state_patch column, and its depth: the number of diffs back to the nearest checkpoint.
- A row without state (a pending detached run) is a checkpoint with a NULL state.
Every N turns (see WithCheckpointInterval) a row is written as a checkpoint, so a read applies at most N patches and each turn writes only its change, however long the session grows. A diff larger than half of the full state is written as a checkpoint instead. Large values need no sharding: PostgreSQL compresses them and stores them out of line, up to 1 GB each.
Rows are keyed by (prefix, snapshot ID), where the prefix is the tenant key from WithSnapshotPathPrefix, or "global". An index on (prefix, session ID, CreatedAt, snapshot ID) answers a session's latest row in one lookup, and a read walks from a row to its checkpoint in one recursive query, so it sees the whole chain at one point in time. State and patches are stored as json, which keeps the exact text the store wrote: jsonb would reject text that contains U+0000, which tool output can carry.
SaveSnapshot runs in one READ COMMITTED transaction, whatever the database's default isolation level, that first takes an advisory lock on the snapshot ID, so concurrent saves of one snapshot wait for each other and fn runs once per save. It rejects, with FAILED_PRECONDITION, a change to the state of a row that another row stores a diff against, since the change would corrupt that row (see aix.SnapshotWriter).
It implements aix.SnapshotSubscriber over LISTEN/NOTIFY: a save that changes the status of an existing row notifies, at commit, every store watching that row, so an abort committed by one process reaches the process running the detached turn at once. Only status changes notify, since PostgreSQL serializes the commits of transactions that do. Each store holds one connection for this while it has subscribers, and a poll re-reads the watched rows as a fallback (see WithPollInterval).
The store creates its table and indexes when the table does not exist, which needs the CREATE privilege on the schema. Where the schema is managed elsewhere, construct the store once with a role that has it; from then on the store needs only SELECT, INSERT, and UPDATE on the table.
func NewPostgresSessionStore ¶
func NewPostgresSessionStore[State any](ctx context.Context, g *genkit.Genkit, opts ...SessionStoreOption) (*PostgresSessionStore[State], error)
NewPostgresSessionStore creates a PostgreSQL-backed snapshot store. It resolves the connection pool from the PostgreSQL plugin registered with g (pass the plugin to genkit.Init before calling this), and creates the snapshot table when it does not exist.
The State type parameter is the user-defined custom-state type carried in aix.SessionState.Custom; it must be JSON-serializable.
func (*PostgresSessionStore[State]) GetLatestSnapshot ¶
func (s *PostgresSessionStore[State]) GetLatestSnapshot(ctx context.Context, sessionID string) (*aix.SessionSnapshot[State], error)
GetLatestSnapshot returns the session's most recently created snapshot regardless of status, per the aix.SnapshotReader.GetLatestSnapshot contract: the greatest CreatedAt, ties broken by the greater snapshot ID.
func (*PostgresSessionStore[State]) GetLatestSnapshotMetadata ¶
func (s *PostgresSessionStore[State]) GetLatestSnapshotMetadata(ctx context.Context, sessionID string) (*aix.SessionSnapshot[State], error)
GetLatestSnapshotMetadata is PostgresSessionStore.GetLatestSnapshot without the state, per aix.SnapshotMetadataReader.
func (*PostgresSessionStore[State]) GetSnapshot ¶
func (s *PostgresSessionStore[State]) GetSnapshot(ctx context.Context, snapshotID string) (*aix.SessionSnapshot[State], error)
GetSnapshot retrieves a snapshot by ID. Returns nil if not found. One query reads the row and its chain back to the checkpoint, at one point in time.
func (*PostgresSessionStore[State]) GetSnapshotMetadata ¶
func (s *PostgresSessionStore[State]) GetSnapshotMetadata(ctx context.Context, snapshotID string) (*aix.SessionSnapshot[State], error)
GetSnapshotMetadata retrieves a snapshot's row without its state, per aix.SnapshotMetadataReader: one row, and none of its state or chain. Returns nil if not found.
func (*PostgresSessionStore[State]) OnSnapshotStatusChange ¶
func (s *PostgresSessionStore[State]) OnSnapshotStatusChange(ctx context.Context, snapshotID string) <-chan aix.SnapshotStatus
OnSnapshotStatusChange returns a channel that yields the snapshot's status at subscription time and on every later change, until ctx is cancelled. Changes arrive over LISTEN/NOTIFY, so one written by any process reaches every process watching the row, and the poll covers one written without a notification. If the snapshot does not exist when the subscription is established, the channel is closed without yielding a value.
Values are level-triggered: the latest status is always delivered, but a slow reader may skip intermediate values. Treat a received value as "the status is now X", not "X happened once".
func (*PostgresSessionStore[State]) SaveSnapshot ¶
func (s *PostgresSessionStore[State]) SaveSnapshot( ctx context.Context, id string, fn func(existing *aix.SessionSnapshot[State]) (*aix.SessionSnapshot[State], error), ) (*aix.SessionSnapshot[State], error)
SaveSnapshot atomically reads the snapshot at id (if any), applies fn, and persists the result. See aix.SnapshotWriter for the full contract. The read-modify-write runs in one transaction under an advisory lock on the snapshot ID, so fn runs once per call and concurrent saves of one snapshot apply one after the other.
type SessionStoreOption ¶
type SessionStoreOption interface {
// contains filtered or unexported methods
}
SessionStoreOption configures a PostgresSessionStore. A later option of the same kind replaces an earlier one.
func WithCheckpointInterval ¶
func WithCheckpointInterval(turns int) SessionStoreOption
WithCheckpointInterval sets the number of turns between full-state checkpoints. Between checkpoints each row stores only the JSON Patch from its parent's state, so a larger value writes less per turn and reads apply more patches; a smaller value does the opposite. A read applies at most this many patches, however long the session. With 1, every row is a checkpoint, and a save does not read or lock the parent row. Must be between 1 and 1000; defaults to 25 when omitted.
func WithPollInterval ¶
func WithPollInterval(d time.Duration) SessionStoreOption
WithPollInterval sets how often a status subscription re-reads the rows it watches. Status changes arrive at once over LISTEN/NOTIFY; the poll is the fallback for a change written without a notification (a manual UPDATE) and for a connection that cannot receive one, such as a connection through a transaction-mode pooler. Defaults to 5 seconds when omitted; a value <= 0 turns the poll off.
func WithSchemaName ¶
func WithSchemaName(name string) SessionStoreOption
WithSchemaName sets the schema the snapshot table lives in. The schema must exist. When omitted, the table name resolves through the connection's search_path, the way an unqualified name does in SQL.
func WithSnapshotPathPrefix ¶
func WithSnapshotPathPrefix(fn func(ctx context.Context) string) SessionStoreOption
WithSnapshotPathPrefix derives a per-call tenant prefix from the operation's context. Every row is keyed by its prefix, so reads, writes, and status subscriptions are isolated per tenant: one tenant can never address another's snapshots, even holding a snapshot ID, because resolving it still requires the matching, auth-derived prefix. A typical fn pulls a stable identity (an authenticated user or org ID) out of ctx. The name matches the other session stores; here the prefix is a column, not a path.
The value must be stable for a given snapshot's lifetime, since every read recomputes it. It must be non-empty: an empty result is rejected at call time, since the way to request the default "global" prefix is to omit this option.
func WithTableName ¶
func WithTableName(name string) SessionStoreOption
WithTableName sets the table session snapshots are stored in. The store creates it, with its indexes, when it does not exist. Defaults to "genkit_session_snapshots" when omitted; at most 51 bytes, so the index names derived from it fit PostgreSQL's identifier limit.
Source Files