Add SchemaLlmTool and TextLlmTool

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<TInput, TOutput> extends LlmTool<TInput, TOutput> {
+  protected abstract readonly outputSchema: ZodSchema<TOutput>;
+
+  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);
+  }
+}

src/core/llm/TextLlmTool.ts

@@ -0,0 +1,8 @@
+import { LlmTool } from './LlmTool';
+import type { RunContext } from './RunContext';
+
+export abstract class TextLlmTool<TInput> extends LlmTool<TInput, string> {
+  protected parseOutput(raw: string, _ctx: RunContext): string {
+    return raw.trim();
+  }
+}

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<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);
+  });
+});

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<Record<string, never>> {
+  protected readonly promptName = 'echo';
+  protected buildVariables(): Record<string, unknown> {
+    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<typeof EchoTool>[0],
+      {} as unknown as ConstructorParameters<typeof EchoTool>[1]
+    ) as unknown as { parseOutput(raw: string, ctx: RunContext): string };
+
+    expect(tool.parseOutput('  hello world  \n', {} as RunContext)).toBe('hello world');
+  });
+});