Building a Type-Safe Gemini Chat Store with Drizzle ORM — Multi-Turn Persistence, Branching Threads, and Vector Search in Production

// Simplified Gemini SDK types
interface Content {
  role: 'user' | 'model';
  parts: Part[];
}
 
type Part =
  | { text: string }
  | { inlineData: { mimeType: string; data: string } }
  | { functionCall: { name: string; args: Record<string, unknown> } }
  | { functionResponse: { name: string; response: Record<string, unknown> } };

// db/schema.ts
import {
  pgTable,
  uuid,
  text,
  timestamp,
  jsonb,
  boolean,
  integer,
  index,
} from 'drizzle-orm/pg-core';
import type { Part } from '@google/genai';
 
export const threads = pgTable(
  'threads',
  {
    id: uuid('id').defaultRandom().primaryKey(),
    userId: uuid('user_id').notNull(),
    title: text('title').notNull().default('New chat'),
    parentMessageId: uuid('parent_message_id'),
    createdAt: timestamp('created_at', { withTimezone: true })
      .notNull()
      .defaultNow(),
    updatedAt: timestamp('updated_at', { withTimezone: true })
      .notNull()
      .defaultNow(),
  },
  (t) => ({
    userIdx: index('threads_user_idx').on(t.userId, t.updatedAt),
  })
);
 
export const messages = pgTable(
  'messages',
  {
    id: uuid('id').defaultRandom().primaryKey(),
    threadId: uuid('thread_id')
      .notNull()
      .references(() => threads.id, { onDelete: 'cascade' }),
    role: text('role', { enum: ['user', 'model', 'system'] }).notNull(),
    parts: jsonb('parts').$type<Part[]>().notNull(),
    // promoted columns for indexing & analytics
    textPreview: text('text_preview'),
    hasImage: boolean('has_image').notNull().default(false),
    hasFunctionCall: boolean('has_function_call').notNull().default(false),
    inputTokens: integer('input_tokens'),
    outputTokens: integer('output_tokens'),
    // streaming state
    status: text('status', {
      enum: ['streaming', 'completed', 'failed', 'cancelled'],
    })
      .notNull()
      .default('completed'),
    finishReason: text('finish_reason'),
    createdAt: timestamp('created_at', { withTimezone: true })
      .notNull()
      .defaultNow(),
    completedAt: timestamp('completed_at', { withTimezone: true }),
  },
  (t) => ({
    threadIdx: index('messages_thread_idx').on(t.threadId, t.createdAt),
    statusIdx: index('messages_status_idx').on(t.status),
  })
);
 
export type Message = typeof messages.$inferSelect;
export type NewMessage = typeof messages.$inferInsert;
export type Thread = typeof threads.$inferSelect;

// app/api/chat/route.ts
import { GoogleGenAI } from '@google/genai';
import { db } from '@/db';
import { messages } from '@/db/schema';
import { eq } from 'drizzle-orm';
 
export async function POST(req: Request) {
  const { threadId, userMessage } = await req.json();
  const ai = new GoogleGenAI({ apiKey: process.env.YOUR_GEMINI_API_KEY! });
 
  // 1) Persist the user turn immediately
  await db.insert(messages).values({
    threadId,
    role: 'user',
    parts: [{ text: userMessage }],
    textPreview: userMessage.slice(0, 200),
    status: 'completed',
    completedAt: new Date(),
  });
 
  // 2) Pre-create the model row in streaming state
  const [modelRow] = await db
    .insert(messages)
    .values({
      threadId,
      role: 'model',
      parts: [],
      status: 'streaming',
    })
    .returning({ id: messages.id });
 
  // 3) Pull completed history into the contents array
  const history = await db.query.messages.findMany({
    where: (m, { eq, and }) =>
      and(eq(m.threadId, threadId), eq(m.status, 'completed')),
    orderBy: (m, { asc }) => asc(m.createdAt),
  });
  const contents = history.map((m) => ({
    role: m.role === 'model' ? 'model' : 'user',
    parts: m.parts as { text: string }[],
  }));
 
  // 4) Stream and persist incrementally
  const stream = await ai.models.generateContentStream({
    model: 'gemini-2.5-pro',
    contents,
  });
 
  const collected: string[] = [];
  let inputTokens = 0;
  let outputTokens = 0;
  let finishReason: string | undefined;
 
  const transformStream = new TransformStream({
    async transform(chunk: any, controller) {
      const text = chunk.text ?? '';
      if (text) {
        collected.push(text);
        controller.enqueue(new TextEncoder().encode(text));
      }
      if (chunk.usageMetadata) {
        inputTokens = chunk.usageMetadata.promptTokenCount ?? inputTokens;
        outputTokens = chunk.usageMetadata.candidatesTokenCount ?? outputTokens;
      }
      if (chunk.candidates?.[0]?.finishReason) {
        finishReason = chunk.candidates[0].finishReason;
      }
      // Mid-flight save every ~20 chunks
      if (collected.length % 20 === 0) {
        await db
          .update(messages)
          .set({ parts: [{ text: collected.join('') }] })
          .where(eq(messages.id, modelRow.id));
      }
    },
    async flush() {
      const fullText = collected.join('');
      await db
        .update(messages)
        .set({
          parts: [{ text: fullText }],
          textPreview: fullText.slice(0, 200),
          inputTokens,
          outputTokens,
          status: 'completed',
          finishReason: finishReason ?? 'STOP',
          completedAt: new Date(),
        })
        .where(eq(messages.id, modelRow.id));
    },
  });
 
  return new Response(
    (stream as any).readable.pipeThrough(transformStream),
    { headers: { 'Content-Type': 'text/plain; charset=utf-8' } }
  );
}

