sugarr/test/core/llm/SchemaLlmTool.test.ts

import { describe, expect, it } from 'vitest';
import { z } from 'zod';
import { extractJson, LlmOutputParseError, SchemaLlmTool } from '../../../src/core/llm/SchemaLlmTool';
import type { RunContext } from '../../../src/core/llm/RunContext';

describe('extractJson', () => {
  it('parses a bare JSON array', () => {
    expect(extractJson('[{"a":1}]')).toEqual([{ a: 1 }]);
  });

  it('strips a ```json fenced block with surrounding prose', () => {
    const content = 'Sure, here you go:\n```json\n[{"a":1}]\n```\nHope that helps!';
    expect(extractJson(content)).toEqual([{ a: 1 }]);
  });

  it('strips a plain ``` fenced block', () => {
    const content = '```\n{"a":1}\n```';
    expect(extractJson(content)).toEqual({ a: 1 });
  });

  it('throws LlmOutputParseError when no JSON is present', () => {
    expect(() => extractJson('no json here')).toThrow(LlmOutputParseError);
  });
});

const outputSchema = z.array(z.object({ title: z.string(), year: z.number() }));

class ListTool extends SchemaLlmTool<Record<string, never>, z.infer<typeof outputSchema>> {
  protected readonly promptName = 'list';
  protected readonly outputSchema = outputSchema;
  protected buildVariables(): Record<string, unknown> {
    return {};
  }
}

function makeTool() {
  return new ListTool(
    {} as unknown as ConstructorParameters<typeof ListTool>[0],
    {} as unknown as ConstructorParameters<typeof ListTool>[1]
  );
}

describe('SchemaLlmTool.parseOutput (via subclass)', () => {
  it('validates parsed JSON against the schema', () => {
    const tool = makeTool() as unknown as { parseOutput(raw: string, ctx: RunContext): unknown };
    const result = tool.parseOutput('[{"title":"Dune","year":2021}]', {} as RunContext);
    expect(result).toEqual([{ title: 'Dune', year: 2021 }]);
  });

  it('throws LlmOutputParseError when schema validation fails', () => {
    const tool = makeTool() as unknown as { parseOutput(raw: string, ctx: RunContext): unknown };
    expect(() => tool.parseOutput('[{"title":"Dune"}]', {} as RunContext)).toThrow(LlmOutputParseError);
  });
});