sugarr/test/core/llm/LlmTool.test.ts

import { describe, expect, it, vi } from 'vitest';
import { LlmTool } from '../../../src/core/llm/LlmTool';
import type { RunContext } from '../../../src/core/llm/RunContext';
import type { PromptService } from '../../../src/core/langfuse/PromptService';
import type { LlmClient } from '../../../src/core/llm/LlmClient';

interface TestInput {
  who: string;
}

class GreetTool extends LlmTool<TestInput, string> {
  protected readonly promptName = 'greet';

  protected buildVariables(input: TestInput): Record<string, unknown> {
    return { who: input.who };
  }

  protected parseOutput(raw: string): string {
    return raw.toUpperCase();
  }
}

function makeCtx() {
  const generationEnd = vi.fn();
  const startGeneration = vi.fn().mockReturnValue({ end: generationEnd });
  const trace = { id: 'trace-1', startGeneration };
  const logger = { info: vi.fn(), error: vi.fn() } as unknown as RunContext['logger'];
  return { ctx: { trace, logger } as RunContext, startGeneration, generationEnd };
}

describe('LlmTool.runAsync', () => {
  it('fetches the compiled prompt, calls the LLM, and parses the output', async () => {
    const getCompiled = vi.fn().mockResolvedValue({
      name: 'greet',
      version: 1,
      config: { model: 'llama3', temperature: 0.5 },
      messages: [{ role: 'user', content: 'hi alice' }],
    });
    const promptService = { getCompiled } as unknown as PromptService;

    const chat = vi.fn().mockResolvedValue({ content: 'hello alice', usage: { totalTokens: 3 } });
    const llmClient = { chat } as unknown as LlmClient;

    const { ctx, startGeneration, generationEnd } = makeCtx();
    const tool = new GreetTool(promptService, llmClient);

    const result = await tool.runAsync({ who: 'alice' }, ctx);

    expect(result).toBe('HELLO ALICE');
    expect(getCompiled).toHaveBeenCalledWith('greet', { who: 'alice' });
    expect(chat).toHaveBeenCalledWith(
      [{ role: 'user', content: 'hi alice' }],
      { model: 'llama3', temperature: 0.5, maxTokens: undefined }
    );
    expect(startGeneration).toHaveBeenCalledWith(
      expect.objectContaining({ name: 'greet', model: 'llama3' })
    );
    expect(generationEnd).toHaveBeenCalledWith(
      expect.objectContaining({ output: 'hello alice' })
    );
  });

  it('falls back to the constructor-provided defaultModel when the prompt config omits model', async () => {
    const getCompiled = vi
      .fn()
      .mockResolvedValue({ name: 'greet', version: 1, config: {}, messages: [{ role: 'user', content: 'hi' }] });
    const promptService = { getCompiled } as unknown as PromptService;
    const chat = vi.fn().mockResolvedValue({ content: 'hi' });
    const llmClient = { chat } as unknown as LlmClient;

    const { ctx } = makeCtx();
    const tool = new GreetTool(promptService, llmClient, 'llama3');

    await tool.runAsync({ who: 'alice' }, ctx);

    expect(chat).toHaveBeenCalledWith(
      [{ role: 'user', content: 'hi' }],
      { model: 'llama3', temperature: undefined, maxTokens: undefined }
    );
  });

  it('falls back to gpt-4o-mini when no defaultModel is given and the prompt config omits model', async () => {
    const getCompiled = vi
      .fn()
      .mockResolvedValue({ name: 'greet', version: 1, config: {}, messages: [{ role: 'user', content: 'hi' }] });
    const promptService = { getCompiled } as unknown as PromptService;
    const chat = vi.fn().mockResolvedValue({ content: 'hi' });
    const llmClient = { chat } as unknown as LlmClient;

    const { ctx } = makeCtx();
    const tool = new GreetTool(promptService, llmClient);

    await tool.runAsync({ who: 'alice' }, ctx);

    expect(chat).toHaveBeenCalledWith(
      [{ role: 'user', content: 'hi' }],
      { model: 'gpt-4o-mini', temperature: undefined, maxTokens: undefined }
    );
  });

  it('ends the generation with error metadata and rethrows when the LLM call fails', async () => {
    const promptService = {
      getCompiled: vi.fn().mockResolvedValue({ name: 'greet', version: 1, config: {}, messages: [] }),
    } as unknown as PromptService;
    const llmClient = { chat: vi.fn().mockRejectedValue(new Error('boom')) } as unknown as LlmClient;

    const { ctx, generationEnd } = makeCtx();
    const tool = new GreetTool(promptService, llmClient);

    await expect(tool.runAsync({ who: 'alice' }, ctx)).rejects.toThrow('boom');
    expect(generationEnd).toHaveBeenCalledWith(
      expect.objectContaining({ output: null })
    );
  });
});