219 lines
6.3 KiB
TypeScript
219 lines
6.3 KiB
TypeScript
|
|
/**
|
||
|
|
* Generic LLM caller for the analyzer pipeline.
|
||
|
|
*
|
||
|
|
* One round-trip is:
|
||
|
|
* 1. Send (system, user) to the chosen model.
|
||
|
|
* 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.
|
||
|
|
*
|
||
|
|
* The system prompt is marked with `cache_control: ephemeral`. Anthropic
|
||
|
|
* silently no-ops caching when the prefix is below the model's minimum
|
||
|
|
* (~2-4K tokens) — for our short stage prompts this often won't fire, which
|
||
|
|
* is fine; cost is unaffected when caching is skipped.
|
||
|
|
*/
|
||
|
|
|
||
|
|
import type Anthropic from '@anthropic-ai/sdk';
|
||
|
|
import type { ZodType } from 'zod';
|
||
|
|
import { getAnthropicClient } from './client';
|
||
|
|
import { estimateCostUsd, type TokenUsage } from './pricing';
|
||
|
|
import { type ModelId, OPUS } from './models';
|
||
|
|
|
||
|
|
export interface LLMCallOptions<T> {
|
||
|
|
model: ModelId;
|
||
|
|
system: string;
|
||
|
|
user: string;
|
||
|
|
schema: ZodType<T>;
|
||
|
|
maxTokens: number;
|
||
|
|
/** Override the 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;
|
||
|
|
}
|
||
|
|
|
||
|
|
/**
|
||
|
|
* Concatenate token usage from two calls (used to track total cost across
|
||
|
|
* the original attempt + retry).
|
||
|
|
*/
|
||
|
|
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 MessagesCreateBody {
|
||
|
|
model: string;
|
||
|
|
max_tokens: number;
|
||
|
|
system: Array<{
|
||
|
|
type: 'text';
|
||
|
|
text: string;
|
||
|
|
cache_control?: { type: 'ephemeral' };
|
||
|
|
}>;
|
||
|
|
messages: Array<{ role: 'user' | 'assistant'; content: string }>;
|
||
|
|
}
|
||
|
|
|
||
|
|
function buildBody(opts: {
|
||
|
|
model: ModelId;
|
||
|
|
system: string;
|
||
|
|
history: Array<{ role: 'user' | 'assistant'; content: string }>;
|
||
|
|
maxTokens: number;
|
||
|
|
}): MessagesCreateBody {
|
||
|
|
return {
|
||
|
|
model: opts.model,
|
||
|
|
max_tokens: opts.maxTokens,
|
||
|
|
system: [
|
||
|
|
{
|
||
|
|
type: 'text',
|
||
|
|
text: opts.system,
|
||
|
|
cache_control: { type: 'ephemeral' },
|
||
|
|
},
|
||
|
|
],
|
||
|
|
messages: opts.history,
|
||
|
|
};
|
||
|
|
}
|
||
|
|
|
||
|
|
function extractText(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 };
|
||
|
|
}
|
||
|
|
|
||
|
|
export async function callLLMStage<T>(
|
||
|
|
opts: LLMCallOptions<T>
|
||
|
|
): Promise<LLMCallResult<T>> {
|
||
|
|
const client = opts.client ?? getAnthropicClient();
|
||
|
|
const history: Array<{ role: 'user' | 'assistant'; content: string }> = [
|
||
|
|
{ role: 'user', content: opts.user },
|
||
|
|
];
|
||
|
|
|
||
|
|
// Opus 4.7 rejects `temperature`, `top_p`, `top_k`. We don't pass any of
|
||
|
|
// them, so the same body shape works on all three models.
|
||
|
|
const firstBody = buildBody({
|
||
|
|
model: opts.model,
|
||
|
|
system: opts.system,
|
||
|
|
history,
|
||
|
|
maxTokens: opts.maxTokens,
|
||
|
|
});
|
||
|
|
void (firstBody satisfies Anthropic.MessageCreateParamsNonStreaming);
|
||
|
|
|
||
|
|
const first = await client.messages.create(firstBody);
|
||
|
|
const firstText = extractText(first);
|
||
|
|
const firstParse = tryParseValidate(firstText, opts.schema);
|
||
|
|
|
||
|
|
if (firstParse.ok) {
|
||
|
|
const usage: TokenUsage = first.usage as TokenUsage;
|
||
|
|
return {
|
||
|
|
data: firstParse.value,
|
||
|
|
usage,
|
||
|
|
estimated_cost_usd: estimateCostUsd(opts.model, usage),
|
||
|
|
attempts: 1,
|
||
|
|
raw_response: firstText,
|
||
|
|
};
|
||
|
|
}
|
||
|
|
|
||
|
|
// Retry once. Append the model's previous (invalid) response and a follow-up
|
||
|
|
// user turn explaining the parse error.
|
||
|
|
history.push({ role: 'assistant', content: firstText });
|
||
|
|
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 secondBody = buildBody({
|
||
|
|
model: opts.model,
|
||
|
|
system: opts.system,
|
||
|
|
history,
|
||
|
|
maxTokens: opts.maxTokens,
|
||
|
|
});
|
||
|
|
void (secondBody satisfies Anthropic.MessageCreateParamsNonStreaming);
|
||
|
|
|
||
|
|
const second = await client.messages.create(secondBody);
|
||
|
|
const secondText = extractText(second);
|
||
|
|
const secondParse = tryParseValidate(secondText, opts.schema);
|
||
|
|
|
||
|
|
const totalUsage = addUsage(
|
||
|
|
first.usage as TokenUsage,
|
||
|
|
second.usage as TokenUsage
|
||
|
|
);
|
||
|
|
|
||
|
|
if (!secondParse.ok) {
|
||
|
|
throw new Error(
|
||
|
|
`LLM stage on ${opts.model} failed twice. First: ${firstParse.error}. Second: ${secondParse.error}. Last raw response: ${secondText.slice(0, 500)}`
|
||
|
|
);
|
||
|
|
}
|
||
|
|
|
||
|
|
return {
|
||
|
|
data: secondParse.value,
|
||
|
|
usage: totalUsage,
|
||
|
|
estimated_cost_usd: estimateCostUsd(opts.model, totalUsage),
|
||
|
|
attempts: 2,
|
||
|
|
raw_response: secondText,
|
||
|
|
};
|
||
|
|
}
|
||
|
|
|
||
|
|
// Re-export OPUS so callers don't need a second import for the common case.
|
||
|
|
export { OPUS };
|