From 87fc3c1091937aef28285d0fa7032436a3210c73 Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Mon, 6 Jul 2026 05:44:14 +0300 Subject: [PATCH] feat: add OpenAI-compatible /v1/embeddings endpoint Add POST /v1/embeddings through gateway auth, operation-aware routing, provider fallback, budgets, usage normalization, per-gateway-key rate limiting, and best-effort ledger recording. Embedding input is never written to the usage ledger; post-flight hard-budget failures still record spend before rejecting. Routing is now operation-aware: chat requires chat/streaming/tools capabilities while embeddings require the embeddings capability. The OpenAI-compatible adapter builds /embeddings requests and serializes embeddings-specific usage. Rebased onto current main (readiness, contracts, Anthropic/Gemini adapters, per-gateway-key rate limiting); embeddings now enforces the same rate limits and threads env into budget evaluation for parity with chat. Bump 0.1.5 -> 0.1.6. --- package.json | 2 +- src/budget.ts | 4 +- src/config.ts | 15 ++- src/gateway.ts | 165 ++++++++++++++++++++++++++++- src/index.ts | 8 +- src/providers/index.ts | 2 +- src/providers/openai-compatible.ts | 56 +++++++++- src/router.ts | 88 ++++++++++----- src/server.ts | 99 ++++++++++++++++- src/types.ts | 38 ++++++- src/usage.ts | 9 +- src/version.ts | 2 +- tests/contracts.test.ts | 2 +- tests/gateway.test.ts | 128 +++++++++++++++++++++- tests/helpers.ts | 20 ++++ tests/provider.test.ts | 45 ++++++++ tests/router.test.ts | 57 ++++++++++ tests/server.test.ts | 95 +++++++++++++++++ tests/usage.test.ts | 11 +- 19 files changed, 794 insertions(+), 52 deletions(-) diff --git a/package.json b/package.json index aae8b58..ff4b350 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "@hasna/gateway", - "version": "0.1.5", + "version": "0.1.6", "description": "Open-source AI model gateway core for one-key, multi-provider routing across Hasna apps and self-hosted deployments", "type": "module", "main": "dist/index.js", diff --git a/src/budget.ts b/src/budget.ts index 7f16831..f8bd4db 100644 --- a/src/budget.ts +++ b/src/budget.ts @@ -1,5 +1,5 @@ import { createHash } from "node:crypto"; -import type { GatewayBudgetConfig, GatewayConfig, GatewayUsage, OpenAIChatCompletionRequest } from "./types"; +import type { GatewayBudgetConfig, GatewayConfig, GatewayRoutableRequest, GatewayUsage } from "./types"; import { GatewayHttpError } from "./errors"; import { hasUsageLedgerBackend, readBudgetLedgerRecords } from "./storage"; @@ -64,7 +64,7 @@ export function fingerprintGatewayKey(value: string | undefined): string | undef } export function budgetContextFromRequest( - request: OpenAIChatCompletionRequest, + request: GatewayRoutableRequest, base: GatewayBudgetContext = {}, ): GatewayBudgetContext { return { diff --git a/src/config.ts b/src/config.ts index 702caa8..7bb9f6c 100644 --- a/src/config.ts +++ b/src/config.ts @@ -792,12 +792,17 @@ function validateProductionRouteReadiness(config: GatewayConfig, env: Record { const model = route.modelAliases?.[0] ?? route.id; - try { - resolveRoute({ config, env }, routeProbeRequest(model)); - return false; - } catch { - return true; + // A route is ready if a keyed provider can satisfy it for either chat or + // embeddings, so embeddings-only routes are not probed as chat requests. + for (const operation of ["chat", "embeddings"] as const) { + try { + resolveRoute({ config, env }, routeProbeRequest(model), { operation }); + return false; + } catch { + // Try the next operation before marking the route unavailable. + } } + return true; }) .map((route) => route.id); diff --git a/src/gateway.ts b/src/gateway.ts index 70f76ff..a5e3315 100644 --- a/src/gateway.ts +++ b/src/gateway.ts @@ -15,10 +15,12 @@ import { transformOpenAICompatibleStream } from "./streaming"; import type { GatewayRouteCandidate, GatewayRouteDecision, + GatewayRoutableRequest, GatewayRuntimeOptions, OpenAIChatCompletionRequest, + OpenAIEmbeddingsRequest, } from "./types"; -import { estimateCostUsd, normalizeUsage, toOpenAIUsage } from "./usage"; +import { estimateCostUsd, normalizeUsage, toOpenAIEmbeddingsUsage, toOpenAIUsage } from "./usage"; type CompletionResult = { body: Record; @@ -59,7 +61,7 @@ function metadataFor( }; } -function includeGatewayMetadata(options: GatewayRuntimeOptions, request: OpenAIChatCompletionRequest): boolean { +function includeGatewayMetadata(options: GatewayRuntimeOptions, request: GatewayRoutableRequest): boolean { if (request.gateway?.strict_openai_compatibility) return false; return request.gateway?.include_gateway_metadata ?? options.config.server.includeGatewayMetadata; } @@ -267,6 +269,31 @@ async function callProvider( }); } +async function callEmbeddingProvider( + options: GatewayRuntimeOptions, + request: OpenAIEmbeddingsRequest, + candidate: GatewayRouteCandidate, +): Promise { + const adapter = adapterForProvider(candidate.provider); + if (!adapter.embed) { + throw new GatewayHttpError({ + status: 400, + type: "gateway_config_error", + code: "provider_embeddings_unsupported", + message: `Provider ${candidate.provider.id} does not support embeddings requests.`, + provider: candidate.provider.id, + }); + } + return adapter.embed({ + provider: candidate.provider, + model: candidate.model, + request, + apiKey: apiKeyFor(candidate, options.env ?? process.env), + timeoutMs: options.config.server.requestTimeoutMs, + fetchImpl: options.fetchImpl, + }); +} + async function openProviderStream( options: GatewayRuntimeOptions, request: OpenAIChatCompletionRequest, @@ -448,6 +475,140 @@ export async function createChatCompletion( }); } +export async function createEmbeddings( + options: GatewayRuntimeOptions, + request: OpenAIEmbeddingsRequest, +): Promise { + const env = options.env ?? process.env; + const requestBudgetContext = budgetContextFromRequest(request, options.budgetContext); + await assertBudgetPreflight(options.config, requestBudgetContext, { env }); + const route = resolveRoute(options, request, { operation: "embeddings" }); + const maxAttempts = Math.min(options.config.server.maxFallbackAttempts, route.candidates.length); + let lastError: GatewayHttpError | undefined; + + for (const candidate of route.candidates.slice(0, maxAttempts)) { + const started = Date.now(); + const budgetContext = { ...requestBudgetContext, selectedModel: candidate.model.id }; + try { + await assertBudgetPreflight(options.config, budgetContext, { env }); + const response = await callEmbeddingProvider(options, request, candidate); + const latencyMs = Date.now() - started; + + if (!response.ok) { + const error = await providerErrorFromResponse(candidate, response); + route.decision.attempts.push({ + provider: candidate.provider.id, + model: candidate.model.id, + providerModel: candidate.model.providerModel, + status: "failed", + reason: error.message, + errorType: error.type, + errorCode: error.code, + retryable: error.retryable, + latencyMs, + }); + lastError = error; + if (error.retryable) continue; + throw error; + } + + route.decision.selected = candidate.model.id; + route.decision.attempts.push({ + provider: candidate.provider.id, + model: candidate.model.id, + providerModel: candidate.model.providerModel, + status: "selected", + latencyMs, + }); + + const providerJson = await parseProviderJson(response); + const rawUsage = providerJson.usage; + if (rawUsage === undefined && options.rateLimit?.requiresStreamingUsage === true) { + throw new GatewayHttpError({ + status: 429, + type: "gateway_rate_limit_error", + code: "gateway_token_usage_missing", + message: "Provider response did not include usage required to enforce a token rate limit.", + raw: { context: budgetContext }, + }); + } + const usage = normalizeUsage(rawUsage); + await options.rateLimit?.onUsage?.(usage); + const estimatedCostUsd = estimateCostUsd(usage, candidate.model); + const budgets = await evaluateBudgetPostflight( + options.config, + budgetContext, + spendFromUsage(usage, estimatedCostUsd), + { env }, + ); + const body: Record = { + ...providerJson, + object: providerJson.object ?? "list", + model: candidate.model.id, + usage: toOpenAIEmbeddingsUsage(usage), + }; + + if (includeGatewayMetadata(options, request)) { + body.gateway = metadataFor(candidate, route.decision, estimatedCostUsd, budgets); + } + + await appendUsageLedgerBestEffort({ + config: options.config, + provider: candidate.provider, + model: candidate.model, + decision: route.decision, + context: budgetContext, + usage, + estimatedCostUsd, + budgets, + status: "success", + }); + + assertBudgetPostflight(budgets); + return { body, status: 200, decision: route.decision }; + } catch (error) { + const gatewayError = + error instanceof GatewayHttpError + ? error + : new GatewayHttpError({ + status: 502, + type: "provider_unavailable", + code: "provider_fetch_failed", + message: error instanceof Error ? error.message : "Provider fetch failed.", + retryable: true, + provider: candidate.provider.id, + }); + + lastError = gatewayError; + if (!route.decision.attempts.some((attempt) => attempt.provider === candidate.provider.id && attempt.status === "failed")) { + route.decision.attempts.push({ + provider: candidate.provider.id, + model: candidate.model.id, + providerModel: candidate.model.providerModel, + status: "failed", + reason: gatewayError.message, + errorType: gatewayError.type, + errorCode: gatewayError.code, + retryable: gatewayError.retryable, + latencyMs: Date.now() - started, + }); + } + + if (gatewayError.retryable) continue; + throw gatewayError; + } + } + + throw new GatewayHttpError({ + status: lastError?.status ?? 502, + type: lastError?.type ?? "gateway_routing_error", + code: lastError?.code ?? "all_routes_failed", + message: lastError?.message ?? `All route attempts failed for model '${request.model}'.`, + retryable: false, + raw: route.decision, + }); +} + export async function createChatCompletionStream( options: GatewayRuntimeOptions, request: OpenAIChatCompletionRequest, diff --git a/src/index.ts b/src/index.ts index e1066fa..12fbc3a 100644 --- a/src/index.ts +++ b/src/index.ts @@ -6,7 +6,7 @@ export { validateRuntimeSecrets, } from "./config"; export { GatewayHttpError, gatewayErrorResponse, jsonError } from "./errors"; -export { createChatCompletion, createChatCompletionStream } from "./gateway"; +export { createChatCompletion, createChatCompletionStream, createEmbeddings } from "./gateway"; export { appendUsageLedger } from "./ledger"; export { toCapabilityCard, toCapabilityCards, toCostEstimate, toDecisionEnvelope } from "./lib/contracts"; export { @@ -53,6 +53,7 @@ export type { GatewayRateLimitConfig, GatewayRequestOptions, GatewayResponseCacheConfig, + GatewayRoutableRequest, GatewayRouteAttempt, GatewayRouteCandidate, GatewayRouteDecision, @@ -65,8 +66,13 @@ export type { GatewayServerConfigInput, GatewayUsage, OpenAIChatCompletionRequest, + OpenAIEmbeddingsInput, + OpenAIEmbeddingsRequest, + OpenAIEmbeddingsUsage, OpenAIUsage, ProviderAdapter, ProviderBuildInput, + ProviderBuildBaseInput, + ProviderEmbeddingsBuildInput, ProviderHttpRequest, } from "./types"; diff --git a/src/providers/index.ts b/src/providers/index.ts index bdc5d9d..2a605ab 100644 --- a/src/providers/index.ts +++ b/src/providers/index.ts @@ -30,4 +30,4 @@ export { toOpenAIChatCompletionResponse, } from "./anthropic"; export { GoogleGeminiAdapter, googleGeminiOpenAIBaseUrl } from "./google-gemini"; -export { OpenAICompatibleAdapter, toProviderChatBody } from "./openai-compatible"; +export { OpenAICompatibleAdapter, toProviderChatBody, toProviderEmbeddingsBody } from "./openai-compatible"; diff --git a/src/providers/openai-compatible.ts b/src/providers/openai-compatible.ts index 40f25db..efbbf76 100644 --- a/src/providers/openai-compatible.ts +++ b/src/providers/openai-compatible.ts @@ -5,8 +5,10 @@ import type { GatewayProviderConfig, GatewayProviderError, OpenAIChatCompletionRequest, + OpenAIEmbeddingsRequest, ProviderAdapter, ProviderBuildInput, + ProviderEmbeddingsBuildInput, ProviderHttpRequest, } from "../types"; @@ -57,6 +59,14 @@ const openRouterProviderFields = new Set([ const vercelGatewayFields = new Set(["models", "order", "only", "caching", "providerTimeouts"]); +const embeddingsForwardedFields = new Set([ + "model", + "input", + "encoding_format", + "dimensions", + "user", +]); + function joinUrl(baseUrl: string, path: string): string { return `${baseUrl.replace(/\/+$/, "")}/${path.replace(/^\/+/, "")}`; } @@ -241,6 +251,18 @@ export function toProviderChatBody( return body; } +export function toProviderEmbeddingsBody(request: OpenAIEmbeddingsRequest, providerModel: string): Record { + const body: Record = {}; + for (const [key, value] of Object.entries(request)) { + if (embeddingsForwardedFields.has(key) && value !== undefined) { + body[key] = value; + } + } + + body.model = providerModel; + return body; +} + function createAbortSignal(timeoutMs: number, signal?: AbortSignal): AbortSignal { if (signal) return signal; return AbortSignal.timeout(timeoutMs); @@ -249,7 +271,7 @@ function createAbortSignal(timeoutMs: number, signal?: AbortSignal): AbortSignal export class OpenAICompatibleAdapter implements ProviderAdapter { readonly id = "openai-compatible"; readonly kind = "openai-compatible"; - readonly supports: GatewayModelCapability[] = ["chat", "streaming", "tools", "json"]; + readonly supports: GatewayModelCapability[] = ["chat", "streaming", "tools", "json", "embeddings"]; buildRequest(input: ProviderBuildInput): ProviderHttpRequest { const baseUrl = providerBaseUrl(input.provider, input.env); @@ -277,6 +299,33 @@ export class OpenAICompatibleAdapter implements ProviderAdapter { }; } + buildEmbeddingsRequest(input: ProviderEmbeddingsBuildInput): ProviderHttpRequest { + if (!input.provider.baseUrl) { + throw new Error(`Provider ${input.provider.id} does not define a baseUrl.`); + } + + const body = toProviderEmbeddingsBody(input.request, input.model.providerModel); + + return { + url: joinUrl(input.provider.baseUrl, "/embeddings"), + init: { + method: "POST", + headers: { + "content-type": "application/json", + authorization: `Bearer ${input.apiKey}`, + ...(input.provider.id === "openrouter" + ? { + "http-referer": "https://github.com/hasna/open-gateway", + "x-title": "Hasna Gateway", + } + : {}), + }, + body: JSON.stringify(body), + signal: createAbortSignal(input.timeoutMs, input.signal), + }, + }; + } + send(input: ProviderBuildInput): Promise { const request = this.buildRequest({ ...input, @@ -299,6 +348,11 @@ export class OpenAICompatibleAdapter implements ProviderAdapter { return (input.fetchImpl ?? fetch)(request.url, request.init); } + embed(input: ProviderEmbeddingsBuildInput): Promise { + const request = this.buildEmbeddingsRequest(input); + return (input.fetchImpl ?? fetch)(request.url, request.init); + } + mapError(response: Response, bodyText?: string): GatewayProviderError { const mapped = mapProviderStatus(response.status); return { diff --git a/src/router.ts b/src/router.ts index b4e3462..acee574 100644 --- a/src/router.ts +++ b/src/router.ts @@ -11,14 +11,22 @@ import type { GatewayModelCapability, GatewayModelConfig, GatewayProviderConfig, + GatewayRoutableRequest, GatewayRouteCandidate, GatewayRouteDecision, GatewayRoutePolicy, GatewayRouteScore, GatewayRuntimeOptions, OpenAIChatCompletionRequest, + OpenAIEmbeddingsRequest, } from "./types"; +type GatewayRouteOperation = "chat" | "embeddings"; + +type ResolveRouteOptions = { + operation?: GatewayRouteOperation; +}; + type EffectivePolicy = { allowedProviders?: string[]; blockedProviders?: string[]; @@ -96,7 +104,7 @@ function strictMax( function mergePolicy( config: GatewayConfig, route: GatewayRoutePolicy | undefined, - request: OpenAIChatCompletionRequest, + request: GatewayRoutableRequest, ): EffectivePolicy { const requestPolicy = request.gateway; const configPolicy = config.policy; @@ -168,7 +176,7 @@ function mergePolicy( }; } -function routeMode(route: GatewayRoutePolicy | undefined, request: OpenAIChatCompletionRequest): GatewayRoutePolicy["mode"] { +function routeMode(route: GatewayRoutePolicy | undefined, request: GatewayRoutableRequest): GatewayRoutePolicy["mode"] { return request.gateway?.routing ?? route?.mode ?? "fallback"; } @@ -192,13 +200,15 @@ function dynamicCapabilitiesForProvider(provider: GatewayProviderConfig): Gatewa return ["chat", "streaming"]; } -function dynamicCandidate(config: GatewayConfig, id: string): GatewayRouteCandidate | undefined { +function dynamicCandidate(config: GatewayConfig, id: string, operation: GatewayRouteOperation): GatewayRouteCandidate | undefined { const slash = id.indexOf("/"); if (slash <= 0) return undefined; const providerId = id.slice(0, slash); const providerModel = id.slice(slash + 1); const provider = providerMap(config).get(providerId); if (!provider) return undefined; + const capabilities: GatewayModelCapability[] = + operation === "embeddings" ? ["embeddings"] : dynamicCapabilitiesForProvider(provider); return { provider, @@ -207,12 +217,16 @@ function dynamicCandidate(config: GatewayConfig, id: string): GatewayRouteCandid providerId, providerModel, aliases: [], - capabilities: dynamicCapabilitiesForProvider(provider), + capabilities, }, }; } -function initialCandidates(config: GatewayConfig, request: OpenAIChatCompletionRequest): GatewayRouteCandidate[] { +function initialCandidates( + config: GatewayConfig, + request: GatewayRoutableRequest, + operation: GatewayRouteOperation, +): GatewayRouteCandidate[] { const providers = providerMap(config); const models = modelMap(config); const route = routeForRequest(config, request.model); @@ -223,7 +237,7 @@ function initialCandidates(config: GatewayConfig, request: OpenAIChatCompletionR return provider ? [{ model: explicit, provider }] : []; } - const dynamic = dynamicCandidate(config, request.model); + const dynamic = dynamicCandidate(config, request.model, operation); if (dynamic) return [dynamic]; if (route?.fallbackModelIds?.length) { @@ -285,9 +299,10 @@ function hasAllowedRegion(provider: GatewayProviderConfig, policy: EffectivePoli function candidateSkipReason( candidate: GatewayRouteCandidate, - request: OpenAIChatCompletionRequest, + request: GatewayRoutableRequest, policy: EffectivePolicy, env: Record, + operation: GatewayRouteOperation, ): string | undefined { const { model, provider } = candidate; @@ -311,10 +326,15 @@ function candidateSkipReason( if (missingHeaderEnvs.length > 0) { return `provider required header env ${missingHeaderEnvs.join(", ")} is not set`; } - if (!model.capabilities.includes("chat")) return "model does not support chat"; - if (request.stream && !model.capabilities.includes("streaming")) return "model does not support streaming"; - if (request.tools && request.tools.length > 0 && !model.capabilities.includes("tools")) { - return "model does not support tools"; + if (operation === "embeddings") { + if (!model.capabilities.includes("embeddings")) return "model does not support embeddings"; + } else { + const chatRequest = request as OpenAIChatCompletionRequest; + if (!model.capabilities.includes("chat")) return "model does not support chat"; + if (chatRequest.stream && !model.capabilities.includes("streaming")) return "model does not support streaming"; + if (chatRequest.tools && chatRequest.tools.length > 0 && !model.capabilities.includes("tools")) { + return "model does not support tools"; + } } if (request.response_format && !model.capabilities.includes("json")) return "model does not support json output"; for (const capability of request.gateway?.required_capabilities ?? []) { @@ -369,22 +389,31 @@ function configuredTokenPrice(candidate: GatewayRouteCandidate): number { return candidate.model.inputUsdPerMillionTokens! + candidate.model.outputUsdPerMillionTokens!; } -function estimateInputTokens(request: OpenAIChatCompletionRequest): number { +function estimateInputTokens(request: GatewayRoutableRequest): number { if (request.gateway?.expected_input_tokens !== undefined) return request.gateway.expected_input_tokens; - const chars = request.messages.reduce((sum, message) => { - if (typeof message.content === "string") return sum + message.content.length; - if (Array.isArray(message.content)) return sum + JSON.stringify(message.content).length; - return sum; - }, 0); + const messages = (request as OpenAIChatCompletionRequest).messages; + let chars = 0; + if (Array.isArray(messages)) { + chars = messages.reduce((sum, message) => { + if (typeof message.content === "string") return sum + message.content.length; + if (Array.isArray(message.content)) return sum + JSON.stringify(message.content).length; + return sum; + }, 0); + } else { + // Embeddings and other non-chat requests carry their content in `input`. + const input = (request as OpenAIEmbeddingsRequest).input; + if (typeof input === "string") chars = input.length; + else if (Array.isArray(input)) chars = JSON.stringify(input).length; + } return Math.max(1, Math.ceil(chars / 4)); } -function estimateOutputTokens(request: OpenAIChatCompletionRequest): number { +function estimateOutputTokens(request: GatewayRoutableRequest): number { const maxTokens = request.max_completion_tokens ?? request.max_tokens; return typeof maxTokens === "number" && maxTokens > 0 ? maxTokens : 512; } -function estimatedRequestCost(candidate: GatewayRouteCandidate, request: OpenAIChatCompletionRequest): number | undefined { +function estimatedRequestCost(candidate: GatewayRouteCandidate, request: GatewayRoutableRequest): number | undefined { if (!candidateHasConfiguredPrice(candidate)) return undefined; return ( (estimateInputTokens(request) / 1_000_000) * candidate.model.inputUsdPerMillionTokens! + @@ -437,13 +466,13 @@ function hashString(value: string): number { return hash >>> 0; } -function stickyTieBreaker(candidate: GatewayRouteCandidate, request: OpenAIChatCompletionRequest): number { +function stickyTieBreaker(candidate: GatewayRouteCandidate, request: GatewayRoutableRequest): number { const sessionId = request.gateway?.sticky_session_id ?? request.gateway?.session_id ?? request.session_id; if (!sessionId) return 0; return hashString(`${sessionId}:${candidate.model.id}`) / 0xffffffff; } -function providerOrderScore(candidate: GatewayRouteCandidate, request: OpenAIChatCompletionRequest): number | undefined { +function providerOrderScore(candidate: GatewayRouteCandidate, request: GatewayRoutableRequest): number | undefined { const order = request.gateway?.provider_order; if (!order?.length) return undefined; const index = order.indexOf(candidate.provider.id); @@ -453,7 +482,7 @@ function providerOrderScore(candidate: GatewayRouteCandidate, request: OpenAICha function weightsForMode( mode: GatewayRoutePolicy["mode"], - request: OpenAIChatCompletionRequest, + request: GatewayRoutableRequest, ): Record<"cost" | "quality" | "latency" | "success" | "throughput" | "providerOrder", number> { if (mode === "lowest-latency") { return { cost: 0.1, quality: 0.1, latency: 0.55, success: 0.2, throughput: 0, providerOrder: 0.05 }; @@ -487,7 +516,7 @@ function weightsForMode( function scoreCandidates( candidates: GatewayRouteCandidate[], mode: GatewayRoutePolicy["mode"], - request: OpenAIChatCompletionRequest, + request: GatewayRoutableRequest, ): GatewayRouteScore[] { const costs = candidates.map((candidate) => estimatedRequestCost(candidate, request)); const latencies = candidates.map((candidate) => candidate.model.averageLatencyMs); @@ -540,7 +569,7 @@ function originalIndexMap(candidates: GatewayRouteCandidate[]): Map @@ -595,12 +624,17 @@ function policyForDecision(policy: EffectivePolicy): GatewayRouteDecision["polic }; } -export function resolveRoute(options: GatewayRuntimeOptions, request: OpenAIChatCompletionRequest): ResolveResult { +export function resolveRoute( + options: GatewayRuntimeOptions, + request: GatewayRoutableRequest, + resolveOptions: ResolveRouteOptions = {}, +): ResolveResult { + const operation = resolveOptions.operation ?? "chat"; const route = routeForRequest(options.config, request.model); const mode = routeMode(route, request); const policy = mergePolicy(options.config, route, request); const env = options.env ?? process.env; - const candidates = initialCandidates(options.config, request); + const candidates = initialCandidates(options.config, request, operation); const decision: GatewayRouteDecision = { requested_model: request.model, resolved_candidates: unique(candidates.map((candidate) => candidate.model.id)), @@ -612,7 +646,7 @@ export function resolveRoute(options: GatewayRuntimeOptions, request: OpenAIChat const eligible: GatewayRouteCandidate[] = []; for (const candidate of candidates) { - const reason = candidateSkipReason(candidate, request, policy, env); + const reason = candidateSkipReason(candidate, request, policy, env, operation); if (reason) { decision.attempts.push({ provider: candidate.provider.id, diff --git a/src/server.ts b/src/server.ts index 6c4c450..060e274 100644 --- a/src/server.ts +++ b/src/server.ts @@ -1,9 +1,16 @@ import { gatewayErrorResponse, GatewayHttpError, jsonError } from "./errors"; import { fingerprintGatewayKey } from "./budget"; import { validateRuntimeSecrets } from "./config"; -import { createChatCompletion, createChatCompletionStream } from "./gateway"; +import { createChatCompletion, createChatCompletionStream, createEmbeddings } from "./gateway"; import { GatewayKeyRateLimiter, gatewayRateLimitKey, type GatewayRateLimitExceeded } from "./rate-limit"; -import type { GatewayConfig, GatewayFetch, GatewayRuntimeOptions, OpenAIChatCompletionRequest } from "./types"; +import type { + GatewayConfig, + GatewayFetch, + GatewayRuntimeOptions, + OpenAIChatCompletionRequest, + OpenAIEmbeddingsInput, + OpenAIEmbeddingsRequest, +} from "./types"; import { gatewayVersion } from "./version"; type ServerOptions = { @@ -150,6 +157,69 @@ function responseCacheBypassRequested(request: Request, config: GatewayConfig): return normalized === "" || !["0", "false", "no", "off"].includes(normalized); } +function isToken(value: unknown): value is number { + return typeof value === "number" && Number.isInteger(value) && value >= 0; +} + +function isTokenArray(value: unknown): value is number[] { + return Array.isArray(value) && value.length > 0 && value.every(isToken); +} + +function isEmbeddingsInput(value: unknown): value is OpenAIEmbeddingsInput { + if (typeof value === "string") return true; + if (!Array.isArray(value) || value.length === 0) return false; + if (value.every((item) => typeof item === "string")) return true; + if (value.every(isToken)) return true; + return value.every(isTokenArray); +} + +function validateEmbeddingsRequest(body: unknown): OpenAIEmbeddingsRequest { + if (!body || typeof body !== "object" || Array.isArray(body)) { + throw new GatewayHttpError({ + status: 400, + type: "gateway_bad_request", + code: "invalid_request", + message: "Embeddings request body must be an object.", + }); + } + + const request = body as OpenAIEmbeddingsRequest; + if (typeof request.model !== "string" || request.model.length === 0) { + throw new GatewayHttpError({ + status: 400, + type: "gateway_bad_request", + code: "missing_model", + message: "Embeddings request requires a model string.", + }); + } + if (!isEmbeddingsInput(request.input)) { + throw new GatewayHttpError({ + status: 400, + type: "gateway_bad_request", + code: "missing_input", + message: "Embeddings request requires input as a string, string array, token array, or token array array.", + }); + } + if (request.encoding_format !== undefined && typeof request.encoding_format !== "string") { + throw new GatewayHttpError({ + status: 400, + type: "gateway_bad_request", + code: "invalid_encoding_format", + message: "Embeddings request encoding_format must be a string when provided.", + }); + } + if (request.dimensions !== undefined && (!Number.isInteger(request.dimensions) || request.dimensions <= 0)) { + throw new GatewayHttpError({ + status: 400, + type: "gateway_bad_request", + code: "invalid_dimensions", + message: "Embeddings request dimensions must be a positive integer when provided.", + }); + } + + return request; +} + function modelsResponse(config: GatewayConfig): Record { const aliasIds = new Set(); for (const model of config.models) { @@ -328,6 +398,31 @@ export function createGatewayHandler(options: ServerOptions): (request: Request) return json(request, options.config, result.body, result.status); } + if (request.method === "POST" && url.pathname === "/v1/embeddings") { + const body = validateEmbeddingsRequest(await parseJsonBody(request, options.config.server.maxRequestBodyBytes)); + const gatewayKeyFingerprint = fingerprintGatewayKey(bearerToken(request)); + const rateLimitKey = gatewayRateLimitKey(gatewayKeyFingerprint); + const rateLimitConfig = options.config.server.rateLimits?.perGatewayKey; + const rateLimitCheck = keyRateLimiter.checkAndConsumeRequest(rateLimitKey, rateLimitConfig); + if (!rateLimitCheck.allowed) { + return rateLimitResponse(request, options.config, rateLimitCheck.exceeded); + } + + const runtimeWithRequestContext: GatewayRuntimeOptions = { + ...runtime, + budgetContext: { + gatewayKey: gatewayKeyFingerprint, + tenant: request.headers.get("x-gateway-tenant") ?? undefined, + }, + rateLimit: { + onUsage: (usage) => keyRateLimiter.recordUsage(rateLimitKey, rateLimitConfig, usage), + requiresStreamingUsage: rateLimitConfig?.tokensPerMinute !== undefined, + }, + }; + const result = await createEmbeddings(runtimeWithRequestContext, body); + return json(request, options.config, result.body, result.status); + } + return jsonError(404, "Endpoint not found.", "gateway_routing_error", "not_found"); } catch (error) { return gatewayErrorResponse(error); diff --git a/src/types.ts b/src/types.ts index dfaaafb..c7c05be 100644 --- a/src/types.ts +++ b/src/types.ts @@ -271,8 +271,14 @@ export type GatewayRequestOptions = { strict_openai_compatibility?: boolean; }; -export type OpenAIChatCompletionRequest = { +export type GatewayRoutableRequest = { model: string; + user?: string; + gateway?: GatewayRequestOptions; + [key: string]: unknown; +}; + +export type OpenAIChatCompletionRequest = GatewayRoutableRequest & { messages: ChatMessage[]; stream?: boolean; tools?: unknown[]; @@ -298,8 +304,6 @@ export type OpenAIChatCompletionRequest = { n?: number; presence_penalty?: number; frequency_penalty?: number; - user?: string; - gateway?: GatewayRequestOptions; provider_options?: Record; providerOptions?: Record; provider?: unknown; @@ -308,6 +312,15 @@ export type OpenAIChatCompletionRequest = { [key: string]: unknown; }; +export type OpenAIEmbeddingsInput = string | string[] | number[] | number[][]; + +export type OpenAIEmbeddingsRequest = GatewayRoutableRequest & { + input: OpenAIEmbeddingsInput; + encoding_format?: "float" | "base64" | string; + dimensions?: number; + provider_options?: Record; +}; + export type GatewayUsage = { inputTokens: number; outputTokens: number; @@ -332,6 +345,12 @@ export type OpenAIUsage = { [key: string]: unknown; }; +export type OpenAIEmbeddingsUsage = { + prompt_tokens: number; + total_tokens: number; + [key: string]: unknown; +}; + export type GatewayRouteAttempt = { provider: string; model: string; @@ -391,15 +410,16 @@ export type ProviderAdapter = { kind: GatewayProviderKind; supports: GatewayModelCapability[]; buildRequest(input: ProviderBuildInput): ProviderHttpRequest; + buildEmbeddingsRequest?(input: ProviderEmbeddingsBuildInput): ProviderHttpRequest; send(input: ProviderBuildInput): Promise; stream(input: ProviderBuildInput): Promise; + embed?(input: ProviderEmbeddingsBuildInput): Promise; mapError(response: Response, bodyText?: string): GatewayProviderError; }; -export type ProviderBuildInput = { +export type ProviderBuildBaseInput = { provider: GatewayProviderConfig; model: GatewayModelConfig; - request: OpenAIChatCompletionRequest; apiKey: string; timeoutMs: number; env?: Record; @@ -407,6 +427,14 @@ export type ProviderBuildInput = { signal?: AbortSignal; }; +export type ProviderBuildInput = ProviderBuildBaseInput & { + request: OpenAIChatCompletionRequest; +}; + +export type ProviderEmbeddingsBuildInput = ProviderBuildBaseInput & { + request: OpenAIEmbeddingsRequest; +}; + export type GatewayProviderError = { message: string; status: number; diff --git a/src/usage.ts b/src/usage.ts index 8c7768a..90bce2d 100644 --- a/src/usage.ts +++ b/src/usage.ts @@ -1,4 +1,4 @@ -import type { GatewayModelConfig, GatewayUsage, OpenAIUsage } from "./types"; +import type { GatewayModelConfig, GatewayUsage, OpenAIEmbeddingsUsage, OpenAIUsage } from "./types"; function numberFrom(value: unknown): number | undefined { return typeof value === "number" && Number.isFinite(value) ? value : undefined; @@ -69,6 +69,13 @@ export function toOpenAIUsage(usage: GatewayUsage): OpenAIUsage { }; } +export function toOpenAIEmbeddingsUsage(usage: GatewayUsage): OpenAIEmbeddingsUsage { + return { + prompt_tokens: usage.inputTokens, + total_tokens: usage.totalTokens, + }; +} + export function estimateCostUsd(usage: GatewayUsage, model: GatewayModelConfig): number | undefined { const inputPrice = model.inputUsdPerMillionTokens; const outputPrice = model.outputUsdPerMillionTokens; diff --git a/src/version.ts b/src/version.ts index 72f1803..68ff9c9 100644 --- a/src/version.ts +++ b/src/version.ts @@ -1 +1 @@ -export const gatewayVersion = "0.1.5"; +export const gatewayVersion = "0.1.6"; diff --git a/tests/contracts.test.ts b/tests/contracts.test.ts index d0cd5b3..ead82f2 100644 --- a/tests/contracts.test.ts +++ b/tests/contracts.test.ts @@ -121,7 +121,7 @@ describe("contract adapters", () => { expect(decision.selected[0]?.externalId).toBe("openai/gpt-4.1-mini"); const cards = toCapabilityCards(config, { createdAt: "2026-07-06T00:00:00.000Z" }); - expect(cards).toHaveLength(2); + expect(cards).toHaveLength(3); expect(cards[0]?.schema).toBe(SCHEMA_IDS.capabilityCard); expect(cards[0]?.kind).toBe("model"); expect(cards[0]?.metadata?.sourcePackage).toBe("@hasna/gateway"); diff --git a/tests/gateway.test.ts b/tests/gateway.test.ts index 2f75e77..3d36929 100644 --- a/tests/gateway.test.ts +++ b/tests/gateway.test.ts @@ -1,6 +1,12 @@ +import { unlink } from "node:fs/promises"; import { describe, expect, test } from "bun:test"; import { GatewayHttpError } from "../src/errors"; -import { createChatCompletion, createChatCompletionStream, providerErrorMessageFromBody } from "../src/gateway"; +import { + createChatCompletion, + createChatCompletionStream, + createEmbeddings, + providerErrorMessageFromBody, +} from "../src/gateway"; import { testConfig, jsonResponse } from "./helpers"; const env = { @@ -859,4 +865,124 @@ describe("chat completion lifecycle", () => { }, ); }); + + test("normalizes embeddings response with gateway metadata and ledger usage", async () => { + const path = `/tmp/hasna-gateway-embeddings-ledger-${crypto.randomUUID()}.jsonl`; + const config = testConfig(); + config.storage.usageLedgerPath = path; + config.budgets = [ + { + id: "embedding-request", + window: "per-request", + mode: "hard", + scope: { modelAlias: "embeddings" }, + maxTotalTokens: 10, + }, + ]; + + const calls: string[] = []; + const fetchImpl = async (url: string | URL | Request, init?: RequestInit): Promise => { + calls.push(String(url)); + const providerBody = JSON.parse(String(init?.body)); + expect(providerBody).toEqual({ + model: "text-embedding-3-small", + input: "do not write embedding input", + }); + return jsonResponse({ + object: "list", + data: [{ object: "embedding", index: 0, embedding: [0.1, 0.2] }], + model: "text-embedding-3-small", + usage: { + prompt_tokens: 4, + total_tokens: 4, + }, + }); + }; + + const result = await createEmbeddings( + { + config, + env: { + GATEWAY_API_KEY: "gateway", + OPENAI_API_KEY: "openai", + }, + fetchImpl, + }, + { + model: "embeddings", + input: "do not write embedding input", + }, + ); + + expect(calls).toEqual(["https://api.openai.test/v1/embeddings"]); + expect(result.body.object).toBe("list"); + expect(result.body.model).toBe("openai/text-embedding-3-small"); + expect(result.body.usage).toEqual({ + prompt_tokens: 4, + total_tokens: 4, + }); + expect((result.body.gateway as Record).provider).toBe("openai"); + expect(((result.body.gateway as Record).budgets as Array<{ remaining: { totalTokens: number } }>)[0]?.remaining.totalTokens).toBe(6); + + const ledgerText = await Bun.file(path).text(); + const record = JSON.parse(ledgerText.trim()); + expect(record.provider).toBe("openai"); + expect(record.model).toBe("openai/text-embedding-3-small"); + expect(record.usage).toEqual({ + inputTokens: 4, + outputTokens: 0, + totalTokens: 4, + }); + expect(ledgerText).not.toContain("do not write embedding input"); + await unlink(path); + }); + + test("rejects over-budget embeddings responses after recording usage", async () => { + const path = `/tmp/hasna-gateway-embeddings-ledger-${crypto.randomUUID()}.jsonl`; + const config = testConfig(); + config.storage.usageLedgerPath = path; + config.budgets = [ + { + id: "tiny-embedding-request", + window: "per-request", + mode: "hard", + scope: { modelAlias: "embeddings" }, + maxTotalTokens: 1, + }, + ]; + + let thrown: unknown; + try { + await createEmbeddings( + { + config, + env: { + GATEWAY_API_KEY: "gateway", + OPENAI_API_KEY: "openai", + }, + fetchImpl: async () => + jsonResponse({ + object: "list", + data: [{ object: "embedding", index: 0, embedding: [0.1, 0.2] }], + usage: { + prompt_tokens: 2, + total_tokens: 2, + }, + }), + }, + { + model: "embeddings", + input: "too many tokens", + }, + ); + } catch (error) { + thrown = error; + } + + expect(thrown).toBeInstanceOf(GatewayHttpError); + expect(thrown).toMatchObject({ status: 402, code: "budget_exceeded" }); + const ledgerText = await Bun.file(path).text(); + expect(ledgerText).toContain('"totalTokens":2'); + await unlink(path); + }); }); diff --git a/tests/helpers.ts b/tests/helpers.ts index b0ec2af..356f486 100644 --- a/tests/helpers.ts +++ b/tests/helpers.ts @@ -70,6 +70,15 @@ export function testConfig(): GatewayConfig { aliases: ["coding", "china-coding"], capabilities: ["chat", "streaming", "tools"], }, + { + id: "openai/text-embedding-3-small", + providerId: "openai", + providerModel: "text-embedding-3-small", + aliases: ["embeddings"], + capabilities: ["embeddings"], + inputUsdPerMillionTokens: 0.02, + outputUsdPerMillionTokens: 0, + }, ], routes: [ { @@ -95,6 +104,17 @@ export function testConfig(): GatewayConfig { allowedRegions: ["cn", "us"], }, }, + { + id: "embeddings", + mode: "fallback", + modelAliases: ["embeddings"], + fallbackModelIds: ["openai/text-embedding-3-small"], + dataPolicy: { + allowTraining: false, + allowChineseProviders: false, + blockedRegions: ["cn"], + }, + }, ], }); } diff --git a/tests/provider.test.ts b/tests/provider.test.ts index 51ed779..d6d549d 100644 --- a/tests/provider.test.ts +++ b/tests/provider.test.ts @@ -7,6 +7,7 @@ import { toAnthropicMessagesBody, toOpenAIChatCompletionResponse, toProviderChatBody, + toProviderEmbeddingsBody, } from "../src/providers"; import { testConfig } from "./helpers"; @@ -218,6 +219,50 @@ describe("OpenAI-compatible provider adapter", () => { expect(body.provider_options).toBeUndefined(); expect((body.providerOptions as Record>).gateway.byok).toBeUndefined(); }); + + test("strips gateway-only fields from embeddings requests", () => { + const body = toProviderEmbeddingsBody( + { + model: "embeddings", + input: ["alpha", "beta"], + encoding_format: "float", + dimensions: 256, + user: "tenant-user", + gateway: { routing: "fallback" }, + provider_options: { ignored: true }, + }, + "text-embedding-3-small", + ); + + expect(body).toEqual({ + model: "text-embedding-3-small", + input: ["alpha", "beta"], + encoding_format: "float", + dimensions: 256, + user: "tenant-user", + }); + }); + + test("builds embeddings bearer request", () => { + const config = testConfig(); + const provider = config.providers.find((candidate) => candidate.id === "openai")!; + const model = config.models.find((candidate) => candidate.id === "openai/text-embedding-3-small")!; + const adapter = new OpenAICompatibleAdapter(); + const request = adapter.buildEmbeddingsRequest({ + provider, + model, + request: { + model: "embeddings", + input: "hello", + }, + apiKey: "secret", + timeoutMs: 1000, + }); + + expect(request.url).toBe("https://api.openai.test/v1/embeddings"); + expect((request.init.headers as Record).authorization).toBe("Bearer secret"); + expect(JSON.parse(String(request.init.body)).model).toBe("text-embedding-3-small"); + }); }); describe("Anthropic Messages provider adapter", () => { diff --git a/tests/router.test.ts b/tests/router.test.ts index 8826724..3ac2c12 100644 --- a/tests/router.test.ts +++ b/tests/router.test.ts @@ -668,4 +668,61 @@ describe("routing policy", () => { expect(result.decision.selected).toBe("openai/first"); expect(result.decision.scores?.length).toBe(2); }); + + test("routes embeddings only to embeddings-capable models", () => { + const config = testConfig(); + config.routes.push({ + id: "mixed-embeddings", + mode: "fallback", + modelAliases: ["mixed-embeddings"], + fallbackModelIds: ["openai/gpt-4.1-mini", "openai/text-embedding-3-small"], + dataPolicy: { + allowTraining: false, + allowChineseProviders: false, + blockedRegions: ["cn"], + }, + }); + + const result = resolveRoute( + { + config, + env: { + GATEWAY_API_KEY: "gateway", + OPENAI_API_KEY: "openai", + }, + }, + { + model: "mixed-embeddings", + input: "hello", + }, + { operation: "embeddings" }, + ); + + expect(result.decision.selected).toBe("openai/text-embedding-3-small"); + expect(result.decision.attempts[0]?.status).toBe("skipped"); + expect(result.decision.attempts[0]?.reason).toBe("model does not support embeddings"); + }); + + test("keeps chat routing scoped to chat-capable models", () => { + const config = testConfig(); + config.routes[0] = { + ...config.routes[0]!, + fallbackModelIds: ["openai/text-embedding-3-small", "openai/gpt-4.1-mini"], + }; + + const result = resolveRoute( + { + config, + env: { + GATEWAY_API_KEY: "gateway", + OPENAI_API_KEY: "openai", + }, + }, + request, + ); + + expect(result.decision.selected).toBe("openai/gpt-4.1-mini"); + expect(result.decision.attempts[0]?.status).toBe("skipped"); + expect(result.decision.attempts[0]?.reason).toBe("model does not support chat"); + }); }); diff --git a/tests/server.test.ts b/tests/server.test.ts index 7ec996b..8afbf47 100644 --- a/tests/server.test.ts +++ b/tests/server.test.ts @@ -746,6 +746,101 @@ describe("HTTP server handler", () => { const text = await response.text(); expect(text).toContain("data: [DONE]"); }); + + test("handles embeddings", async () => { + const handler = createGatewayHandler({ + config: testConfig(), + env: { GATEWAY_API_KEY: "gateway", OPENAI_API_KEY: "openai" }, + fetchImpl: async (url, init) => { + expect(String(url)).toBe("https://api.openai.test/v1/embeddings"); + expect(JSON.parse(String(init?.body))).toMatchObject({ + model: "text-embedding-3-small", + input: ["alpha", "beta"], + }); + return jsonResponse({ + object: "list", + data: [ + { object: "embedding", index: 0, embedding: [0.1] }, + { object: "embedding", index: 1, embedding: [0.2] }, + ], + usage: { prompt_tokens: 4, total_tokens: 4 }, + }); + }, + }); + const response = await handler( + new Request("http://localhost/v1/embeddings", { + method: "POST", + headers: { + authorization: "Bearer gateway", + "content-type": "application/json", + }, + body: JSON.stringify({ model: "embeddings", input: ["alpha", "beta"] }), + }), + ); + const body = await response.json(); + expect(response.status).toBe(200); + expect(body.object).toBe("list"); + expect(body.model).toBe("openai/text-embedding-3-small"); + expect(body.usage).toEqual({ prompt_tokens: 4, total_tokens: 4 }); + expect(body.gateway.provider).toBe("openai"); + }); + + test("rejects malformed embeddings requests", async () => { + const handler = createGatewayHandler({ + config: testConfig(), + env: { GATEWAY_API_KEY: "gateway", OPENAI_API_KEY: "openai" }, + }); + const response = await handler( + new Request("http://localhost/v1/embeddings", { + method: "POST", + headers: { + authorization: "Bearer gateway", + "content-type": "application/json", + }, + body: JSON.stringify({ model: "embeddings" }), + }), + ); + const body = await response.json(); + expect(response.status).toBe(400); + expect(body.error.code).toBe("missing_input"); + }); + + test("enforces per-gateway-key RPM on embeddings before provider fetch", async () => { + const config = testConfig(); + config.server.rateLimits = { perGatewayKey: { requestsPerMinute: 1 } }; + let providerCalls = 0; + const handler = createGatewayHandler({ + config, + env: { GATEWAY_API_KEY: "gateway", OPENAI_API_KEY: "openai" }, + fetchImpl: async () => { + providerCalls += 1; + return jsonResponse({ + object: "list", + data: [{ object: "embedding", index: 0, embedding: [0.1] }], + usage: { prompt_tokens: 4, total_tokens: 4 }, + }); + }, + }); + + const embeddingsRequest = () => + new Request("http://localhost/v1/embeddings", { + method: "POST", + headers: { + authorization: "Bearer gateway", + "content-type": "application/json", + }, + body: JSON.stringify({ model: "embeddings", input: "hello" }), + }); + + const allowed = await handler(embeddingsRequest()); + const rejected = await handler(embeddingsRequest()); + const body = await rejected.json(); + + expect(allowed.status).toBe(200); + expect(rejected.status).toBe(429); + expect(body.error.type).toBe("gateway_rate_limit_error"); + expect(providerCalls).toBe(1); + }); }); function chatRequest(token: string, overrides: Record = {}): Request { diff --git a/tests/usage.test.ts b/tests/usage.test.ts index 8eb89cb..67be4a7 100644 --- a/tests/usage.test.ts +++ b/tests/usage.test.ts @@ -1,5 +1,5 @@ import { describe, expect, test } from "bun:test"; -import { estimateCostUsd, normalizeUsage, toOpenAIUsage } from "../src/usage"; +import { estimateCostUsd, normalizeUsage, toOpenAIEmbeddingsUsage, toOpenAIUsage } from "../src/usage"; describe("usage normalization", () => { test("normalizes OpenAI usage details", () => { @@ -101,6 +101,15 @@ describe("usage normalization", () => { ).toBe(0.000111); }); + test("serializes embeddings usage without chat completion fields", () => { + const usage = normalizeUsage({ prompt_tokens: 12, total_tokens: 12 }); + expect(toOpenAIEmbeddingsUsage(usage)).toEqual({ + prompt_tokens: 12, + total_tokens: 12, + }); + expect(toOpenAIEmbeddingsUsage(usage)).not.toHaveProperty("completion_tokens"); + }); + test("returns unknown cost when a used token side has no configured price", () => { const usage = normalizeUsage({ prompt_tokens: 1, completion_tokens: 1_000_000, total_tokens: 1_000_001 }); expect(