test: 暂存代码

2026-08-17 19:47:08 +08:00 · 2026-08-06 20:29:25 +08:00
parent 307d247660
commit ad4b3a8074
43 changed files with 9978 additions and 512 deletions
@@ -0,0 +1,209 @@
+import { mkdtempSync } from 'fs'
+import { rm } from 'fs/promises'
+import { tmpdir } from 'os'
+import { join } from 'path'
+import { performance } from 'perf_hooks'
+import { afterAll, describe, expect, it } from 'vitest'
+import {
+  DEFAULT_KNOWLEDGE_CHUNKER,
+  type KnowledgeEvidence,
+  type KnowledgeFtsConfig,
+  type KnowledgeSourceMessage
+} from '../../src/shared/knowledge'
+import { KnowledgeStore } from '../../src/main/knowledge/knowledge-store'
+import {
+  createRealisticKnowledgeFixture,
+  type RealisticBenchmarkCase,
+  type RealisticBenchmarkCategory
+} from '../fixtures/knowledge-realistic'
+
+const root = mkdtempSync(join(tmpdir(), 'wxe-realistic-fts-'))
+
+afterAll(async () => {
+  await rm(root, { recursive: true, force: true })
+})
+
+type Retrieval = { messageIds: string[] }
+type Metrics = {
+  recallAt5: number
+  recallAt10: number
+  mrr: number
+  evidenceAccuracy: number
+  findSuccessAt10: number
+  p50LatencyMs: number
+  p95LatencyMs: number
+}
+
+function percentile(values: number[], ratio: number): number {
+  if (!values.length) return 0
+  const sorted = values.slice().sort((left, right) => left - right)
+  return sorted[Math.min(sorted.length - 1, Math.ceil(sorted.length * ratio) - 1)]
+}
+
+function oldSearch(
+  messages: KnowledgeSourceMessage[],
+  testCase: RealisticBenchmarkCase
+): Retrieval[] {
+  return messages
+    .map((message) => {
+      const text =
+        `${message.text || ''}\n${message.voiceTranscript || ''}\n${message.attachment?.name || ''}`.toLowerCase()
+      const score = testCase.searchTerms.reduce(
+        (total, term) => total + (text.includes(term.toLowerCase()) ? 1 : 0),
+        0
+      )
+      return { messageIds: [message.messageId], score }
+    })
+    .filter((item) => item.score > 0)
+    .sort((left, right) => right.score - left.score)
+}
+
+function score(
+  cases: RealisticBenchmarkCase[],
+  search: (testCase: RealisticBenchmarkCase) => Retrieval[]
+): Metrics {
+  const latency: number[] = []
+  let recallAt5 = 0
+  let recallAt10 = 0
+  let reciprocalRank = 0
+  let evidenceAccuracy = 0
+  let findSuccessAt10 = 0
+  for (const testCase of cases) {
+    const started = performance.now()
+    const results = search(testCase)
+    latency.push(performance.now() - started)
+    const expected = new Set(testCase.expectedMessageIds)
+    const firstTen = results.slice(0, 10)
+    const hitIndex = firstTen.findIndex((item) => item.messageIds.some((id) => expected.has(id)))
+    if (results.slice(0, 5).some((item) => item.messageIds.some((id) => expected.has(id)))) {
+      recallAt5 += 1
+    }
+    if (hitIndex >= 0) {
+      recallAt10 += 1
+      findSuccessAt10 += 1
+      reciprocalRank += 1 / (hitIndex + 1)
+    }
+    const firstFive = results.slice(0, 5)
+    if (firstFive.length) {
+      evidenceAccuracy +=
+        firstFive.filter((item) => item.messageIds.some((id) => expected.has(id))).length /
+        firstFive.length
+    }
+  }
+  return {
+    recallAt5: recallAt5 / cases.length,
+    recallAt10: recallAt10 / cases.length,
+    mrr: reciprocalRank / cases.length,
+    evidenceAccuracy: evidenceAccuracy / cases.length,
+    findSuccessAt10: findSuccessAt10 / cases.length,
+    p50LatencyMs: percentile(latency, 0.5),
+    p95LatencyMs: percentile(latency, 0.95)
+  }
+}
+
+function groupedByCategory<T>(
+  cases: RealisticBenchmarkCase[],
+  evaluate: (items: RealisticBenchmarkCase[]) => T
+): Record<RealisticBenchmarkCategory, T> {
+  const groups = new Map<RealisticBenchmarkCategory, RealisticBenchmarkCase[]>()
+  for (const item of cases) groups.set(item.category, [...(groups.get(item.category) || []), item])
+  return Object.fromEntries(
+    Array.from(groups.entries()).map(([category, items]) => [category, evaluate(items)])
+  ) as Record<RealisticBenchmarkCategory, T>
+}
+
+const profiles: KnowledgeFtsConfig[] = [
+  {
+    profileId: 'unicode61-external-full-columnsize',
+    tokenizer: 'unicode61',
+    contentMode: 'external',
+    detail: 'full',
+    columnsize: 1
+  },
+  {
+    profileId: 'trigram-external-full-columnsize',
+    tokenizer: 'trigram',
+    contentMode: 'external',
+    detail: 'full',
+    columnsize: 1
+  },
+  {
+    profileId: 'trigram-external-column-no-columnsize',
+    tokenizer: 'trigram',
+    contentMode: 'external',
+    detail: 'column',
+    columnsize: 0
+  }
+]
+
+describe('realistic desensitized WeChat FTS5 benchmark', () => {
+  it('compares unicode61 and trigram by recall quality before Task 3 chooses a runtime profile', async () => {
+    const fixture = createRealisticKnowledgeFixture()
+    const messages = fixture.conversations.flatMap((conversation) => conversation.messages)
+    const oldMetrics = score(fixture.cases, (testCase) => oldSearch(messages, testCase))
+    const comparisons: Array<{
+      profile: KnowledgeFtsConfig
+      metrics: Metrics
+      categoryMetrics: Record<RealisticBenchmarkCategory, Metrics>
+      databaseBytes: number
+    }> = []
+    for (const profile of profiles) {
+      const store = new KnowledgeStore(
+        join(root, profile.profileId),
+        messages[0].accountId,
+        profile
+      )
+      await store.index({
+        conversations: fixture.conversations,
+        chunker: DEFAULT_KNOWLEDGE_CHUNKER
+      })
+      const search = (testCase: RealisticBenchmarkCase): Retrieval[] =>
+        store
+          .search({
+            accountId: messages[0].accountId,
+            text: testCase.question,
+            terms: testCase.searchTerms,
+            limit: 10
+          })
+          .map((item: KnowledgeEvidence) => ({ messageIds: item.messageIds }))
+      const metrics = score(fixture.cases, search)
+      store.checkpoint()
+      comparisons.push({
+        profile,
+        metrics,
+        categoryMetrics: groupedByCategory(fixture.cases, (items) => score(items, search)),
+        databaseBytes: store.getStorageStats().databaseBytes
+      })
+      store.close()
+    }
+    console.log(
+      `KNOWLEDGE_REALISTIC_FTS_REPORT=${JSON.stringify(
+        {
+          fixture: 'realistic-desensitized-wechat-v1',
+          questionCount: fixture.cases.length,
+          categories: Array.from(new Set(fixture.cases.map((item) => item.category))),
+          oldSearch: oldMetrics,
+          fts5Comparisons: comparisons
+        },
+        null,
+        2
+      )}`
+    )
+    expect(fixture.cases).toHaveLength(14)
+    expect(new Set(fixture.cases.map((item) => item.category))).toEqual(
+      new Set([
+        'chinese-continuous',
+        'chinese-short',
+        'person-name',
+        'mixed-language',
+        'url',
+        'file-name',
+        'technical-term',
+        'number-email-path',
+        'short-message',
+        'long-voice'
+      ])
+    )
+    expect(comparisons).toHaveLength(profiles.length)
+  })
+})