import { createAssistantMessageEventStream } from "@earendil-works/pi-ai"; import type { Api, AssistantMessage, AssistantMessageEventStream, Model, ProviderStreams, StreamOptions } from "@earendil-works/pi-ai"; import { asRecord } from "./values.ts"; import type { ExplabsModel } from "./types.ts"; export interface GatewayDetails { requestId?: string; provider?: string; routeDepth?: string; routeReason?: string; zdr?: string; replayRepair?: string; chargedUsd?: number; isByok?: boolean; upstreamUsd?: number; ignoredParameters?: string[]; omittedSampling?: string[]; } export type GatewayMessage = AssistantMessage & { explabs?: GatewayDetails }; function nonnegativeNumber(value: unknown): number | undefined { return typeof value === "number" && Number.isFinite(value) && value >= 0 ? value : undefined; } export function gatewayHeaders(headers: Record): GatewayDetails { const lower = Object.fromEntries(Object.entries(headers).map(([key, value]) => [key.toLowerCase(), value])); return { requestId: lower["x-request-id"], provider: lower["x-gateway-provider"], routeDepth: lower["x-gateway-route-depth"], routeReason: lower["x-gateway-route-reason"], zdr: lower["x-gateway-zdr"], replayRepair: lower["x-gateway-replay-repair"], }; } export function gatewayEvent(data: unknown): GatewayDetails { const event = asRecord(data); if (!event) return {}; // Responses wraps the final reply in response; Messages reports usage on message_delta. const response = asRecord(event.response) ?? event; const finalUsage = event.type === "response.completed" || event.type === "response.incomplete" || event.type === "response.failed" || event.type === "message_delta" || event.type === undefined; const usage = finalUsage ? asRecord(response.usage) : undefined; const details = asRecord(usage?.cost_details); const ignored = response["x-experiential-ignored-parameters"]; const chargedUsd = nonnegativeNumber(usage?.cost); const upstreamUsd = nonnegativeNumber(details?.upstream_inference_cost); return { ...(typeof response.provider === "string" ? { provider: response.provider } : {}), ...(chargedUsd !== undefined ? { chargedUsd } : {}), ...(typeof usage?.is_byok === "boolean" ? { isByok: usage.is_byok } : {}), ...(upstreamUsd !== undefined ? { upstreamUsd } : {}), ...(Array.isArray(ignored) ? { ignoredParameters: ignored.filter((item): item is string => typeof item === "string") } : {}), }; } export function rewriteSampling(payload: unknown, model: Model & Pick): unknown { const record = asRecord(payload); const capabilities = model.samplingCapabilities; if (!record || !capabilities) return payload; const reasoning = asRecord(record.reasoning); const thinking = asRecord(record.thinking); const outputConfig = asRecord(record.output_config); const effort = reasoning?.effort ?? record.reasoning_effort ?? outputConfig?.effort; const reasoningEnabled = effort !== undefined ? effort !== "none" && effort !== "off" : thinking !== undefined ? thinking.type !== "disabled" : model.reasoning && model.thinkingLevelMap?.off === null; const samplingConflictsWithReasoning = capabilities.sampling_requires_reasoning_none && reasoningEnabled; const next = { ...record }; function outsideRange(value: unknown, min: number | undefined, max: number | undefined): boolean { return typeof value === "number" && ((min !== undefined && value < min) || (max !== undefined && value > max)); } if (capabilities.supports_temperature === false || samplingConflictsWithReasoning || outsideRange(next.temperature, capabilities.minimum_temperature, capabilities.maximum_temperature)) { delete next.temperature; } if (capabilities.supports_top_p === false || samplingConflictsWithReasoning || outsideRange(next.top_p, capabilities.minimum_top_p, capabilities.maximum_top_p)) { delete next.top_p; } return next; } export interface GatewayOptions { trace?: boolean; traceName?: string; userId?: string; tags?: Record; } export function readGatewayOptions(env: NodeJS.ProcessEnv = process.env): GatewayOptions { let tags: Record | undefined; if (env.EXPLABS_TAGS) { const parsed = asRecord(JSON.parse(env.EXPLABS_TAGS)); if (!parsed || Object.values(parsed).some((value) => typeof value !== "string")) { throw new Error("EXPLABS_TAGS must be a JSON object with string values"); } if (Object.keys(parsed).length > 16 || Buffer.byteLength(JSON.stringify(parsed)) > 8192) { throw new Error("EXPLABS_TAGS must contain at most 16 entries and 8 KiB of JSON"); } tags = parsed as Record; } return { trace: env.EXPLABS_TRACE === "1", traceName: env.EXPLABS_TRACE_NAME, userId: env.EXPLABS_USER_ID, tags }; } export function withGatewayInstrumentation(api: ProviderStreams, config: GatewayOptions = {}): ProviderStreams { function instrument(options: T | undefined) { let details: GatewayDetails = {}; let omittedSampling: string[] = []; const next = { ...options, headers: { ...(config.trace && options?.sessionId ? { "X-Explabs-Session-Id": options.sessionId } : {}), ...(config.traceName ? { "X-Explabs-Trace-Name": config.traceName } : {}), ...(config.userId ? { "X-Explabs-User-Id": config.userId } : {}), ...(config.tags ? { "X-Explabs-Tags": JSON.stringify(config.tags) } : {}), ...options?.headers, }, onPayload: async (payload: unknown, model: Model) => { const rewritten = (await options?.onPayload?.(payload, model)) ?? payload; const filtered = rewriteSampling(rewritten, model); const before = asRecord(rewritten); const after = asRecord(filtered); omittedSampling = ["temperature", "top_p"].filter((key) => before && key in before && after && !(key in after)); return filtered; }, onResponse: async (response, model) => { // A retry starts a new attempt; never retain another attempt's charge or route. details = gatewayHeaders(response.headers); await options?.onResponse?.(response, model); }, onProviderStreamEvent: async (event: unknown, model: Model) => { details = { ...details, ...gatewayEvent(event) }; await options?.onProviderStreamEvent?.(event, model); }, } satisfies StreamOptions; return { options: next as T, attach(message: GatewayMessage) { if (Object.values(details).some((value) => value !== undefined) || omittedSampling.length) { message.explabs = { ...details, ...(omittedSampling.length ? { omittedSampling: [...omittedSampling] } : {}) }; } }, }; } function forward(source: AssistantMessageEventStream, attach: (message: GatewayMessage) => void, model: Model): AssistantMessageEventStream { const target = createAssistantMessageEventStream(); const started = Date.now(); void (async () => { for await (const event of source) { if (event.type === "done") attach(event.message); else if (event.type === "error") attach(event.error); target.push(event); } const result = await source.result(); attach(result); target.end(result); })().catch((error: unknown) => { const message: GatewayMessage = { role: "assistant", api: model.api, provider: model.provider, model: model.id, content: [], timestamp: started, stopReason: "error", errorMessage: error instanceof Error ? error.message : String(error), usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 } }, }; attach(message); target.push({ type: "error", reason: "error", error: message }); target.end(message); }); return target; } return { stream(model, context, options) { const request = instrument(options); return forward(api.stream(model, context, request.options), request.attach, model); }, streamSimple(model, context, options) { const request = instrument(options); return forward(api.streamSimple(model, context, request.options), request.attach, model); }, fetchDeferred: api.fetchDeferred ? (model, handle, options) => { const request = instrument(options); return forward(api.fetchDeferred!(model, handle, request.options), request.attach, model); } : undefined, cancelDeferred: api.cancelDeferred ? (model, handle, options) => api.cancelDeferred!(model, handle, options) : undefined, }; }