connectai/tests/largeInputMapReduce.test.ts

/**
 * Phase 2 — large-input map-reduce core.
 *
 * Pure orchestration with an injected `callLLM`, so no network / SDK is touched.
 */

import {
    runMapReduce,
    shouldMapReduce,
    chunkCharBudget,
    inputBudgetTokens,
    type MapReduceConfig,
    type MapReduceDeps,
} from '../src/agent/handlePrompt/largeInputMapReduce';
import type { ChatMessage } from '../src/agent';

const estimateTokens = (s: string) => Math.ceil((s || '').length / 4);

const cfg: MapReduceConfig = {
    enabled: true,
    triggerRatio: 0.6,
    concurrency: 2,
    maxDepth: 3,
    showProvenance: false,
};

function isExtract(messages: ChatMessage[]): boolean {
    return /추출기/.test(messages[0]?.content ?? '');
}
function chunkLabel(messages: ChatMessage[]): string {
    const m = (messages[1]?.content ?? '').match(/자료 조각 (\d+)\/(\d+)/);
    return m ? m[1] : '?';
}

// ~12 short markdown sections → forces multiple chunks under a small window.
const bigContent = Array.from({ length: 12 }, (_, i) =>
    `## 섹션 ${i + 1}\n안건 ${i + 1}: 결정사항과 수치 ${i * 10}. ` + '내용 '.repeat(40)
).join('\n\n');

describe('shouldMapReduce', () => {
    test('triggers only above window * triggerRatio and when enabled', () => {
        expect(shouldMapReduce(6200, 10000, cfg)).toBe(true);  // > 6000
        expect(shouldMapReduce(5000, 10000, cfg)).toBe(false); // < 6000
        expect(shouldMapReduce(99999, 10000, { ...cfg, enabled: false })).toBe(false);
        expect(shouldMapReduce(100, 0, cfg)).toBe(false);      // unknown window
    });
});

describe('budget helpers', () => {
    test('inputBudgetTokens reserves output + safety', () => {
        // 10000 - sys(500) - max(2048, 1000)=2048 - safety(512) = 6940
        expect(inputBudgetTokens(10000, 500, 512)).toBe(6940);
    });
    test('chunkCharBudget is positive and scales with the window', () => {
        const small = chunkCharBudget(4000, 200, 512);
        const big = chunkCharBudget(16000, 200, 512);
        expect(small).toBeGreaterThan(0);
        expect(big).toBeGreaterThan(small);
    });
});

describe('runMapReduce', () => {
    function deps(callLLM: MapReduceDeps['callLLM']): MapReduceDeps {
        return { callLLM, estimateTokens };
    }
    const params = {
        intent: '회의록을 안건별로 정리해줘',
        largeContent: bigContent,
        windowTokens: 4000,
        systemTokens: 200,
        safetyMargin: 512,
        cfg,
    };

    test('extracts relevant facts per chunk and condenses them', async () => {
        const seen: string[] = [];
        const r = await runMapReduce(
            deps(async (messages) => {
                expect(isExtract(messages)).toBe(true);
                const k = chunkLabel(messages);
                seen.push(k);
                return `추출-${k}`;
            }),
            params,
        );
        expect(r.allIrrelevant).toBe(false);
        expect(r.chunkCount).toBeGreaterThan(1);
        expect(r.relevantCount).toBe(r.chunkCount);
        expect(r.condensedContext).toContain('추출-1');
        // every chunk was visited
        expect(seen.length).toBe(r.chunkCount);
    });

    test('all-irrelevant chunks → allIrrelevant with empty context', async () => {
        const r = await runMapReduce(
            deps(async () => '(관련 없음)'),
            params,
        );
        expect(r.allIrrelevant).toBe(true);
        expect(r.relevantCount).toBe(0);
        expect(r.condensedContext).toBe('');
    });

    test('respects concurrency limit', async () => {
        let active = 0;
        let peak = 0;
        await runMapReduce(
            deps(async (messages) => {
                active++;
                peak = Math.max(peak, active);
                await new Promise((res) => setTimeout(res, 5));
                active--;
                return `x-${chunkLabel(messages)}`;
            }),
            params,
        );
        expect(peak).toBeLessThanOrEqual(cfg.concurrency);
    });

    test('a failing chunk extraction falls back to truncated raw (not a crash)', async () => {
        let call = 0;
        const r = await runMapReduce(
            deps(async (messages) => {
                if (isExtract(messages) && ++call === 1) throw new Error('boom');
                return `ok-${chunkLabel(messages)}`;
            }),
            params,
        );
        expect(r.allIrrelevant).toBe(false);
        // The failed chunk still contributed (raw fallback), so relevantCount === chunkCount.
        expect(r.relevantCount).toBe(r.chunkCount);
    });

    test('tags provenance when showProvenance is on', async () => {
        const r = await runMapReduce(
            deps(async (messages) => `발췌-${chunkLabel(messages)}`),
            { ...params, cfg: { ...cfg, showProvenance: true } },
        );
        expect(r.condensedContext).toMatch(/\[조각 \d+\]/);
    });

    test('hierarchical reduce kicks in when extractions overflow the context ceiling', async () => {
        // Tiny window so even a few extractions exceed the ceiling → reduce rounds run.
        let reduceCalls = 0;
        const r = await runMapReduce(
            deps(async (messages) => {
                if (isExtract(messages)) {
                    return '관련 사실 '.repeat(60); // big extraction per chunk
                }
                reduceCalls++;
                return '통합본'; // reduce collapses to something small
            }),
            { ...params, windowTokens: 2200 },
        );
        expect(reduceCalls).toBeGreaterThan(0);
        expect(r.reduceDepth).toBeGreaterThan(0);
        expect(r.allIrrelevant).toBe(false);
    });
});