diff --git a/src/gateway.ts b/src/gateway.ts new file mode 100644 index 0000000..25e9467 --- /dev/null +++ b/src/gateway.ts @@ -0,0 +1,191 @@ +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, + }; +}

Add automatic protocols and diagnostics; release 0.3.0.

Changes