diff --git a/src/core/llm/LlmPipeline.ts b/src/core/llm/LlmPipeline.ts new file mode 100644 index 0000000..c700013 --- /dev/null +++ b/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 { + constructor( + protected readonly tracer: Tracer, + protected readonly logger: Logger + ) {} + + protected abstract readonly name: string; + protected abstract runInternal(input: TInput, ctx: RunContext): Promise; + + async runAsync(input: TInput): Promise { + 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( + tool: LlmTool, + input: TToolInput, + ctx: RunContext + ): Promise { + return tool.runAsync(input, ctx); + } +} diff --git a/test/core/llm/LlmPipeline.test.ts b/test/core/llm/LlmPipeline.test.ts new file mode 100644 index 0000000..622d0e4 --- /dev/null +++ b/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 { + protected readonly name = 'sum-pipeline'; + protected async runInternal(input: number[], ctx: RunContext): Promise { + expect(ctx.trace.id).toBe('trace-1'); + return input.reduce((a, b) => a + b, 0); + } +} + +class ThrowingPipeline extends LlmPipeline { + protected readonly name = 'throwing-pipeline'; + protected async runInternal(): Promise { + 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 { + protected readonly name = 'tool-runner'; + protected async runInternal(_input: void, ctx: RunContext): Promise { + const tool = { runAsync: vi.fn().mockResolvedValue('tool-output') } as unknown as LlmTool< + Record, + string + >; + return this.runTool(tool, {}, ctx); + } + } + + const pipeline = new ToolRunnerPipeline(tracer, logger); + const result = await pipeline.runAsync(); + expect(result).toBe('tool-output'); + }); +});