// lib/forkThread.ts
import { db } from '@/db';
import { threads, messages } from '@/db/schema';
 
export async function forkThread(params: {
  sourceThreadId: string;
  forkAtMessageId: string;
  userId: string;
}) {
  const sourceMessage = await db.query.messages.findFirst({
    where: (m, { eq }) => eq(m.id, params.forkAtMessageId),
  });
  if (!sourceMessage) throw new Error('Source message not found');
 
  const [newThread] = await db
    .insert(threads)
    .values({
      userId: params.userId,
      title: 'Forked chat',
      parentMessageId: params.forkAtMessageId,
    })
    .returning();
 
  const ancestors = await db.query.messages.findMany({
    where: (m, { and, eq, lte }) =>
      and(
        eq(m.threadId, params.sourceThreadId),
        lte(m.createdAt, sourceMessage.createdAt)
      ),
    orderBy: (m, { asc }) => asc(m.createdAt),
  });
 
  if (ancestors.length > 0) {
    await db.insert(messages).values(
      ancestors.map((m) => ({
        threadId: newThread.id,
        role: m.role,
        parts: m.parts,
        textPreview: m.textPreview,
        hasImage: m.hasImage,
        hasFunctionCall: m.hasFunctionCall,
        inputTokens: m.inputTokens,
        outputTokens: m.outputTokens,
        status: m.status,
        finishReason: m.finishReason,
        completedAt: m.completedAt,
      }))
    );
  }
 
  return newThread;
}

// lib/replayHistory.ts
import type { Content, Part } from '@google/genai';
 
export function rebuildContents(
  messages: { role: string; parts: Part[] }[]
): Content[] {
  return messages.map((m) => {
    const hasFunctionResponse = m.parts.some((p) => 'functionResponse' in p);
    const role = hasFunctionResponse
      ? 'user'
      : m.role === 'model'
        ? 'model'
        : 'user';
    return { role, parts: m.parts };
  });
}

// db/schema-vector.ts
import { pgTable, uuid, text, index } from 'drizzle-orm/pg-core';
import { vector } from 'drizzle-orm/pg-core';
import { messages } from './schema';
 
export const messageEmbeddings = pgTable(
  'message_embeddings',
  {
    messageId: uuid('message_id')
      .primaryKey()
      .references(() => messages.id, { onDelete: 'cascade' }),
    embedding: vector('embedding', { dimensions: 768 }).notNull(),
    model: text('model').notNull().default('gemini-embedding-001'),
  },
  (t) => ({
    embeddingIdx: index('embedding_hnsw_idx').using(
      'hnsw',
      t.embedding.op('vector_cosine_ops')
    ),
  })
);

// lib/embedMessage.ts
import { GoogleGenAI } from '@google/genai';
import { db } from '@/db';
import { messageEmbeddings } from '@/db/schema-vector';
 
const ai = new GoogleGenAI({ apiKey: process.env.YOUR_GEMINI_API_KEY! });
 
export async function embedAndStore(messageId: string, text: string) {
  if (!text.trim()) return;
  try {
    const result = await ai.models.embedContent({
      model: 'gemini-embedding-001',
      contents: text.slice(0, 8000),
    });
    const embedding = result.embeddings?.[0]?.values;
    if (!embedding) return;
    await db
      .insert(messageEmbeddings)
      .values({ messageId, embedding })
      .onConflictDoUpdate({
        target: messageEmbeddings.messageId,
        set: { embedding },
      });
  } catch (err) {
    // Embedding failures must not block the main response
    console.error('Embedding failed:', messageId, err);
  }
}

import { sql, eq } from 'drizzle-orm';
import { messages, threads } from '@/db/schema';
 
export async function searchSimilarMessages(
  queryEmbedding: number[],
  userId: string,
  limit = 10
) {
  return db
    .select({
      messageId: messageEmbeddings.messageId,
      threadId: messages.threadId,
      preview: messages.textPreview,
      similarity: sql<number>`1 - (${messageEmbeddings.embedding} <=> ${queryEmbedding}::vector)`,
    })
    .from(messageEmbeddings)
    .innerJoin(messages, eq(messages.id, messageEmbeddings.messageId))
    .innerJoin(threads, eq(threads.id, messages.threadId))
    .where(eq(threads.userId, userId))
    .orderBy(sql`${messageEmbeddings.embedding} <=> ${queryEmbedding}::vector`)
    .limit(limit);
}

