diff --git a/src/core/llm/SchemaLlmTool.ts b/src/core/llm/SchemaLlmTool.ts new file mode 100644 index 0000000..4a37a17 --- /dev/null +++ b/src/core/llm/SchemaLlmTool.ts @@ -0,0 +1,60 @@ +import type { ZodSchema } from 'zod'; +import { LlmTool } from './LlmTool'; +import type { RunContext } from './RunContext'; + +export class LlmOutputParseError extends Error { + constructor( + message: string, + readonly raw: string + ) { + super(message); + this.name = 'LlmOutputParseError'; + } +} + +export abstract class SchemaLlmTool extends LlmTool { + protected abstract readonly outputSchema: ZodSchema; + + protected parseOutput(raw: string, _ctx: RunContext): TOutput { + const json = extractJson(raw); + const result = this.outputSchema.safeParse(json); + if (!result.success) { + throw new LlmOutputParseError(`Output failed schema validation: ${result.error.message}`, raw); + } + return result.data; + } +} + +export function extractJson(content: string): unknown { + let text = content.trim(); + + const fenceStart = text.indexOf('```'); + if (fenceStart !== -1) { + text = text.slice(fenceStart); + text = text.replace(/^```json/, '').replace(/^```/, ''); + const fenceEnd = text.indexOf('```'); + if (fenceEnd !== -1) { + text = text.slice(0, fenceEnd); + } + text = text.trim(); + } + + const arrayStart = text.indexOf('['); + const objectStart = text.indexOf('{'); + const starts = [arrayStart, objectStart].filter((i) => i !== -1); + if (starts.length === 0) { + throw new LlmOutputParseError(`No JSON found in LLM response: ${content}`, content); + } + const start = Math.min(...starts); + const isArray = start === arrayStart; + const end = isArray ? text.lastIndexOf(']') : text.lastIndexOf('}'); + if (end === -1 || end <= start) { + throw new LlmOutputParseError(`No JSON found in LLM response: ${content}`, content); + } + + try { + return JSON.parse(text.slice(start, end + 1)); + } catch (err) { + throw new LlmOutputParseError(`Failed to parse JSON from LLM response: ${String(err)}`, content); + } +} diff --git a/src/core/llm/TextLlmTool.ts b/src/core/llm/TextLlmTool.ts new file mode 100644 index 0000000..2a6f479 --- /dev/null +++ b/src/core/llm/TextLlmTool.ts @@ -0,0 +1,8 @@ +import { LlmTool } from './LlmTool'; +import type { RunContext } from './RunContext'; + +export abstract class TextLlmTool extends LlmTool { + protected parseOutput(raw: string, _ctx: RunContext): string { + return raw.trim(); + } +} diff --git a/test/core/llm/SchemaLlmTool.test.ts b/test/core/llm/SchemaLlmTool.test.ts new file mode 100644 index 0000000..53f7ec8 --- /dev/null +++ b/test/core/llm/SchemaLlmTool.test.ts @@ -0,0 +1,54 @@ +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, z.infer> { + protected readonly promptName = 'list'; + protected readonly outputSchema = outputSchema; + protected buildVariables(): Record { + return {}; + } +} + +function makeTool() { + return new ListTool( + {} as unknown as ConstructorParameters[0], + {} as unknown as ConstructorParameters[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); + }); +}); diff --git a/test/core/llm/TextLlmTool.test.ts b/test/core/llm/TextLlmTool.test.ts new file mode 100644 index 0000000..1a15331 --- /dev/null +++ b/test/core/llm/TextLlmTool.test.ts @@ -0,0 +1,21 @@ +import { describe, expect, it } from 'vitest'; +import { TextLlmTool } from '../../../src/core/llm/TextLlmTool'; +import type { RunContext } from '../../../src/core/llm/RunContext'; + +class EchoTool extends TextLlmTool> { + protected readonly promptName = 'echo'; + protected buildVariables(): Record { + return {}; + } +} + +describe('TextLlmTool.parseOutput (via subclass)', () => { + it('trims the raw LLM text and returns it as-is', () => { + const tool = new EchoTool( + {} as unknown as ConstructorParameters[0], + {} as unknown as ConstructorParameters[1] + ) as unknown as { parseOutput(raw: string, ctx: RunContext): string }; + + expect(tool.parseOutput(' hello world \n', {} as RunContext)).toBe('hello world'); + }); +});