Skip to content
Sign up

flagon-io/g1t

public
0
589 linesCodeBlameRaw
1/**
2 * The semantic index: pages and projects' docs files as passages
3 * (src/chunks.ts) in D1 (`doc_chunks`, with full text in
4 * `doc_chunks_fts`) and as vectors in the vector store (src/vectors.ts).
5 *
6 * A page is indexed a little after its Markdown is saved (the room,
7 * src/room.ts, in the background, so editing never waits on it); a
8 * project's docs files after they are read (src/repo-spaces.ts). Only
9 * passages whose text changed are embedded; one that only moved keeps its
10 * vector. Embedding is capped per workspace and hour (EMBED_PER_HOUR):
11 * past the cap passages are kept for words, and a catch-up run on the
12 * queue embeds them later. A failure leaves them for the next save or run.
13 * Nothing here throws to its caller.
14 *
15 * The backfill indexes a workspace's existing pages and files in batches
16 * on the queue (`docs.index` jobs, on the artifacts service's own events queue).
17 */
18import { chunkId, chunkMarkdown, embedText, repoFileId, textHash } from "./chunks.ts";
19import { overFreeBudget } from "./embed-billing.ts";
20import { scopesOf, type FolioRow, FOLIO_COLUMNS } from "./folios/access-store.ts";
21import { kindModel } from "./kinds/index.ts";
22import { cloudflareEmbedder, cloudflareVectors, type Embedder, type VectorMetadata, type VectorStore } from "./vectors.ts";
23
24/** Passages embedded per workspace per hour at most. Logged when reached. */
25export const EMBED_PER_HOUR = 200;
26/** Documents one backfill job indexes before handing on to the next. */
27const BACKFILL_BATCH = 20;
28/** A backfill not finished after this long is taken to have died and may start again. */
29const BACKFILL_STALE_MS = 2 * 60 * 60 * 1000;
30/** How long a run waits when the hour's cap is reached. */
31const CAP_DELAY_SECONDS = 3600;
32
33/**
34 * A job on the queue: index more of a workspace's docs and folios; or
35 * bring a large folio subtree's access, rooms and index scope up to date
36 * after a move or a sharing change (src/folios/service.ts).
37 */
38export type DocsJob = { type: "docs.index"; workspace_id: string } | { type: "folios.reacl"; folio_id: string };
39
40export type IndexEnv = {
41 DB: D1Database;
42 AI?: Ai;
43 VECTORS?: Vectorize;
44 /** Folios' semantic index, Vectorize `g1t-folios`. Optional: without it folios match words only. */
45 FOLIO_VECTORS?: Vectorize;
46 JOBS?: Queue<DocsJob>;
47};
48
49/** The embedder and store this deployment has; null without them (words only). */
50export function adapters(env: IndexEnv): { embedder: Embedder | null; store: VectorStore | null } {
51 return env.AI && env.VECTORS ? { embedder: cloudflareEmbedder(env.AI), store: cloudflareVectors(env.VECTORS) } : { embedder: null, store: null };
52}
53
54/** One document to index: a page, or a project's docs file. */
55export type IndexDoc = {
56 kind: "page" | "repo_file";
57 /** The page's id, or the file's `rf_` id. */
58 doc_id: string;
59 workspace_id: string;
60 space_id: string;
61 title: string;
62 markdown: string;
63 repo_id?: string | null;
64 path?: string | null;
65};
66
67type ChunkRow = { id: string; seq: number; space_id: string; hash: string; vector_hash: string | null };
68
69const hourOf = (now: Date) => now.toISOString().slice(0, 13);
70
71function metadataOf(doc: IndexDoc): VectorMetadata {
72 return doc.kind === "page"
73 ? { workspace_id: doc.workspace_id, space_id: doc.space_id, kind: "page", page_id: doc.doc_id }
74 : { workspace_id: doc.workspace_id, space_id: doc.space_id, kind: "repo_file", repo_file_id: doc.doc_id, repo_id: doc.repo_id ?? "" };
75}
76
77/**
78 * Passages the workspace may still embed this hour; none once a workspace
79 * without the plan has used its free month (src/embed-billing.ts).
80 */
81async function allowance(db: D1Database, workspaceId: string, now: Date): Promise<number> {
82 if (await overFreeBudget(db, workspaceId, now)) {
83 console.log("embedding budget of a free workspace used for the month", workspaceId);
84 return 0;
85 }
86 const used = await db.prepare("SELECT chunks FROM doc_embed_usage WHERE workspace_id = ? AND hour = ?").bind(workspaceId, hourOf(now)).first<{ chunks: number }>();
87 return Math.max(0, EMBED_PER_HOUR - (used?.chunks ?? 0));
88}
89
90async function meter(db: D1Database, workspaceId: string, chunks: number, chars: number, now: Date): Promise<void> {
91 if (!chunks) return;
92 await db
93 .prepare(
94 "INSERT INTO doc_embed_usage (workspace_id, hour, chunks, tokens) VALUES (?1, ?2, ?3, ?4) ON CONFLICT (workspace_id, hour) DO UPDATE SET chunks = chunks + ?3, tokens = tokens + ?4",
95 )
96 .bind(workspaceId, hourOf(now), chunks, Math.ceil(chars / 4))
97 .run();
98}
99
100/**
101 * Brings one document's passages up to date: rows and full text in D1,
102 * vectors for those that changed (within the hour's cap). `capped` when
103 * some were left for later.
104 */
105export async function indexDoc(
106 db: D1Database,
107 embedder: Embedder | null,
108 store: VectorStore | null,
109 doc: IndexDoc,
110 now = new Date(),
111): Promise<{ chunks: number; embedded: number; capped: boolean }> {
112 const at = now.toISOString();
113 const chunks = chunkMarkdown(doc.markdown, doc.title).map((c) => {
114 const embed = embedText(doc.title, c);
115 return { ...c, id: chunkId(doc.doc_id, c.seq), embed, hash: textHash(embed) };
116 });
117 const column = doc.kind === "page" ? "page_id" : "repo_file_id";
118 const old = (await db.prepare(`SELECT id, seq, space_id, hash, vector_hash FROM doc_chunks WHERE ${column} = ?`).bind(doc.doc_id).all<ChunkRow>()).results;
119 const byId = new Map(old.map((r) => [r.id, r]));
120 const statements: D1PreparedStatement[] = [];
121 for (const c of chunks) {
122 const was = byId.get(c.id);
123 if (was && was.hash === c.hash && was.space_id === doc.space_id) continue;
124 statements.push(
125 db
126 .prepare(
127 `INSERT INTO doc_chunks (id, workspace_id, space_id, page_id, repo_file_id, repo_id, path, seq, heading, text, hash, vector_hash, updated_at)
128 VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, NULL, ?)
129 ON CONFLICT (id) DO UPDATE SET space_id = excluded.space_id, repo_id = excluded.repo_id, path = excluded.path, heading = excluded.heading, text = excluded.text, hash = excluded.hash, updated_at = excluded.updated_at`,
130 )
131 .bind(
132 c.id,
133 doc.workspace_id,
134 doc.space_id,
135 doc.kind === "page" ? doc.doc_id : null,
136 doc.kind === "repo_file" ? doc.doc_id : null,
137 doc.repo_id ?? null,
138 doc.path ?? null,
139 c.seq,
140 c.heading,
141 c.text,
142 c.hash,
143 at,
144 ),
145 db.prepare("DELETE FROM doc_chunks_fts WHERE chunk_id = ?").bind(c.id),
146 db.prepare("INSERT INTO doc_chunks_fts (chunk_id, space_id, doc_id, heading, text) VALUES (?, ?, ?, ?, ?)").bind(c.id, doc.space_id, doc.doc_id, c.heading ?? "", c.text),
147 );
148 }
149 const keep = new Set(chunks.map((c) => c.id));
150 const gone = old.filter((r) => !keep.has(r.id));
151 for (const r of gone) {
152 statements.push(db.prepare("DELETE FROM doc_chunks WHERE id = ?").bind(r.id), db.prepare("DELETE FROM doc_chunks_fts WHERE chunk_id = ?").bind(r.id));
153 }
154 for (let i = 0; i < statements.length; i += 60) await db.batch(statements.slice(i, i + 60));
155
156 if (!store || !embedder) return { chunks: chunks.length, embedded: 0, capped: false };
157 const result = await embedChanged(db, embedder, store, doc, chunks, old, byId, now);
158 // Last, so a passage that moved from a place now gone could take its vector first.
159 try {
160 if (gone.some((r) => r.vector_hash)) await store.delete(gone.filter((r) => r.vector_hash).map((r) => r.id));
161 } catch (error) {
162 console.error("docs could not drop old passages from the index", doc.doc_id, String(error));
163 }
164 return result;
165}
166
167type IndexedChunk = { id: string; seq: number; heading: string | null; text: string; embed: string; hash: string };
168
169/** Vectors for the passages that need one: reused when the index already holds their text, embedded otherwise (within the cap). */
170async function embedChanged(
171 db: D1Database,
172 embedder: Embedder,
173 store: VectorStore,
174 doc: IndexDoc,
175 chunks: IndexedChunk[],
176 old: ChunkRow[],
177 byId: Map<string, ChunkRow>,
178 now: Date,
179): Promise<{ chunks: number; embedded: number; capped: boolean }> {
180 // What needs a vector: a passage whose text the index doesn't hold under its id, or one whose space changed.
181 const stale = chunks.filter((c) => {
182 const was = byId.get(c.id);
183 return !was || was.vector_hash !== c.hash || was.space_id !== doc.space_id;
184 });
185 if (!stale.length) return { chunks: chunks.length, embedded: 0, capped: false };
186 // A passage that only moved (or changed space) has its vector already, under another id or this one.
187 const held = new Map<string, string>();
188 for (const r of old) if (r.vector_hash) held.set(r.vector_hash, r.id);
189 const reuse = stale.filter((c) => held.has(c.hash));
190 const fresh = stale.filter((c) => !held.has(c.hash));
191 const budget = fresh.length ? await allowance(db, doc.workspace_id, now) : 0;
192 const embedNow = fresh.slice(0, budget);
193 const capped = embedNow.length < fresh.length;
194 if (capped) console.log("docs embedding cap reached", doc.workspace_id, `${fresh.length - embedNow.length} passages wait for the next hour`);
195 const metadata = metadataOf(doc);
196 const done: { id: string; hash: string }[] = [];
197 try {
198 const vectors: { id: string; values: number[]; metadata: VectorMetadata }[] = [];
199 if (reuse.length) {
200 const values = new Map((await store.get([...new Set(reuse.map((c) => held.get(c.hash)!))])).map((v) => [v.id, v.values]));
201 for (const c of reuse) {
202 const v = values.get(held.get(c.hash)!);
203 if (v) vectors.push({ id: c.id, values: v, metadata });
204 }
205 }
206 if (embedNow.length) {
207 const embedded = await embedder.embed(embedNow.map((c) => c.embed));
208 embedNow.forEach((c, i) => vectors.push({ id: c.id, values: embedded[i]!, metadata }));
209 await meter(
210 db,
211 doc.workspace_id,
212 embedNow.length,
213 embedNow.reduce((n, c) => n + c.embed.length, 0),
214 now,
215 );
216 }
217 if (vectors.length) await store.upsert(vectors);
218 const hashOf = new Map(chunks.map((c) => [c.id, c.hash]));
219 for (const v of vectors) done.push({ id: v.id, hash: hashOf.get(v.id)! });
220 } catch (error) {
221 console.error("docs could not embed passages; the next save tries again", doc.doc_id, String(error));
222 }
223 if (done.length) {
224 const updates = done.map((d) => db.prepare("UPDATE doc_chunks SET vector_hash = ? WHERE id = ?").bind(d.hash, d.id));
225 for (let i = 0; i < updates.length; i += 60) await db.batch(updates.slice(i, i + 60));
226 }
227 return { chunks: chunks.length, embedded: embedNow.length, capped };
228}
229
230/** Takes passages out of D1 and the index: by page ids, file ids, or a whole space. */
231export async function forgetDocs(env: IndexEnv, which: { page_ids?: string[]; repo_file_ids?: string[]; space_id?: string }): Promise<void> {
232 const db = env.DB;
233 const { store } = adapters(env);
234 try {
235 const found: string[] = [];
236 const select = async (column: string, values: string[]) => {
237 for (let i = 0; i < values.length; i += 90) {
238 const part = values.slice(i, i + 90);
239 const rows = await db.prepare(`SELECT id FROM doc_chunks WHERE ${column} IN (${part.map(() => "?").join(",")})`).bind(...part).all<{ id: string }>();
240 found.push(...rows.results.map((r) => r.id));
241 }
242 };
243 if (which.page_ids?.length) await select("page_id", which.page_ids);
244 if (which.repo_file_ids?.length) await select("repo_file_id", which.repo_file_ids);
245 if (which.space_id) await select("space_id", [which.space_id]);
246 if (!found.length) return;
247 if (store) {
248 try {
249 await store.delete(found);
250 } catch (error) {
251 console.error("docs could not drop passages from the index", found.length, String(error));
252 }
253 }
254 const statements = found.flatMap((id) => [db.prepare("DELETE FROM doc_chunks WHERE id = ?").bind(id), db.prepare("DELETE FROM doc_chunks_fts WHERE chunk_id = ?").bind(id)]);
255 for (let i = 0; i < statements.length; i += 80) await db.batch(statements.slice(i, i + 80));
256 } catch (error) {
257 console.error("docs could not forget passages", String(error));
258 }
259}
260
261type PageForIndex = { id: string; workspace_id: string; space_id: string; title: string; markdown: string; archived_at: string | null };
262
263/** Indexes one page as it is saved now; forgets it when it is gone or in the trash. Never throws. */
264export async function indexPage(env: IndexEnv, pageId: string, now = new Date()): Promise<{ capped: boolean }> {
265 try {
266 const page = await env.DB.prepare("SELECT id, workspace_id, space_id, title, markdown, archived_at FROM pages WHERE id = ?").bind(pageId).first<PageForIndex>();
267 if (!page || page.archived_at) {
268 await forgetDocs(env, { page_ids: [pageId] });
269 return { capped: false };
270 }
271 const { embedder, store } = adapters(env);
272 const result = await indexDoc(env.DB, embedder, store, { kind: "page", doc_id: page.id, workspace_id: page.workspace_id, space_id: page.space_id, title: page.title, markdown: page.markdown }, now);
273 if (result.capped) await startBackfill(env, page.workspace_id, { delaySeconds: CAP_DELAY_SECONDS });
274 return { capped: result.capped };
275 } catch (error) {
276 console.error("docs could not index a page", pageId, String(error));
277 return { capped: false };
278 }
279}
280
281type RepoFileForIndex = { space_id: string; path: string; title: string; markdown: string; workspace_id: string; repo_id: string };
282
283function repoDoc(row: RepoFileForIndex): IndexDoc {
284 return { kind: "repo_file", doc_id: repoFileId(row.space_id, row.path), workspace_id: row.workspace_id, space_id: row.space_id, title: row.title, markdown: row.markdown, repo_id: row.repo_id, path: row.path };
285}
286
287/**
288 * Indexes a project's docs files after they were read: `changed` paths
289 * again, `gone` paths forgotten. Never throws.
290 */
291export async function indexRepoFiles(env: IndexEnv, spaceId: string, changed: string[], gone: string[], now = new Date()): Promise<void> {
292 try {
293 if (gone.length) await forgetDocs(env, { repo_file_ids: gone.map((p) => repoFileId(spaceId, p)) });
294 if (!changed.length) return;
295 const { embedder, store } = adapters(env);
296 let capped = false;
297 let workspaceId: string | null = null;
298 for (let i = 0; i < changed.length; i += 50) {
299 const part = changed.slice(i, i + 50);
300 const rows = (
301 await env.DB.prepare(
302 `SELECT f.space_id, f.path, f.title, f.markdown, s.workspace_id, s.repo_id FROM repo_files f JOIN repo_spaces s ON s.id = f.space_id WHERE f.space_id = ? AND f.path IN (${part.map(() => "?").join(",")})`,
303 )
304 .bind(spaceId, ...part)
305 .all<RepoFileForIndex>()
306 ).results;
307 for (const row of rows) {
308 workspaceId = row.workspace_id;
309 // Past the cap, the rest are only kept for words until the catch-up run.
310 const result = await indexDoc(env.DB, capped ? null : embedder, capped ? null : store, repoDoc(row), now);
311 capped ||= result.capped;
312 }
313 }
314 if (capped && workspaceId) await startBackfill(env, workspaceId, { delaySeconds: CAP_DELAY_SECONDS });
315 } catch (error) {
316 console.error("docs could not index a project's docs", spaceId, String(error));
317 }
318}
319
320/**
321 * Starts (or, with `force`, restarts) indexing a workspace's existing
322 * pages and projects' docs on the queue. Does nothing while a run is
323 * going, unless it has been going so long it must have died. Without a
324 * queue, runs one batch now. True when a run was started.
325 */
326export async function startBackfill(env: IndexEnv, workspaceId: string, options: { force?: boolean; delaySeconds?: number } = {}): Promise<boolean> {
327 const at = new Date();
328 const staleBefore = new Date(at.getTime() - BACKFILL_STALE_MS).toISOString();
329 const started = await env.DB.prepare(
330 `INSERT INTO doc_index_runs (workspace_id, started_at, finished_at, cursor, pages, files) VALUES (?1, ?2, NULL, NULL, 0, 0)
331 ON CONFLICT (workspace_id) DO UPDATE SET started_at = ?2, finished_at = NULL, cursor = NULL, pages = 0, files = 0
332 WHERE ?3 OR doc_index_runs.finished_at IS NOT NULL OR doc_index_runs.started_at < ?4
333 RETURNING workspace_id`,
334 )
335 .bind(workspaceId, at.toISOString(), options.force ? 1 : 0, staleBefore)
336 .first<{ workspace_id: string }>();
337 if (!started) return false;
338 await enqueue(env, workspaceId, options.delaySeconds ?? 0);
339 return true;
340}
341
342async function enqueue(env: IndexEnv, workspaceId: string, delaySeconds: number): Promise<void> {
343 if (env.JOBS) {
344 await env.JOBS.send({ type: "docs.index", workspace_id: workspaceId }, delaySeconds ? { delaySeconds } : undefined);
345 return;
346 }
347 // No queue (self-hosted without one): one batch now; saves and later recalls carry on from there.
348 if (!delaySeconds) await runBackfill(env, workspaceId, false);
349}
350
351/**
352 * One batch of a workspace's backfill: pages by id, then projects' docs
353 * files by space and path, from the run's cursor. Hands on to the next
354 * batch (or waits out the hour's cap) through the queue.
355 */
356export async function runBackfill(env: IndexEnv, workspaceId: string, chain = true): Promise<{ done: boolean }> {
357 const run = await env.DB.prepare("SELECT cursor, finished_at, pages, files FROM doc_index_runs WHERE workspace_id = ?").bind(workspaceId).first<{ cursor: string | null; finished_at: string | null; pages: number; files: number }>();
358 if (!run || run.finished_at) return { done: true };
359 const { embedder, store } = adapters(env);
360 const cursor = run.cursor ?? "p:";
361 let next: string | null = cursor;
362 let pages = 0;
363 let files = 0;
364 let capped = false;
365 if (cursor.startsWith("o:")) {
366 // Folios, after pages and before projects' docs.
367 const rows = (
368 await env.DB.prepare(`SELECT ${FOLIO_COLUMNS} FROM folios WHERE workspace_id = ? AND trashed_at IS NULL AND id > ? ORDER BY id LIMIT ?`)
369 .bind(workspaceId, cursor.slice(2), BACKFILL_BATCH)
370 .all<FolioRow>()
371 ).results;
372 for (const row of rows) {
373 const result = await indexFolio(env, row.id);
374 pages++;
375 if (result.capped) {
376 capped = true;
377 break;
378 }
379 next = `o:${row.id}`;
380 }
381 if (!capped && rows.length < BACKFILL_BATCH) next = "f:";
382 } else if (cursor.startsWith("p:")) {
383 const rows = (
384 await env.DB.prepare("SELECT id, workspace_id, space_id, title, markdown, archived_at FROM pages WHERE workspace_id = ? AND archived_at IS NULL AND id > ? ORDER BY id LIMIT ?")
385 .bind(workspaceId, cursor.slice(2), BACKFILL_BATCH)
386 .all<PageForIndex>()
387 ).results;
388 for (const page of rows) {
389 const result = await indexDoc(env.DB, embedder, store, { kind: "page", doc_id: page.id, workspace_id: page.workspace_id, space_id: page.space_id, title: page.title, markdown: page.markdown });
390 pages++;
391 if (result.capped) {
392 // This page again next time, after the hour: the cursor stays before it.
393 capped = true;
394 break;
395 }
396 next = `p:${page.id}`;
397 }
398 if (!capped && rows.length < BACKFILL_BATCH) next = "o:";
399 } else {
400 const [space, path] = splitFileCursor(cursor.slice(2));
401 const rows = (
402 await env.DB.prepare(
403 `SELECT f.space_id, f.path, f.title, f.markdown, s.workspace_id, s.repo_id FROM repo_files f JOIN repo_spaces s ON s.id = f.space_id
404 WHERE s.workspace_id = ? AND (f.space_id > ? OR (f.space_id = ? AND f.path > ?)) ORDER BY f.space_id, f.path LIMIT ?`,
405 )
406 .bind(workspaceId, space, space, path, BACKFILL_BATCH)
407 .all<RepoFileForIndex>()
408 ).results;
409 for (const row of rows) {
410 const result = await indexDoc(env.DB, embedder, store, repoDoc(row));
411 files++;
412 if (result.capped) {
413 capped = true;
414 break;
415 }
416 next = `f:${row.space_id}\n${row.path}`;
417 }
418 if (!capped && rows.length < BACKFILL_BATCH) next = null;
419 }
420 await env.DB.prepare("UPDATE doc_index_runs SET cursor = ?, pages = pages + ?, files = files + ?, finished_at = ? WHERE workspace_id = ?")
421 .bind(next, pages, files, next === null ? new Date().toISOString() : null, workspaceId)
422 .run();
423 if (next === null) return { done: true };
424 if (chain && env.JOBS) await enqueue(env, workspaceId, capped ? CAP_DELAY_SECONDS : 0);
425 return { done: false };
426}
427
428function splitFileCursor(cursor: string): [string, string] {
429 const at = cursor.indexOf("\n");
430 // Empty: from the start.
431 return at < 0 ? ["", ""] : [cursor.slice(0, at), cursor.slice(at + 1)];
432}
433
434// ── Folios (Artifacts mode) ─────────────────────────────────────────────
435
436let saidNoFolioIndex = false;
437
438/**
439 * The embedder and folio index this deployment has. Without Workers AI or
440 * the `g1t-folios` index (FOLIO_VECTORS, which is made by hand before a
441 * deploy that uses it), folios keep their passages in D1 and recall and
442 * search match words; that is said once per isolate.
443 */
444export function folioAdapters(env: IndexEnv): { embedder: Embedder | null; store: VectorStore | null } {
445 if (env.AI && env.FOLIO_VECTORS) return { embedder: cloudflareEmbedder(env.AI), store: cloudflareVectors(env.FOLIO_VECTORS) };
446 if (!saidNoFolioIndex) {
447 saidNoFolioIndex = true;
448 console.log(`folios: no ${env.AI ? "FOLIO_VECTORS (Vectorize g1t-folios)" : "AI"} binding, so artifacts are searched and recalled by their words only`);
449 }
450 return { embedder: null, store: null };
451}
452
453type FolioChunkRow = { id: string; seq: number; scope: string; hash: string; vector_hash: string | null };
454
455/**
456 * Brings one folio's passages up to date, as `indexDoc` does for pages:
457 * rows and full text in D1, under the folio's index scope, and vectors
458 * for passages whose text or scope changed (a passage whose text the
459 * index already holds keeps its vector; only its metadata is written
460 * again). Forgets it when it is gone or in the trash. Never throws.
461 */
462export async function indexFolio(env: IndexEnv, folioId: string, now = new Date()): Promise<{ capped: boolean }> {
463 try {
464 const db = env.DB;
465 const row = await db.prepare(`SELECT ${FOLIO_COLUMNS.replace("'' AS text", "text")} FROM folios WHERE id = ?`).bind(folioId).first<FolioRow>();
466 if (!row || row.trashed_at) {
467 await forgetFolios(env, [folioId]);
468 return { capped: false };
469 }
470 const scope = (await scopesOf(db, [row])).get(row.id) ?? `folio:${row.acl_root}`;
471 const model = kindModel(row.kind);
472 const at = now.toISOString();
473 const chunks = (model ? model.chunks(row.text, row.title) : chunkMarkdown(row.text, row.title)).map((c) => {
474 const embed = embedText(row.title, c);
475 return { ...c, id: chunkId(row.id, c.seq), embed, hash: textHash(embed) };
476 });
477 const old = (await db.prepare("SELECT id, seq, scope, hash, vector_hash FROM folio_chunks WHERE folio_id = ?").bind(row.id).all<FolioChunkRow>()).results;
478 const byId = new Map(old.map((r) => [r.id, r]));
479 const statements: D1PreparedStatement[] = [];
480 for (const c of chunks) {
481 const was = byId.get(c.id);
482 if (was && was.hash === c.hash && was.scope === scope) continue;
483 statements.push(
484 db
485 .prepare(
486 `INSERT INTO folio_chunks (id, workspace_id, folio_id, kind, scope, seq, heading, text, hash, vector_hash, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, NULL, ?)
487 ON CONFLICT (id) DO UPDATE SET kind = excluded.kind, scope = excluded.scope, heading = excluded.heading, text = excluded.text, hash = excluded.hash, updated_at = excluded.updated_at`,
488 )
489 .bind(c.id, row.workspace_id, row.id, row.kind, scope, c.seq, c.heading, c.text, c.hash, at),
490 db.prepare("DELETE FROM folio_chunks_fts WHERE chunk_id = ?").bind(c.id),
491 db.prepare("INSERT INTO folio_chunks_fts (chunk_id, scope, folio_id, heading, text) VALUES (?, ?, ?, ?, ?)").bind(c.id, scope, row.id, c.heading ?? "", c.text),
492 );
493 }
494 const keep = new Set(chunks.map((c) => c.id));
495 const gone = old.filter((r) => !keep.has(r.id));
496 for (const r of gone) statements.push(db.prepare("DELETE FROM folio_chunks WHERE id = ?").bind(r.id), db.prepare("DELETE FROM folio_chunks_fts WHERE chunk_id = ?").bind(r.id));
497 for (let i = 0; i < statements.length; i += 60) await db.batch(statements.slice(i, i + 60));
498
499 const { embedder, store } = folioAdapters(env);
500 if (!embedder || !store) return { capped: false };
501 // What needs a vector written: text the index doesn't hold under this id, or a new scope.
502 const stale = chunks.filter((c) => {
503 const was = byId.get(c.id);
504 return !was || was.vector_hash !== c.hash || was.scope !== scope;
505 });
506 let capped = false;
507 if (stale.length) {
508 const held = new Map<string, string>();
509 for (const r of old) if (r.vector_hash) held.set(r.vector_hash, r.id);
510 const reuse = stale.filter((c) => held.has(c.hash));
511 const fresh = stale.filter((c) => !held.has(c.hash));
512 const budget = fresh.length ? await allowance(db, row.workspace_id, now) : 0;
513 const embedNow = fresh.slice(0, budget);
514 capped = embedNow.length < fresh.length;
515 if (capped) console.log("folios embedding cap reached", row.workspace_id, `${fresh.length - embedNow.length} passages wait for the next hour`);
516 const metadata: VectorMetadata = { workspace_id: row.workspace_id, scope, kind: row.kind, folio_id: row.id };
517 const done: { id: string; hash: string }[] = [];
518 try {
519 const vectors: { id: string; values: number[]; metadata: VectorMetadata }[] = [];
520 if (reuse.length) {
521 const values = new Map((await store.get([...new Set(reuse.map((c) => held.get(c.hash)!))])).map((v) => [v.id, v.values]));
522 for (const c of reuse) {
523 const v = values.get(held.get(c.hash)!);
524 if (v) vectors.push({ id: c.id, values: v, metadata });
525 }
526 }
527 if (embedNow.length) {
528 const embedded = await embedder.embed(embedNow.map((c) => c.embed));
529 embedNow.forEach((c, i) => vectors.push({ id: c.id, values: embedded[i]!, metadata }));
530 await meter(
531 db,
532 row.workspace_id,
533 embedNow.length,
534 embedNow.reduce((n, c) => n + c.embed.length, 0),
535 now,
536 );
537 }
538 if (vectors.length) await store.upsert(vectors);
539 const hashOf = new Map(chunks.map((c) => [c.id, c.hash]));
540 for (const v of vectors) done.push({ id: v.id, hash: hashOf.get(v.id)! });
541 } catch (error) {
542 console.error("folios could not embed passages; the next save tries again", row.id, String(error));
543 }
544 if (done.length) {
545 const updates = done.map((d) => db.prepare("UPDATE folio_chunks SET vector_hash = ? WHERE id = ?").bind(d.hash, d.id));
546 for (let i = 0; i < updates.length; i += 60) await db.batch(updates.slice(i, i + 60));
547 }
548 }
549 try {
550 if (gone.some((r) => r.vector_hash)) await store.delete(gone.filter((r) => r.vector_hash).map((r) => r.id));
551 } catch (error) {
552 console.error("folios could not drop old passages from the index", row.id, String(error));
553 }
554 if (capped) await startBackfill(env, row.workspace_id, { delaySeconds: CAP_DELAY_SECONDS });
555 return { capped };
556 } catch (error) {
557 console.error("folios could not index", folioId, String(error));
558 return { capped: false };
559 }
560}
561
562/** Takes folios' passages out of D1 and the index. Never throws. */
563export async function forgetFolios(env: IndexEnv, ids: readonly string[]): Promise<void> {
564 if (!ids.length) return;
565 const db = env.DB;
566 try {
567 const found: string[] = [];
568 for (let i = 0; i < ids.length; i += 500) {
569 const rows = await db
570 .prepare("SELECT id FROM folio_chunks WHERE folio_id IN (SELECT value FROM json_each(?))")
571 .bind(JSON.stringify(ids.slice(i, i + 500)))
572 .all<{ id: string }>();
573 found.push(...rows.results.map((r) => r.id));
574 }
575 if (!found.length) return;
576 const { store } = folioAdapters(env);
577 if (store) {
578 try {
579 await store.delete(found);
580 } catch (error) {
581 console.error("folios could not drop passages from the index", found.length, String(error));
582 }
583 }
584 const statements = found.flatMap((id) => [db.prepare("DELETE FROM folio_chunks WHERE id = ?").bind(id), db.prepare("DELETE FROM folio_chunks_fts WHERE chunk_id = ?").bind(id)]);
585 for (let i = 0; i < statements.length; i += 80) await db.batch(statements.slice(i, i + 80));
586 } catch (error) {
587 console.error("folios could not forget passages", String(error));
588 }
589}