app/lib/search/ask.server.ts

Functions

Imports

It imports app/db/publications.ts, app/lib/error-message.mjs, app/lib/health/verdicts.mjs, app/lib/publications/paths.mjs, app/lib/search/ask-guard.server.ts, app/lib/search/ask-keys.mjs, app/lib/search/ask-prompt.mjs, app/lib/search/ask-twins.mjs, app/lib/search/records.mjs, app/lib/search/search.server.ts, app/lib/search/sse.mjs, app/lib/search/visibility.mjs, app/lib/timing.ts. Imported by app/lib/admin/posts-actions.server.ts, app/lib/editor/publish.server.ts, app/lib/health/checks.server.ts, app/lib/operator/sync-tools.server.ts, app/lib/publications/save.server.ts, app/routes/admin._index.tsx, app/routes/admin.posts._index.tsx, app/routes/admin.tsx, app/routes/search.ask.ts, app/routes/search.tsx.

Source

1// With AI_SEARCH unbound no Ask affordance renders and /search is pure D1: that off switch is a requirement.
2
3import { createContext } from "react-router";
4
5import {
6  dropCachedDrift,
7  invalidateAnswerCache,
8  readCachedDrift,
9  writeCachedDrift,
10} from "./ask-guard.server";
11import { KEY_SEPARATOR, keyForUrl, ownsAskKey } from "./ask-keys.mjs";
12import { uploadTwins, withRetry } from "./ask-twins.mjs";
13import { answerDelta, takeSseFrames } from "./sse.mjs";
14import { setDrift } from "~/lib/health/verdicts.mjs";
15import { isPubliclyVisible, statusForDraft } from "./visibility.mjs";
16import { askCorpusRecords, askExpectedUrls } from "./search.server";
17import { listPublishedPublicationTwins } from "~/db/publications";
18import { paperPath } from "~/lib/publications/paths.mjs";
19import { timed, type Timings } from "~/lib/timing";
20import { recordsForPosts } from "./records.mjs";
21import { errorMessage } from "~/lib/error-message.mjs";
22import {
23  answerLeaksPrompt,
24  askMessages,
25  citedSlugs,
26  guardAnswerStream,
27  replayFrames,
28} from "./ask-prompt.mjs";
29
30export { answerLeaksPrompt, citedSlugs, guardAnswerStream };
31
32const ASK_MODEL = "@cf/meta/llama-3.3-70b-instruct-fp8-fast";
33
34/** Bounds the answer, not the corpus: past it, retrieval widens the net without widening the answer. */
35const MAX_CHUNKS = 6;
36
37/** A property check, not a try/catch: only a removed binding, not an erroring instance, should hide Ask. */
38export function askAvailable(env: Env): boolean {
39  return Boolean(env.AI_SEARCH);
40}
41
42/** The raw upstream SSE is handed through untouched: the client already parses SSE. */
43export async function askStream(env: Env, question: string): Promise<ReadableStream> {
44  return env.AI_SEARCH.chatCompletions({
45    // The last message is the retrieval query; decorating it costs the whole search.
46    messages: askMessages(question),
47    model: ASK_MODEL,
48    stream: true,
49    ai_search_options: {
50      max_num_results: MAX_CHUNKS,
51    },
52  });
53}
54
55/** The cache copy is parsed inside waitUntil, so caching costs time-to-first-token nothing. */
56export function teeForCache(upstream: ReadableStream): {
57  toReader: ReadableStream;
58  captured: Promise<{ answer: string; chunks: unknown[] } | null>;
59} {
60  const [toReader, toCache] = upstream.tee();
61  return { toReader, captured: parseSseAnswer(toCache) };
62}
63
64/** Same frame handling as the client, or a cached replay would not match what the reader first saw. */
65async function parseSseAnswer(
66  stream: ReadableStream,
67): Promise<{ answer: string; chunks: unknown[] } | null> {
68  const reader = stream.getReader();
69  const decoder = new TextDecoder();
70  let buffer = "";
71  let answer = "";
72  let chunks: unknown[] = [];
73
74  try {
75    for (;;) {
76      const { done, value } = await reader.read();
77      if (done) break;
78      buffer += decoder.decode(value, { stream: true });
79      const taken = takeSseFrames(buffer);
80      buffer = taken.rest;
81      for (const { event, data } of taken.frames) {
82        if (event === "chunks") {
83          if (Array.isArray(data)) chunks = data;
84          continue;
85        }
86        answer += answerDelta(data) ?? "";
87      }
88    }
89  } catch {
90    return null;
91  }
92
93  return answer.trim().length > 0 ? { answer, chunks } : null;
94}
95
96export function replayCachedAnswer(cached: {
97  answer: string;
98  chunks: unknown[];
99}): ReadableStream {
100  const encoder = new TextEncoder();
101  return new ReadableStream({
102    start(controller) {
103      for (const frame of replayFrames(cached)) controller.enqueue(encoder.encode(frame));
104      controller.close();
105    },
106  });
107}
108
109interface CorpusSyncResult {
110  uploaded: number;
111  keys: string[];
112  /** Non-empty means the sync did NOT converge. */
113  failed: Array<{ key: string; error: string }>;
114  cacheDropped: number;
115}
116
117/** items.list() is paged: a bare call returns only the first page, and a prune cannot delete what it cannot see. */
118async function listAllAskItems(env: Env, timings?: Timings) {
119  const all: Awaited<ReturnType<typeof env.AI_SEARCH.items.list>>["result"] = [];
120  // 50 is the API maximum.
121  const perPage = 50;
122  // One mark per page, so round trips are counted rather than inferred from the index size.
123  const listPage = (page: number) =>
124    timed(timings, "ask_list_page", () => env.AI_SEARCH.items.list({ page, per_page: perPage }));
125
126  // The first page says how many items there are, so the rest are fetched together instead of one round trip after
127  // another. The loop below still runs after them: it ends at once when they were all of it, and carries on page by
128  // page when the count was wrong, so a listing is never taken for complete on the count alone.
129  let page = 1;
130  const first = await listPage(1);
131  all.push(...(first.result ?? []));
132  const total = first.result_info?.total_count;
133  let last = first;
134  if (typeof total === "number" && (first.result ?? []).length === perPage && total > perPage) {
135    const lastPage = Math.ceil(total / perPage);
136    const rest = await Promise.all(Array.from({ length: lastPage - 1 }, (_, i) => listPage(i + 2)));
137    for (const listed of rest) all.push(...(listed.result ?? []));
138    page = lastPage;
139    last = rest[rest.length - 1] ?? first;
140  }
141  for (;;) {
142    const batch = last.result ?? [];
143    if (batch.length < perPage) break;
144    if (typeof total === "number" && all.length >= total) break;
145    // Backstop against a server that never shrinks a page. Thrown, not a break: a truncated listing
146    // would make every unlisted item look absent to the prune and the drift count.
147    if (page > 200) {
148      throw new Error(
149        `AI Search listing did not end after ${page} pages of ${perPage} (${all.length} items); ` +
150          `refusing to act on a listing that may be truncated.`,
151      );
152    }
153    page += 1;
154    last = await listPage(page);
155    all.push(...(last.result ?? []));
156  }
157  return all;
158}
159
160/** AI Search has no per-item status to filter on, so nothing unpublished may be uploaded at all. */
161function publishableForAsk<T extends { draft?: boolean; publishAt?: string | null }>(
162  posts: readonly T[],
163): T[] {
164  const now = Date.now();
165  return posts.filter((p) =>
166    isPubliclyVisible({ status: statusForDraft(p.draft), publishAt: p.publishAt }, now),
167  );
168}
169
170function askPublishable(post: { draft?: boolean; publishAt?: string | null }): boolean {
171  return publishableForAsk([post]).length === 1;
172}
173
174/**
175 * One search record as an Ask item. Fails closed: a url containing the separator would produce a
176 * key that resolves to the wrong URL. The heading is in the text: the body alone loses what the
177 * section is about.
178 */
179function askItemFor(record: Record<string, any>) {
180  if (record.url.includes(KEY_SEPARATOR)) {
181    throw new Error(
182      `record url contains the key separator "${KEY_SEPARATOR}" and cannot be ` +
183        `uploaded safely: ${record.url}`,
184    );
185  }
186  return { key: keyForUrl(record.url), body: `# ${record.title}\n\n${record.body}\n` };
187}
188
189/** Lists every item, deletes the ones whose key matches, and returns their keys. */
190async function deleteAskItemsWhere(env: Env, matches: (key: string) => boolean) {
191  const removed: string[] = [];
192  for (const item of await listAllAskItems(env)) {
193    if (!matches(item.key)) continue;
194    await env.AI_SEARCH.items.delete(item.id);
195    removed.push(item.key);
196  }
197  return removed;
198}
199
200/** Uploaded rather than crawled: the crawler would index the apex (the legacy site) and lose heading granularity. */
201export async function syncAskCorpus(env: Env): Promise<CorpusSyncResult> {
202  // From D1 with visibilityClause composed in the SQL, so drafts and future posts never enter the index.
203  const records = await askCorpusRecords(env);
204  const keys: string[] = [];
205
206  for (const record of records) {
207    const { key, body } = askItemFor(record);
208    await env.AI_SEARCH.items.upload(key, body);
209    keys.push(key);
210  }
211
212  // Papers from their twins, not search_docs: uploading both would let the shorter win questions the longer answers.
213  const twins = await uploadPaperTwins(env);
214  keys.push(...twins.keys);
215
216  // The corpus changed, so cached answers may now be untrue.
217  const dropped = await invalidateAnswerCache(env);
218  await dropCachedDrift(env);
219
220  return {
221    uploaded: keys.length - twins.failed.length,
222    keys,
223    failed: twins.failed,
224    cacheDropped: dropped,
225  };
226}
227
228/** The published papers' Ask keys, from D1: a draft paper is not public, so it is not asked about. */
229async function paperItemKeys(env: Env): Promise<string[]> {
230  return (await listPublishedPublicationTwins(env)).map(({ slug }) => keyForUrl(paperPath(slug)));
231}
232
233/** Each published twin from D1, the row the twin route serves, so this indexes exactly what the site says. */
234async function uploadPaperTwins(env: Env) {
235  const rows = await listPublishedPublicationTwins(env);
236  const bySlug = new Map(rows.map((row) => [row.slug, row.markdown]));
237  const twins = rows.map(({ slug }) => ({ key: keyForUrl(paperPath(slug)), path: slug }));
238  return uploadTwins(twins, {
239    fetchText: async (slug: string) => {
240      const markdown = bySlug.get(slug);
241      if (markdown === undefined) throw new Error(`twin ${slug} has no published row`);
242      return markdown;
243    },
244    upload: (key: string, body: string) => env.AI_SEARCH.items.upload(key, body),
245  });
246}
247
248/** Scoped to this post's keys, which is sound because sections are a pure function of one post's markdown. */
249export async function syncAskPost(
250  env: Env,
251  post: Parameters<typeof recordsForPosts>[0][number],
252): Promise<{ uploaded: number; removed: number; failed: string[] }> {
253  // A draft uploads nothing, and the empty live set makes the prune below remove what it had.
254  const records = askPublishable(post) ? recordsForPosts([post]) : [];
255  const live = new Set<string>();
256
257  // A failed key still joins live: the prune deletes every key not in it, which would delete the indexed copy.
258  const failed: string[] = [];
259
260  for (const record of records) {
261    // Outside the retry: a bad url would produce the same wrong key every time.
262    const { key, body } = askItemFor(record);
263    try {
264      await withRetry(() => env.AI_SEARCH.items.upload(key, body), { attempts: 2 });
265    } catch (error) {
266      console.error(JSON.stringify({ alert: "ask-upload-failed-twice", key, detail: errorMessage(error) }));
267      failed.push(key);
268    }
269    live.add(key);
270  }
271
272  const ours = ownsAskKey(post.slug);
273  const removed = await deleteAskItemsWhere(env, (key) => ours(key) && !live.has(key));
274
275  await invalidateAnswerCache(env);
276  await dropCachedDrift(env);
277  return { uploaded: records.length - failed.length, removed: removed.length, failed };
278}
279
280export interface AskIndexStatus {
281  expected: number;
282  present: number;
283  missing: string[];
284  stale: string[];
285}
286
287/** Both directions: an item the corpus does not know is as much a defect as a missing one. */
288export async function askIndexStatus(env: Env, timings?: Timings): Promise<AskIndexStatus> {
289  // publishableForAsk must stay in step with the SQL predicate here; check:ask-guards binds the two.
290  // The papers are read with the rest, not after it: without them the badge would report them stale and the repair
291  // button would delete them, and a read that waited for the listing was a round trip of pure waiting.
292  const [expectedUrls, listed, paperKeys] = await Promise.all([
293    askExpectedUrls(env),
294    listAllAskItems(env, timings),
295    paperItemKeys(env),
296  ]);
297  const expected = new Set([...expectedUrls.map((u) => keyForUrl(u)), ...paperKeys]);
298  const present = new Set(listed.map((item) => item.key));
299
300  const { extra, ...drift } = setDrift(expected, present);
301  return { ...drift, stale: extra };
302}
303
304/** On the context because a parent cannot read a child's loader data; a getter, so a route that never asks never pays. */
305type AskStatusReader = () => Promise<AskIndexStatus | null>;
306
307export const askStatusContext = createContext<AskStatusReader>();
308
309/** Null rather than throwing: an enhancement must not take an admin page down. */
310export function askStatusReader(env: Env, timings?: Timings): AskStatusReader {
311  let pending: Promise<AskIndexStatus | null> | undefined;
312  return () => {
313    pending ??= (async () => {
314      if (!askAvailable(env)) return null;
315      try {
316        return await askIndexStatus(env, timings);
317      } catch (error) {
318        console.error(JSON.stringify({ alert: "ask-index-status-failed", detail: errorMessage(error) }));
319        return null;
320      }
321    })();
322    return pending;
323  };
324}
325
326/**
327 * On a hit nothing touches AI Search: keep the early return before any index use. Null, not zero, on
328 * failure. A late listing is handed to waitUntil, since a floating write never lands after the response.
329 */
330export async function askDriftCount(
331  env: Env,
332  ctx: ExecutionContext,
333  timings?: Timings,
334): Promise<number | null> {
335  if (!askAvailable(env)) return null;
336
337  const cached = await timed(timings, "drift_cache_read", () => readCachedDrift(env));
338  if (cached !== null) return cached;
339
340  const listing = askIndexStatus(env, timings).catch((error) => {
341    console.error(JSON.stringify({ alert: "ask-drift-count-failed", detail: errorMessage(error) }));
342    return null;
343  });
344
345  const budget = new Promise<undefined>((resolve) => {
346    setTimeout(() => resolve(undefined), DRIFT_BUDGET_MS);
347  });
348
349  const status = await Promise.race([listing, budget]);
350
351  if (status === undefined) {
352    // The catch stays: a rejecting waitUntil is no better than a floating rejection.
353    ctx.waitUntil(
354      listing
355        .then((late) =>
356          late ? writeCachedDrift(env, late.missing.length + late.stale.length) : null,
357        )
358        .catch((error) =>
359          console.error(JSON.stringify({ alert: "ask-drift-cache-write-failed", detail: errorMessage(error) })),
360        ),
361    );
362    return null;
363  }
364
365  if (!status) return null;
366
367  const drift = status.missing.length + status.stale.length;
368  await writeCachedDrift(env, drift);
369  return drift;
370}
371
372/** Measured: admits all but the one pathological listing this budget exists for. */
373const DRIFT_BUDGET_MS = 1000;
374
375export async function removeAskPost(env: Env, slug: string): Promise<number> {
376  const removed = await deleteAskItemsWhere(env, ownsAskKey(slug));
377  await invalidateAnswerCache(env);
378  await dropCachedDrift(env);
379  return removed.length;
380}
381
382/** Upload is an upsert, so a renamed or deleted post leaves its old item answering forever. */
383export async function pruneAskCorpus(env: Env, liveKeys: string[]): Promise<string[]> {
384  const live = new Set(liveKeys);
385  const removed = await deleteAskItemsWhere(env, (key) => !live.has(key));
386  // Only on an actual removal; dropping the key otherwise would cost the next reader a listing.
387  if (removed.length > 0) await dropCachedDrift(env);
388  return removed;
389}
390
391/**
392 * One paper's twin in or out of the Ask index: uploaded when `markdown` is given, removed when it is null
393 * (a draft, or a paper that is gone). Scoped to this paper's key, as syncAskPost is to a post's.
394 */
395export async function syncAskPaper(
396  env: Env,
397  slug: string,
398  markdown: string | null,
399): Promise<{ uploaded: number; removed: number }> {
400  const key = keyForUrl(paperPath(slug));
401  let uploaded = 0;
402  let removed = 0;
403  if (markdown !== null) {
404    await withRetry(() => env.AI_SEARCH.items.upload(key, markdown), { attempts: 2 });
405    uploaded = 1;
406  } else {
407    removed = (await deleteAskItemsWhere(env, (itemKey) => itemKey === key)).length;
408  }
409  await invalidateAnswerCache(env);
410  await dropCachedDrift(env);
411  return { uploaded, removed };
412}