Tutorial: Build Semantic Notes with Postgres, pgvector, Drizzle, and Redis

npm pkg set scripts.db:generate="drizzle-kit generate" scripts.db:migrate="drizzle-kit migrate" scripts.notes="node --env-file=.env src/cli.ts" scripts.test="node --test" scripts.typecheck="tsc"

import { sql } from "drizzle-orm";
import { check, index, pgTable, text, timestamp, uuid, vector } from "drizzle-orm/pg-core";
 
export const EMBEDDING_DIMENSIONS = 256;
 
export const users = pgTable("users", {
  id: uuid("id").primaryKey().defaultRandom(),
  email: text("email").notNull().unique(),
  createdAt: timestamp("created_at", { withTimezone: true }).notNull().defaultNow(),
});
 
export const notes = pgTable(
  "notes",
  {
    id: uuid("id").primaryKey().defaultRandom(),
    userId: uuid("user_id")
      .notNull()
      .references(() => users.id, { onDelete: "cascade" }),
    title: text("title").notNull(),
    body: text("body").notNull(),
    // Must match the embedding function's output size exactly.
    embedding: vector("embedding", { dimensions: EMBEDDING_DIMENSIONS }).notNull(),
    createdAt: timestamp("created_at", { withTimezone: true }).notNull().defaultNow(),
  },
  (t) => [
    index("notes_user_idx").on(t.userId, t.createdAt),
    index("notes_embedding_idx").using("hnsw", t.embedding.op("vector_cosine_ops")),
    // The database enforces the rule too, so no code path can store an invalid title.
    check("notes_title_length", sql`char_length(${t.title}) between 1 and 120`),
  ],
);

import { defineConfig } from "drizzle-kit";
 
process.loadEnvFile(); // reads .env, so drizzle-kit sees DATABASE_URL
 
export default defineConfig({
  schema: "./src/db/schema.ts",
  out: "./drizzle",
  dialect: "postgresql",
  dbCredentials: { url: process.env.DATABASE_URL! },
});

npm run db:generate
npm run db:migrate

CREATE TABLE "notes" (
	"id" uuid PRIMARY KEY DEFAULT gen_random_uuid() NOT NULL,
	"user_id" uuid NOT NULL,
	"title" text NOT NULL,
	"body" text NOT NULL,
	"embedding" vector(256) NOT NULL,
	"created_at" timestamp with time zone DEFAULT now() NOT NULL,
	CONSTRAINT "notes_title_length" CHECK (char_length("notes"."title") between 1 and 120)
);
CREATE INDEX "notes_user_idx" ON "notes" USING btree ("user_id","created_at");
CREATE INDEX "notes_embedding_idx" ON "notes" USING hnsw ("embedding" vector_cosine_ops);

import { createHash } from "node:crypto";
import { setTimeout as sleep } from "node:timers/promises";
import { EMBEDDING_DIMENSIONS } from "./db/schema.ts";
 
const STOP_WORDS = new Set(["the", "and", "for", "with", "this", "that", "from", "are", "was", "how", "what", "our", "your", "you", "keeps", "into"]);
 
/** Lowercase words, minus very common ones, with a crude plural/-ing/-ed strip. */
export function words(text: string): string[] {
  return (text.toLowerCase().match(/[a-z0-9]+/g) ?? [])
    .filter((w) => w.length > 2 && !STOP_WORDS.has(w))
    .map((w) => (w.length > 4 ? w.replace(/(ing|ed|s)$/, "") : w));
}
 
/**
 * A toy embedding: the "hashing trick". Each word is hashed to one of 256 dimensions, and the
 * vector is normalized to length 1. Texts that share words point in similar directions.
 * Real embedding models also match synonyms and meaning; this one only matches shared words.
 */
export function embedLocal(text: string): number[] {
  const vector = new Array<number>(EMBEDDING_DIMENSIONS).fill(0);
  for (const word of words(text)) {
    const hash = createHash("sha256").update(word).digest();
    vector[hash.readUInt16BE(0) % EMBEDDING_DIMENSIONS]! += hash[2]! & 1 ? 1 : -1;
  }
  const length = Math.hypot(...vector);
  return length === 0 ? vector : vector.map((x) => x / length);
}
 
/** Stand-in for an embedding API call: slow, and in real life billed per token. */
export async function embed(text: string): Promise<number[]> {
  await sleep(300);
  return embedLocal(text);
}

import { Redis } from "ioredis";
 
export const redis = new Redis(process.env.REDIS_URL ?? "redis://localhost:6379");
 
/** Fixed-window limiter: at most `limit` calls per `windowSeconds` for this key. */
export async function rateLimit(key: string, limit: number, windowSeconds: number) {
  const window = Math.floor(Date.now() / 1000 / windowSeconds);
  const bucket = `ratelimit:${key}:${window}`;
  // MULTI runs both commands atomically, so the counter always gets an expiry.
  const results = await redis.multi().incr(bucket).expire(bucket, windowSeconds).exec();
  const count = Number(results?.[0]?.[1] ?? 0);
  return { allowed: count <= limit, remaining: Math.max(0, limit - count) };
}
 
