Memory

Drizzle

Implement MemoryStore with Drizzle ORM.

This example uses Drizzle with PostgreSQL. Store messages as jsonb, then load them by sessionId in insertion order.

Tables

import { index, integer, jsonb, pgTable, text, timestamp, bigserial } from "drizzle-orm/pg-core";
import type { Message } from "@anvia/core";

export const agentMemoryMessages = pgTable(
  "agent_memory_messages",
  {
    id: bigserial("id", { mode: "number" }).primaryKey(),
    sessionId: text("session_id").notNull(),
    userId: text("user_id"),
    runId: text("run_id").notNull(),
    turn: integer("turn").notNull(),
    message: jsonb("message").$type<Message>().notNull(),
    metadata: jsonb("metadata"),
    createdAt: timestamp("created_at", { withTimezone: true }).notNull().defaultNow(),
  },
  (table) => ({
    sessionOrderIdx: index("agent_memory_messages_session_id_id_idx").on(
      table.sessionId,
      table.id,
    ),
  }),
);

export const agentMemoryErrors = pgTable(
  "agent_memory_errors",
  {
    id: bigserial("id", { mode: "number" }).primaryKey(),
    sessionId: text("session_id").notNull(),
    userId: text("user_id"),
    runId: text("run_id").notNull(),
    error: jsonb("error").notNull(),
    messages: jsonb("messages").$type<Message[]>().notNull(),
    createdAt: timestamp("created_at", { withTimezone: true }).notNull().defaultNow(),
  },
  (table) => ({
    sessionCreatedIdx: index("agent_memory_errors_session_id_created_at_idx").on(
      table.sessionId,
      table.createdAt,
    ),
  }),
);

Store

import type {
  MemoryAppendInput,
  MemoryContext,
  MemoryErrorInput,
  MemoryStore,
  Message,
} from "@anvia/core";
import { asc, eq } from "drizzle-orm";
import type { NodePgDatabase } from "drizzle-orm/node-postgres";
import { agentMemoryErrors, agentMemoryMessages } from "./schema";

export class DrizzleMemoryStore implements MemoryStore {
  constructor(private readonly db: NodePgDatabase) {}

  async load(context: MemoryContext): Promise<Message[]> {
    const rows = await this.db
      .select({ message: agentMemoryMessages.message })
      .from(agentMemoryMessages)
      .where(eq(agentMemoryMessages.sessionId, context.sessionId))
      .orderBy(asc(agentMemoryMessages.id));

    return rows.map((row) => row.message);
  }

  async append(input: MemoryAppendInput): Promise<void> {
    if (input.messages.length === 0) {
      return;
    }

    await this.db.insert(agentMemoryMessages).values(
      input.messages.map((message) => ({
        sessionId: input.context.sessionId,
        userId: input.context.userId,
        runId: input.runId,
        turn: input.turn,
        message,
        metadata: input.context.metadata,
      })),
    );
  }

  async clear(context: MemoryContext): Promise<void> {
    await this.db
      .delete(agentMemoryMessages)
      .where(eq(agentMemoryMessages.sessionId, context.sessionId));
  }

  async recordError(input: MemoryErrorInput): Promise<void> {
    await this.db.insert(agentMemoryErrors).values({
      sessionId: input.context.sessionId,
      userId: input.context.userId,
      runId: input.runId,
      error: serializeError(input.error),
      messages: input.messages,
    });
  }
}

function serializeError(error: unknown): Record<string, unknown> {
  if (error instanceof Error) {
    return {
      name: error.name,
      message: error.message,
      stack: error.stack,
    };
  }
  return { message: String(error) };
}

Use It

const memory = new DrizzleMemoryStore(db);

const agent = new AgentBuilder("support", model)
  .memory(memory)
  .build();

await agent.session("thread_123", { userId: "user_456" }).prompt("Hello").send();