Arcrun/kbdb/src/embed.ts

// KBDB optional embed module (issue #7 / mira-dissolve SDD T2.4).
//
// 鐵律對齊：
//   - embedding 屬 **base 的 optional 模組**（非 graph/ingest）。CF 內建（Vectorize+AI），程式薄。
//   - **不拆 repo，binding 開/關**：有 env.VECTORIZE + env.AI 才啟用；沒有 → base 維持 LIKE keyword，API 不變。
//   - 不動三表結構（只標既有 entries.is_embedded / content_hash bookkeeping 欄；那些 base 從不讀，embed 才寫）。
//   - 不對每個 block 地毯式 embed（精耕，非 RAG 一股腦灌）：只 embed「被標記為 embeddable」的 entry
//     （wiki 段落 + graph node gloss）。標記方式＝寫入時 metadata_json.embed === true（caller 顯式標）。
//
// 為何用 metadata flag 而非 entry_type 白名單：base 不該寫死「哪些 entry_type 該 embed」（那是上游語意，
// 會讓 base 知道 wiki/graph 概念，破壞解耦）。改由 caller（wiki/gloss 寫入端）顯式標 embed:true，
// base 只認這個通用旗標 → base 維持對內容語意無知。

import type { Bindings, Entry } from './types';

const EMBED_MODEL = '@cf/baai/bge-base-en-v1.5'; // 768-dim，與 Vectorize index dimensions=768 對齊

/** embed 模組是否啟用（binding 都在才算開）。base 一切 embed 動作先過這關。 */
export function embedEnabled(env: Bindings): boolean {
  return !!(env.VECTORIZE && env.AI);
}

/** 一段文字 → 768 維向量（Workers AI bge）。空字串回 null（不 embed）。 */
async function embedText(env: Bindings, text: string): Promise<number[] | null> {
  const t = (text ?? '').trim();
  if (!t || !env.AI) return null;
  const res = (await env.AI.run(EMBED_MODEL, { text: [t] })) as { data: number[][] };
  return res?.data?.[0] ?? null;
}

/**
 * 寫入時選擇性 embed（embed-on-write，#5 第4點併入此）。
 *   - 模組未開 → no-op（base 輕量）。
 *   - 只 embed 被標 embeddable 的 entry（metadata_json.embed === true）。其餘略過（非地毯式）。
 * 失敗不致命（fire-and-forget 由 caller 用 waitUntil 包；這裡只負責「能 embed 就 embed」）。
 * 回傳是否真的 embed 了（讓 caller 決定要不要標 is_embedded）。
 */
export async function embedOnWrite(env: Bindings, entry: Entry): Promise<boolean> {
  if (!embedEnabled(env)) return false;
  if (!isEmbeddable(entry)) return false;
  const vec = await embedText(env, entry.content ?? '');
  if (!vec) return false;
  await env.VECTORIZE!.upsert([
    {
      id: entry.id,
      values: vec,
      // metadata 走 indexed 範圍：owner_id（租戶隔離）、entry_type、source（#5.1 過濾與語義共用）。
      metadata: {
        owner_id: entry.owner_id ?? '',
        entry_type: entry.entry_type,
        source: readSource(entry) ?? '',
      },
    },
  ]);
  // 標記 bookkeeping（既有欄，base 不讀、僅供「已 embed」可查）。不動表結構。
  await env.DB.prepare('UPDATE entries SET is_embedded = 1 WHERE id = ?').bind(entry.id).run();
  return true;
}

/** entry 是否該被 embed：caller 在 metadata_json 標 embed:true（精耕，非地毯式）。 */
function isEmbeddable(entry: Entry): boolean {
  const meta = parseMeta(entry.metadata_json);
  return meta?.embed === true;
}

function readSource(entry: Entry): string | null {
  const meta = parseMeta(entry.metadata_json);
  const s = meta?.source;
  return typeof s === 'string' ? s : null;
}

function parseMeta(json: string | null): Record<string, unknown> | null {
  if (!json) return null;
  try {
    const p = JSON.parse(json);
    return p && typeof p === 'object' ? (p as Record<string, unknown>) : null;
  } catch {
    return null;
  }
}