/** Return the cached value for `key`, or compute it, store it with a TTL, and return it. */
export async function cached<T>(key: string, ttlSeconds: number, compute: () => Promise<T>): Promise<{ value: T; hit: boolean }> {
  const hit = await redis.get(key);
  if (hit !== null) return { value: JSON.parse(hit) as T, hit: true };
  const value = await compute();
  await redis.set(key, JSON.stringify(value), "EX", ttlSeconds);
  return { value, hit: false };
}

import { drizzle } from "drizzle-orm/node-postgres";
import { Pool } from "pg";
import * as schema from "./schema.ts";
 
// One pool per process. Each connection costs memory on the database server.
export const pool = new Pool({ connectionString: process.env.DATABASE_URL, max: 10 });
export const db = drizzle({ client: pool, schema });

import { createHash } from "node:crypto";
import { and, cosineDistance, count, desc, eq, lt, sql } from "drizzle-orm";
import { db } from "./db/client.ts";
import { notes, users } from "./db/schema.ts";
import { embed, words } from "./embed.ts";
import { cached } from "./redis.ts";
 
export interface NoteInput {
  title: string;
  body: string;
}
 
/** Find a user by email, creating them on first use. */
export async function getOrCreateUser(email: string) {
  const [user] = await db
    .insert(users)
    .values({ email })
    .onConflictDoUpdate({ target: users.email, set: { email } }) // a no-op update, so RETURNING gives the existing row
    .returning();
  return user!;
}
 
/** Embeddings never change for the same text and model, so cache them by a hash of the text. */
async function embedCached(text: string) {
  const key = `embed:v1:${createHash("sha256").update(text).digest("hex")}`;
  return cached(key, 60 * 60 * 24 * 30, () => embed(text));
}
 
/** Import several notes atomically: either every note is stored, or none is. */
export async function importNotes(userId: string, input: NoteInput[]) {
  // Slow external calls happen before the transaction, so it holds its connection for milliseconds only.
  const embeddings = await Promise.all(input.map((n) => embedCached(`${n.title}\n${n.body}`)));
  return db.transaction(async (tx) => {
    const rows = [];
    for (const [i, note] of input.entries()) {
      const [row] = await tx.insert(notes).values({ userId, ...note, embedding: embeddings[i]!.value }).returning({ id: notes.id });
      rows.push(row!);
    }
    return rows;
  });
}
 
/** The notes closest in meaning to `query`, from this user's notes only. */
export async function searchNotes(userId: string, query: string, { limit = 3, minSimilarity = 0.1 } = {}) {
  // A query of only common words embeds to the zero vector, which has no direction to compare.
  if (words(query).length === 0) return { results: [], cacheHit: false };
 
  const { value: queryEmbedding, hit } = await embedCached(query);
  const distance = cosineDistance(notes.embedding, queryEmbedding);
  const results = await db
    .select({ title: notes.title, similarity: sql<number>`1 - (${distance})` })
    .from(notes)
    .where(and(
      eq(notes.userId, userId), // the ownership filter: never search other users' notes
      lt(distance, 1 - minSimilarity), // drop matches too weak to be useful
    ))
    .orderBy(distance)
    .limit(limit);
  return { results, cacheHit: hit };
}
 
export async function noteCounts() {
  return db
    .select({ email: users.email, notes: count(notes.id) })
    .from(users)
    .leftJoin(notes, eq(notes.userId, users.id))
    .groupBy(users.email)
    .orderBy(desc(count(notes.id)));
}

import type { NoteInput } from "./notes.ts";
 
export const SEED: Record<string, NoteInput[]> = {
  "ada@example.com": [
    { title: "Postgres connection pool limits", body: "Too many clients: create one pool per process and cap max connections at 10." },
    { title: "pgvector index choice", body: "Use an HNSW index with cosine distance for embedding similarity search." },
    { title: "Redis rate limiting", body: "A fixed window counter with INCR and EXPIRE protects the model budget." },
    { title: "Prompt caching savings", body: "Cache the long system prompt; cached input tokens cost a fraction of the price." },
    { title: "Weekly team sync", body: "Demo the streaming chat on Friday and collect feedback on latency." },
  ],
  "grace@example.com": [
    { title: "Database connection errors in staging", body: "Staging Postgres refused connections: the pool max was set to 200." },
    { title: "Sourdough schedule", body: "Feed the starter at 8, mix at noon, bake the next morning." },
  ],
};

import { pool } from "./db/client.ts";
import { getOrCreateUser, importNotes, noteCounts, searchNotes } from "./notes.ts";
import { rateLimit, redis } from "./redis.ts";
import { SEED } from "./seed-data.ts";
 
