bilin/tests/main/ai-context-builder.test.ts

import { describe, it, expect, beforeEach, afterEach } from 'vitest'
import { DatabaseSync } from 'node:sqlite'
import { mkdtempSync, rmSync } from 'fs'
import { join } from 'path'
import { tmpdir } from 'os'
import { migrate } from '../../src/main/db/migrate'
import { ChapterRepository } from '../../src/main/db/repositories/chapter.repo'
import { VolumeRepository } from '../../src/main/db/repositories/volume.repo'
import { SettingRepository } from '../../src/main/db/repositories/setting.repo'
import { KnowledgeRepository } from '../../src/main/db/repositories/knowledge.repo'
import { BookRegistryService } from '../../src/main/services/book-registry'
import { aiContextBuilder } from '../../src/main/services/ai-context-builder.service'
import type { SqliteDb } from '../../src/main/db/types'

describe('AiContextBuilderService', () => {
  let userDir: string
  let registry: BookRegistryService
  let bookId: string
  let volId: string
  let db: SqliteDb

  beforeEach(() => {
    userDir = mkdtempSync(join(tmpdir(), 'bilin-ctx-'))
    registry = new BookRegistryService(userDir)
    const meta = registry.create({ name: '测试书', category: '玄幻' })
    bookId = meta.id
    db = registry.getDb(bookId)
    volId = new VolumeRepository(db).create('卷一', 0).id
  })

  afterEach(() => {
    db?.close()
    rmSync(userDir, { recursive: true, force: true })
  })

  it('truncates long chapter content to 2000 chars per block', () => {
    const chapters = new ChapterRepository(db)
    const longHtml = `<p>${'章'.repeat(3000)}</p>`
    const chapter = chapters.create(volId, '长章节', 0, longHtml)

    const result = aiContextBuilder.build(
      { chapterIds: [chapter.id], outlineIds: [], settingIds: [], inspirationIds: [], knowledgeIds: [] },
      registry,
      bookId,
      '作者'
    )

    expect(result.preview).toHaveLength(1)
    expect(result.preview[0].excerpt.length).toBeLessThanOrEqual(2001)
    expect(result.systemPrompt.length).toBeLessThanOrEqual(16 * 1024 + 500)
  })

  it('keeps total context within 16KB', () => {
    const chapters = new ChapterRepository(db)
    const ids: string[] = []
    for (let i = 0; i < 12; i++) {
      const ch = chapters.create(volId, `章节${i}`, i, `<p>${'内容'.repeat(900)}</p>`)
      ids.push(ch.id)
    }

    const result = aiContextBuilder.build(
      { chapterIds: ids, outlineIds: [], settingIds: [], inspirationIds: [], knowledgeIds: [] },
      registry,
      bookId,
      '作者'
    )

    const blocksText = result.preview.map((p) => p.excerpt).join('')
    expect(blocksText.length).toBeLessThanOrEqual(16 * 1024)
    expect(result.preview.length).toBeLessThan(ids.length)
  })

  it('includes setting in preview', () => {
    const settings = new SettingRepository(db)
    const setting = settings.create('character', '主角林远', '少年修士')

    const result = aiContextBuilder.build(
      { chapterIds: [], outlineIds: [], settingIds: [setting.id], inspirationIds: [], knowledgeIds: [] },
      registry,
      bookId,
      '作者'
    )

    expect(result.preview[0].title).toBe('主角林远')
    expect(result.systemPrompt).toContain('主角林远')
  })

  it('UT-CTX-02: build includes knowledge preview items', () => {
    const know = new KnowledgeRepository(db)
    const entry = know.create({
      type: 'foreshadow',
      title: '测灵石',
      content: '石头发热',
      status: 'approved'
    })
    const result = aiContextBuilder.build(
      {
        chapterIds: [],
        outlineIds: [],
        settingIds: [],
        inspirationIds: [],
        knowledgeIds: [entry.id]
      },
      registry,
      bookId,
      '作者'
    )
    expect(result.preview.some((p) => p.kind === 'knowledge')).toBe(true)
    expect(result.systemPrompt).toContain('【知识·伏笔】')
    expect(result.systemPrompt).toContain('测灵石')
  })

  it('UT-CTX-03: caps knowledge section at 4000 chars', () => {
    const know = new KnowledgeRepository(db)
    const ids: string[] = []
    for (let i = 0; i < 20; i++) {
      const e = know.create({
        type: 'fact',
        title: `条目${i}`,
        content: '长'.repeat(400),
        importance: 5,
        status: 'approved'
      })
      ids.push(e.id)
    }
    const result = aiContextBuilder.build(
      {
        chapterIds: [],
        outlineIds: [],
        settingIds: [],
        inspirationIds: [],
        knowledgeIds: ids
      },
      registry,
      bookId,
      '作者'
    )
    const knowledgeChars = result.preview
      .filter((p) => p.kind === 'knowledge')
      .reduce((s, p) => s + p.charCount, 0)
    expect(knowledgeChars).toBeLessThanOrEqual(4000)
  })
})