Neon-Desk/lib/server/ai.test.ts

import { beforeEach, describe, expect, it, mock } from 'bun:test';
import {
  __resetAiWarningsForTests,
  getAiConfig,
  getExtractionAiConfig,
  runAiAnalysis
} from './ai';

const CODING_API_BASE_URL = 'https://api.z.ai/api/coding/paas/v4';

describe('ai config and runtime', () => {
  beforeEach(() => {
    __resetAiWarningsForTests();
  });

  it('uses coding endpoint defaults when optional env values are missing', () => {
    const config = getAiConfig({
      env: {
        ZHIPU_API_KEY: 'key'
      },
      warn: () => {}
    });

    expect(config.provider).toBe('zhipu');
    expect(config.apiKey).toBe('key');
    expect(config.baseUrl).toBe(CODING_API_BASE_URL);
    expect(config.model).toBe('glm-5');
    expect(config.temperature).toBe(0.2);
  });

  it('ignores ZHIPU_BASE_URL and keeps the hardcoded coding endpoint', () => {
    const config = getAiConfig({
      env: {
        ZHIPU_API_KEY: 'key',
        ZHIPU_BASE_URL: 'https://api.z.ai/api/paas/v4'
      },
      warn: () => {}
    });

    expect(config.baseUrl).toBe(CODING_API_BASE_URL);
  });

  it('clamps report temperature into [0, 2]', () => {
    const negative = getAiConfig({
      env: {
        ZHIPU_API_KEY: 'key',
        AI_TEMPERATURE: '-2'
      },
      warn: () => {}
    });
    expect(negative.temperature).toBe(0);

    const high = getAiConfig({
      env: {
        ZHIPU_API_KEY: 'key',
        AI_TEMPERATURE: '9'
      },
      warn: () => {}
    });
    expect(high.temperature).toBe(2);

    const invalid = getAiConfig({
      env: {
        ZHIPU_API_KEY: 'key',
        AI_TEMPERATURE: 'not-a-number'
      },
      warn: () => {}
    });
    expect(invalid.temperature).toBe(0.2);
  });

  it('uses extraction workload with zhipu config and zero temperature', async () => {
    const createModel = mock((config: {
      provider: string;
      apiKey?: string;
      model: string;
      baseUrl: string;
      temperature: number;
    }) => {
      expect(config.provider).toBe('zhipu');
      expect(config.apiKey).toBe('new-key');
      expect(config.baseUrl).toBe(CODING_API_BASE_URL);
      expect(config.model).toBe('glm-5');
      expect(config.temperature).toBe(0);
      return { modelId: config.model };
    });
    const generate = mock(async (input: {
      model: unknown;
      system?: string;
      prompt: string;
      temperature: number;
      maxRetries?: number;
    }) => {
      expect(input.system).toBe('Return strict JSON only.');
      expect(input.prompt).toBe('Extract this filing');
      expect(input.temperature).toBe(0);
      expect(input.maxRetries).toBe(0);
      return { text: '{"summary":"ok"}' };
    });

    const result = await runAiAnalysis('Extract this filing', 'Return strict JSON only.', {
      env: {
        ZHIPU_API_KEY: 'new-key'
      },
      warn: () => {},
      workload: 'extraction',
      createModel,
      generate
    });

    expect(result.provider).toBe('zhipu');
    expect(result.model).toBe('glm-5');
    expect(result.text).toBe('{"summary":"ok"}');
    expect(createModel).toHaveBeenCalledTimes(1);
    expect(generate).toHaveBeenCalledTimes(1);
  });

  it('warns once when ZHIPU_BASE_URL is set because coding endpoint is hardcoded', () => {
    const warn = mock((_message: string) => {});

    const env = {
      ZHIPU_API_KEY: 'new-key',
      ZHIPU_BASE_URL: 'https://api.z.ai/api/paas/v4'
    };

    getAiConfig({ env, warn });
    getAiConfig({ env, warn });

    expect(warn).toHaveBeenCalledTimes(1);
  });

  it('uses configured ZHIPU values and injected generator when API key exists', async () => {
    const createModel = mock((config: {
      provider: string;
      apiKey?: string;
      model: string;
      baseUrl: string;
      temperature: number;
    }) => {
      expect(config.provider).toBe('zhipu');
      expect(config.apiKey).toBe('new-key');
      expect(config.baseUrl).toBe(CODING_API_BASE_URL);
      expect(config.model).toBe('glm-4-plus');
      expect(config.temperature).toBe(0.4);
      return { modelId: config.model };
    });
    const generate = mock(async (input: {
      model: unknown;
      system?: string;
      prompt: string;
      temperature: number;
      maxRetries?: number;
    }) => {
      expect(input.system).toBe('Use concise style');
      expect(input.prompt).toBe('Analyze this filing');
      expect(input.temperature).toBe(0.4);
      expect(input.maxRetries).toBe(0);
      return { text: '  Generated insight  ' };
    });

    const result = await runAiAnalysis('Analyze this filing', 'Use concise style', {
      env: {
        ZHIPU_API_KEY: 'new-key',
        ZHIPU_MODEL: 'glm-4-plus',
        ZHIPU_BASE_URL: 'https://api.z.ai/api/paas/v4',
        AI_TEMPERATURE: '0.4'
      },
      warn: () => {},
      createModel,
      generate
    });

    expect(createModel).toHaveBeenCalledTimes(1);
    expect(generate).toHaveBeenCalledTimes(1);
    expect(result.provider).toBe('zhipu');
    expect(result.model).toBe('glm-4-plus');
    expect(result.text).toBe('Generated insight');
  });

  it('throws when report workload runs without ZHIPU_API_KEY', async () => {
    await expect(
      runAiAnalysis('Analyze this filing', undefined, {
        env: {},
        warn: () => {},
        createModel: () => ({}),
        generate: async () => ({ text: 'should-not-be-used' })
      })
    ).rejects.toThrow('ZHIPU_API_KEY is required for AI workloads');
  });

  it('throws when extraction workload runs without ZHIPU_API_KEY', async () => {
    await expect(
      runAiAnalysis('Extract this filing', 'Return strict JSON only.', {
        env: {},
        warn: () => {},
        workload: 'extraction',
        createModel: () => ({}),
        generate: async () => ({ text: 'should-not-be-used' })
      })
    ).rejects.toThrow('ZHIPU_API_KEY is required for AI workloads');
  });

  it('throws when AI generation returns an empty response', async () => {
    await expect(
      runAiAnalysis('Analyze this filing', undefined, {
        env: { ZHIPU_API_KEY: 'new-key' },
        warn: () => {},
        createModel: () => ({}),
        generate: async () => ({ text: '   ' })
      })
    ).rejects.toThrow('AI SDK returned an empty response');
  });

  it('keeps throwing unknown provider errors', async () => {
    await expect(
      runAiAnalysis('Analyze this filing', 'Use concise style', {
        env: {
          ZHIPU_API_KEY: 'new-key'
        },
        warn: () => {},
        createModel: () => ({}),
        generate: async () => {
          throw new Error('unexpected schema mismatch');
        }
      })
    ).rejects.toThrow('unexpected schema mismatch');
  });

  it('returns extraction config with same zhipu model and zero temperature', () => {
    const config = getExtractionAiConfig({
      env: {
        ZHIPU_API_KEY: 'new-key',
        ZHIPU_MODEL: 'glm-4-plus',
        AI_TEMPERATURE: '0.9'
      },
      warn: () => {}
    });

    expect(config.provider).toBe('zhipu');
    expect(config.apiKey).toBe('new-key');
    expect(config.baseUrl).toBe(CODING_API_BASE_URL);
    expect(config.model).toBe('glm-4-plus');
    expect(config.temperature).toBe(0);
  });
});