// SQL predicate for "an entry that SHOULD be embedded but isn't yet".
//   - isEmbeddable 契約 = metadata_json.embed === true（base 通用旗標，對內容語意無知，不寫死 entry_type）。
//     SQLite json_extract 對 JSON boolean true 回整數 1 → `= 1` 精確對齊 TS 的 `=== true`。
//   - is_embedded = 0：尚未（對「當前」index）補嵌的 bookkeeping。
//   - content 非空：空字串 embedText 會回 null，排除以免變成永遠清不掉的殘留候選。
const BACKFILL_PREDICATE =
  "is_embedded = 0 AND content IS NOT NULL AND content <> '' AND json_extract(metadata_json, '$.embed') = 1";

export interface BackfillResult {
  enabled: boolean;   // 模組是否開（false → 什麼都沒做，caller 該誠實回錯，不假裝）。
  processed: number;  // 本次真的嵌進 Vectorize 並標 is_embedded=1 的筆數。
  skipped: number;    // 掃到但沒嵌（例如 embedText 回 null）的筆數。
  remaining: number;  // 本次之後仍待補嵌的筆數（可重複呼叫直到 0）。
  scanned: number;    // 本批掃出的候選筆數（受 limit 限制）。
}

/**
 * Backfill（回填）：對「開 Vectorize 之前就寫入、或 embed-on-write 當時漏掉」的既有 entry 批次補嵌。
 * 冪等（重跑已補嵌的不會重複算，upsert 同 id 冪等）、分批（單次 limit 上限，避開 subrequest/CPU/timeout）、
 * 回傳處理筆數 + 剩餘筆數（caller 重複呼叫直到 remaining=0）。
 *   - 模組未開（無 VECTORIZE+AI）→ 誠實回 { enabled:false }，不假裝成功（mindset §7 禁假綠）。
 *   - 只補「isEmbeddable（metadata.embed===true）且 is_embedded=0」的 entry——與 embedOnWrite 同一契約，
 *     base 維持對內容語意無知（不知 triplet/wiki，只認通用 embed 旗標）。
 *   - 效率：整批用「單次 AI.run（陣列輸入）+ 單次 VECTORIZE.upsert（陣列）+ 單次 UPDATE ... IN(...)」，
 *     一批 ≈ 3 個 subrequest，不隨 limit 線性增長 → free/paid tier 都安全。
 */
export async function backfillEmbeddings(
  env: Bindings,
  opts: { limit?: number; owner_id?: string; source?: string; reindex?: boolean; offset?: number } = {},
): Promise<BackfillResult> {
  if (!embedEnabled(env)) return { enabled: false, processed: 0, skipped: 0, remaining: 0, scanned: 0 };
  const limit = Math.min(Math.max(opts.limit ?? 25, 1), 100);
  const offset = Math.max(opts.offset ?? 0, 0);

  // reindex（Arcrun#11 根因修復）：對「既有已嵌」向量原樣重嵌重推 upsert，讓它們被『事後才建立』的
  //   Vectorize metadata index（owner_id/entry_type/source）收錄。Vectorize 只索引「metadata index
  //   建立之後 upsert」的向量 → 既有向量不重推就永遠 filter 不到（= 本 bug）。upsert 同 id 冪等。
  //   非 reindex（預設）＝原行為：只補 is_embedded=0 的漏網。
  const basePredicate = opts.reindex
    ? "content IS NOT NULL AND content <> '' AND json_extract(metadata_json, '$.embed') = 1"
    : BACKFILL_PREDICATE;

  const conds = [basePredicate];
  const params: unknown[] = [];
  if (opts.owner_id) { conds.push('owner_id = ?'); params.push(opts.owner_id); }
  if (opts.source) { conds.push("json_extract(metadata_json, '$.source') = ?"); params.push(opts.source); }
  const where = conds.join(' AND ');

  const res = await env.DB
    .prepare(`SELECT * FROM entries WHERE ${where} ORDER BY created_at ASC LIMIT ? OFFSET ?`)
    .bind(...params, limit, offset)
    .all<Entry>();
  const rows = res.results ?? [];
  const scanned = rows.length;

  let processed = 0;
  const embeddable = rows.filter((e) => (e.content ?? '').trim().length > 0);
  if (embeddable.length > 0 && env.AI && env.VECTORIZE) {
    const texts = embeddable.map((e) => (e.content ?? '').trim());
    const out = (await env.AI.run(EMBED_MODEL, { text: texts })) as { data: number[][] };
    const data = out?.data ?? [];
    const vectors = embeddable
      .map((e, i) => ({ e, vec: data[i] }))
      .filter((x): x is { e: Entry; vec: number[] } => Array.isArray(x.vec) && x.vec.length > 0)
      .map((x) => ({
        id: x.e.id,
        values: x.vec,
        metadata: {
          owner_id: x.e.owner_id ?? '',
          entry_type: x.e.entry_type,
          source: readSource(x.e) ?? '',
        },
      }));
    if (vectors.length > 0) {
      await env.VECTORIZE.upsert(vectors);
      const ids = vectors.map((v) => v.id);
      const placeholders = ids.map(() => '?').join(',');
      await env.DB.prepare(`UPDATE entries SET is_embedded = 1 WHERE id IN (${placeholders})`).bind(...ids).run();
      processed = vectors.length;
    }
  }

  const remRow = await env.DB
    .prepare(`SELECT COUNT(*) as c FROM entries WHERE ${where}`)
    .bind(...params)
    .first<{ c: number }>();
  const totalMatching = remRow?.c ?? 0;
  // 非 reindex：predicate 含 is_embedded=0，處理後該筆變 1 → COUNT 自然遞減（重呼直到 0）。
  // reindex：predicate 不含 is_embedded，COUNT 恆等於總數 → 改用 offset 分頁計 remaining（否則永不終止）。
  const remaining = opts.reindex ? Math.max(0, totalMatching - (offset + scanned)) : totalMatching;
  return { enabled: true, processed, skipped: scanned - processed, remaining, scanned };
}

