@neutrome/open-ai-router 0.1.0

package/src/router/execute.ts

@@ -0,0 +1,735 @@
+import {
+  emitChatCompletionsResponse,
+  emitChatCompletionsStreamChunk,
+  emitProviderResponse,
+  emitProviderRequest,
+  emitProviderStreamChunk,
+  getModel,
+  isStreaming,
+  parseChatCompletionsRequest,
+  parseProviderRequest,
+  parseProviderResponse,
+  parseProviderStreamChunk,
+  setModel,
+  usageObject,
+  type Program,
+} from "@neutrome/lil-engine";
+import {
+  createExecutionRuntime,
+  ExecutionError,
+  type ExecutionRuntime,
+  type ExecutionRuntimeOptions,
+  type Executor,
+  type ProgramTransform,
+  type ProviderInvoker,
+} from "@neutrome/lil-engine";
+import type { RequestContext, TargetAuthContext } from "../auth/types.ts";
+import { cloneForwardHeaders, isSse } from "../util/headers.ts";
+import { eventDataLines, splitSseEvents } from "../util/sse.ts";
+import { resolveInvocationTarget, type ResolvedTarget } from "./resolve.ts";
+import type { RouterRuntime } from "./runtime.ts";
+
+const encoder = new TextEncoder();
+const decoder = new TextDecoder();
+const DEFAULT_UPSTREAM_TIMEOUT_MS = 120_000;
+
+export type UsageReport = {
+  model: string;
+  tokensInput: number;
+  tokensOutput: number;
+  cachedTokensInput: number;
+};
+
+export type RouterExecutionOptions = Pick<
+  ExecutionRuntimeOptions,
+  "observe" | "requestIdFactory" | "executionIdFactory" | "resolveExecutor"
+> & {
+  executors?: Readonly<Record<string, Executor>>;
+  transforms?: Readonly<Record<string, ProgramTransform>>;
+  upstreamTimeoutMs?: number;
+  onUsage?: (report: UsageReport) => void;
+};
+
+export type HandleChatCompletionsOptions = RouterExecutionOptions & {
+  runtime: RouterRuntime;
+  ctx: RequestContext;
+};
+
+export type HandleResponsesOptions = RouterExecutionOptions & {
+  runtime: RouterRuntime;
+  ctx: RequestContext;
+};
+
+export class RouterError extends Error {
+  constructor(
+    message: string,
+    readonly status: number,
+    readonly code: string,
+    readonly type = "invalid_request_error",
+  ) {
+    super(message);
+  }
+}
+
+export class ProviderHttpError extends Error {
+  constructor(
+    message: string,
+    readonly status: number,
+    readonly body: Uint8Array,
+    readonly contentType: string,
+  ) {
+    super(message);
+  }
+}
+
+export function createRouterExecutionRuntime(
+  runtime: RouterRuntime,
+  ctx: RequestContext,
+  options: RouterExecutionOptions = {},
+): ExecutionRuntime {
+  const runtimeOptions: ExecutionRuntimeOptions = {
+    signal: ctx.request.signal,
+    resolveTarget(request) {
+      return resolveRequestTarget(runtime, request).target;
+    },
+    providerInvoker: createFetchProviderInvoker(
+      runtime,
+      ctx,
+      options.upstreamTimeoutMs,
+    ),
+  };
+
+  if (options.executors) {
+    runtimeOptions.executors = options.executors;
+  }
+  if (options.resolveExecutor) {
+    runtimeOptions.resolveExecutor = options.resolveExecutor;
+  }
+  if (options.transforms) {
+    runtimeOptions.transforms = options.transforms;
+  }
+  if (options.observe) {
+    runtimeOptions.observe = options.observe;
+  }
+  if (options.requestIdFactory) {
+    runtimeOptions.requestIdFactory = options.requestIdFactory;
+  }
+  if (options.executionIdFactory) {
+    runtimeOptions.executionIdFactory = options.executionIdFactory;
+  }
+
+  return createExecutionRuntime(runtimeOptions);
+}
+
+export async function handleChatCompletions(
+  options: HandleChatCompletionsOptions,
+): Promise<Response> {
+  try {
+    const requestBody = new Uint8Array(await options.ctx.request.arrayBuffer());
+    const request = parseChatCompletionsRequest(requestBody);
+    const resolved = resolveRequestTarget(options.runtime, request);
+    const execution = createRouterExecutionRuntime(
+      options.runtime,
+      options.ctx,
+      options,
+    );
+
+    const startTime = Date.now();
+    if (isStreaming(request)) {
+      const userWantsUsage = requestedStreamUsage(requestBody);
+      const chunks = execution.stream(request, { target: resolved.target });
+      return await streamExecutionResponse(
+        "chat-completions",
+        chunks,
+        resolved,
+        startTime,
+        options.onUsage,
+        userWantsUsage,
+      );
+    }
+
+    const response = await execution.execute(request, {
+      target: resolved.target,
+    });
+    const userTime = Date.now() - startTime;
+    reportUsage(response, resolved, options.onUsage);
+    return jsonExecutionResponse(
+      emitChatCompletionsResponse(response),
+      resolved,
+      userTime,
+    );
+  } catch (error) {
+    return errorResponse(error);
+  }
+}
+
+export async function handleResponses(
+  options: HandleResponsesOptions,
+): Promise<Response> {
+  try {
+    const requestBody = new Uint8Array(await options.ctx.request.arrayBuffer());
+    const request = parseProviderRequest("responses", requestBody);
+    const resolved = resolveRequestTarget(options.runtime, request);
+    const execution = createRouterExecutionRuntime(
+      options.runtime,
+      options.ctx,
+      options,
+    );
+
+    const startTime = Date.now();
+    if (isStreaming(request)) {
+      const chunks = execution.stream(request, { target: resolved.target });
+      return await streamExecutionResponse(
+        "responses",
+        chunks,
+        resolved,
+        startTime,
+        options.onUsage,
+      );
+    }
+
+    const response = await execution.execute(request, {
+      target: resolved.target,
+    });
+    const userTime = Date.now() - startTime;
+    reportUsage(response, resolved, options.onUsage);
+    return jsonExecutionResponse(
+      emitProviderResponse("responses", response),
+      resolved,
+      userTime,
+    );
+  } catch (error) {
+    return errorResponse(error);
+  }
+}
+
+export function resolveRequestTarget(
+  runtime: RouterRuntime,
+  request: Program,
+): ResolvedTarget {
+  const model = getModel(request);
+  if (!model) {
+    throw new RouterError("Request is missing `model`", 400, "missing_model");
+  }
+
+  try {
+    return resolveInvocationTarget(runtime, model);
+  } catch (error) {
+    throw new RouterError(
+      error instanceof Error ? error.message : "Model resolution failed",
+      404,
+      "model_not_found",
+    );
+  }
+}
+
+export function createFetchProviderInvoker(
+  runtime: RouterRuntime,
+  ctx: RequestContext,
+  upstreamTimeoutMs = DEFAULT_UPSTREAM_TIMEOUT_MS,
+): ProviderInvoker {
+  return {
+    async execute(request, providerCtx) {
+      const provider = getProviderOrThrow(runtime, providerCtx.target.provider, providerCtx.target.model);
+      const timed = createUpstreamSignal(providerCtx.signal, upstreamTimeoutMs);
+
+      try {
+        const upstream = await fetchProvider(
+          runtime,
+          ctx,
+          provider.name,
+          provider.allowAnonymous,
+          provider.headers,
+          provider.apiBaseUrl,
+          provider.endpointPath,
+          provider.style,
+          setModel(request, providerCtx.target.model),
+          timed.signal,
+        );
+        if (!upstream.ok) {
+          throw await toProviderError(provider.name, upstream);
+        }
+
+        const bytes = new Uint8Array(await upstream.arrayBuffer());
+        return parseProviderResponse(provider.style, bytes);
+      } catch (error) {
+        throw normalizeUpstreamAbort(error, timed.signal);
+      } finally {
+        timed.cleanup();
+      }
+    },
+
+    async *stream(request, providerCtx) {
+      const provider = getProviderOrThrow(runtime, providerCtx.target.provider, providerCtx.target.model);
+      const timed = createUpstreamSignal(providerCtx.signal, upstreamTimeoutMs);
+      let reader: ReadableStreamDefaultReader<Uint8Array> | null = null;
+      let buffer = "";
+
+      try {
+        const upstream = await fetchProvider(
+          runtime,
+          ctx,
+          provider.name,
+          provider.allowAnonymous,
+          provider.headers,
+          provider.apiBaseUrl,
+          provider.endpointPath,
+          provider.style,
+          setModel(request, providerCtx.target.model),
+          timed.signal,
+        );
+        if (!upstream.ok) {
+          throw await toProviderError(provider.name, upstream);
+        }
+        if (!isSse(upstream.headers) || !upstream.body) {
+          throw new RouterError(
+            `Provider ${provider.name} did not return an event stream`,
+            502,
+            "invalid_upstream_stream",
+          );
+        }
+
+        reader = upstream.body.getReader();
+        while (true) {
+          const { done, value } = await reader.read();
+          if (done) {
+            break;
+          }
+          buffer += decoder.decode(value, { stream: true });
+          const split = splitSseEvents(buffer);
+          buffer = split.remainder;
+          for (const event of split.events) {
+            for (const data of eventDataLines(event)) {
+              if (data === "[DONE]") {
+                return;
+              }
+              if (data[0] !== "{") continue;
+              yield parseProviderStreamChunk(
+                provider.style,
+                encoder.encode(data),
+              );
+            }
+          }
+        }
+
+        if (buffer.length > 0) {
+          for (const data of eventDataLines(buffer)) {
+            if (data === "[DONE]") {
+              return;
+            }
+            if (data[0] !== "{") continue;
+            yield parseProviderStreamChunk(
+              provider.style,
+              encoder.encode(data),
+            );
+          }
+        }
+      } catch (error) {
+        throw normalizeUpstreamAbort(error, timed.signal);
+      } finally {
+        await reader?.cancel().catch(() => {});
+        reader?.releaseLock();
+        timed.cleanup();
+      }
+    },
+  };
+}
+
+async function streamExecutionResponse(
+  responseStyle: "chat-completions" | "responses",
+  chunks: AsyncIterable<Program>,
+  resolved: ResolvedTarget,
+  startTime: number,
+  onUsage?: (report: UsageReport) => void,
+  emitUsageToClient = true,
+): Promise<Response> {
+  const iterator = chunks[Symbol.asyncIterator]();
+
+  const first = await iterator.next();
+  const userTime = Date.now() - startTime;
+
+  const stream = new ReadableStream<Uint8Array>({
+    async start(controller) {
+      try {
+        let lastChunkWithUsage: Program | null = null;
+
+        function emitChunk(chunk: Program): void {
+          const hasUsage = usageObject(chunk) !== undefined;
+          if (hasUsage) lastChunkWithUsage = chunk;
+          if (hasUsage && !emitUsageToClient) return;
+          const bytes =
+            responseStyle === "chat-completions"
+              ? emitChatCompletionsStreamChunk(chunk)
+              : emitProviderStreamChunk("responses", chunk);
+          controller.enqueue(
+            encoder.encode(`data: ${decoder.decode(bytes)}\n\n`),
+          );
+        }
+
+        if (!first.done) {
+          emitChunk(first.value);
+        }
+
+        while (true) {
+          const { done, value } = await iterator.next();
+          if (done) {
+            break;
+          }
+          emitChunk(value);
+        }
+        if (responseStyle === "chat-completions") {
+          controller.enqueue(encoder.encode("data: [DONE]\n\n"));
+        }
+        if (lastChunkWithUsage) {
+          reportUsage(lastChunkWithUsage, resolved, onUsage);
+        }
+        controller.close();
+      } catch (error) {
+        controller.enqueue(
+          encoder.encode(
+            `event: error\ndata: ${JSON.stringify({ message: error instanceof Error ? error.message : String(error) })}\n\n`,
+          ),
+        );
+        controller.close();
+      }
+    },
+    async cancel() {
+      await iterator.return?.();
+    },
+  });
+
+  return new Response(stream, {
+    headers: {
+      "content-type": "text/event-stream; charset=utf-8",
+      ...routingHeaders(resolved),
+      ...timingHeaders(userTime),
+    },
+  });
+}
+
+function jsonExecutionResponse(
+  body: Uint8Array,
+  resolved: ResolvedTarget,
+  userTime: number,
+): Response {
+  return new Response(toBody(body), {
+    headers: {
+      "content-type": "application/json; charset=utf-8",
+      ...routingHeaders(resolved),
+      ...timingHeaders(userTime),
+    },
+  });
+}
+
+function routingHeaders(resolved: ResolvedTarget): Record<string, string> {
+  if (resolved.target.kind === "provider") {
+    return {
+      "x-openairouter-source": resolved.source,
+      "x-openairouter-namespace": resolved.namespace,
+      "x-openairouter-provider-id": resolved.target.provider ?? "",
+      "x-openairouter-model-id": resolved.target.model,
+    };
+  }
+
+  return {
+    "x-openairouter-engine": "open-ai-router",
+    "x-openairouter-source": resolved.source,
+    "x-openairouter-namespace": resolved.namespace,
+    "x-openairouter-executor-id": resolved.target.executor,
+    "x-openairouter-model-id": resolved.target.alias,
+  };
+}
+
+function timingHeaders(userTimeMs: number): Record<string, string> {
+  return {
+    "x-openairouter-user-time": String(userTimeMs),
+  };
+}
+
+async function fetchProvider(
+  runtime: RouterRuntime,
+  ctx: RequestContext,
+  providerName: string,
+  allowAnonymous: boolean,
+  providerHeaders: Readonly<Record<string, string>>,
+  apiBaseUrl: string,
+  endpointPath: string,
+  providerStyle: RouterRuntime["providers"] extends ReadonlyMap<string, infer T>
+    ? T extends { style: infer S }
+      ? S
+      : never
+    : never,
+  request: Program,
+  signal: AbortSignal,
+): Promise<Response> {
+  const headers = cloneForwardHeaders(ctx.request.headers);
+  headers.set("content-type", "application/json");
+  for (const [key, value] of Object.entries(providerHeaders)) {
+    headers.set(key, value);
+  }
+
+  let hasAuth = false;
+  for (const driver of runtime.authChain) {
+    const targetCtx: TargetAuthContext = { ...ctx, providerName };
+    const auth = await driver.collectTarget(targetCtx);
+    if (!auth) {
+      continue;
+    }
+    new Headers(auth.headers).forEach((value, key) => headers.set(key, value));
+    hasAuth = true;
+    break;
+  }
+
+  if (!hasAuth && !allowAnonymous) {
+    throw new RouterError(
+      `No auth available for provider ${providerName}`,
+      401,
+      "provider_auth_missing",
+    );
+  }
+
+  let body = emitProviderRequest(providerStyle, request);
+  if (providerStyle === "chat-completions" && isStreaming(request)) {
+    body = injectStreamUsage(body);
+  }
+  return runtime.fetchImpl(`${apiBaseUrl}${endpointPath}`, {
+    method: "POST",
+    headers,
+    body: toBody(body),
+    signal,
+  });
+}
+
+function getProviderOrThrow(runtime: RouterRuntime, providerName: string | undefined, model?: string) {
+  if (providerName) {
+    const provider = runtime.providers.get(providerName);
+    if (!provider) {
+      throw new RouterError(
+        `Provider ${providerName} is not configured`,
+        500,
+        "provider_not_configured",
+      );
+    }
+    return provider;
+  }
+
+  for (const name of runtime.providerOrder) {
+    const provider = runtime.providers.get(name)!;
+    if (!provider.noPrefix) continue;
+    if (provider.exportsMatcher(model ?? "")) {
+      return provider;
+    }
+  }
+
+  throw new RouterError(
+    `No provider found for model \`${model ?? "(unknown)"}\``,
+    404,
+    "provider_not_found",
+  );
+}
+
+async function toProviderError(
+  providerName: string,
+  response: Response,
+): Promise<ProviderHttpError> {
+  const body = new Uint8Array(await response.arrayBuffer());
+  return new ProviderHttpError(
+    `Provider ${providerName} returned ${response.status}`,
+    response.status,
+    body,
+    response.headers.get("content-type") ?? "application/json; charset=utf-8",
+  );
+}
+
+function errorResponse(error: unknown): Response {
+  if (error instanceof ProviderHttpError) {
+    return new Response(toBody(error.body), {
+      status: error.status,
+      headers: {
+        "content-type": error.contentType,
+      },
+    });
+  }
+
+  if (error instanceof RouterError) {
+    return new Response(
+      JSON.stringify({
+        error: {
+          message: error.message,
+          type: error.type,
+          param: null,
+          code: error.code,
+        },
+      }),
+      {
+        status: error.status,
+        headers: { "content-type": "application/json; charset=utf-8" },
+      },
+    );
+  }
+
+  if (error instanceof ExecutionError) {
+    const providerError = unwrapCause(error.cause, ProviderHttpError);
+    if (providerError) {
+      return errorResponse(providerError);
+    }
+    const nestedRouterError = unwrapCause(error.cause, RouterError);
+    if (nestedRouterError) {
+      return errorResponse(nestedRouterError);
+    }
+    return errorResponse(routerErrorFromExecutionError(error));
+  }
+
+  console.error("[open-ai-router] unhandled error:", error);
+  return new Response(
+    JSON.stringify({
+      error: {
+        message: "Internal router error",
+        type: "invalid_request_error",
+        param: null,
+        code: "internal_error",
+      },
+    }),
+    {
+      status: 500,
+      headers: { "content-type": "application/json; charset=utf-8" },
+    },
+  );
+}
+
+function routerErrorFromExecutionError(error: ExecutionError): RouterError {
+  switch (error.kind) {
+    case "validation":
+      return new RouterError(error.message, 400, "invalid_program");
+    case "resolution":
+      return new RouterError(error.message, 400, "target_resolution_failed");
+    case "cancellation":
+      return new RouterError(error.message, 499, "request_cancelled");
+    case "provider":
+      return new RouterError(error.message, 502, "provider_execution_failed");
+    case "transform":
+      return new RouterError(error.message, 500, "transform_failed");
+    case "executor":
+      return new RouterError(error.message, 500, "executor_failed");
+    case "stream":
+      return new RouterError(error.message, 500, "stream_failed");
+    case "recursion_limit":
+      return new RouterError(error.message, 500, "recursion_limit");
+  }
+}
+
+function unwrapCause<T extends Error>(
+  error: unknown,
+  ctor: new (...args: any[]) => T,
+): T | null {
+  let current = error;
+  while (current && typeof current === "object") {
+    if (current instanceof ctor) {
+      return current;
+    }
+    if (!("cause" in current)) {
+      return null;
+    }
+    current = (current as { cause?: unknown }).cause;
+  }
+  return null;
+}
+
+function toBody(bytes: Uint8Array): ArrayBuffer {
+  return bytes.buffer.slice(
+    bytes.byteOffset,
+    bytes.byteOffset + bytes.byteLength,
+  ) as ArrayBuffer;
+}
+
+function createUpstreamSignal(
+  signal: AbortSignal,
+  timeoutMs: number,
+): {
+  signal: AbortSignal;
+  cleanup(): void;
+} {
+  if (timeoutMs <= 0) {
+    return {
+      signal,
+      cleanup() {},
+    };
+  }
+
+  const controller = new AbortController();
+  const onAbort = () => controller.abort(signal.reason);
+  if (signal.aborted) {
+    onAbort();
+  } else {
+    signal.addEventListener("abort", onAbort, { once: true });
+  }
+
+  const timer = setTimeout(() => {
+    controller.abort(
+      new RouterError(
+        `Upstream provider timed out after ${timeoutMs}ms`,
+        504,
+        "upstream_timeout",
+      ),
+    );
+  }, timeoutMs);
+
+  return {
+    signal: controller.signal,
+    cleanup() {
+      clearTimeout(timer);
+      signal.removeEventListener("abort", onAbort);
+    },
+  };
+}
+
+function reportUsage(
+  response: Program,
+  resolved: ResolvedTarget,
+  onUsage?: (report: UsageReport) => void,
+): void {
+  if (!onUsage) return;
+  const usage = usageObject(response);
+  if (!usage || typeof usage !== "object") return;
+  const u = usage as Record<string, unknown>;
+  const tokensInput = asNum(u.prompt_tokens) || asNum(u.input_tokens);
+  const tokensOutput = asNum(u.completion_tokens) || asNum(u.output_tokens);
+  if (tokensInput === 0 && tokensOutput === 0) return;
+  const details = u.prompt_tokens_details as Record<string, unknown> | undefined;
+  const cachedTokensInput = details ? asNum(details.cached_tokens) : 0;
+  const model = resolved.target.kind === "provider" ? resolved.target.model : resolved.target.alias;
+  try {
+    onUsage({ model, tokensInput, tokensOutput, cachedTokensInput });
+  } catch {}
+}
+
+function asNum(v: unknown): number {
+  return typeof v === "number" && v >= 0 ? v : 0;
+}
+
+function requestedStreamUsage(body: Uint8Array): boolean {
+  try {
+    const obj = JSON.parse(decoder.decode(body)) as Record<string, unknown>;
+    const opts = obj.stream_options as Record<string, unknown> | undefined;
+    return opts?.include_usage === true;
+  } catch {
+    return false;
+  }
+}
+
+function injectStreamUsage(body: Uint8Array): Uint8Array {
+  try {
+    const obj = JSON.parse(decoder.decode(body)) as Record<string, unknown>;
+    obj.stream_options = { include_usage: true };
+    return encoder.encode(JSON.stringify(obj));
+  } catch {
+    return body;
+  }
+}
+
+function normalizeUpstreamAbort(error: unknown, signal: AbortSignal): unknown {
+  if (signal.aborted && signal.reason instanceof RouterError) {
+    return signal.reason;
+  }
+  return error;
+}