Skip to content
Sign up

flagon-io/g1t

public
00

Where people and agents ship software together. The open-source git platform for the whole job: issues, agents, checks and deploys to the edge.

gitagentscloudflarerustforgemcp
688 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
87 .prepare("SELECT chunks FROM doc_embed_usage WHERE workspace_id = ? AND hour = ?")
88 .bind(workspaceId, hourOf(now))
89 .first<{ chunks: number }>();
90 return Math.max(0, EMBED_PER_HOUR - (used?.chunks ?? 0));
91}
92
93async function meter(db: D1Database, workspaceId: string, chunks: number, chars: number, now: Date): Promise<void> {
94 if (!chunks) return;
95 await db
96 .prepare(
97 "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",
98 )
99 .bind(workspaceId, hourOf(now), chunks, Math.ceil(chars / 4))
100 .run();
101}
102
103/**
104 * Brings one document's passages up to date: rows and full text in D1,
105 * vectors for those that changed (within the hour's cap). `capped` when
106 * some were left for later.
107 */
108export async function indexDoc(
109 db: D1Database,
110 embedder: Embedder | null,
111 store: VectorStore | null,
112 doc: IndexDoc,
113 now = new Date(),
114): Promise<{ chunks: number; embedded: number; capped: boolean }> {
115 const at = now.toISOString();
116 const chunks = chunkMarkdown(doc.markdown, doc.title).map((c) => {
117 const embed = embedText(doc.title, c);
118 return { ...c, id: chunkId(doc.doc_id, c.seq), embed, hash: textHash(embed) };
119 });
120 const column = doc.kind === "page" ? "page_id" : "repo_file_id";
121 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;
122 const byId = new Map(old.map((r) => [r.id, r]));
123 const statements: D1PreparedStatement[] = [];
124 for (const c of chunks) {
125 const was = byId.get(c.id);
126 if (was && was.hash === c.hash && was.space_id === doc.space_id) continue;
127 statements.push(
128 db
129 .prepare(
130 `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)
131 VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, NULL, ?)
132 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`,
133 )
134 .bind(
135 c.id,
136 doc.workspace_id,
137 doc.space_id,
138 doc.kind === "page" ? doc.doc_id : null,
139 doc.kind === "repo_file" ? doc.doc_id : null,
140 doc.repo_id ?? null,
141 doc.path ?? null,
142 c.seq,
143 c.heading,
144 c.text,
145 c.hash,
146 at,
147 ),
148 db.prepare("DELETE FROM doc_chunks_fts WHERE chunk_id = ?").bind(c.id),
149 // Only beside a chunk row the database kept: an erased document's insert was ignored (migration 0010).
150 db
151 .prepare(
152 "INSERT INTO doc_chunks_fts (chunk_id, space_id, doc_id, heading, text) SELECT ?1, ?2, ?3, ?4, ?5 WHERE EXISTS (SELECT 1 FROM doc_chunks WHERE id = ?1)",
153 )
154 .bind(c.id, doc.space_id, doc.doc_id, c.heading ?? "", c.text),
155 );
156 }
157 const keep = new Set(chunks.map((c) => c.id));
158 const gone = old.filter((r) => !keep.has(r.id));
159 for (const r of gone) {
160 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));
161 }
162 for (let i = 0; i < statements.length; i += 60) await db.batch(statements.slice(i, i + 60));
163
164 if (!store || !embedder) return { chunks: chunks.length, embedded: 0, capped: false };
165 if (await erasedSource(db, doc.kind === "page" ? "page" : null, doc.doc_id, doc.workspace_id)) return { chunks: 0, embedded: 0, capped: false };
166 const result = await embedChanged(db, embedder, store, doc, chunks, old, byId, now);
167 await dropIfErased(
168 db,
169 store,
170 doc.kind === "page" ? "page" : null,
171 doc.doc_id,
172 doc.workspace_id,
173 chunks.map((c) => c.id),
174 );
175 // Last, so a passage that moved from a place now gone could take its vector first.
176 try {
177 if (gone.some((r) => r.vector_hash)) await store.delete(gone.filter((r) => r.vector_hash).map((r) => r.id));
178 } catch (error) {
179 console.error("docs could not drop old passages from the index", doc.doc_id, String(error));
180 }
181 return result;
182}
183
184type IndexedChunk = { id: string; seq: number; heading: string | null; text: string; embed: string; hash: string };
185
186/** Vectors for the passages that need one: reused when the index already holds their text, embedded otherwise (within the cap). */
187async function embedChanged(
188 db: D1Database,
189 embedder: Embedder,
190 store: VectorStore,
191 doc: IndexDoc,
192 chunks: IndexedChunk[],
193 old: ChunkRow[],
194 byId: Map<string, ChunkRow>,
195 now: Date,
196): Promise<{ chunks: number; embedded: number; capped: boolean }> {
197 // What needs a vector: a passage whose text the index doesn't hold under its id, or one whose space changed.
198 const stale = chunks.filter((c) => {
199 const was = byId.get(c.id);
200 return !was || was.vector_hash !== c.hash || was.space_id !== doc.space_id;
201 });
202 if (!stale.length) return { chunks: chunks.length, embedded: 0, capped: false };
203 // A passage that only moved (or changed space) has its vector already, under another id or this one.
204 const held = new Map<string, string>();
205 for (const r of old) if (r.vector_hash) held.set(r.vector_hash, r.id);
206 const reuse = stale.filter((c) => held.has(c.hash));
207 const fresh = stale.filter((c) => !held.has(c.hash));
208 const budget = fresh.length ? await allowance(db, doc.workspace_id, now) : 0;
209 const embedNow = fresh.slice(0, budget);
210 const capped = embedNow.length < fresh.length;
211 if (capped) console.log("docs embedding cap reached", doc.workspace_id, `${fresh.length - embedNow.length} passages wait for the next hour`);
212 const metadata = metadataOf(doc);
213 const done: { id: string; hash: string }[] = [];
214 try {
215 const vectors: { id: string; values: number[]; metadata: VectorMetadata }[] = [];
216 if (reuse.length) {
217 const values = new Map((await store.get([...new Set(reuse.map((c) => held.get(c.hash)!))])).map((v) => [v.id, v.values]));
218 for (const c of reuse) {
219 const v = values.get(held.get(c.hash)!);
220 if (v) vectors.push({ id: c.id, values: v, metadata });
221 }
222 }
223 if (embedNow.length) {
224 const embedded = await embedder.embed(embedNow.map((c) => c.embed));
225 embedNow.forEach((c, i) => vectors.push({ id: c.id, values: embedded[i]!, metadata }));
226 await meter(
227 db,
228 doc.workspace_id,
229 embedNow.length,
230 embedNow.reduce((n, c) => n + c.embed.length, 0),
231 now,
232 );
233 }
234 if (vectors.length) await store.upsert(vectors);
235 const hashOf = new Map(chunks.map((c) => [c.id, c.hash]));
236 for (const v of vectors) done.push({ id: v.id, hash: hashOf.get(v.id)! });
237 } catch (error) {
238 console.error("docs could not embed passages; the next save tries again", doc.doc_id, String(error));
239 }
240 if (done.length) {
241 const updates = done.map((d) => db.prepare("UPDATE doc_chunks SET vector_hash = ? WHERE id = ?").bind(d.hash, d.id));
242 for (let i = 0; i < updates.length; i += 60) await db.batch(updates.slice(i, i + 60));
243 }
244 return { chunks: chunks.length, embedded: embedNow.length, capped };
245}
246
247/** Takes passages out of D1 and the index: by page ids, file ids, or a whole space. */
248export async function forgetDocs(env: IndexEnv, which: { page_ids?: string[]; repo_file_ids?: string[]; space_id?: string }): Promise<void> {
249 const db = env.DB;
250 const { store } = adapters(env);
251 try {
252 const found: string[] = [];
253 const select = async (column: string, values: string[]) => {
254 for (let i = 0; i < values.length; i += 90) {
255 const part = values.slice(i, i + 90);
256 const rows = await db
257 .prepare(`SELECT id FROM doc_chunks WHERE ${column} IN (${part.map(() => "?").join(",")})`)
258 .bind(...part)
259 .all<{ id: string }>();
260 found.push(...rows.results.map((r) => r.id));
261 }
262 };
263 if (which.page_ids?.length) await select("page_id", which.page_ids);
264 if (which.repo_file_ids?.length) await select("repo_file_id", which.repo_file_ids);
265 if (which.space_id) await select("space_id", [which.space_id]);
266 if (!found.length) return;
267 if (store) {
268 try {
269 await store.delete(found);
270 } catch (error) {
271 console.error("docs could not drop passages from the index", found.length, String(error));
272 }
273 }
274 const statements = found.flatMap((id) => [
275 db.prepare("DELETE FROM doc_chunks WHERE id = ?").bind(id),
276 db.prepare("DELETE FROM doc_chunks_fts WHERE chunk_id = ?").bind(id),
277 ]);
278 for (let i = 0; i < statements.length; i += 80) await db.batch(statements.slice(i, i + 80));
279 } catch (error) {
280 console.error("docs could not forget passages", String(error));
281 }
282}
283
284type PageForIndex = { id: string; workspace_id: string; space_id: string; title: string; markdown: string; archived_at: string | null };
285
286/** Indexes one page as it is saved now; forgets it when it is gone or in the trash. Never throws. */
287export async function indexPage(env: IndexEnv, pageId: string, now = new Date()): Promise<{ capped: boolean }> {
288 try {
289 const page = await env.DB.prepare("SELECT id, workspace_id, space_id, title, markdown, archived_at FROM pages WHERE id = ?")
290 .bind(pageId)
291 .first<PageForIndex>();
292 if (!page || page.archived_at) {
293 await forgetDocs(env, { page_ids: [pageId] });
294 return { capped: false };
295 }
296 const { embedder, store } = adapters(env);
297 const result = await indexDoc(
298 env.DB,
299 embedder,
300 store,
301 { kind: "page", doc_id: page.id, workspace_id: page.workspace_id, space_id: page.space_id, title: page.title, markdown: page.markdown },
302 now,
303 );
304 if (result.capped) await startBackfill(env, page.workspace_id, { delaySeconds: CAP_DELAY_SECONDS });
305 return { capped: result.capped };
306 } catch (error) {
307 console.error("docs could not index a page", pageId, String(error));
308 return { capped: false };
309 }
310}
311
312type RepoFileForIndex = { space_id: string; path: string; title: string; markdown: string; workspace_id: string; repo_id: string };
313
314function repoDoc(row: RepoFileForIndex): IndexDoc {
315 return {
316 kind: "repo_file",
317 doc_id: repoFileId(row.space_id, row.path),
318 workspace_id: row.workspace_id,
319 space_id: row.space_id,
320 title: row.title,
321 markdown: row.markdown,
322 repo_id: row.repo_id,
323 path: row.path,
324 };
325}
326
327/**
328 * Indexes a project's docs files after they were read: `changed` paths
329 * again, `gone` paths forgotten. Never throws.
330 */
331export async function indexRepoFiles(env: IndexEnv, spaceId: string, changed: string[], gone: string[], now = new Date()): Promise<void> {
332 try {
333 if (gone.length) await forgetDocs(env, { repo_file_ids: gone.map((p) => repoFileId(spaceId, p)) });
334 if (!changed.length) return;
335 const { embedder, store } = adapters(env);
336 let capped = false;
337 let workspaceId: string | null = null;
338 for (let i = 0; i < changed.length; i += 50) {
339 const part = changed.slice(i, i + 50);
340 const rows = (
341 await env.DB.prepare(
342 `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(",")})`,
343 )
344 .bind(spaceId, ...part)
345 .all<RepoFileForIndex>()
346 ).results;
347 for (const row of rows) {
348 workspaceId = row.workspace_id;
349 // Past the cap, the rest are only kept for words until the catch-up run.
350 const result = await indexDoc(env.DB, capped ? null : embedder, capped ? null : store, repoDoc(row), now);
351 capped ||= result.capped;
352 }
353 }
354 if (capped && workspaceId) await startBackfill(env, workspaceId, { delaySeconds: CAP_DELAY_SECONDS });
355 } catch (error) {
356 console.error("docs could not index a project's docs", spaceId, String(error));
357 }
358}
359
360/**
361 * Starts (or, with `force`, restarts) indexing a workspace's existing
362 * pages and projects' docs on the queue. Does nothing while a run is
363 * going, unless it has been going so long it must have died. Without a
364 * queue, runs one batch now. True when a run was started.
365 */
366export async function startBackfill(env: IndexEnv, workspaceId: string, options: { force?: boolean; delaySeconds?: number } = {}): Promise<boolean> {
367 const at = new Date();
368 const staleBefore = new Date(at.getTime() - BACKFILL_STALE_MS).toISOString();
369 const started = await env.DB.prepare(
370 `INSERT INTO doc_index_runs (workspace_id, started_at, finished_at, cursor, pages, files) VALUES (?1, ?2, NULL, NULL, 0, 0)
371 ON CONFLICT (workspace_id) DO UPDATE SET started_at = ?2, finished_at = NULL, cursor = NULL, pages = 0, files = 0
372 WHERE ?3 OR doc_index_runs.finished_at IS NOT NULL OR doc_index_runs.started_at < ?4
373 RETURNING workspace_id`,
374 )
375 .bind(workspaceId, at.toISOString(), options.force ? 1 : 0, staleBefore)
376 .first<{ workspace_id: string }>();
377 if (!started) return false;
378 await enqueue(env, workspaceId, options.delaySeconds ?? 0);
379 return true;
380}
381
382async function enqueue(env: IndexEnv, workspaceId: string, delaySeconds: number): Promise<void> {
383 if (env.JOBS) {
384 await env.JOBS.send({ type: "docs.index", workspace_id: workspaceId }, delaySeconds ? { delaySeconds } : undefined);
385 return;
386 }
387 // No queue (self-hosted without one): one batch now; saves and later recalls carry on from there.
388 if (!delaySeconds) await runBackfill(env, workspaceId, false);
389}
390
391/**
392 * One batch of a workspace's backfill: pages by id, then projects' docs
393 * files by space and path, from the run's cursor. Hands on to the next
394 * batch (or waits out the hour's cap) through the queue.
395 */
396export async function runBackfill(env: IndexEnv, workspaceId: string, chain = true): Promise<{ done: boolean }> {
397 const run = await env.DB.prepare("SELECT cursor, finished_at, pages, files FROM doc_index_runs WHERE workspace_id = ?")
398 .bind(workspaceId)
399 .first<{ cursor: string | null; finished_at: string | null; pages: number; files: number }>();
400 if (!run || run.finished_at) return { done: true };
401 const { embedder, store } = adapters(env);
402 const cursor = run.cursor ?? "p:";
403 let next: string | null = cursor;
404 let pages = 0;
405 let files = 0;
406 let capped = false;
407 if (cursor.startsWith("o:")) {
408 // Folios, after pages and before projects' docs.
409 const rows = (
410 await env.DB.prepare(`SELECT ${FOLIO_COLUMNS} FROM folios WHERE workspace_id = ? AND trashed_at IS NULL AND id > ? ORDER BY id LIMIT ?`)
411 .bind(workspaceId, cursor.slice(2), BACKFILL_BATCH)
412 .all<FolioRow>()
413 ).results;
414 for (const row of rows) {
415 const result = await indexFolio(env, row.id);
416 pages++;
417 if (result.capped) {
418 capped = true;
419 break;
420 }
421 next = `o:${row.id}`;
422 }
423 if (!capped && rows.length < BACKFILL_BATCH) next = "f:";
424 } else if (cursor.startsWith("p:")) {
425 const rows = (
426 await env.DB.prepare(
427 "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 ?",
428 )
429 .bind(workspaceId, cursor.slice(2), BACKFILL_BATCH)
430 .all<PageForIndex>()
431 ).results;
432 for (const page of rows) {
433 const result = await indexDoc(env.DB, embedder, store, {
434 kind: "page",
435 doc_id: page.id,
436 workspace_id: page.workspace_id,
437 space_id: page.space_id,
438 title: page.title,
439 markdown: page.markdown,
440 });
441 pages++;
442 if (result.capped) {
443 // This page again next time, after the hour: the cursor stays before it.
444 capped = true;
445 break;
446 }
447 next = `p:${page.id}`;
448 }
449 if (!capped && rows.length < BACKFILL_BATCH) next = "o:";
450 } else {
451 const [space, path] = splitFileCursor(cursor.slice(2));
452 const rows = (
453 await env.DB.prepare(
454 `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
455 WHERE s.workspace_id = ? AND (f.space_id > ? OR (f.space_id = ? AND f.path > ?)) ORDER BY f.space_id, f.path LIMIT ?`,
456 )
457 .bind(workspaceId, space, space, path, BACKFILL_BATCH)
458 .all<RepoFileForIndex>()
459 ).results;
460 for (const row of rows) {
461 const result = await indexDoc(env.DB, embedder, store, repoDoc(row));
462 files++;
463 if (result.capped) {
464 capped = true;
465 break;
466 }
467 next = `f:${row.space_id}\n${row.path}`;
468 }
469 if (!capped && rows.length < BACKFILL_BATCH) next = null;
470 }
471 await env.DB.prepare("UPDATE doc_index_runs SET cursor = ?, pages = pages + ?, files = files + ?, finished_at = ? WHERE workspace_id = ?")
472 .bind(next, pages, files, next === null ? new Date().toISOString() : null, workspaceId)
473 .run();
474 if (next === null) return { done: true };
475 if (chain && env.JOBS) await enqueue(env, workspaceId, capped ? CAP_DELAY_SECONDS : 0);
476 return { done: false };
477}
478
479function splitFileCursor(cursor: string): [string, string] {
480 const at = cursor.indexOf("\n");
481 // Empty: from the start.
482 return at < 0 ? ["", ""] : [cursor.slice(0, at), cursor.slice(at + 1)];
483}
484
485// ── Folios (Artifacts mode) ─────────────────────────────────────────────
486
487let saidNoFolioIndex = false;
488
489/**
490 * The embedder and folio index this deployment has. Without Workers AI or
491 * the `g1t-folios` index (FOLIO_VECTORS, which is made by hand before a
492 * deploy that uses it), folios keep their passages in D1 and recall and
493 * search match words; that is said once per isolate.
494 */
495export function folioAdapters(env: IndexEnv): { embedder: Embedder | null; store: VectorStore | null } {
496 if (env.AI && env.FOLIO_VECTORS) return { embedder: cloudflareEmbedder(env.AI), store: cloudflareVectors(env.FOLIO_VECTORS) };
497 if (!saidNoFolioIndex) {
498 saidNoFolioIndex = true;
499 console.log(`folios: no ${env.AI ? "FOLIO_VECTORS (Vectorize g1t-folios)" : "AI"} binding, so artifacts are searched and recalled by their words only`);
500 }
501 return { embedder: null, store: null };
502}
503
504type FolioChunkRow = { id: string; seq: number; scope: string; hash: string; vector_hash: string | null };
505
506/**
507 * Brings one folio's passages up to date, as `indexDoc` does for pages:
508 * rows and full text in D1, under the folio's index scope, and vectors
509 * for passages whose text or scope changed (a passage whose text the
510 * index already holds keeps its vector; only its metadata is written
511 * again). Forgets it when it is gone or in the trash. Never throws.
512 */
513export async function indexFolio(env: IndexEnv, folioId: string, now = new Date()): Promise<{ capped: boolean }> {
514 try {
515 const db = env.DB;
516 const row = await db
517 .prepare(`SELECT ${FOLIO_COLUMNS.replace("'' AS text", "text")} FROM folios WHERE id = ?`)
518 .bind(folioId)
519 .first<FolioRow>();
520 if (!row || row.trashed_at) {
521 await forgetFolios(env, [folioId]);
522 return { capped: false };
523 }
524 const scope = (await scopesOf(db, [row])).get(row.id) ?? `folio:${row.acl_root}`;
525 const model = kindModel(row.kind);
526 const at = now.toISOString();
527 const chunks = (model ? model.chunks(row.text, row.title) : chunkMarkdown(row.text, row.title)).map((c) => {
528 const embed = embedText(row.title, c);
529 return { ...c, id: chunkId(row.id, c.seq), embed, hash: textHash(embed) };
530 });
531 const old = (await db.prepare("SELECT id, seq, scope, hash, vector_hash FROM folio_chunks WHERE folio_id = ?").bind(row.id).all<FolioChunkRow>()).results;
532 const byId = new Map(old.map((r) => [r.id, r]));
533 const statements: D1PreparedStatement[] = [];
534 for (const c of chunks) {
535 const was = byId.get(c.id);
536 if (was && was.hash === c.hash && was.scope === scope) continue;
537 statements.push(
538 db
539 .prepare(
540 `INSERT INTO folio_chunks (id, workspace_id, folio_id, kind, scope, seq, heading, text, hash, vector_hash, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, NULL, ?)
541 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`,
542 )
543 .bind(c.id, row.workspace_id, row.id, row.kind, scope, c.seq, c.heading, c.text, c.hash, at),
544 db.prepare("DELETE FROM folio_chunks_fts WHERE chunk_id = ?").bind(c.id),
545 db
546 .prepare(
547 "INSERT INTO folio_chunks_fts (chunk_id, scope, folio_id, heading, text) SELECT ?1, ?2, ?3, ?4, ?5 WHERE EXISTS (SELECT 1 FROM folio_chunks WHERE id = ?1)",
548 )
549 .bind(c.id, scope, row.id, c.heading ?? "", c.text),
550 );
551 }
552 const keep = new Set(chunks.map((c) => c.id));
553 const gone = old.filter((r) => !keep.has(r.id));
554 for (const r of gone)
555 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));
556 for (let i = 0; i < statements.length; i += 60) await db.batch(statements.slice(i, i + 60));
557
558 const { embedder, store } = folioAdapters(env);
559 if (!embedder || !store) return { capped: false };
560 if (await erasedSource(db, "folio", row.id, row.workspace_id)) return { capped: false };
561 // What needs a vector written: text the index doesn't hold under this id, or a new scope.
562 const stale = chunks.filter((c) => {
563 const was = byId.get(c.id);
564 return !was || was.vector_hash !== c.hash || was.scope !== scope;
565 });
566 let capped = false;
567 if (stale.length) {
568 const held = new Map<string, string>();
569 for (const r of old) if (r.vector_hash) held.set(r.vector_hash, r.id);
570 const reuse = stale.filter((c) => held.has(c.hash));
571 const fresh = stale.filter((c) => !held.has(c.hash));
572 const budget = fresh.length ? await allowance(db, row.workspace_id, now) : 0;
573 const embedNow = fresh.slice(0, budget);
574 capped = embedNow.length < fresh.length;
575 if (capped) console.log("folios embedding cap reached", row.workspace_id, `${fresh.length - embedNow.length} passages wait for the next hour`);
576 const metadata: VectorMetadata = { workspace_id: row.workspace_id, scope, kind: row.kind, folio_id: row.id };
577 const done: { id: string; hash: string }[] = [];
578 try {
579 const vectors: { id: string; values: number[]; metadata: VectorMetadata }[] = [];
580 if (reuse.length) {
581 const values = new Map((await store.get([...new Set(reuse.map((c) => held.get(c.hash)!))])).map((v) => [v.id, v.values]));
582 for (const c of reuse) {
583 const v = values.get(held.get(c.hash)!);
584 if (v) vectors.push({ id: c.id, values: v, metadata });
585 }
586 }
587 if (embedNow.length) {
588 const embedded = await embedder.embed(embedNow.map((c) => c.embed));
589 embedNow.forEach((c, i) => vectors.push({ id: c.id, values: embedded[i]!, metadata }));
590 await meter(
591 db,
592 row.workspace_id,
593 embedNow.length,
594 embedNow.reduce((n, c) => n + c.embed.length, 0),
595 now,
596 );
597 }
598 if (vectors.length) await store.upsert(vectors);
599 if (vectors.length)
600 await dropIfErased(
601 db,
602 store,
603 "folio",
604 row.id,
605 row.workspace_id,
606 vectors.map((v) => v.id),
607 );
608 const hashOf = new Map(chunks.map((c) => [c.id, c.hash]));
609 for (const v of vectors) done.push({ id: v.id, hash: hashOf.get(v.id)! });
610 } catch (error) {
611 console.error("folios could not embed passages; the next save tries again", row.id, String(error));
612 }
613 if (done.length) {
614 const updates = done.map((d) => db.prepare("UPDATE folio_chunks SET vector_hash = ? WHERE id = ?").bind(d.hash, d.id));
615 for (let i = 0; i < updates.length; i += 60) await db.batch(updates.slice(i, i + 60));
616 }
617 }
618 try {
619 if (gone.some((r) => r.vector_hash)) await store.delete(gone.filter((r) => r.vector_hash).map((r) => r.id));
620 } catch (error) {
621 console.error("folios could not drop old passages from the index", row.id, String(error));
622 }
623 if (capped) await startBackfill(env, row.workspace_id, { delaySeconds: CAP_DELAY_SECONDS });
624 return { capped };
625 } catch (error) {
626 console.error("folios could not index", folioId, String(error));
627 return { capped: false };
628 }
629}
630
631/** Takes folios' passages out of D1 and the index. Never throws. */
632export async function forgetFolios(env: IndexEnv, ids: readonly string[]): Promise<void> {
633 if (!ids.length) return;
634 const db = env.DB;
635 try {
636 const found: string[] = [];
637 for (let i = 0; i < ids.length; i += 500) {
638 const rows = await db
639 .prepare("SELECT id FROM folio_chunks WHERE folio_id IN (SELECT value FROM json_each(?))")
640 .bind(JSON.stringify(ids.slice(i, i + 500)))
641 .all<{ id: string }>();
642 found.push(...rows.results.map((r) => r.id));
643 }
644 if (!found.length) return;
645 const { store } = folioAdapters(env);
646 if (store) {
647 try {
648 await store.delete(found);
649 } catch (error) {
650 console.error("folios could not drop passages from the index", found.length, String(error));
651 }
652 }
653 const statements = found.flatMap((id) => [
654 db.prepare("DELETE FROM folio_chunks WHERE id = ?").bind(id),
655 db.prepare("DELETE FROM folio_chunks_fts WHERE chunk_id = ?").bind(id),
656 ]);
657 for (let i = 0; i < statements.length; i += 80) await db.batch(statements.slice(i, i + 80));
658 } catch (error) {
659 console.error("folios could not forget passages", String(error));
660 }
661}
662
663/** Whether a document, or its workspace, was erased (src/erasure.ts): nothing of it is indexed again. */
664export async function erasedSource(db: D1Database, kind: "page" | "folio" | null, id: string, workspaceId: string): Promise<boolean> {
665 const row = await db
666 .prepare("SELECT 1 FROM erased_subjects WHERE id = ?1 AND kind = 'workspace' UNION ALL SELECT 1 FROM erased_documents WHERE id = ?2 AND kind = ?3 LIMIT 1")
667 .bind(workspaceId, id, kind ?? "")
668 .first();
669 return row != null;
670}
671
672/**
673 * Checked again after vectors were written: an erasure that ran while they
674 * were being embedded found no chunk rows for them, so they are taken back
675 * here. Throws when they cannot be, so the caller does not record them.
676 */
677export async function dropIfErased(
678 db: D1Database,
679 store: VectorStore,
680 kind: "page" | "folio" | null,
681 id: string,
682 workspaceId: string,
683 ids: string[],
684): Promise<void> {
685 if (!ids.length || !(await erasedSource(db, kind, id, workspaceId))) return;
686 await store.delete(ids);
687 await db.batch([db.prepare(`DELETE FROM ${kind === "folio" ? "folio_chunks_fts WHERE folio_id" : "doc_chunks_fts WHERE doc_id"} = ?`).bind(id)]);
688}