Add core Tracer/RunContext types and LlmTool base class

This commit is contained in:
2026-08-12 20:53:27 +02:00
parent 3ca58194e6
commit ee4aba83e8
4 changed files with 169 additions and 0 deletions

59
src/core/llm/LlmTool.ts Normal file
View File

@@ -0,0 +1,59 @@
import type { PromptService } from '../langfuse/PromptService';
import type { ChatMessage, LlmClient } from './LlmClient';
import type { RunContext } from './RunContext';
export abstract class LlmTool<TInput, TOutput> {
protected abstract readonly promptName: string;
constructor(
protected readonly promptService: PromptService,
protected readonly llmClient: LlmClient
) {}
protected abstract buildVariables(input: TInput, ctx: RunContext): Record<string, unknown>;
protected async preProcess(input: TInput, _ctx: RunContext): Promise<TInput> {
return input;
}
protected async postProcess(output: TOutput, _input: TInput, _ctx: RunContext): Promise<TOutput> {
return output;
}
protected abstract parseOutput(raw: string, ctx: RunContext): TOutput;
async runAsync(input: TInput, ctx: RunContext): Promise<TOutput> {
const processedInput = await this.preProcess(input, ctx);
const variables = this.buildVariables(processedInput, ctx);
const compiled = await this.promptService.getCompiled(this.promptName, variables);
const model = String(compiled.config.model ?? 'gpt-4o-mini');
const temperature =
typeof compiled.config.temperature === 'number' ? compiled.config.temperature : undefined;
const maxTokens =
typeof compiled.config.maxTokens === 'number' ? compiled.config.maxTokens : undefined;
const messages: ChatMessage[] = compiled.messages;
const generation = ctx.trace.startGeneration({
name: this.promptName,
model,
input: messages,
metadata: { promptVersion: compiled.version },
});
let raw: string;
try {
const result = await this.llmClient.chat(messages, { model, temperature, maxTokens });
raw = result.content;
generation.end({ output: raw, usage: result.usage });
} catch (err) {
generation.end({ output: null, metadata: { error: String(err) } });
throw err;
}
const parsed = this.parseOutput(raw, ctx);
const output = await this.postProcess(parsed, processedInput, ctx);
ctx.logger.info({ tool: this.promptName }, 'llm tool completed');
return output;
}
}

View File

@@ -0,0 +1,7 @@
import type { Logger } from '../logging/logger';
import type { TraceHandle } from './Tracer';
export interface RunContext {
trace: TraceHandle;
logger: Logger;
}

26
src/core/llm/Tracer.ts Normal file
View File

@@ -0,0 +1,26 @@
export interface GenerationInput {
name: string;
model: string;
input: unknown;
metadata?: Record<string, unknown>;
}
export interface GenerationOutput {
output: unknown;
usage?: { promptTokens?: number; completionTokens?: number; totalTokens?: number };
metadata?: Record<string, unknown>;
}
export interface GenerationHandle {
end(output: GenerationOutput): void;
}
export interface TraceHandle {
readonly id: string;
startGeneration(input: GenerationInput): GenerationHandle;
}
export abstract class Tracer {
abstract startTrace(name: string, metadata?: Record<string, unknown>): TraceHandle;
abstract flush(): Promise<void>;
}

View File

@@ -0,0 +1,77 @@
import { describe, expect, it, vi } from 'vitest';
import { LlmTool } from '../../../src/core/llm/LlmTool';
import type { RunContext } from '../../../src/core/llm/RunContext';
import type { PromptService } from '../../../src/core/langfuse/PromptService';
import type { LlmClient } from '../../../src/core/llm/LlmClient';
interface TestInput {
who: string;
}
class GreetTool extends LlmTool<TestInput, string> {
protected readonly promptName = 'greet';
protected buildVariables(input: TestInput): Record<string, unknown> {
return { who: input.who };
}
protected parseOutput(raw: string): string {
return raw.toUpperCase();
}
}
function makeCtx() {
const generationEnd = vi.fn();
const startGeneration = vi.fn().mockReturnValue({ end: generationEnd });
const trace = { id: 'trace-1', startGeneration };
const logger = { info: vi.fn(), error: vi.fn() } as unknown as RunContext['logger'];
return { ctx: { trace, logger } as RunContext, startGeneration, generationEnd };
}
describe('LlmTool.runAsync', () => {
it('fetches the compiled prompt, calls the LLM, and parses the output', async () => {
const getCompiled = vi.fn().mockResolvedValue({
name: 'greet',
version: 1,
config: { model: 'llama3', temperature: 0.5 },
messages: [{ role: 'user', content: 'hi alice' }],
});
const promptService = { getCompiled } as unknown as PromptService;
const chat = vi.fn().mockResolvedValue({ content: 'hello alice', usage: { totalTokens: 3 } });
const llmClient = { chat } as unknown as LlmClient;
const { ctx, startGeneration, generationEnd } = makeCtx();
const tool = new GreetTool(promptService, llmClient);
const result = await tool.runAsync({ who: 'alice' }, ctx);
expect(result).toBe('HELLO ALICE');
expect(getCompiled).toHaveBeenCalledWith('greet', { who: 'alice' });
expect(chat).toHaveBeenCalledWith(
[{ role: 'user', content: 'hi alice' }],
{ model: 'llama3', temperature: 0.5, maxTokens: undefined }
);
expect(startGeneration).toHaveBeenCalledWith(
expect.objectContaining({ name: 'greet', model: 'llama3' })
);
expect(generationEnd).toHaveBeenCalledWith(
expect.objectContaining({ output: 'hello alice' })
);
});
it('ends the generation with error metadata and rethrows when the LLM call fails', async () => {
const promptService = {
getCompiled: vi.fn().mockResolvedValue({ name: 'greet', version: 1, config: {}, messages: [] }),
} as unknown as PromptService;
const llmClient = { chat: vi.fn().mockRejectedValue(new Error('boom')) } as unknown as LlmClient;
const { ctx, generationEnd } = makeCtx();
const tool = new GreetTool(promptService, llmClient);
await expect(tool.runAsync({ who: 'alice' }, ctx)).rejects.toThrow('boom');
expect(generationEnd).toHaveBeenCalledWith(
expect.objectContaining({ output: null })
);
});
});