// lib/usage-report.ts
import { sql, gte, eq, and } from 'drizzle-orm';
import { db } from '@/db';
import { messages, threads } from '@/db/schema';
 
const PRICE_PER_M_TOKENS = { input: 1.25, output: 5.0 };
 
export async function userMonthlyCost(userId: string, monthStart: Date) {
  const result = await db
    .select({
      inputTokens: sql<number>`SUM(${messages.inputTokens})::int`,
      outputTokens: sql<number>`SUM(${messages.outputTokens})::int`,
      messageCount: sql<number>`COUNT(*)::int`,
    })
    .from(messages)
    .innerJoin(threads, eq(threads.id, messages.threadId))
    .where(
      and(
        eq(threads.userId, userId),
        eq(messages.role, 'model'),
        gte(messages.createdAt, monthStart)
      )
    );
  const row = result[0];
  const inputUsd = (row.inputTokens / 1_000_000) * PRICE_PER_M_TOKENS.input;
  const outputUsd = (row.outputTokens / 1_000_000) * PRICE_PER_M_TOKENS.output;
  return {
    inputTokens: row.inputTokens ?? 0,
    outputTokens: row.outputTokens ?? 0,
    messageCount: row.messageCount ?? 0,
    estimatedUsd: Number((inputUsd + outputUsd).toFixed(4)),
  };
}

// __tests__/thread-fork.test.ts
import { describe, it, expect, beforeEach } from 'vitest';
import { db } from '@/db';
import { threads, messages } from '@/db/schema';
import { forkThread } from '@/lib/forkThread';
 
describe('forkThread', () => {
  beforeEach(async () => {
    await db.delete(messages);
    await db.delete(threads);
  });
 
  it('copies ancestor messages up to the fork point', async () => {
    const userId = '00000000-0000-0000-0000-000000000001';
    const [thread] = await db
      .insert(threads)
      .values({ userId, title: 'Source' })
      .returning();
    const inserted = await db
      .insert(messages)
      .values([
        { threadId: thread.id, role: 'user', parts: [{ text: 'Q1' }] },
        { threadId: thread.id, role: 'model', parts: [{ text: 'A1' }] },
        { threadId: thread.id, role: 'user', parts: [{ text: 'Q2' }] },
        { threadId: thread.id, role: 'model', parts: [{ text: 'A2' }] },
      ])
      .returning();
 
    const fork = await forkThread({
      sourceThreadId: thread.id,
      forkAtMessageId: inserted[1].id,
      userId,
    });
 
    const forkMessages = await db.query.messages.findMany({
      where: (m, { eq }) => eq(m.threadId, fork.id),
      orderBy: (m, { asc }) => asc(m.createdAt),
    });
    expect(forkMessages.length).toBe(2);
    expect(fork.parentMessageId).toBe(inserted[1].id);
  });
});

ALTER TABLE threads ENABLE ROW LEVEL SECURITY;
ALTER TABLE messages ENABLE ROW LEVEL SECURITY;
 
CREATE POLICY threads_owner_only ON threads
  FOR ALL
  USING (user_id = current_setting('app.current_user_id', true)::uuid);
 
CREATE POLICY messages_via_thread ON messages
  FOR ALL
  USING (
    thread_id IN (
      SELECT id FROM threads
      WHERE user_id = current_setting('app.current_user_id', true)::uuid
    )
  );

// lib/db-with-user.ts
import { db } from '@/db';
import { sql } from 'drizzle-orm';
 
export async function withUser<T>(userId: string, fn: () => Promise<T>): Promise<T> {
  await db.execute(sql`SET LOCAL app.current_user_id = ${userId}`);
  return fn();
}

import { db } from '@/db';
import { sql } from 'drizzle-orm';
 
await db.execute(sql`
  UPDATE messages
  SET has_video = true
  WHERE has_video = false
    AND parts::text LIKE '%"mimeType":"video/%'
    AND created_at >= NOW() - INTERVAL '30 days'
`);

import { and, eq, lte } from 'drizzle-orm';
await db
  .update(messages)
  .set({ status: 'failed', finishReason: 'TIMEOUT' })
  .where(
    and(
      eq(messages.status, 'streaming'),
      lte(messages.createdAt, new Date(Date.now() - 10 * 60 * 1000))
    )
  );

const threadsWithLastMessage = await db.execute(sql`
  SELECT t.*, m.text_preview AS last_preview, m.created_at AS last_message_at
  FROM threads t
  LEFT JOIN LATERAL (
    SELECT text_preview, created_at
    FROM messages
    WHERE thread_id = t.id AND status = 'completed'
    ORDER BY created_at DESC
    LIMIT 1
  ) m ON true
  WHERE t.user_id = ${userId}
  ORDER BY t.updated_at DESC
  LIMIT 50
`);