0pt/tests/document-workbench/litellm-gateway-client.test.ts

import { describe, expect, it } from 'vitest';
import { LiteLlmGatewayClient } from '../../src/document-workbench/litellm-gateway-client';
import { Inf0OutcomeUnknownError } from '../../src/document-workbench/inf0-gateway-client';
import { createInf0Gateway } from '../../src/document-workbench/inf0-gateway';

const context = {
  tenantId: 't_0pt',
  projectId: 'prj_alpha',
  actorId: 'user_1',
  capability: '0pt.document.structure.v1',
  correlationId: 'corr-1',
};
const admission = {
  tenantId: 't_0pt',
  projectId: 'prj_alpha',
  jobId: 'job-1',
  modelId: 'inf0-standard',
  dataClass: 'low',
  maxInputTokens: 1000,
  maxOutputTokens: 200,
  maximumCredits: 5,
};
const policy = {
  tenantId: 't_0pt',
  projectId: 'prj_alpha',
  allowedModelIds: ['inf0-standard'],
  allowedDataClasses: ['low'],
  destination: 'external' as const,
  externalTransferAllowed: true,
  maxInputTokens: 1000,
  maxOutputTokens: 200,
};
const messages = [
  { role: 'system' as const, content: 'system' },
  { role: 'user' as const, content: 'data' },
];

function gateway(fetchImpl: typeof fetch, providerKind: 'local' | 'external' = 'external') {
  return new LiteLlmGatewayClient({
    baseUrl: 'https://ai-gateway.siax.io/v1',
    apiKey: 'sk-test',
    providerKind,
    fetchImpl,
  });
}

function ok(payload: unknown): Response {
  return new Response(JSON.stringify(payload), { status: 200 });
}

describe('LiteLLM INF0 adapter', () => {
  it('maps an OpenAI-compatible response into the INF0 receipt', async () => {
    const client = gateway(async () =>
      ok({
        model: 'inf0-standard',
        choices: [{ message: { content: '{"ok":true}' }, finish_reason: 'stop' }],
        usage: { prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 },
      }),
    );
    await expect(
      client.complete({
        context,
        admission,
        policy,
        reservation: {
          tenantId: 't_0pt',
          projectId: 'prj_alpha',
          jobId: 'job-1',
          status: 'reserved',
          credits: 5,
        },
        inputTokens: 10,
        messages,
      }),
    ).resolves.toMatchObject({
      status: 'model_stopped',
      content: '{"ok":true}',
      modelExecuted: 'inf0-standard',
      providerKind: 'external',
      totalTokens: 15,
    });
  });

  it('denies a model outside the allow-list before any request', async () => {
    let called = false;
    const client = gateway(async () => {
      called = true;
      return ok({});
    });
    await expect(
      client.complete({
        context,
        admission: { ...admission, modelId: 'inf0-frontier' },
        policy,
        reservation: undefined,
        inputTokens: 10,
        messages,
      }),
    ).rejects.toThrow(/denied/i);
    expect(called).toBe(false);
  });

  it('refuses an external provider when the policy requires internal', async () => {
    const client = gateway(
      async () =>
        ok({
          model: 'inf0-standard',
          choices: [{ message: { content: 'x' }, finish_reason: 'stop' }],
          usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 },
        }),
      'external',
    );
    await expect(
      client.complete({
        context,
        admission,
        policy: { ...policy, destination: 'internal' },
        reservation: {
          tenantId: 't_0pt',
          projectId: 'prj_alpha',
          jobId: 'job-1',
          status: 'reserved',
          credits: 5,
        },
        inputTokens: 1,
        messages,
      }),
    ).rejects.toBeInstanceOf(Inf0OutcomeUnknownError);
  });

  it('treats malformed or inconsistent usage as an unknown outcome', async () => {
    const client = gateway(async () =>
      ok({
        model: 'inf0-standard',
        choices: [{ message: { content: 'x' }, finish_reason: 'stop' }],
        usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 99 },
      }),
    );
    await expect(
      client.complete({
        context,
        admission,
        policy,
        reservation: {
          tenantId: 't_0pt',
          projectId: 'prj_alpha',
          jobId: 'job-1',
          status: 'reserved',
          credits: 5,
        },
        inputTokens: 1,
        messages,
      }),
    ).rejects.toBeInstanceOf(Inf0OutcomeUnknownError);
  });

  it('prefers the live LiteLLM key and falls back to the native surface', () => {
    expect(createInf0Gateway({} as unknown as NodeJS.ProcessEnv)).toBeUndefined();
    expect(
      createInf0Gateway({
        INF0_GATEWAY_LITELLM_KEY: 'sk-test',
      } as unknown as NodeJS.ProcessEnv),
    ).toBeInstanceOf(LiteLlmGatewayClient);
    expect(
      createInf0Gateway({
        INF0_URL: 'https://api.inf0.siax.io',
        INF0_API_TOKEN: 'token',
      } as unknown as NodeJS.ProcessEnv),
    ).toBeDefined();
  });
});