promptfoo--promptfoo/test/assertions/traceSpanCount.test.ts

import { describe, expect, it } from 'vitest';
import { handleTraceSpanCount } from '../../src/assertions/traceSpanCount';
import { createMockProvider, createProviderResponse } from '../factories/provider';

import type { AssertionParams, AtomicTestCase } from '../../src/types/index';
import type { TraceData } from '../../src/types/tracing';

const mockProvider = createMockProvider({
  id: 'mock',
  response: createProviderResponse({ output: 'mock' }),
});

const mockTraceData: TraceData = {
  traceId: 'test-trace-id',
  evaluationId: 'test-evaluation-id',
  testCaseId: 'test-test-case-id',
  metadata: { test: 'value' },
  spans: [
    {
      spanId: 'span-1',
      name: 'llm.completion',
      startTime: 1000,
      endTime: 1500,
    },
    {
      spanId: 'span-2',
      name: 'llm.chat',
      startTime: 1600,
      endTime: 2000,
    },
    {
      spanId: 'span-3',
      name: 'database.query',
      startTime: 2100,
      endTime: 2200,
    },
    {
      spanId: 'span-4',
      name: 'retrieval.search',
      startTime: 2300,
      endTime: 2400,
    },
    {
      spanId: 'span-5',
      name: 'api.external_call',
      startTime: 2500,
      endTime: 2600,
    },
  ],
};

const defaultParams = {
  baseType: 'trace-span-count' as const,
  assertionValueContext: {
    vars: {},
    test: {} as AtomicTestCase,
    prompt: 'test prompt',
    logProbs: undefined,
    provider: mockProvider,
    providerResponse: { output: 'test output' },
  },
  output: 'test output',
  outputString: 'test output',
  providerResponse: { output: 'test output' },
  test: {} as AtomicTestCase,
  inverse: false,
};

