From d34a5c1f0c7f83246570aa4dd31799181c898cef Mon Sep 17 00:00:00 2001 From: Lucas Winther Date: Fri, 30 Oct 2026 20:28:23 +0100 Subject: [PATCH] feat(ai): citations column on ai_messages (contract request 5) New packages/feature-ai/migrations/0001_ai_message_citations.sql adds ai_messages.citations_json and folds v1 'sources' rows into the assistant row before them. The feature writes and reads the column directly. Co-Authored-By: Claude Opus 5.5 --- .../migrations/0001_ai_message_citations.sql | 10 +++++ packages/feature-ai/scripts/dev.ts | 3 +- packages/feature-ai/src/feature.ts | 32 +++++++-------- packages/feature-ai/src/index.ts | 3 +- packages/feature-ai/src/migrations.ts | 7 ++++ packages/feature-ai/test/feature.test.ts | 7 +++- packages/feature-ai/test/helpers.ts | 11 +++++ packages/feature-ai/test/settings.test.ts | 2 +- packages/feature-ai/test/upgrade.test.ts | 41 +++++++++++++++++++ 9 files changed, 94 insertions(+), 22 deletions(-) create mode 100644 packages/feature-ai/migrations/0001_ai_message_citations.sql create mode 100644 packages/feature-ai/src/migrations.ts create mode 100644 packages/feature-ai/test/upgrade.test.ts diff --git a/packages/feature-ai/migrations/0001_ai_message_citations.sql b/packages/feature-ai/migrations/0001_ai_message_citations.sql new file mode 100644 index 0000000..4e04b03 --- /dev/null +++ b/packages/feature-ai/migrations/0001_ai_message_citations.sql @@ -0,0 +1,10 @@ +-- v1.1 (contract request 5): an assistant message's citations live in its own column. +-- v1 stored them as an extra row with role 'sources' right after the assistant row (ids sort in write order). +-- Fold each such row into the assistant row before it, then drop every 'sources' row. +ALTER TABLE ai_messages ADD COLUMN citations_json TEXT; +UPDATE ai_messages SET citations_json = ( + SELECT s.content FROM ai_messages s + WHERE s.conversation_id = ai_messages.conversation_id AND s.role = 'sources' + AND s.id = (SELECT MIN(n.id) FROM ai_messages n WHERE n.conversation_id = ai_messages.conversation_id AND n.id > ai_messages.id) +) WHERE role = 'assistant'; +DELETE FROM ai_messages WHERE role = 'sources'; diff --git a/packages/feature-ai/scripts/dev.ts b/packages/feature-ai/scripts/dev.ts index a31fe02..55e0af2 100644 --- a/packages/feature-ai/scripts/dev.ts +++ b/packages/feature-ai/scripts/dev.ts @@ -7,6 +7,7 @@ import { createCourseRepository } from "@stc/course-format"; import { connect, migrate } from "@stc/db"; import indexHtml from "../../../apps/web/index.html"; import { aiFeature } from "../src/feature"; +import { migrationsDirs } from "../src/migrations"; const config = loadConfig(); const who = process.env.STC_AI_DEV_USER ?? "owner"; @@ -32,7 +33,7 @@ const fake: AiProvider = { }; const db = await connect(config.databaseUrl); -await migrate(db); +await migrate(db, { "feature-ai": migrationsDirs }); const courses = await createCourseRepository(config.courseDirs); const ctx = { config, db, courses, auth } as ServerContext; const feature = aiFeature(ctx, process.env.STC_AI_DEV_FAKE ? { providerFactory: () => fake } : {}); diff --git a/packages/feature-ai/src/feature.ts b/packages/feature-ai/src/feature.ts index 13b5180..e283fc1 100644 --- a/packages/feature-ai/src/feature.ts +++ b/packages/feature-ai/src/feature.ts @@ -198,22 +198,23 @@ export function aiFeature(ctx: ServerContext, opts: AiFeatureOptions = {}): Feat const rows = await ctx.db.sql`SELECT id, user_id, course_id, title, created_at FROM ai_conversations WHERE id = ${id} AND user_id = ${userId}`; return rows[0] as { id: string; user_id: string; course_id: string; title: string | null; created_at: string } | undefined; }; - // Roles: "user", "assistant", and "sources" (JSON citations for the assistant message before it). - const addMessage = async (conversationId: string, role: "user" | "assistant" | "sources", content: string) => { - await ctx.db.sql`INSERT INTO ai_messages (id, conversation_id, role, content, created_at) VALUES (${Bun.randomUUIDv7()}, ${conversationId}, ${role}, ${content}, ${new Date().toISOString()})`; + /** Roles "user" and "assistant"; an assistant message's citations go in `citations_json` (contract request 5). */ + const addMessage = async (conversationId: string, role: "user" | "assistant", content: string, citations?: unknown[]) => { + const cites = citations ? JSON.stringify(citations) : null; + await ctx.db.sql`INSERT INTO ai_messages (id, conversation_id, role, content, citations_json, created_at) VALUES (${Bun.randomUUIDv7()}, ${conversationId}, ${role}, ${content}, ${cites}, ${new Date().toISOString()})`; }; - const messagesOf = async (conversationId: string) => - (await ctx.db.sql`SELECT id, role, content, created_at FROM ai_messages WHERE conversation_id = ${conversationId}`) - .map((r: { id: string; role: string; content: string; created_at: string }) => ({ id: r.id, role: r.role, content: r.content, createdAt: r.created_at })) - .sort((a: { id: string }, b: { id: string }) => (a.id < b.id ? -1 : 1)) as { id: string; role: "user" | "assistant" | "sources"; content: string; createdAt: string }[]; - /** Chat turns with each assistant message's stored citations attached. */ + /** Chat turns in write order (UUIDv7 ids), each assistant message with its stored citations. */ const turnsOf = async (conversationId: string) => { + const rows = (await ctx.db.sql`SELECT id, role, content, citations_json, created_at FROM ai_messages WHERE conversation_id = ${conversationId}`) as + { id: string; role: string; content: string; citations_json: string | null; created_at: string }[]; const out: { id: string; role: "user" | "assistant"; content: string; createdAt: string; citations?: unknown[] }[] = []; - for (const m of await messagesOf(conversationId)) { - if (m.role === "sources") { - const last = out[out.length - 1]; - if (last?.role === "assistant") { try { last.citations = JSON.parse(m.content); } catch {} } - } else out.push({ ...m, role: m.role }); + for (const r of rows.sort((a, b) => (a.id < b.id ? -1 : 1))) { + if (r.role !== "user" && r.role !== "assistant") continue; + const turn: (typeof out)[number] = { id: r.id, role: r.role, content: r.content, createdAt: r.created_at }; + if (r.role === "assistant" && r.citations_json) { + try { const c = JSON.parse(r.citations_json); if (Array.isArray(c)) turn.citations = c; } catch {} + } + out.push(turn); } return out; }; @@ -282,10 +283,7 @@ export function aiFeature(ctx: ServerContext, opts: AiFeatureOptions = {}): Feat send({ type: "done", text: ev.text, citations: cited }); } else send(ev); } - if (answer) { - await addMessage(convId, "assistant", answer); - await addMessage(convId, "sources", JSON.stringify(cited)); - } + if (answer) await addMessage(convId, "assistant", answer, cited); } catch (e) { const m = (e as Error).message; if (resolved.source === "byok") log(`chat error for user ${user.id} (own key): ${m}`); diff --git a/packages/feature-ai/src/index.ts b/packages/feature-ai/src/index.ts index 405c90e..66c642d 100644 --- a/packages/feature-ai/src/index.ts +++ b/packages/feature-ai/src/index.ts @@ -1,7 +1,6 @@ export { createSecretBox, type SecretBox } from "./crypto"; -/** Migration folders for this package. Empty: the AI package uses the ai_* tables from 0001_init.sql. */ -export const migrationsDirs: string[] = []; +export { migrationsDirs } from "./migrations"; export * from "./settings"; export * from "./providers"; export * from "./retrieval/course-index"; diff --git a/packages/feature-ai/src/migrations.ts b/packages/feature-ai/src/migrations.ts new file mode 100644 index 0000000..4c24074 --- /dev/null +++ b/packages/feature-ai/src/migrations.ts @@ -0,0 +1,7 @@ +import { join } from "node:path"; + +/** + * Migration folders for this package; the server passes them to core `migrate()`. The files upgrade the v1 ai_* + * tables of core's 0001_init.sql. Data that portable SQL cannot move runs in `upgradeAiData` (src/upgrade.ts). + */ +export const migrationsDirs: string[] = [join(import.meta.dir, "..", "migrations")]; diff --git a/packages/feature-ai/test/feature.test.ts b/packages/feature-ai/test/feature.test.ts index 84579dc..5042444 100644 --- a/packages/feature-ai/test/feature.test.ts +++ b/packages/feature-ai/test/feature.test.ts @@ -6,7 +6,7 @@ import type { AppConfig, AuthUser, JobRecord, ServerContext } from "@stc/contrac import { createApp, createJobQueue, resolveConfig } from "@stc/core"; import { createCourseRepository } from "@stc/course-format"; import type { Db } from "@stc/db"; -import { forEachDb } from "../../db/test/helpers"; +import { forEachAiDb as forEachDb } from "./helpers"; import { aiFeature, type AiFeatureOptions } from "../src/feature"; import { fakeProvider } from "./helpers"; @@ -104,6 +104,11 @@ forEachDb("chat cites course files, applies mentor rules, and stores the convers expect(conv.messages.map((m: { role: string }) => m.role)).toEqual(["user", "assistant"]); expect(conv.messages[1].citations.some((c: { path: string }) => c.path === "lessons/04-exercise.md")).toBe(true); expect((await call("GET", `/api/ai/conversations/${convs[0].id}`, "bob")).status).toBe(404); + // citations live in the assistant row's column, not in a separate "sources" row + const stored = await db.sql`SELECT role, citations_json FROM ai_messages WHERE conversation_id = ${convs[0].id}`; + expect(stored.map((m: { role: string }) => m.role).sort()).toEqual(["assistant", "user"]); + const assistantRow = stored.find((m: { role: string }) => m.role === "assistant"); + expect(JSON.parse(assistantRow.citations_json).some((c: { path: string }) => c.path === "lessons/04-exercise.md")).toBe(true); // Learner switches mentor mode off for this course: rules leave the prompt. expect((await call("PUT", "/api/ai/courses/sample/mentor", "ann", { off: true })).status).toBe(200); diff --git a/packages/feature-ai/test/helpers.ts b/packages/feature-ai/test/helpers.ts index d54648f..cc3846a 100644 --- a/packages/feature-ai/test/helpers.ts +++ b/packages/feature-ai/test/helpers.ts @@ -1,4 +1,7 @@ import type { AiProvider, ChatMessage, ToolSpec } from "@stc/contracts"; +import { migrate, type Db } from "@stc/db"; +import { forEachDb } from "../../db/test/helpers"; +import { migrationsDirs } from "../src/migrations"; /** SSE body from a list of [event, data] pairs, split into awkward chunks to exercise buffering. */ export function sseBody(events: [string | null, unknown][], chunk = 7): ReadableStream { @@ -51,3 +54,11 @@ export function fakeProvider(turns: ((messages: ChatMessage[], tools?: ToolSpec[ }; return { provider, seen }; } + +/** forEachDb with the AI package's migrations applied on top of core (the schema the server runs). */ +export function forEachAiDb(name: string, fn: (db: Db) => Promise): void { + forEachDb(name, async (db) => { + await migrate(db, { "feature-ai": migrationsDirs }); + await fn(db); + }); +} diff --git a/packages/feature-ai/test/settings.test.ts b/packages/feature-ai/test/settings.test.ts index 6dec14a..3febe6c 100644 --- a/packages/feature-ai/test/settings.test.ts +++ b/packages/feature-ai/test/settings.test.ts @@ -1,5 +1,5 @@ import { expect } from "bun:test"; -import { forEachDb } from "../../db/test/helpers"; +import { forEachAiDb as forEachDb } from "./helpers"; import { createSecretBox } from "../src/crypto"; import { aiSettingsRepo } from "../src/settings"; diff --git a/packages/feature-ai/test/upgrade.test.ts b/packages/feature-ai/test/upgrade.test.ts new file mode 100644 index 0000000..56fe1a4 --- /dev/null +++ b/packages/feature-ai/test/upgrade.test.ts @@ -0,0 +1,41 @@ +import { expect } from "bun:test"; +import { migrate, type Db } from "@stc/db"; +// Core migrations only: the ai_* tables exactly as v1 shipped them. +import { forEachDb as forEachV1Db } from "../../db/test/helpers"; +import { migrationsDirs } from "../src/migrations"; + +const T = "2026-10-01T00:00:00.000Z"; +const aiMigrations = { "feature-ai": migrationsDirs }; + +async function seedMessage(db: Db, id: string, conv: string, role: string, content: string) { + await db.sql`INSERT INTO ai_messages (id, conversation_id, role, content, created_at) VALUES (${id}, ${conv}, ${role}, ${content}, ${T})`; +} + +forEachV1Db("v1 upgrade: sources rows fold into citations_json on the assistant row before them", async (db) => { + const cites = JSON.stringify([{ path: "lessons/04-exercise.md", title: "Adding numbers", lessonId: "m2-ex-add" }]); + await db.sql`INSERT INTO ai_conversations (id, user_id, course_id, title, created_at) VALUES ('c1', 'u1', 'sample', 'Hi', ${T})`; + await db.sql`INSERT INTO ai_conversations (id, user_id, course_id, title, created_at) VALUES ('c2', 'u1', 'sample', 'Two', ${T})`; + // ids sort in the order v1 wrote them (UUIDv7); c1: user, assistant, sources, user, assistant (no sources) + await seedMessage(db, "0001-a", "c1", "user", "How do I add?"); + await seedMessage(db, "0001-b", "c1", "assistant", "Which operator?"); + await seedMessage(db, "0001-c", "c1", "sources", cites); + await seedMessage(db, "0001-d", "c1", "user", "And then?"); + await seedMessage(db, "0001-e", "c1", "assistant", "Try it."); + // c2: an orphan sources row (no assistant before it) and an assistant whose next row belongs to c1's ids range + await seedMessage(db, "0002-a", "c2", "sources", "[]"); + await seedMessage(db, "0002-b", "c2", "assistant", "Alone."); + + const applied = await migrate(db, aiMigrations); + expect(applied).toContain("0001_ai_message_citations.sql"); + + const rows = await db.sql`SELECT id, role, content, citations_json FROM ai_messages ORDER BY id`; + expect(rows.map((r: { id: string; role: string }) => `${r.id}:${r.role}`)).toEqual([ + "0001-a:user", "0001-b:assistant", "0001-d:user", "0001-e:assistant", "0002-b:assistant", + ]); + const byId = Object.fromEntries(rows.map((r: { id: string; citations_json: string | null }) => [r.id, r.citations_json])); + expect(JSON.parse(byId["0001-b"]!)).toEqual(JSON.parse(cites)); + expect(byId["0001-e"]).toBeNull(); + expect(byId["0002-b"]).toBeNull(); + expect(byId["0001-a"]).toBeNull(); + expect(await migrate(db, aiMigrations)).toEqual([]); +});