Arcrun/kbdb/tests/embed-backfill.test.ts

import { describe, it, expect } from 'vitest';
import { backfillEmbeddings, backfillStatus, embedEnabled } from '../src/embed';
import type { Bindings, Entry } from '../src/types';

// ── Minimal in-memory fakes (no Workers runtime) ─────────────────────────────
// The fake DB interprets only the 3 statement shapes backfill issues, by keyword:
//   SELECT * ... LIMIT ? OFFSET ?  → candidate rows (embeddable & non-empty content;
//                                     +is_embedded=0 for normal backfill, any for reindex)
//   UPDATE ... IN (...)            → flip is_embedded=1 for the bound ids
//   SELECT COUNT(*)                → count of matching candidates
// embeddable = metadata.embed===true & non-empty content（reindex predicate）。
function isEmbeddable(e: Entry): boolean {
  if (!e.content || e.content.trim() === '') return false;
  try {
    const m = JSON.parse(e.metadata_json ?? 'null');
    return m?.embed === true;
  } catch {
    return false;
  }
}
// normal backfill 額外要求 is_embedded=0（漏網補嵌）。
function isCandidate(e: Entry): boolean {
  return e.is_embedded === 0 && isEmbeddable(e);
}

function makeFakeDB(store: Entry[]) {
  const prepare = (sql: string) => {
    // reindex predicate 不含 "is_embedded = 0" → 依 SQL 判斷該用哪個 filter（對齊 embed.ts）。
    const pred = /is_embedded = 0/.test(sql) ? isCandidate : isEmbeddable;
    let bound: unknown[] = [];
    const stmt = {
      bind(...args: unknown[]) { bound = args; return stmt; },
      async all<T>() {
        // SELECT * ... LIMIT ? OFFSET ?  (bound tail = [..., limit, offset])
        const offset = Number(bound[bound.length - 1]);
        const limit = Number(bound[bound.length - 2]);
        const results = store.filter(pred).slice(offset, offset + limit) as unknown as T[];
        return { results };
      },
      async first<T>() {
        // SELECT COUNT(*) as c ...
        const c = store.filter(pred).length;
        return { c } as unknown as T;
      },
      async run() {
        // UPDATE entries SET is_embedded = 1 WHERE id IN (...)  → bound = ids
        const ids = new Set(bound.map(String));
        for (const e of store) if (ids.has(e.id)) e.is_embedded = 1;
        return { success: true };
      },
    };
    return stmt;
  };
  return { prepare } as unknown as D1Database;
}

function mkEntry(id: string, content: string | null, embed: boolean, is_embedded = 0): Entry {
  return {
    id, content, entry_type: 'workflow', owner_id: 'leo', parent_id: null, page_name: null,
    refs_json: '[]', tags_json: '[]', task_status: null, content_hash: null, is_embedded,
    confidence: null, metadata_json: JSON.stringify({ embed }), created_at: 1, updated_at: 1,
  };
}

function makeEnv(store: Entry[], withBindings: boolean): Bindings {
  const upserts: { id: string }[] = [];
  const aiCalls: string[][] = [];
  const env = {
    DB: makeFakeDB(store),
    ENVIRONMENT: 'test',
    ...(withBindings
      ? {
          AI: { async run(_m: string, i: { text: string[] }) { aiCalls.push(i.text); return { data: i.text.map(() => [0.1, 0.2, 0.3]) }; } },
          VECTORIZE: { async upsert(v: { id: string }[]) { upserts.push(...v); return { count: v.length }; } },
        }
      : {}),
  } as unknown as Bindings;
  (env as unknown as { __upserts: unknown[]; __ai: unknown[] }).__upserts = upserts;
  (env as unknown as { __upserts: unknown[]; __ai: unknown[] }).__ai = aiCalls;
  return env;
}

