import { ZodError } from "zod";
import { setTimeout as delay } from "node:timers/promises";
import { setDefaultAutoSelectFamilyAttemptTimeout } from "node:net";
import {
acquireModel,
cancelModelRequests,
ModelError,
providerCooldown,
retryAfterMs,
} from "./model-runtime";
import type { Draft, Source, ExtractionMetrics } from "../src/types";
import { extractLocal } from "./context";
import { compactInput, parseCompactDraft, COMPACT_PROMPT } from "./compact";
type ModelConnection = {
endpoint: string;
key: string;
model: string;
auth: string;
};
// Node gives IPv6 250 ms before trying IPv4; api.meta.ai's IPv6 address can stall longer than that.
setDefaultAutoSelectFamilyAttemptTimeout(2_000);
let sessionConnection: ModelConnection | null = null;
let lastVerifiedAt: string | null = null;
let revision = 0;
const cache = new Map<string, { expires: number; draft: Draft }>();
const CACHE_TTL = 3 * 60 * 1000;
export const modelVerification = () => lastVerifiedAt;
export function configureAzureSession(connection: ModelConnection): void {
const url = new URL(connection.endpoint);
if (
url.protocol !== "https:" ||
url.username ||
url.password ||
url.port ||
!/^[a-z0-9-]+\.(services\.ai|openai)\.azure\.com$/.test(url.hostname) ||
url.pathname !== "/openai/v1/chat/completions" ||
url.search ||
url.hash
) {
throw new Error(
"Use the Azure deployment HTTPS chat-completions endpoint shown in Foundry.",
);
}
if (
!connection.key.trim() ||
!/^[a-zA-Z0-9._-]{1,150}$/.test(connection.model) ||
!["api-key", "bearer"].includes(connection.auth)
) {
throw new Error(
"The model connection fields are incomplete or unsupported.",
);
}
cancelModelRequests();
sessionConnection = { ...connection, key: connection.key.trim() };
lastVerifiedAt = null;
revision++;
cache.clear();
}
export const modelConfig = () =>
sessionConnection ?? {
endpoint: process.env.CONTEXT_MODEL_ENDPOINT,
key: process.env.CONTEXT_MODEL_API_KEY,
model: process.env.CONTEXT_MODEL_ID,
auth: process.env.CONTEXT_MODEL_AUTH || "bearer",
};
export const modelReady = () => {
const c = modelConfig();
return Boolean(c.endpoint && c.key && c.model);
};
// Meta Model API (api.meta.ai) serves Muse Spark, a reasoning model: its thinking counts against the output cap.
export const modelProvider = (): "meta" | "azure" | "other" => {
const endpoint = modelConfig().endpoint;
if (!endpoint) return "other";
const host = new URL(endpoint).hostname;
return host === "api.meta.ai" ? "meta" : host.endsWith(".azure.com") ? "azure" : "other";
};
const MODEL_NAMES = { meta: "Meta Muse Spark", azure: "Azure model", other: "Model" };
export const modelLabel = () => (modelReady() ? `${MODEL_NAMES[modelProvider()]} (${modelConfig().model})` : null);
export async function extract(
sources: Source[],
useModel: boolean,
freshModel = false,
signal?: AbortSignal,
): Promise<Draft> {
signal?.throwIfAborted();
const started = performance.now(),
draft = extractLocal(sources);
const metrics: ExtractionMetrics = {
providerCalls: 0,
inputCharacters: 0,
outputCharacters: 0,
inputTokens: null,
outputTokens: null,
cacheHit: false,
};
const narrative = compactInput(sources);
if (!useModel || narrative.length === 0)
return { ...draft, strategy: "structured", attempts: 0, metrics };
const c = modelConfig();
if (!modelReady())
throw new Error(
"No model endpoint is configured. Choose the local structured reader.",
);
if (!/^https:\/\//.test(c.endpoint!))
throw new Error("The model endpoint must use HTTPS.");
const key = revision + ":" + c.model + ":" + draft.fingerprint;
for (const [id, entry] of cache)
if (entry.expires <= Date.now()) cache.delete(id);
const hit = cache.get(key);
if (!freshModel && hit)
return {
...structuredClone(hit.draft),
strategy: "cached",
durationMs: Math.round(performance.now() - started),
attempts: 0,
metrics: { ...metrics, cacheHit: true },
};
const slot = acquireModel(signal);
try {
const input = JSON.stringify(narrative);
metrics.inputCharacters = input.length;
let usageReports = 0;
const meta = modelProvider() === "meta";
const deadline = performance.now() + (meta ? 60_000 : 22_000);
let retryReason = "format";
for (let attempt = 1; attempt <= 2; attempt++) {
slot.signal.throwIfAborted();
const remaining = Math.floor(deadline - performance.now());
if (remaining < 500)
throw new Error(
"Extraction exceeded its time budget. No partial model response was used. Retry when the provider is available.",
);
let payload: {
choices?: { finish_reason?: string; message?: { content?: unknown } }[];
usage?: { prompt_tokens?: number; completion_tokens?: number };
};
try {
metrics.providerCalls++;
const response = await fetch(c.endpoint!, {
method: "POST",
signal: AbortSignal.any([
slot.signal,
AbortSignal.timeout(Math.min(meta ? 45_000 : 12_000, remaining)),
]),
headers: {
"Content-Type": "application/json",
...(c.auth === "api-key"
? { "api-key": c.key! }
: { Authorization: "Bearer " + c.key }),
},
body: JSON.stringify({
model: c.model,
temperature: 0,
...(meta
? { max_completion_tokens: 12_000, reasoning_effort: "low" }
: {
max_tokens: Math.min(
3000,
Math.max(500, draft.unparsed.length * 190),
),
}),
response_format: { type: "json_object" },
messages: [
{
role: "system",
content:
COMPACT_PROMPT +
(attempt === 2 && retryReason === "format"
? "\nReturn the exact JSON schema, with numeric references and one normalized string per fact. No prose."
: ""),
},
{ role: "user", content: input },
],
}),
});
if (
[408, 429, 500, 502, 503, 504].includes(response.status) &&
attempt === 1
) {
const waitMs = retryAfterMs(response.headers.get("retry-after"));
await response.body?.cancel();
if (response.status === 429) providerCooldown(waitMs || 1000);
if (waitMs > 2000)
throw new ModelError(
"The model provider is busy. Retry after the indicated delay; no model result was used.",
503,
Math.ceil(waitMs / 1000),
);
await delay(Math.max(150, waitMs), undefined, {
signal: slot.signal,
});
retryReason = "transport";
continue;
}
if (!response.ok) {
const waitMs = retryAfterMs(response.headers.get("retry-after"));
await response.body?.cancel();
if (response.status === 429) providerCooldown(waitMs || 1000);
throw new ModelError(
response.status === 401 || response.status === 403
? "The provider rejected this connection. Check the model connection; no model result was used."
: "The provider rejected the request or is busy. No model result was used. Retry later.",
502,
waitMs ? Math.ceil(waitMs / 1000) : undefined,
);
}
payload = (await response.json()) as typeof payload;
} catch (error) {
if (slot.signal.aborted)
throw new ModelError(
"Extraction cancelled. No model result was used.",
499,
);
if (error instanceof ModelError) throw error;
if (attempt === 1) {
retryReason = "transport";
continue;
}
throw new Error(
"The model request failed or timed out within the bounded retry budget. No model result was used.",
);
}
slot.signal.throwIfAborted();
if (
!payload ||
typeof payload !== "object" ||
!Array.isArray(payload.choices)
)
throw new ModelError(
"The provider returned an invalid response. No model result was used.",
502,
);
const usage = payload.usage;
if (
usage &&
Number.isSafeInteger(usage.prompt_tokens) &&
Number.isSafeInteger(usage.completion_tokens) &&
usage.prompt_tokens! >= 0 &&
usage.completion_tokens! >= 0
) {
usageReports++;
metrics.inputTokens = (metrics.inputTokens ?? 0) + usage.prompt_tokens!;
metrics.outputTokens =
(metrics.outputTokens ?? 0) + usage.completion_tokens!;
}
const choice = payload.choices?.[0];
if (choice?.finish_reason === "content_filter")
throw new Error(
"The provider's content filter blocked this input or response. No model result was used.",
);
if (choice?.finish_reason === "length")
throw new Error(
"The model reached its output limit. No model result was used. Split long records into smaller cases.",
);
const content = choice?.message?.content;
if (typeof content !== "string" || !content.trim())
throw new Error(
"The provider returned no text. No model result was used.",
);
metrics.outputCharacters += content.length;
let parsed: unknown;
try {
parsed = JSON.parse(
content.trim().replace(/^```(?:json)?\s*([\s\S]*?)\s*```$/i, "$1"),
);
} catch {
if (attempt === 1) {
retryReason = "format";
continue;
}
throw new Error(
"The provider returned malformed JSON on both attempts. No model result was used.",
);
}
let validated: Draft;
try {
validated = parseCompactDraft(
parsed,
sources,
c.model!,
Math.round(performance.now() - started),
);
} catch (error) {
if (error instanceof ZodError && attempt === 1) {
retryReason = "format";
continue;
}
const reason =
error instanceof ZodError
? "Unsupported model response structure."
: error instanceof Error
? error.message
: "Unsupported response.";
throw new Error(
"Model source checks failed. " +
reason +
" No model result was used.",
);
}
lastVerifiedAt = new Date().toISOString();
if (usageReports !== metrics.providerCalls) {
metrics.inputTokens = null;
metrics.outputTokens = null;
}
const result = { ...validated, attempts: attempt, metrics };
if (cache.size >= 20) cache.delete(cache.keys().next().value!);
cache.set(key, {
expires: Date.now() + CACHE_TTL,
draft: structuredClone(result),
});
return result;
}
throw new Error("No model result was used.");
} catch (error) {
if (slot.signal.aborted)
throw new ModelError(
"Extraction cancelled. No model result was used.",
499,
);
if (error instanceof ModelError) throw error;
throw new ModelError(
error instanceof Error
? error.message
: "Model extraction failed. No result was used.",
502,
);
} finally {
slot.release();
}
}
const TRANSCRIBE_PROMPT =
"Transcribe this synthetic medical document exactly as written, one output line per line on the page. " +
"Keep the writer's words, numbers, units and abbreviations. Do not correct, interpret, summarise or add anything. " +
"Write [illegible] for any word you cannot read. Return plain text only.";
/** Muse Spark reads a photo or PDF of a record into its lines; the reader then cites those lines like any typed record. */
export async function transcribe(data: Buffer, contentType: string, signal?: AbortSignal): Promise<string> {
if (modelProvider() !== "meta") throw new ModelError("Photos and PDFs are read by Meta Muse Spark; connect it first.", 400);
const c = modelConfig();
const file =
contentType === "application/pdf"
? { type: "file", file: { filename: "record.pdf", file_data: `data:application/pdf;base64,${data.toString("base64")}` } }
: { type: "image_url", image_url: { url: `data:${contentType};base64,${data.toString("base64")}` } };
const slot = acquireModel(signal);
try {
const response = await fetch(c.endpoint!, {
method: "POST",
signal: AbortSignal.any([slot.signal, AbortSignal.timeout(60_000)]),
headers: { "Content-Type": "application/json", Authorization: "Bearer " + c.key },
body: JSON.stringify({
model: c.model,
temperature: 0,
max_completion_tokens: 8_000,
reasoning_effort: "low",
messages: [
{ role: "system", content: TRANSCRIBE_PROMPT },
{ role: "user", content: [{ type: "text", text: "Transcribe this record." }, file] },
],
}),
});
if (!response.ok) {
await response.body?.cancel();
throw new ModelError("Muse Spark could not read this file. No reading was used.", 502);
}
const payload = (await response.json()) as { choices?: { message?: { content?: unknown } }[] };
const text = payload.choices?.[0]?.message?.content;
if (typeof text !== "string" || !text.trim()) throw new ModelError("Muse Spark returned no text for this file. No reading was used.", 502);
return text.replace(/\r\n/g, "\n").replace(/^```\w*\n?|```$/g, "").trim();
} finally {
slot.release();
}
}