From ee4aba83e83b3c8398bcaed0817469408a91b229 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fran=20Jurmanovi=C4=87?= Date: Wed, 12 Aug 2026 20:53:27 +0200 Subject: [PATCH] Add core Tracer/RunContext types and LlmTool base class --- src/core/llm/LlmTool.ts | 59 +++++++++++++++++++++++++++ src/core/llm/RunContext.ts | 7 ++++ src/core/llm/Tracer.ts | 26 ++++++++++++ test/core/llm/LlmTool.test.ts | 77 +++++++++++++++++++++++++++++++++++ 4 files changed, 169 insertions(+) create mode 100644 src/core/llm/LlmTool.ts create mode 100644 src/core/llm/RunContext.ts create mode 100644 src/core/llm/Tracer.ts create mode 100644 test/core/llm/LlmTool.test.ts diff --git a/src/core/llm/LlmTool.ts b/src/core/llm/LlmTool.ts new file mode 100644 index 0000000..5abef38 --- /dev/null +++ b/src/core/llm/LlmTool.ts @@ -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 { + protected abstract readonly promptName: string; + + constructor( + protected readonly promptService: PromptService, + protected readonly llmClient: LlmClient + ) {} + + protected abstract buildVariables(input: TInput, ctx: RunContext): Record; + + protected async preProcess(input: TInput, _ctx: RunContext): Promise { + return input; + } + + protected async postProcess(output: TOutput, _input: TInput, _ctx: RunContext): Promise { + return output; + } + + protected abstract parseOutput(raw: string, ctx: RunContext): TOutput; + + async runAsync(input: TInput, ctx: RunContext): Promise { + 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; + } +} diff --git a/src/core/llm/RunContext.ts b/src/core/llm/RunContext.ts new file mode 100644 index 0000000..a1f52f0 --- /dev/null +++ b/src/core/llm/RunContext.ts @@ -0,0 +1,7 @@ +import type { Logger } from '../logging/logger'; +import type { TraceHandle } from './Tracer'; + +export interface RunContext { + trace: TraceHandle; + logger: Logger; +} diff --git a/src/core/llm/Tracer.ts b/src/core/llm/Tracer.ts new file mode 100644 index 0000000..eff5113 --- /dev/null +++ b/src/core/llm/Tracer.ts @@ -0,0 +1,26 @@ +export interface GenerationInput { + name: string; + model: string; + input: unknown; + metadata?: Record; +} + +export interface GenerationOutput { + output: unknown; + usage?: { promptTokens?: number; completionTokens?: number; totalTokens?: number }; + metadata?: Record; +} + +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): TraceHandle; + abstract flush(): Promise; +} diff --git a/test/core/llm/LlmTool.test.ts b/test/core/llm/LlmTool.test.ts new file mode 100644 index 0000000..cd44628 --- /dev/null +++ b/test/core/llm/LlmTool.test.ts @@ -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 { + protected readonly promptName = 'greet'; + + protected buildVariables(input: TestInput): Record { + 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 }) + ); + }); +});