Antigravity × Cloudflare Vectorize：エッジで動く本番 RAG パイプラインを構築する

// src/lib/chunker.ts
type Section = { heading: string; level: number; body: string };
 
export function parseMarkdown(md: string): Section[] {
  const lines = md.split("\n");
  const sections: Section[] = [];
  let current: Section = { heading: "", level: 1, body: "" };
 
  for (const line of lines) {
    const headingMatch = line.match(/^(#{1,6})\s+(.+)$/);
    if (headingMatch) {
      if (current.body.trim() || current.heading) sections.push(current);
      current = {
        heading: headingMatch[2].trim(),
        level: headingMatch[1].length,
        body: "",
      };
    } else {
      current.body += line + "\n";
    }
  }
  if (current.body.trim() || current.heading) sections.push(current);
  return sections;
}
 
export function chunkSections(
  sections: Section[],
  target = 800,
  overlap = 100
): { text: string; heading: string }[] {
  const out: { text: string; heading: string }[] = [];
 
  for (const section of sections) {
    const fullText = section.heading
      ? `## ${section.heading}\n${section.body.trim()}`
      : section.body.trim();
 
    if (fullText.length <= target) {
      out.push({ text: fullText, heading: section.heading });
      continue;
    }
 
    // 段落単位で再分割
    const paragraphs = fullText.split(/\n{2,}/);
    let buf = "";
    for (const p of paragraphs) {
      if ((buf + "\n\n" + p).length > target && buf) {
        out.push({ text: buf.trim(), heading: section.heading });
        const tail = buf.slice(Math.max(0, buf.length - overlap));
        buf = tail + "\n\n" + p;
      } else {
        buf = buf ? `${buf}\n\n${p}` : p;
      }
    }
    if (buf.trim()) out.push({ text: buf.trim(), heading: section.heading });
  }
  return out;
}

// src/lib/hybrid-search.ts
type HybridResult = {
  id: string;
  score: number;
  source: "vector" | "fts" | "both";
  text: string;
  title: string;
};
 
export async function hybridSearch(
  env: { VECTORIZE_INDEX: VectorizeIndex; DB: D1Database; AI: Ai },
  query: string,
  topK = 10
): Promise<HybridResult[]> {
  // 並列実行で待ち時間を最小化
  const [vectorRes, ftsRes] = await Promise.all([
    (async () => {
      const { data } = await env.AI.run("@cf/baai/bge-m3", { text: [query] });
      return env.VECTORIZE_INDEX.query(data[0], {
        topK,
        returnMetadata: "all",
      });
    })(),
    env.DB.prepare(
      `SELECT id, title, text, rank
       FROM chunks_fts
       WHERE chunks_fts MATCH ?1
       ORDER BY rank LIMIT ?2`
    )
      .bind(query, topK)
      .all<{ id: string; title: string; text: string; rank: number }>(),
  ]);
 
  // RRF（Reciprocal Rank Fusion）でスコアを統合
  const map = new Map<string, HybridResult>();
  vectorRes.matches.forEach((m, i) => {
    map.set(m.id, {
      id: m.id,
      score: 1 / (60 + i),
      source: "vector",
      text: String(m.metadata?.text ?? ""),
      title: String(m.metadata?.title ?? ""),
    });
  });
  (ftsRes.results ?? []).forEach((r, i) => {
    const existing = map.get(r.id);
    const ftsScore = 1 / (60 + i);
    if (existing) {
      existing.score += ftsScore;
      existing.source = "both";
    } else {
      map.set(r.id, {
        id: r.id,
        score: ftsScore,
        source: "fts",
        text: r.text,
        title: r.title,
      });
    }
  });
 
  return Array.from(map.values()).sort((a, b) => b.score - a.score);
}

-- migrations/0001_chunks_fts.sql
CREATE VIRTUAL TABLE chunks_fts USING fts5(
  id UNINDEXED,
  title,
  text,
  tokenize='unicode61 remove_diacritics 2'
);

// src/lib/rerank.ts
export async function rerank(
  env: { AI: Ai },
  query: string,
  candidates: { id: string; text: string }[]
): Promise<{ id: string; score: number }[]> {
  const result = await env.AI.run("@cf/baai/bge-reranker-base", {
    query,
    contexts: candidates.map((c) => ({ text: c.text })),
  });
 
  return result.response
    .map((r: { index: number; score: number }) => ({
      id: candidates[r.index].id,
      score: r.score,
    }))
    .sort((a, b) => b.score - a.score);
}

// src/lib/cost-logger.ts
type UsageEvent = {
  type: "embedding" | "llm" | "rerank" | "vectorize_query";
  inputTokens?: number;
  outputTokens?: number;
  vectorCount?: number;
  durationMs: number;
  ts: number;
};
 
export class UsageLogger {
  private events: UsageEvent[] = [];
  constructor(private analytics: AnalyticsEngineDataset) {}
 
  log(event: UsageEvent) {
    this.events.push(event);
    this.analytics.writeDataPoint({
      indexes: [event.type],
      blobs: [JSON.stringify(event)],
      doubles: [event.durationMs, event.inputTokens ?? 0],
    });
  }
 
  summary(): Record<string, { calls: number; tokens: number; ms: number }> {
    return this.events.reduce((acc, e) => {
      const k = e.type;
      acc[k] ??= { calls: 0, tokens: 0, ms: 0 };
      acc[k].calls += 1;
      acc[k].tokens += (e.inputTokens ?? 0) + (e.outputTokens ?? 0);
      acc[k].ms += e.durationMs;
      return acc;
    }, {} as Record<string, { calls: number; tokens: number; ms: number }>);
  }
}

// src/lib/query-cache.ts
import { createHash } from "node:crypto";
 
export async function getCachedAnswer(
  kv: KVNamespace,
  query: string
): Promise<string | null> {
  const key = `q:${createHash("sha256").update(query).digest("hex").slice(0, 32)}`;
  return await kv.get(key);
}
 
export async function setCachedAnswer(
  kv: KVNamespace,
  query: string,
  answer: string,
  ttlSec = 86400
) {
  const key = `q:${createHash("sha256").update(query).digest("hex").slice(0, 32)}`;
  await kv.put(key, answer, { expirationTtl: ttlSec });
}

// src/lib/health-check.ts
export async function searchHealthCheck(env: Bindings) {
  const goldenQueries = [
    { q: "Antigravity の使い方", expectedDocId: "antigravity-getting-started" },
    { q: "Cloudflare Vectorize の費用", expectedDocId: "antigravity-cloudflare-vectorize-rag-production-guide" },
  ];
 
  const failures: { q: string; got: string[] }[] = [];
  for (const { q, expectedDocId } of goldenQueries) {
    const results = await hybridSearch(env, q, 5);
    const ids = results.map((r) => r.id.split("#")[0]);
    if (!ids.includes(expectedDocId)) failures.push({ q, got: ids });
  }
  return { passed: failures.length === 0, failures };
}

# こうしないと filter: { docId: "xxx" } が効かない
wrangler vectorize create-metadata-index blog-knowledge \
  --property-name=docId --type=string

async function reindexDoc(env: Bindings, docId: string, body: string) {
  // 既存チャンク数を取得
  const oldCount = parseInt((await env.KV.get(`doc:${docId}:chunkCount`)) ?? "0");
  if (oldCount > 0) {
    const ids = Array.from({ length: oldCount }, (_, i) => `${docId}#${i}`);
    await env.VECTORIZE_INDEX.deleteByIds(ids);
  }
  const chunks = chunk(body);
  const vectors = await embedAll(env, chunks, docId);
  await env.VECTORIZE_INDEX.upsert(vectors);
  await env.KV.put(`doc:${docId}:chunkCount`, String(chunks.length));
}

async function ingestIfChanged(env: Bindings, docId: string, body: string) {
  const newHash = await sha256(body);
  const oldHash = await env.KV.get(`doc:${docId}:hash`);
  if (oldHash === newHash) {
    return { changed: false };
  }
  await reindexDoc(env, docId, body);
  await env.KV.put(`doc:${docId}:hash`, newHash);
  return { changed: true };
}

// src/lib/eval.ts
export type EvalCase = { question: string; expectedDocId: string; tags?: string[] };
 
export async function runEval(env: Bindings, cases: EvalCase[]) {
  let recall1 = 0;
  let recall5 = 0;
  const failures: { case: EvalCase; got: string[] }[] = [];
 
  for (const c of cases) {
    const results = await hybridSearch(env, c.question, 10);
    const ids = results.map((r) => r.id.split("#")[0]);
    if (ids[0] === c.expectedDocId) recall1 += 1;
    if (ids.slice(0, 5).includes(c.expectedDocId)) recall5 += 1;
    else failures.push({ case: c, got: ids.slice(0, 5) });
  }
 
  return {
    recall1: recall1 / cases.length,
    recall5: recall5 / cases.length,
    failures,
  };
}