describe('backfillEmbeddings', () => {
  it('module off → enabled:false, no-op (誠實不假綠)', async () => {
    const store = [mkEntry('e1', 'hello', true)];
    const env = makeEnv(store, false);
    expect(embedEnabled(env)).toBe(false);
    const r = await backfillEmbeddings(env);
    expect(r).toEqual({ enabled: false, processed: 0, skipped: 0, remaining: 0, scanned: 0 });
    expect(store[0].is_embedded).toBe(0); // untouched
  });

  it('embeds embeddable+is_embedded=0 entries, marks is_embedded=1, batches AI+upsert', async () => {
    const store = [
      mkEntry('e1', 'doorbell workflow', true),
      mkEntry('e2', 'notify workflow', true),
      mkEntry('e3', 'not tagged', false),         // embed:false → not a candidate
      mkEntry('e4', 'already done', true, 1),      // is_embedded=1 → not a candidate
      mkEntry('e5', '   ', true),                  // empty content → not embeddable
    ];
    const env = makeEnv(store, true);
    const r = await backfillEmbeddings(env, { limit: 100 });
    expect(r.enabled).toBe(true);
    expect(r.processed).toBe(2);       // only e1,e2
    expect(r.remaining).toBe(0);       // nothing left embeddable
    expect(store.find((e) => e.id === 'e1')!.is_embedded).toBe(1);
    expect(store.find((e) => e.id === 'e2')!.is_embedded).toBe(1);
    expect(store.find((e) => e.id === 'e3')!.is_embedded).toBe(0);
    const upserts = (env as unknown as { __upserts: { id: string }[] }).__upserts;
    expect(upserts.map((u) => u.id).sort()).toEqual(['e1', 'e2']);
    const ai = (env as unknown as { __ai: string[][] }).__ai;
    expect(ai.length).toBe(1);         // single batched AI.run for the whole batch
    expect(ai[0].length).toBe(2);
  });

  it('idempotent: re-run after all embedded processes nothing', async () => {
    const store = [mkEntry('e1', 'x', true)];
    const env = makeEnv(store, true);
    await backfillEmbeddings(env);
    const r2 = await backfillEmbeddings(env);
    expect(r2.processed).toBe(0);
    expect(r2.remaining).toBe(0);
  });

  it('batches via limit → remaining reported so caller can loop to zero', async () => {
    const store = [mkEntry('a', 'x', true), mkEntry('b', 'y', true), mkEntry('c', 'z', true)];
    const env = makeEnv(store, true);
    const r1 = await backfillEmbeddings(env, { limit: 2 });
    expect(r1.processed).toBe(2);
    expect(r1.remaining).toBe(1);
    const r2 = await backfillEmbeddings(env, { limit: 2 });
    expect(r2.processed).toBe(1);
    expect(r2.remaining).toBe(0);
  });

  it('reindex: 重推所有 embeddable（含 is_embedded=1），offset 分頁到 remaining=0（Arcrun#11）', async () => {
    // 三筆皆已 is_embedded=1（既有向量）：正常 backfill 不會碰（pending=0），reindex 要全部重推
    // 讓事後建立的 Vectorize metadata index 收錄。
    const store = [
      mkEntry('a', 'x', true, 1), mkEntry('b', 'y', true, 1), mkEntry('c', 'z', true, 1),
    ];
    const env = makeEnv(store, true);
    // 正常 backfill：沒有 is_embedded=0 → 什麼都不做（證明「不重推就補不到」）。
    const normal = await backfillEmbeddings(env, { limit: 100 });
    expect(normal.processed).toBe(0);
    // reindex 分頁：第一批 2 筆、remaining=1；第二批 1 筆、remaining=0。
    const r1 = await backfillEmbeddings(env, { reindex: true, limit: 2, offset: 0 });
    expect(r1.processed).toBe(2);
    expect(r1.remaining).toBe(1);
    const r2 = await backfillEmbeddings(env, { reindex: true, limit: 2, offset: 2 });
    expect(r2.processed).toBe(1);
    expect(r2.remaining).toBe(0);
    const upserts = (env as unknown as { __upserts: { id: string }[] }).__upserts;
    expect(upserts.map((u) => u.id).sort()).toEqual(['a', 'b', 'c']);
  });

  it('status reports pending/embedded counts', async () => {
    const store = [mkEntry('e1', 'x', true), mkEntry('e2', 'y', true, 1)];
    const env = makeEnv(store, true);
    const s = await backfillStatus(env);
    // fake first() returns candidate count for pending; embedded query also runs through
    // the same COUNT fake, so this asserts the call path works (enabled:true).
    expect(s.enabled).toBe(true);
    expect(typeof s.pending).toBe('number');
  });
});