Add core LlmPipeline base class

src/core/llm/LlmPipeline.ts

@@ -0,0 +1,32 @@
+import type { Logger } from '../logging/logger';
+import type { LlmTool } from './LlmTool';
+import type { RunContext } from './RunContext';
+import type { Tracer } from './Tracer';
+
+export abstract class LlmPipeline<TInput, TOutput> {
+  constructor(
+    protected readonly tracer: Tracer,
+    protected readonly logger: Logger
+  ) {}
+
+  protected abstract readonly name: string;
+  protected abstract runInternal(input: TInput, ctx: RunContext): Promise<TOutput>;
+
+  async runAsync(input: TInput): Promise<TOutput> {
+    const trace = this.tracer.startTrace(this.name);
+    const ctx: RunContext = { trace, logger: this.logger };
+    try {
+      return await this.runInternal(input, ctx);
+    } finally {
+      await this.tracer.flush();
+    }
+  }
+
+  protected runTool<TToolInput, TToolOutput>(
+    tool: LlmTool<TToolInput, TToolOutput>,
+    input: TToolInput,
+    ctx: RunContext
+  ): Promise<TToolOutput> {
+    return tool.runAsync(input, ctx);
+  }
+}

test/core/llm/LlmPipeline.test.ts

@@ -0,0 +1,74 @@
+import { describe, expect, it, vi } from 'vitest';
+import { LlmPipeline } from '../../../src/core/llm/LlmPipeline';
+import type { RunContext } from '../../../src/core/llm/RunContext';
+import type { Tracer, TraceHandle } from '../../../src/core/llm/Tracer';
+import type { LlmTool } from '../../../src/core/llm/LlmTool';
+import type { Logger } from '../../../src/core/logging/logger';
+
+class SumPipeline extends LlmPipeline<number[], number> {
+  protected readonly name = 'sum-pipeline';
+  protected async runInternal(input: number[], ctx: RunContext): Promise<number> {
+    expect(ctx.trace.id).toBe('trace-1');
+    return input.reduce((a, b) => a + b, 0);
+  }
+}
+
+class ThrowingPipeline extends LlmPipeline<void, void> {
+  protected readonly name = 'throwing-pipeline';
+  protected async runInternal(): Promise<void> {
+    throw new Error('pipeline failed');
+  }
+}
+
+function makeTracer(trace: TraceHandle) {
+  const flush = vi.fn().mockResolvedValue(undefined);
+  const startTrace = vi.fn().mockReturnValue(trace);
+  const tracer = { startTrace, flush } as unknown as Tracer;
+  return { tracer, startTrace, flush };
+}
+
+describe('LlmPipeline', () => {
+  it('opens one trace, runs runInternal with it, and flushes afterward', async () => {
+    const trace = { id: 'trace-1', startGeneration: vi.fn() } as unknown as TraceHandle;
+    const { tracer, startTrace, flush } = makeTracer(trace);
+    const logger = {} as unknown as Logger;
+
+    const pipeline = new SumPipeline(tracer, logger);
+    const result = await pipeline.runAsync([1, 2, 3]);
+
+    expect(result).toBe(6);
+    expect(startTrace).toHaveBeenCalledWith('sum-pipeline');
+    expect(flush).toHaveBeenCalledTimes(1);
+  });
+
+  it('still flushes the trace when runInternal throws', async () => {
+    const trace = { id: 'trace-2', startGeneration: vi.fn() } as unknown as TraceHandle;
+    const { tracer, flush } = makeTracer(trace);
+    const logger = {} as unknown as Logger;
+
+    const pipeline = new ThrowingPipeline(tracer, logger);
+    await expect(pipeline.runAsync()).rejects.toThrow('pipeline failed');
+    expect(flush).toHaveBeenCalledTimes(1);
+  });
+
+  it('runTool delegates to the given tool.runAsync with the same context', async () => {
+    const trace = { id: 'trace-3', startGeneration: vi.fn() } as unknown as TraceHandle;
+    const { tracer } = makeTracer(trace);
+    const logger = {} as unknown as Logger;
+
+    class ToolRunnerPipeline extends LlmPipeline<void, string> {
+      protected readonly name = 'tool-runner';
+      protected async runInternal(_input: void, ctx: RunContext): Promise<string> {
+        const tool = { runAsync: vi.fn().mockResolvedValue('tool-output') } as unknown as LlmTool<
+          Record<string, never>,
+          string
+        >;
+        return this.runTool(tool, {}, ctx);
+      }
+    }
+
+    const pipeline = new ToolRunnerPipeline(tracer, logger);
+    const result = await pipeline.runAsync();
+    expect(result).toBe('tool-output');
+  });
+});