/** 補嵌進度統計（回報用；模組未開仍可查 pending 數，誠實標 enabled:false）。 */
export async function backfillStatus(
  env: Bindings,
  opts: { owner_id?: string; source?: string } = {},
): Promise<{ enabled: boolean; pending: number; embedded: number }> {
  const conds: string[] = [];
  const params: unknown[] = [];
  if (opts.owner_id) { conds.push('owner_id = ?'); params.push(opts.owner_id); }
  if (opts.source) { conds.push("json_extract(metadata_json, '$.source') = ?"); params.push(opts.source); }
  const extra = conds.length ? ` AND ${conds.join(' AND ')}` : '';
  const pendingRow = await env.DB
    .prepare(`SELECT COUNT(*) as c FROM entries WHERE ${BACKFILL_PREDICATE}${extra}`)
    .bind(...params)
    .first<{ c: number }>();
  const embeddedRow = await env.DB
    .prepare(`SELECT COUNT(*) as c FROM entries WHERE is_embedded = 1 AND json_extract(metadata_json, '$.embed') = 1${extra}`)
    .bind(...params)
    .first<{ c: number }>();
  return { enabled: embedEnabled(env), pending: pendingRow?.c ?? 0, embedded: embeddedRow?.c ?? 0 };
}

export interface SemanticHit {
  id: string;
  score: number;
  owner_id?: string;
  entry_type?: string;
  source?: string;
}

/**
 * 語義搜尋（mode:'semantic'）。模組未開 → 回 null（caller 降級 keyword + 告知缺能力）。
 * owner_id / source / entry_type 過濾走 Vectorize metadata filter（entry_type 已 index，見上 upsert metadata）。
 * entry_type 是 base 通用 filter（caller 傳任意 type，base 不寫死語意）。
 */
export async function semanticSearch(
  env: Bindings,
  q: string,
  opts: { owner_id?: string; source?: string; entry_type?: string; topK?: number } = {},
): Promise<SemanticHit[] | null> {
  if (!embedEnabled(env)) return null;
  const vec = await embedText(env, q);
  if (!vec) return [];
  const filter: Record<string, string> = {};
  if (opts.owner_id) filter.owner_id = opts.owner_id;
  if (opts.source) filter.source = opts.source;
  if (opts.entry_type) filter.entry_type = opts.entry_type;
  const res = await env.VECTORIZE!.query(vec, {
    topK: Math.min(opts.topK ?? 20, 100),
    returnMetadata: 'indexed',
    ...(Object.keys(filter).length ? { filter } : {}),
  });
  return (res.matches ?? []).map((m) => ({
    id: m.id,
    score: m.score,
    owner_id: m.metadata?.owner_id as string | undefined,
    entry_type: m.metadata?.entry_type as string | undefined,
    source: m.metadata?.source as string | undefined,
  }));
}