tests/generator_core.test.mjs

import { describe, it, expect, vi } from 'vitest';
import { runGenerator } from '../src/generator/generator_core.mjs';

describe('generator_core.mjs (thinking generator)', () => {
  it('includes question and context in the prompt and parses JSON output', async () => {
    const fakeContext = [
      { content: 'First context chunk' },
      { content: 'Second context chunk' },
    ];

    const provider = {
      generate: vi.fn(async (prompt) => {
        // Prompt should contain the question
        expect(prompt).toContain('What is love?');
        // And the context chunks (we expect {{CONTEXT}} is wired up)
        expect(prompt).toContain('First context chunk');
        expect(prompt).toContain('Second context chunk');

        // Return JSON output
        return JSON.stringify({
          reasoning: ['step A', 'step B'],
          answer: 'Love is the recognition of shared being.',
          confidence: 0.92,
          evidence: ['quote (para #1)'],
        });
      }),
    };

    const result = await runGenerator('What is love?', fakeContext, provider);

    expect(provider.generate).toHaveBeenCalledOnce();
    expect(result.question).toBe('What is love?');
    expect(result.context).toHaveLength(2);
    expect(result.raw).toContain('step A');
    expect(result.answer).toBe('Love is the recognition of shared being.');
    expect(result.thought).toContain('step A');
    expect(result.confidence).toBeCloseTo(0.92);
  });

  it('extracts thought and answer correctly when <think> block is present', async () => {
    const fakeContext = [{ content: 'ctx' }];

    const provider = {
      generate: vi.fn(async () => {
        return `<think>
Step 1: Read the context carefully.
Step 2: Identify the relevant statements.
Step 3: Synthesize an answer.
</think>
The final answer derived from the context.`;
      }),
    };

    const result = await runGenerator(
      'Test question?',
      fakeContext,
      provider,
    );

    expect(result.raw).toContain('<think>');
    expect(result.thought).toContain('Step 1:');
    expect(result.thought).toContain('Step 3:');
    expect(result.answer).toBe('The final answer derived from the context.');
  });

  it('handles output without <think> block gracefully', async () => {
    const fakeContext = [{ content: 'ctx' }];

    const provider = {
      generate: vi.fn(async () => {
        // No <think> tags at all
        return 'Just a direct answer with no visible reasoning.';
      }),
    };

    const result = await runGenerator(
      'Another question?',
      fakeContext,
      provider,
    );

    expect(result.raw).toBe('Just a direct answer with no visible reasoning.');
    // No JSON or think tags means thought falls back to raw
    expect(result.thought).toBe('Just a direct answer with no visible reasoning.');
    expect(result.answer).toBe('Just a direct answer with no visible reasoning.');
  });

  it('parses Qwen answer block and preserves thinking object', async () => {
    const fakeContext = [{ content: 'ctx' }];
    const provider = {
      generate: vi.fn(async () => ({
        response: `<|thought|>step1<|end_of_thought|>\n<|answer|>\nConfidence: High\nAnswer: Supported answer\nEvidence: ["quote1 (loc1)", "quote2 (loc2)"]\nLimitations: None\n<|end_of_answer|>`,
        thinking: { steps: ['t1', 't2'] },
      })),
    };

    const result = await runGenerator('Test?', fakeContext, provider);

    expect(result.thought).toEqual({ steps: ['t1', 't2'] });
    expect(result.answer).toBe('Supported answer');
    expect(result.confidence).toBe('High');
    expect(result.evidence).toEqual(['quote1 (loc1)', 'quote2 (loc2)']);
    expect(result.limitations).toBe('None');
  });

  it('parses legacy reasoning tags without answer block', async () => {
    const fakeContext = [{ content: 'ctx' }];
    const provider = {
      generate: vi.fn(async () =>
        `<understanding>step A</understanding>\n<reasoning_chain>step B</reasoning_chain>\nConfidence: Medium\nAnswer: Legacy answer\nEvidence: ["ev1 (loc)"]\nLimitations: None`
      ),
    };

    const result = await runGenerator('Test?', fakeContext, provider);

    expect(typeof result.thought).toBe('string');
    expect(result.answer).toBe('Legacy answer');
    expect(result.confidence).toBe('Medium');
    expect(result.evidence).toEqual(['ev1 (loc)']);
    expect(result.limitations).toBe('None');
  });
});