sugarr/src/core/llm/LlmTool.ts

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 readonly defaultModel: string = 'gpt-4o-mini'
  ) {}

  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 ?? this.defaultModel);
    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;
  }
}