describe('handleTraceSpanCount', () => {
  it('should pass when span count is within max limit', () => {
    const params: AssertionParams = {
      ...defaultParams,
      assertion: {
        type: 'trace-span-count',
        value: { pattern: '*llm*', max: 3 },
      },
      renderedValue: { pattern: '*llm*', max: 3 },
      assertionValueContext: {
        ...defaultParams.assertionValueContext,
        trace: {
          ...mockTraceData,
          evaluationId: 'test-evaluation-id',
          testCaseId: 'test-test-case-id',
          metadata: { test: 'value' },
        },
      },
    };

    const result = handleTraceSpanCount(params);
    expect(result).toEqual({
      pass: true,
      score: 1,
      reason:
        'Found 2 spans matching pattern "*llm*" (expected at most 3). Matched spans: llm.completion, llm.chat',
      assertion: params.assertion,
    });
  });

  it('should fail when span count exceeds max limit', () => {
    const params: AssertionParams = {
      ...defaultParams,
      assertion: {
        type: 'trace-span-count',
        value: { pattern: '*llm*', max: 1 },
      },
      renderedValue: { pattern: '*llm*', max: 1 },
      assertionValueContext: {
        ...defaultParams.assertionValueContext,
        trace: {
          ...mockTraceData,
          evaluationId: 'test-evaluation-id',
          testCaseId: 'test-test-case-id',
          metadata: { test: 'value' },
        },
      },
    };

    const result = handleTraceSpanCount(params);
    expect(result).toEqual({
      pass: false,
      score: 0,
      reason:
        'Found 2 spans matching pattern "*llm*", expected at most 1. Matched spans: llm.completion, llm.chat',
      assertion: params.assertion,
    });
  });

  it('should pass when span count meets min requirement', () => {
    const params: AssertionParams = {
      ...defaultParams,
      assertion: {
        type: 'trace-span-count',
        value: { pattern: '*retrieval*', min: 1 },
      },
      renderedValue: { pattern: '*retrieval*', min: 1 },
      assertionValueContext: {
        ...defaultParams.assertionValueContext,
        trace: {
          ...mockTraceData,
          evaluationId: 'test-evaluation-id',
          testCaseId: 'test-test-case-id',
          metadata: { test: 'value' },
        },
      },
    };

    const result = handleTraceSpanCount(params);
    expect(result).toEqual({
      pass: true,
      score: 1,
      reason:
        'Found 1 spans matching pattern "*retrieval*" (expected at least 1). Matched spans: retrieval.search',
      assertion: params.assertion,
    });
  });

  it('should fail when span count is below min requirement', () => {
    const params: AssertionParams = {
      ...defaultParams,
      assertion: {
        type: 'trace-span-count',
        value: { pattern: '*retrieval*', min: 2 },
      },
      renderedValue: { pattern: '*retrieval*', min: 2 },
      assertionValueContext: {
        ...defaultParams.assertionValueContext,
        trace: {
          ...mockTraceData,
          evaluationId: 'test-evaluation-id',
          testCaseId: 'test-test-case-id',
          metadata: { test: 'value' },
        },
      },
    };

    const result = handleTraceSpanCount(params);
    expect(result).toEqual({
      pass: false,
      score: 0,
      reason:
        'Found 1 spans matching pattern "*retrieval*", expected at least 2. Matched spans: retrieval.search',
      assertion: params.assertion,
    });
  });

  it('should pass when span count is within min and max range', () => {
    const params: AssertionParams = {
      ...defaultParams,
      assertion: {
        type: 'trace-span-count',
        value: { pattern: '*', min: 3, max: 10 },
      },
      renderedValue: { pattern: '*', min: 3, max: 10 },
      assertionValueContext: {
        ...defaultParams.assertionValueContext,
        trace: {
          ...mockTraceData,
          evaluationId: 'test-evaluation-id',
          testCaseId: 'test-test-case-id',
          metadata: { test: 'value' },
        },
      },
    };

    const result = handleTraceSpanCount(params);
    expect(result).toEqual({
      pass: true,
      score: 1,
      reason:
        'Found 5 spans matching pattern "*" (expected 3-10). Matched spans: llm.completion, llm.chat, database.query, retrieval.search, api.external_call',
      assertion: params.assertion,
    });
  });

  it('should handle patterns with ? wildcard', () => {
    const params: AssertionParams = {
      ...defaultParams,
      assertion: {
        type: 'trace-span-count',
        value: { pattern: 'llm.c?at', max: 1 },
      },
      renderedValue: { pattern: 'llm.c?at', max: 1 },
      assertionValueContext: {
        ...defaultParams.assertionValueContext,
        trace: mockTraceData,
      },
    };

    const result = handleTraceSpanCount(params);
    expect(result).toEqual({
      pass: true,
      score: 1,
      reason:
        'Found 1 spans matching pattern "llm.c?at" (expected at most 1). Matched spans: llm.chat',
      assertion: params.assertion,
    });
  });

  it('should handle case-insensitive matching', () => {
    const params: AssertionParams = {
      ...defaultParams,
      assertion: {
        type: 'trace-span-count',
        value: { pattern: '*LLM*', max: 5 },
      },
      renderedValue: { pattern: '*LLM*', max: 5 },
      assertionValueContext: {
        ...defaultParams.assertionValueContext,
        trace: {
          ...mockTraceData,
          evaluationId: 'test-evaluation-id',
          testCaseId: 'test-test-case-id',
          metadata: { test: 'value' },
        },
      },
    };

    const result = handleTraceSpanCount(params);
    expect(result).toEqual({
      pass: true,
      score: 1,
      reason:
        'Found 2 spans matching pattern "*LLM*" (expected at most 5). Matched spans: llm.completion, llm.chat',
      assertion: params.assertion,
    });
  });

  it('should fail when no trace data is available', () => {
    const params: AssertionParams = {
      ...defaultParams,
      assertion: {
        type: 'trace-span-count',
        value: { pattern: '*', min: 1 },
      },
      renderedValue: { pattern: '*', min: 1 },
    };

    expect(() => handleTraceSpanCount(params)).toThrow(
      'No trace data available for trace-span-count assertion',
    );
  });

  it('should handle empty trace spans', () => {
    const params: AssertionParams = {
      ...defaultParams,
      assertion: {
        type: 'trace-span-count',
        value: { pattern: '*', max: 0 },
      },
      renderedValue: { pattern: '*', max: 0 },
      assertionValueContext: {
        ...defaultParams.assertionValueContext,
        trace: {
          traceId: 'empty-trace',
          evaluationId: 'test-evaluation-id',
          testCaseId: 'test-test-case-id',
          metadata: { test: 'value' },
          spans: [],
        },
      },
    };

    const result = handleTraceSpanCount(params);
    expect(result).toEqual({
      pass: true,
      score: 1,
      reason: 'Found 0 spans matching pattern "*" (expected at most 0)',
      assertion: params.assertion,
    });
  });

  it('should throw error for invalid value format', () => {
    const params: AssertionParams = {
      ...defaultParams,
      assertion: { type: 'trace-span-count', value: 'invalid' },
      renderedValue: 'invalid',
      assertionValueContext: {
        ...defaultParams.assertionValueContext,
        trace: mockTraceData,
      },
    };

    expect(() => handleTraceSpanCount(params)).toThrow(
      'trace-span-count assertion must have a value object with pattern property',
    );
  });

  it('should throw error for missing pattern', () => {
    const params: AssertionParams = {
      ...defaultParams,
      assertion: { type: 'trace-span-count', value: { max: 5 } },
      renderedValue: { max: 5 },
      assertionValueContext: {
        ...defaultParams.assertionValueContext,
        trace: {
          ...mockTraceData,
          evaluationId: 'test-evaluation-id',
          testCaseId: 'test-test-case-id',
          metadata: { test: 'value' },
        },
      },
    };

    expect(() => handleTraceSpanCount(params)).toThrow(
      'trace-span-count assertion must have a value object with pattern property',
    );
  });
});