wulf-pulse/lib/services/llm/call.ts

/**
 * Generic LLM caller for the analyzer pipeline.
 *
 * One round-trip is:
 *   1. Send (system, user) to the chosen model (Anthropic or OpenRouter).
 *   2. Extract the assistant's text content.
 *   3. JSON.parse + Zod-validate against the caller's schema.
 *   4. On failure: retry ONCE with the prior raw response + parse error in a
 *      follow-up user turn, then validate again.
 *   5. After two failures: throw.
 *
 * Provider dispatch:
 *   - claude-* → Anthropic SDK (with prompt-cache hint on the system prefix)
 *   - <vendor>/<model> → OpenRouter chat-completions (OpenAI-compatible)
 *
 * The retry logic, schema validation, and result shape are identical across
 * providers so callers stay provider-agnostic.
 */

import type Anthropic from '@anthropic-ai/sdk';
import type { ZodType } from 'zod';
import { getAnthropicClient } from './client';
import { callOpenRouterChat } from './openrouter-call';
import { estimateCostUsd, type TokenUsage } from './pricing';
import {
  type ModelId,
  OPUS,
  providerForModel,
} from './models';

export interface LLMCallOptions<T> {
  model: ModelId;
  system: string;
  user: string;
  schema: ZodType<T>;
  maxTokens: number;
  /** Override the Anthropic singleton (test injection). */
  client?: Anthropic;
}

export interface LLMCallResult<T> {
  data: T;
  usage: TokenUsage;
  estimated_cost_usd: number;
  /** Number of attempts made (1 = first try succeeded; 2 = retry succeeded). */
  attempts: 1 | 2;
  /** Raw response text for debugging / model_traces. */
  raw_response: string;
}

interface RoundTripResult {
  text: string;
  usage: TokenUsage;
}

type RoundTripFn = (
  history: Array<{ role: 'user' | 'assistant'; content: string }>
) => Promise<RoundTripResult>;

function addUsage(a: TokenUsage, b: TokenUsage): TokenUsage {
  return {
    input_tokens: a.input_tokens + b.input_tokens,
    output_tokens: a.output_tokens + b.output_tokens,
    cache_creation_input_tokens:
      (a.cache_creation_input_tokens ?? 0) + (b.cache_creation_input_tokens ?? 0),
    cache_read_input_tokens:
      (a.cache_read_input_tokens ?? 0) + (b.cache_read_input_tokens ?? 0),
  };
}

interface AnthropicMessagesCreateBody {
  model: string;
  max_tokens: number;
  system: Array<{
    type: 'text';
    text: string;
    cache_control?: { type: 'ephemeral' };
  }>;
  messages: Array<{ role: 'user' | 'assistant'; content: string }>;
}

function buildAnthropicBody(opts: {
  model: ModelId;
  system: string;
  history: Array<{ role: 'user' | 'assistant'; content: string }>;
  maxTokens: number;
}): AnthropicMessagesCreateBody {
  return {
    model: opts.model,
    max_tokens: opts.maxTokens,
    system: [
      {
        type: 'text',
        text: opts.system,
        cache_control: { type: 'ephemeral' },
      },
    ],
    messages: opts.history,
  };
}

function extractAnthropicText(response: Anthropic.Message): string {
  const parts: string[] = [];
  for (const block of response.content) {
    if (block.type === 'text') parts.push(block.text);
  }
  return parts.join('').trim();
}

/**
 * Some models occasionally wrap JSON in code fences despite "respond ONLY with
 * JSON" instructions. Strip a single ```json ... ``` fence if present so the
 * happy path doesn't bounce into the retry just for that.
 */
function unwrapFences(text: string): string {
  const trimmed = text.trim();
  if (trimmed.startsWith('```')) {
    const stripped = trimmed
      .replace(/^```(?:json)?\s*/i, '')
      .replace(/```\s*$/, '')
      .trim();
    return stripped;
  }
  return trimmed;
}

function tryParseValidate<T>(
  text: string,
  schema: ZodType<T>
): { ok: true; value: T } | { ok: false; error: string } {
  let parsed: unknown;
  try {
    parsed = JSON.parse(unwrapFences(text));
  } catch (err) {
    return {
      ok: false,
      error: `JSON.parse failed: ${err instanceof Error ? err.message : String(err)}`,
    };
  }
  const result = schema.safeParse(parsed);
  if (!result.success) {
    return {
      ok: false,
      error: `Zod validation failed: ${JSON.stringify(result.error.issues, null, 2)}`,
    };
  }
  return { ok: true, value: result.data };
}

function makeAnthropicRoundTrip(
  opts: LLMCallOptions<unknown>,
  client: Anthropic
): RoundTripFn {
  return async (history) => {
    const body = buildAnthropicBody({
      model: opts.model,
      system: opts.system,
      history,
      maxTokens: opts.maxTokens,
    });
    void (body satisfies Anthropic.MessageCreateParamsNonStreaming);
    const response = await client.messages.create(body);
    return {
      text: extractAnthropicText(response),
      usage: response.usage as TokenUsage,
    };
  };
}

function makeOpenRouterRoundTrip(opts: LLMCallOptions<unknown>): RoundTripFn {
  return async (history) => {
    const result = await callOpenRouterChat({
      model: opts.model,
      system: opts.system,
      history,
      maxTokens: opts.maxTokens,
    });
    return { text: result.text, usage: result.usage };
  };
}

export async function callLLMStage<T>(
  opts: LLMCallOptions<T>
): Promise<LLMCallResult<T>> {
  const provider = providerForModel(opts.model);
  const roundTrip: RoundTripFn =
    provider === 'anthropic'
      ? makeAnthropicRoundTrip(opts, opts.client ?? getAnthropicClient())
      : makeOpenRouterRoundTrip(opts);

  const history: Array<{ role: 'user' | 'assistant'; content: string }> = [
    { role: 'user', content: opts.user },
  ];

  const first = await roundTrip(history);
  const firstParse = tryParseValidate(first.text, opts.schema);

  if (firstParse.ok) {
    return {
      data: firstParse.value,
      usage: first.usage,
      estimated_cost_usd: estimateCostUsd(opts.model, first.usage),
      attempts: 1,
      raw_response: first.text,
    };
  }

  // Retry once. Append the model's previous (invalid) response and a follow-up
  // user turn explaining the parse error.
  history.push({ role: 'assistant', content: first.text });
  history.push({
    role: 'user',
    content: [
      'Your previous response could not be parsed. Error:',
      firstParse.error,
      '',
      'Re-emit the SAME response, corrected to be valid JSON that conforms to the requested schema.',
      'Respond ONLY with JSON. No prose, no code fences.',
    ].join('\n'),
  });

  const second = await roundTrip(history);
  const secondParse = tryParseValidate(second.text, opts.schema);

  const totalUsage = addUsage(first.usage, second.usage);

  if (!secondParse.ok) {
    throw new Error(
      `LLM stage on ${opts.model} failed twice. First: ${firstParse.error}. Second: ${secondParse.error}. Last raw response: ${second.text.slice(0, 500)}`
    );
  }

  return {
    data: secondParse.value,
    usage: totalUsage,
    estimated_cost_usd: estimateCostUsd(opts.model, totalUsage),
    attempts: 2,
    raw_response: second.text,
  };
}

// Re-export OPUS so callers don't need a second import for the common case.
export { OPUS };