sugarr/test/core/langfuse/LangfuseTracer.test.ts

import { describe, expect, it, vi } from 'vitest';

const generationEnd = vi.fn();
const generation = vi.fn().mockReturnValue({ end: generationEnd });
const trace = vi.fn().mockReturnValue({ id: 'lf-trace-1', generation });
const flushAsync = vi.fn().mockResolvedValue(undefined);

vi.mock('langfuse', () => ({
  Langfuse: vi.fn().mockImplementation(() => ({ trace, flushAsync })),
}));

import { LangfuseTracer } from '../../../src/core/langfuse/LangfuseTracer';

describe('LangfuseTracer', () => {
  it('starts a Langfuse trace and exposes its id', () => {
    const tracer = new LangfuseTracer({ publicKey: 'pk', secretKey: 'sk', baseUrl: 'http://lf' });
    const handle = tracer.startTrace('my-pipeline', { foo: 'bar' });

    expect(trace).toHaveBeenCalledWith({ name: 'my-pipeline', metadata: { foo: 'bar' } });
    expect(handle.id).toBe('lf-trace-1');
  });

  it('startGeneration creates a Langfuse generation and end() forwards output/usage', () => {
    const tracer = new LangfuseTracer({ publicKey: 'pk', secretKey: 'sk', baseUrl: 'http://lf' });
    const handle = tracer.startTrace('my-pipeline');
    const gen = handle.startGeneration({ name: 'step-1', model: 'llama3', input: ['hi'] });
    gen.end({ output: 'done', usage: { totalTokens: 5 } });

    expect(generation).toHaveBeenCalledWith(
      expect.objectContaining({ name: 'step-1', model: 'llama3', input: ['hi'] })
    );
    expect(generationEnd).toHaveBeenCalledWith(
      expect.objectContaining({ output: 'done', usage: { totalTokens: 5 } })
    );
  });

  it('flush() calls flushAsync on the underlying client', async () => {
    const tracer = new LangfuseTracer({ publicKey: 'pk', secretKey: 'sk', baseUrl: 'http://lf' });
    await tracer.flush();
    expect(flushAsync).toHaveBeenCalledTimes(1);
  });
});