const [command, email, ...rest] = process.argv.slice(2);
 
async function main() {
  switch (command) {
    case "seed": {
      for (const [address, input] of Object.entries(SEED)) {
        const user = await getOrCreateUser(address);
        const rows = await importNotes(user.id, input);
        console.log(`${address}: imported ${rows.length} notes`);
      }
      break;
    }
    case "add": {
      const [title, body] = rest;
      if (!email || !title || !body) throw new Error('Usage: add <email> "<title>" "<body>"');
      const user = await getOrCreateUser(email);
      await importNotes(user.id, [{ title, body }]);
      console.log(`added "${title}" for ${email}`);
      break;
    }
    case "search": {
      const query = rest.join(" ");
      if (!email || !query) throw new Error('Usage: search <email> "<query>"');
      const user = await getOrCreateUser(email);
 
      const limit = await rateLimit(`search:${user.id}`, 5, 60);
      if (!limit.allowed) {
        console.log("Rate limited: at most 5 searches per minute. Try again shortly.");
        process.exitCode = 1;
        return;
      }
 
      const started = performance.now();
      const { results, cacheHit } = await searchNotes(user.id, query);
      const ms = Math.round(performance.now() - started);
      console.log(`Top matches for ${email} (${ms} ms, embedding cache ${cacheHit ? "hit" : "miss"}, ${limit.remaining} searches left this minute):`);
      for (const r of results) console.log(`  ${r.similarity.toFixed(2)}  ${r.title}`);
      if (results.length === 0) console.log("  (no notes match closely enough)");
      break;
    }
    case "import-bad": {
      if (!email) throw new Error("Usage: import-bad <email>");
      const user = await getOrCreateUser(email);
      const before = (await noteCounts()).find((r) => r.email === email)?.notes ?? 0;
      try {
        await importNotes(user.id, [
          { title: "A perfectly fine note", body: "This one is valid." },
          { title: "x".repeat(200), body: "This title breaks the 120-character check constraint." },
        ]);
      } catch (error) {
        const cause = (error as { cause?: { message?: string } }).cause;
        console.log(`import failed: ${cause?.message ?? (error as Error).message}`);
      }
      const after = (await noteCounts()).find((r) => r.email === email)?.notes ?? 0;
      console.log(`notes for ${email}: ${before} before, ${after} after (the valid note was rolled back too)`);
      break;
    }
    case "stats": {
      for (const row of await noteCounts()) console.log(`${row.email.padEnd(20)} ${row.notes} notes`);
      break;
    }
    default:
      console.log("Commands: seed | add <email> <title> <body> | search <email> <query> | import-bad <email> | stats");
  }
}
 
try {
  await main();
} finally {
  // Close connections, or the process stays alive waiting on them.
  await pool.end();
  redis.disconnect();
}

npm run notes -- seed
npm run notes -- search ada@example.com "my database keeps running out of connections"

ada@example.com: imported 5 notes
grace@example.com: imported 2 notes
 
Top matches for ada@example.com (318 ms, embedding cache miss, 4 searches left this minute):
  0.23  Postgres connection pool limits

Top matches for ada@example.com (4 ms, embedding cache hit, 3 searches left this minute):
  0.23  Postgres connection pool limits

npm run notes -- search ada@example.com "how do I stop users from burning my model budget"
#   0.27  Redis rate limiting

npm run notes -- search grace@example.com "database connection"
#   0.53  Database connection errors in staging
npm run notes -- search ada@example.com "staging errors"
#   (no notes match closely enough)

npm run notes -- import-bad ada@example.com

import failed: new row for relation "notes" violates check constraint "notes_title_length"
notes for ada@example.com: 5 before, 5 after (the valid note was rolled back too)

import assert from "node:assert/strict";
import { test } from "node:test";
import { embedLocal, words } from "../src/embed.ts";
 
const cosine = (a: number[], b: number[]) => a.reduce((sum, x, i) => sum + x * b[i]!, 0);
 
test("words lowercases, drops short and common words, and strips plurals", () => {
  assert.deepEqual(words("The Connections keep failing in Postgres"), ["connection", "keep", "fail", "postgre"]);
});
 
test("embeddings have the column's dimension and unit length", () => {
  const v = embedLocal("pgvector index for embeddings");
  assert.equal(v.length, 256);
  assert.ok(Math.abs(Math.hypot(...v) - 1) < 1e-9);
});
 
test("texts that share words are more similar than texts that do not", () => {
  const query = embedLocal("postgres connection pool");
  const related = embedLocal("Postgres refused connections: the pool max was too high");
  const unrelated = embedLocal("Feed the sourdough starter in the morning");
  assert.ok(cosine(query, related) > cosine(query, unrelated));
});
 
test("text with no searchable words embeds to the zero vector", () => {
  assert.ok(embedLocal("the and for").every((x) => x === 0));
});

npm test          # 4 passing
npm run typecheck