code:
/* This file is fucking horrendous, sorry */
import { Request, Response } from "express";
import * as http from "http";
import util from "util";
import zlib from "zlib";
import { enqueue, trackWaitTime } from "../../queue";
import { HttpError } from "../../../shared/errors";
import { keyPool } from "../../../shared/key-management";
import { getOpenAIModelFamily } from "../../../shared/models";
import { countTokens } from "../../../shared/tokenization";
import {
  incrementPromptCount,
  incrementTokenCount,
} from "../../../shared/users/user-store";
import { assertNever } from "../../../shared/utils";
import { refundLastAttempt } from "../../rate-limit";
import {
  getCompletionFromBody,
  isImageGenerationRequest,
  isTextGenerationRequest,
  sendProxyError,
} from "../common";
import { handleStreamedResponse } from "./handle-streamed-response";
import { logPrompt } from "./log-prompt";
import { saveImage } from "./save-image";

const DECODER_MAP = {
  gzip: util.promisify(zlib.gunzip),
  deflate: util.promisify(zlib.inflate),
  br: util.promisify(zlib.brotliDecompress),
};

const isSupportedContentEncoding = (
  contentEncoding: string
): contentEncoding is keyof typeof DECODER_MAP => {
  return contentEncoding in DECODER_MAP;
};

export class RetryableError extends Error {
  constructor(message: string) {
    super(message);
    this.name = "RetryableError";
  }
}

/**
 * Either decodes or streams the entire response body and then passes it as the
 * last argument to the rest of the middleware stack.
 */
export type RawResponseBodyHandler = (
  proxyRes: http.IncomingMessage,
  req: Request,
  res: Response
) => Promise<string | Record<string, any>>;
export type ProxyResHandlerWithBody = (
  proxyRes: http.IncomingMessage,
  req: Request,
  res: Response,
  /**
   * This will be an object if the response content-type is application/json,
   * or if the response is a streaming response. Otherwise it will be a string.
   */
  body: string | Record<string, any>
) => Promise<void>;
export type ProxyResMiddleware = ProxyResHandlerWithBody[];

/**
 * Returns a on.proxyRes handler that executes the given middleware stack after
 * the common proxy response handlers have processed the response and decoded
 * the body.  Custom middleware won't execute if the response is determined to
 * be an error from the upstream service as the response will be taken over by
 * the common error handler.
 *
 * For streaming responses, the handleStream middleware will block remaining
 * middleware from executing as it consumes the stream and forwards events to
 * the client. Once the stream is closed, the finalized body will be attached
 * to res.body and the remaining middleware will execute.
 */
export const createOnProxyResHandler = (apiMiddleware: ProxyResMiddleware) => {
  return async (
    proxyRes: http.IncomingMessage,
    req: Request,
    res: Response
  ) => {
    const initialHandler = req.isStreaming
      ? handleStreamedResponse
      : decodeResponseBody;

    let lastMiddleware = initialHandler.name;

    try {
      const body = await initialHandler(proxyRes, req, res);

      const middlewareStack: ProxyResMiddleware = [];

      if (req.isStreaming) {
        // `handleStreamedResponse` writes to the response and ends it, so
        // we can only execute middleware that doesn't write to the response.
        middlewareStack.push(
          trackRateLimit,
          countResponseTokens,
          incrementUsage,
          logPrompt
        );
      } else {
        middlewareStack.push(
          trackRateLimit,
          handleUpstreamErrors,
          countResponseTokens,
          incrementUsage,
          copyHttpHeaders,
          saveImage,
          logPrompt,
          ...apiMiddleware
        );
      }

      for (const middleware of middlewareStack) {
        lastMiddleware = middleware.name;
        await middleware(proxyRes, req, res, body);
      }

      trackWaitTime(req);
    } catch (error) {
      // Hack: if the error is a retryable rate-limit error, the request has
      // been re-enqueued and we can just return without doing anything else.
      if (error instanceof RetryableError) {
        return;
      }

      // Already logged and responded to the client by handleUpstreamErrors
      if (error instanceof HttpError) {
        if (!res.writableEnded) res.end();
        return;
      }

      const { stack, message } = error;
      const info = { stack, lastMiddleware, key: req.key?.hash };
      const description = `Error while executing proxy response middleware: ${lastMiddleware} (${message})`;

      if (res.headersSent) {
        req.log.error(info, description);
        if (!res.writableEnded) res.end();
        return;
      } else {
        req.log.error(info, description);
        res
          .status(500)
          .json({ error: "Internal server error", proxy_note: description });
      }
    }
  };
};

async function reenqueueRequest(req: Request) {
  req.log.info(
    { key: req.key?.hash, retryCount: req.retryCount },
    `Re-enqueueing request due to retryable error`
  );
  req.retryCount++;
  await enqueue(req);
}

/**
 * Handles the response from the upstream service and decodes the body if
 * necessary.  If the response is JSON, it will be parsed and returned as an
 * object.  Otherwise, it will be returned as a string.
 * @throws {Error} Unsupported content-encoding or invalid application/json body
 */
export const decodeResponseBody: RawResponseBodyHandler = async (
  proxyRes,
  req,
  res
) => {
  if (req.isStreaming) {
    const err = new Error("decodeResponseBody called for a streaming request.");
    req.log.error({ stack: err.stack, api: req.inboundApi }, err.message);
    throw err;
  }

  return new Promise<string>((resolve, reject) => {
    let chunks: Buffer[] = [];
    proxyRes.on("data", (chunk) => chunks.push(chunk));
    proxyRes.on("end", async () => {
      let body = Buffer.concat(chunks);

      const contentEncoding = proxyRes.headers["content-encoding"];
      if (contentEncoding) {
        if (isSupportedContentEncoding(contentEncoding)) {
          const decoder = DECODER_MAP[contentEncoding];
          // @ts-ignore - started failing after upgrading TypeScript, don't care
          // as it was never a problem.
          body = await decoder(body);
        } else {
          const error = `Proxy received response with unsupported content-encoding: ${contentEncoding}`;
          req.log.warn({ contentEncoding, key: req.key?.hash }, error);
          sendProxyError(req, res, 500, "Internal Server Error", {
            error,
            contentEncoding,
          });
          return reject(error);
        }
      }

      try {
        if (proxyRes.headers["content-type"]?.includes("application/json")) {
          const json = JSON.parse(body.toString());
          return resolve(json);
        }
        return resolve(body.toString());
      } catch (e) {
        const msg = `Proxy received response with invalid JSON: ${e.message}`;
        req.log.warn({ error: e.stack, key: req.key?.hash }, msg);
        sendProxyError(req, res, 500, "Internal Server Error", { error: msg });
        return reject(msg);
      }
    });
  });
};

type ProxiedErrorPayload = {
  error?: Record<string, any>;
  message?: string;
  proxy_note?: string;
};

/**
 * Handles non-2xx responses from the upstream service.  If the proxied response
 * is an error, this will respond to the client with an error payload and throw
 * an error to stop the middleware stack.
 * On 429 errors, if request queueing is enabled, the request will be silently
 * re-enqueued.  Otherwise, the request will be rejected with an error payload.
 * @throws {HttpError} On HTTP error status code from upstream service
 */
const handleUpstreamErrors: ProxyResHandlerWithBody = async (
  proxyRes,
  req,
  res,
  body
) => {
  const statusCode = proxyRes.statusCode || 500;
  const statusMessage = proxyRes.statusMessage || "Internal Server Error";

  if (statusCode < 400) {
    return;
  }

  let errorPayload: ProxiedErrorPayload;
  const tryAgainMessage = keyPool.available(req.body?.model)
    ? `There may be more keys available for this model; try again in a few seconds.`
    : "There are no more keys available for this model.";

  try {
    assertJsonResponse(body);
    errorPayload = body;
  } catch (parseError) {
    // Likely Bad Gateway or Gateway Timeout from upstream's reverse proxy
    const hash = req.key?.hash;
    req.log.warn({ statusCode, statusMessage, key: hash }, parseError.message);

    const errorObject = {
      error: parseError.message,
      status: statusCode,
      statusMessage,
      proxy_note: `Proxy got back an error, but it was not in JSON format. This is likely a temporary problem with the upstream service.`,
    };

    sendProxyError(req, res, statusCode, statusMessage, errorObject);
    throw new HttpError(statusCode, parseError.message);
  }

  const errorType =
    errorPayload.error?.code ||
    errorPayload.error?.type ||
    getAwsErrorType(proxyRes.headers["x-amzn-errortype"]);

  req.log.warn(
    { statusCode, type: errorType, errorPayload, key: req.key?.hash },
    `Received error response from upstream. (${proxyRes.statusMessage})`
  );

  const service = req.key!.service;
  if (service === "aws") {
    // Try to standardize the error format for AWS
    errorPayload.error = { message: errorPayload.message, type: errorType };
    delete errorPayload.message;
  }

  if (statusCode === 400) {
    // Bad request. For OpenAI, this is usually due to prompt length.
    // For Anthropic, this is usually due to missing preamble.
    switch (service) {
      case "openai":
      case "google-ai":
      case "mistral-ai":
      case "azure":
        const filteredCodes = ["content_policy_violation", "content_filter"];
        if (filteredCodes.includes(errorPayload.error?.code)) {
          errorPayload.proxy_note = `Request was filtered by the upstream API's content moderation system. Modify your prompt and try again.`;
          refundLastAttempt(req);
        } else if (errorPayload.error?.code === "billing_hard_limit_reached") {
          // For some reason, some models return this 400 error instead of the
          // same 429 billing error that other models return.
          await handleOpenAIRateLimitError(req, tryAgainMessage, errorPayload);
        } else {
          errorPayload.proxy_note = `The upstream API rejected the request. Your prompt may be too long for ${req.body?.model}.`;
        }
        break;
      case "anthropic":
      case "aws":
        await handleAnthropicBadRequestError(req, errorPayload);
        break;
      default:
        assertNever(service);
    }
  } else if (statusCode === 401) {
    // Key is invalid or was revoked
    keyPool.disable(req.key!, "revoked");
    errorPayload.proxy_note = `API key is invalid or revoked. ${tryAgainMessage}`;
  } else if (statusCode === 403) {
    if (service === "anthropic") {
      keyPool.disable(req.key!, "revoked");
      errorPayload.proxy_note = `API key is invalid or revoked. ${tryAgainMessage}`;
      return;
    }
    switch (errorType) {
      case "UnrecognizedClientException":
        // Key is invalid.
        keyPool.disable(req.key!, "revoked");
        errorPayload.proxy_note = `API key is invalid or revoked. ${tryAgainMessage}`;
        break;
      case "AccessDeniedException":
        const isModelAccessError =
          errorPayload.error?.message?.includes(`specified model ID`);
        if (!isModelAccessError) {
          req.log.error(
            { key: req.key?.hash, model: req.body?.model },
            "Disabling key due to AccessDeniedException when invoking model. If credentials are valid, check IAM permissions."
          );
          keyPool.disable(req.key!, "revoked");
        }
        errorPayload.proxy_note = `API key doesn't have access to the requested resource. Model ID: ${req.body?.model}`;
        break;
      default:
        errorPayload.proxy_note = `Received 403 error. Key may be invalid.`;
    }
  } else if (statusCode === 429) {
    switch (service) {
      case "openai":
        await handleOpenAIRateLimitError(req, tryAgainMessage, errorPayload);
        break;
      case "anthropic":
        await handleAnthropicRateLimitError(req, errorPayload);
        break;
      case "aws":
        await handleAwsRateLimitError(req, errorPayload);
        break;
      case "azure":
      case "mistral-ai":
        await handleAzureRateLimitError(req, errorPayload);
        break;
      case "google-ai":
        await handleGoogleAIRateLimitError(req, errorPayload);
        break;
      default:
        assertNever(service);
    }
  } else if (statusCode === 404) {
    // Most likely model not found
    switch (service) {
      case "openai":
        if (errorPayload.error?.code === "model_not_found") {
          const requestedModel = req.body.model;
          const modelFamily = getOpenAIModelFamily(requestedModel);
          errorPayload.proxy_note = `The key assigned to your prompt does not support the requested model (${requestedModel}, family: ${modelFamily}).`;
          req.log.error(
            { key: req.key?.hash, model: requestedModel, modelFamily },
            "Prompt was routed to a key that does not support the requested model."
          );
        }
        break;
      case "anthropic":
        errorPayload.proxy_note = `The requested Claude model might not exist, or the key might not be provisioned for it.`;
        break;
      case "google-ai":
        errorPayload.proxy_note = `The requested Google AI model might not exist, or the key might not be provisioned for it.`;
        break;
      case "mistral-ai":
        errorPayload.proxy_note = `The requested Mistral AI model might not exist, or the key might not be provisioned for it.`;
        break;
      case "aws":
        errorPayload.proxy_note = `The requested AWS resource might not exist, or the key might not have access to it.`;
        break;
      case "azure":
        errorPayload.proxy_note = `The assigned Azure deployment does not support the requested model.`;
        break;
      default:
        assertNever(service);
    }
  } else {
    errorPayload.proxy_note = `Unrecognized error from upstream service.`;
  }

  // Some OAI errors contain the organization ID, which we don't want to reveal.
  if (errorPayload.error?.message) {
    errorPayload.error.message = errorPayload.error.message.replace(
      /org-.{24}/gm,
      "org-xxxxxxxxxxxxxxxxxxx"
    );
  }

  sendProxyError(req, res, statusCode, statusMessage, errorPayload);
  // This is bubbled up to onProxyRes's handler for logging but will not trigger
  // a write to the response as `sendProxyError` has just done that.
  throw new HttpError(statusCode, errorPayload.error?.message);
};

async function handleAnthropicBadRequestError(
  req: Request,
  errorPayload: ProxiedErrorPayload
) {
  const { error } = errorPayload;
  const isMissingPreamble = error?.message.startsWith(
    `prompt must start with "\n\nHuman:" turn`
  );

  // Some keys mandate a \n\nHuman: preamble, which we can add and retry
  if (isMissingPreamble) {
    req.log.warn(
      { key: req.key?.hash },
      "Request failed due to missing preamble. Key will be marked as such for subsequent requests."
    );
    keyPool.update(req.key!, { requiresPreamble: true });
    await reenqueueRequest(req);
    throw new RetryableError("Claude request re-enqueued to add preamble.");
  }

  // {"type":"error","error":{"type":"invalid_request_error","message":"Usage blocked until 2024-03-01T00:00:00+00:00 due to user specified spend limits."}}
  // {"type":"error","error":{"type":"invalid_request_error","message":"Your credit balance is too low to access the Claude API. Please go to Plans & Billing to upgrade or purchase credits."}}
  const isOverQuota =
    error?.message?.match(/usage blocked until/i) ||
    error?.message?.match(/credit balance is too low/i);
  if (isOverQuota) {
    req.log.warn(
      { key: req.key?.hash, message: error?.message },
      "Anthropic key has hit spending limit and will be disabled."
    );
    keyPool.disable(req.key!, "quota");
    errorPayload.proxy_note = `Assigned key has hit its spending limit. ${error?.message}`;
    return;
  }

  const isDisabled = error?.message?.match(/organization has been disabled/i);
  if (isDisabled) {
    req.log.warn(
      { key: req.key?.hash, message: error?.message },
      "Anthropic key has been disabled."
    );
    keyPool.disable(req.key!, "revoked");
    errorPayload.proxy_note = `Assigned key has been disabled. ${error?.message}`;
    return;
  }

  errorPayload.proxy_note = `Unrecognized error from the API. (${error?.message})`;
}

async function handleAnthropicRateLimitError(
  req: Request,
  errorPayload: ProxiedErrorPayload
) {
  if (errorPayload.error?.type === "rate_limit_error") {
    keyPool.markRateLimited(req.key!);
    await reenqueueRequest(req);
    throw new RetryableError("Claude rate-limited request re-enqueued.");
  } else {
    errorPayload.proxy_note = `Unrecognized 429 Too Many Requests error from the API.`;
  }
}

async function handleAwsRateLimitError(
  req: Request,
  errorPayload: ProxiedErrorPayload
) {
  const errorType = errorPayload.error?.type;
  switch (errorType) {
    case "ThrottlingException":
      keyPool.markRateLimited(req.key!);
      await reenqueueRequest(req);
      throw new RetryableError("AWS rate-limited request re-enqueued.");
    case "ModelNotReadyException":
      errorPayload.proxy_note = `The requested model is overloaded. Try again in a few seconds.`;
      break;
    default:
      errorPayload.proxy_note = `Unrecognized rate limit error from AWS. (${errorType})`;
  }
}

async function handleOpenAIRateLimitError(
  req: Request,
  tryAgainMessage: string,
  errorPayload: ProxiedErrorPayload
): Promise<Record<string, any>> {
  const type = errorPayload.error?.type;
  switch (type) {
    case "insufficient_quota":
    case "invalid_request_error": // this is the billing_hard_limit_reached error seen in some cases
      // Billing quota exceeded (key is dead, disable it)
      keyPool.disable(req.key!, "quota");
      errorPayload.proxy_note = `Assigned key's quota has been exceeded. ${tryAgainMessage}`;
      break;
    case "access_terminated":
      // Account banned (key is dead, disable it)
      keyPool.disable(req.key!, "revoked");
      errorPayload.proxy_note = `Assigned key has been banned by OpenAI for policy violations. ${tryAgainMessage}`;
      break;
    case "billing_not_active":
      // Key valid but account billing is delinquent
      keyPool.disable(req.key!, "quota");
      errorPayload.proxy_note = `Assigned key has been disabled due to delinquent billing. ${tryAgainMessage}`;
      break;
    case "requests":
    case "tokens":
      keyPool.markRateLimited(req.key!);
      if (errorPayload.error?.message?.match(/on requests per day/)) {
        // This key has a very low rate limit, so we can't re-enqueue it.
        errorPayload.proxy_note = `Assigned key has reached its per-day request limit for this model. Try another model.`;
        break;
      }

      // Per-minute request or token rate limit is exceeded, which we can retry
      await reenqueueRequest(req);
      throw new RetryableError("Rate-limited request re-enqueued.");
    // WIP/nonfunctional
    // case "tokens_usage_based":
    //   // Weird new rate limit type that seems limited to preview models.
    //   // Distinct from `tokens` type. Can be per-minute or per-day.
    //
    //   // I've seen reports of this error for 500k tokens/day and 10k tokens/min.
    //   // 10k tokens per minute is problematic, because this is much less than
    //   // GPT4-Turbo's max context size for a single prompt and is effectively a
    //   // cap on the max context size for just that key+model, which the app is
    //   // not able to deal with.
    //
    //   // Similarly if there is a 500k tokens per day limit and 450k tokens have
    //   // been used today, the max context for that key becomes 50k tokens until
    //   // the next day and becomes progressively smaller as more tokens are used.
    //
    //   // To work around these keys we will first retry the request a few times.
    //   // After that we will reject the request, and if it's a per-day limit we
    //   // will also disable the key.
    //
    //   // "Rate limit reached for gpt-4-1106-preview in organization org-xxxxxxxxxxxxxxxxxxx on tokens_usage_based per day: Limit 500000, Used 460000, Requested 50000"
    //   // "Rate limit reached for gpt-4-1106-preview in organization org-xxxxxxxxxxxxxxxxxxx on tokens_usage_based per min: Limit 10000, Requested 40000"
    //
    //   const regex =
    //     /Rate limit reached for .+ in organization .+ on \w+ per (day|min): Limit (\d+)(?:, Used (\d+))?, Requested (\d+)/;
    //   const [, period, limit, used, requested] =
    //     errorPayload.error?.message?.match(regex) || [];
    //
    //   req.log.warn(
    //     { key: req.key?.hash, period, limit, used, requested },
    //     "Received `tokens_usage_based` rate limit error from OpenAI."
    //   );
    //
    //   if (!period || !limit || !requested) {
    //     errorPayload.proxy_note = `Unrecognized rate limit error from OpenAI. (${errorPayload.error?.message})`;
    //     break;
    //   }
    //
    //   if (req.retryCount < 2) {
    //     await reenqueueRequest(req);
    //     throw new RetryableError("Rate-limited request re-enqueued.");
    //   }
    //
    //   if (period === "min") {
    //     errorPayload.proxy_note = `Assigned key can't be used for prompts longer than ${limit} tokens, and no other keys are available right now. Reduce the length of your prompt or try again in a few minutes.`;
    //   } else {
    //     errorPayload.proxy_note = `Assigned key  has reached its per-day request limit for this model. Try another model.`;
    //   }
    //
    //   keyPool.markRateLimited(req.key!);
    //   break;
    default:
      errorPayload.proxy_note = `This is likely a temporary error with OpenAI. Try again in a few seconds.`;
      break;
  }
  return errorPayload;
}

async function handleAzureRateLimitError(
  req: Request,
  errorPayload: ProxiedErrorPayload
) {
  const code = errorPayload.error?.code;
  switch (code) {
    case "429":
      keyPool.markRateLimited(req.key!);
      await reenqueueRequest(req);
      throw new RetryableError("Rate-limited request re-enqueued.");
    default:
      errorPayload.proxy_note = `Unrecognized rate limit error from Azure (${code}). Please report this.`;
      break;
  }
}

//{"error":{"code":429,"message":"Resource has been exhausted (e.g. check quota).","status":"RESOURCE_EXHAUSTED"}
async function handleGoogleAIRateLimitError(
  req: Request,
  errorPayload: ProxiedErrorPayload
) {
  const status = errorPayload.error?.status;
  switch (status) {
    case "RESOURCE_EXHAUSTED":
      keyPool.markRateLimited(req.key!);
      await reenqueueRequest(req);
      throw new RetryableError("Rate-limited request re-enqueued.");
    default:
      errorPayload.proxy_note = `Unrecognized rate limit error from Google AI (${status}). Please report this.`;
      break;
  }
}

const incrementUsage: ProxyResHandlerWithBody = async (_proxyRes, req) => {
  if (isTextGenerationRequest(req) || isImageGenerationRequest(req)) {
    const model = req.body.model;
    const tokensUsed = req.promptTokens! + req.outputTokens!;
    req.log.debug(
      {
        model,
        tokensUsed,
        promptTokens: req.promptTokens,
        outputTokens: req.outputTokens,
      },
      `Incrementing usage for model`
    );
    keyPool.incrementUsage(req.key!, model, tokensUsed);
    if (req.user) {
      incrementPromptCount(req.user.token);
      incrementTokenCount(req.user.token, model, req.outboundApi, tokensUsed);
    }
  }
};

const countResponseTokens: ProxyResHandlerWithBody = async (
  _proxyRes,
  req,
  _res,
  body
) => {
  if (req.outboundApi === "openai-image") {
    req.outputTokens = req.promptTokens;
    req.promptTokens = 0;
    return;
  }

  // This function is prone to breaking if the upstream API makes even minor
  // changes to the response format, especially for SSE responses. If you're
  // seeing errors in this function, check the reassembled response body from
  // handleStreamedResponse to see if the upstream API has changed.
  try {
    assertJsonResponse(body);
    const service = req.outboundApi;
    const completion = getCompletionFromBody(req, body);
    const tokens = await countTokens({ req, completion, service });

    req.log.debug(
      { service, tokens, prevOutputTokens: req.outputTokens },
      `Counted tokens for completion`
    );
    if (req.tokenizerInfo) {
      req.tokenizerInfo.completion_tokens = tokens;
    }

    req.outputTokens = tokens.token_count;
  } catch (error) {
    req.log.warn(
      error,
      "Error while counting completion tokens; assuming `max_output_tokens`"
    );
    // req.outputTokens will already be set to `max_output_tokens` from the
    // prompt counting middleware, so we don't need to do anything here.
  }
};

const trackRateLimit: ProxyResHandlerWithBody = async (proxyRes, req) => {
  keyPool.updateRateLimits(req.key!, proxyRes.headers);
};

const copyHttpHeaders: ProxyResHandlerWithBody = async (
  proxyRes,
  _req,
  res
) => {
  Object.keys(proxyRes.headers).forEach((key) => {
    // Omit content-encoding because we will always decode the response body
    if (key === "content-encoding") {
      return;
    }
    // We're usually using res.json() to send the response, which causes express
    // to set content-length. That's not valid for chunked responses and some
    // clients will reject it so we need to omit it.
    if (key === "transfer-encoding") {
      return;
    }
    res.setHeader(key, proxyRes.headers[key] as string);
  });
};

function getAwsErrorType(header: string | string[] | undefined) {
  const val = String(header).match(/^(\w+):?/)?.[1];
  return val || String(header);
}

function assertJsonResponse(body: any): asserts body is Record<string, any> {
  if (typeof body !== "object") {
    throw new Error("Expected response to be an object");
  }
}


code:
import { Request } from "express";
import { config } from "../../../config";
import { logQueue } from "../../../shared/prompt-logging";
import {
  getCompletionFromBody,
  getModelFromBody,
  isImageGenerationRequest,
  isTextGenerationRequest,
} from "../common";
import { ProxyResHandlerWithBody } from ".";
import { assertNever } from "../../../shared/utils";
import {
  AnthropicChatMessage,
  flattenAnthropicMessages,
  MistralAIChatMessage,
  OpenAIChatMessage,
} from "../../../shared/api-schemas";
import { APIFormat } from "../../../shared/key-management";

/** If prompt logging is enabled, enqueues the prompt for logging. */
export const logPrompt: ProxyResHandlerWithBody = async (
  _proxyRes,
  req,
  _res,
  responseBody
) => {
  if (!config.promptLogging) {
    return;
  }
  if (typeof responseBody !== "object") {
    throw new Error("Expected body to be an object");
  }

  const loggable =
    isTextGenerationRequest(req) || isImageGenerationRequest(req);
  if (!loggable) return;

  const promptPayload = getPromptForRequest(req, responseBody);
  const promptFlattened = flattenMessages(promptPayload, req.outboundApi);
  const response = getCompletionFromBody(req, responseBody);
  const model = getModelFromBody(req, responseBody);

  logQueue.enqueue({
    endpoint: req.inboundApi,
    promptRaw: JSON.stringify(promptPayload),
    promptFlattened,
    model,
    response,
  });
};

type OaiImageResult = {
  prompt: string;
  size: string;
  style: string;
  quality: string;
  revisedPrompt?: string;
};

const getPromptForRequest = (
  req: Request,
  responseBody: Record<string, any>
):
  | string
  | OpenAIChatMessage[]
  | AnthropicChatMessage[]
  | MistralAIChatMessage[]
  | OaiImageResult => {
  // Since the prompt logger only runs after the request has been proxied, we
  // can assume the body has already been transformed to the target API's
  // format.
  switch (req.outboundApi) {
    case "openai":
    case "mistral-ai":
    case "anthropic-chat":
      return req.body.messages;
    case "openai-text":
      return req.body.prompt;
    case "openai-image":
      return {
        prompt: req.body.prompt,
        size: req.body.size,
        style: req.body.style,
        quality: req.body.quality,
        revisedPrompt: responseBody.data[0].revised_prompt,
      };
    case "anthropic-text":
      return req.body.prompt;
    case "google-ai":
      return req.body.prompt.text;
    default:
      assertNever(req.outboundApi);
  }
};

const flattenMessages = (
  val:
    | string
    | OaiImageResult
    | OpenAIChatMessage[]
    | AnthropicChatMessage[]
    | MistralAIChatMessage[],
  format: APIFormat
): string => {
  if (typeof val === "string") {
    return val.trim();
  }
  if (format === "anthropic-chat") {
    return flattenAnthropicMessages(val as AnthropicChatMessage[]);
  }
  if (Array.isArray(val)) {
    return val
      .map(({ content, role }) => {
        const text = Array.isArray(content)
          ? content
              .map((c) => {
                if ("text" in c) return c.text;
                if ("image_url" in c) return "(( Attached Image ))";
                if ("source" in c) return "(( Attached Image ))";
                return "(( Unsupported Content ))";
              })
              .join("\n")
          : content;
        return `${role}: ${text}`;
      })
      .join("\n");
  }
  return val.prompt.trim();
};

repo_name: Cucumber/oai

license: unknown


code:
import { ProxyResHandlerWithBody } from "./index";
import { mirrorGeneratedImage, OpenAIImageGenerationResult } from "../../../shared/file-storage/mirror-generated-image";

export const saveImage: ProxyResHandlerWithBody = async (
  _proxyRes,
  req,
  _res,
  body,
) => {
  if (req.outboundApi !== "openai-image") {
    return;
  }

  if (typeof body !== "object") {
    throw new Error("Expected body to be an object");
  }

  if (body.data) {
    const baseUrl = req.protocol + "://" + req.get("host");
    const prompt = body.data[0].revised_prompt ?? req.body.prompt;
    await mirrorGeneratedImage(
      baseUrl,
      prompt,
      body as OpenAIImageGenerationResult
    );
  }
};


code:
import { OpenAIChatCompletionStreamEvent } from "../index";

export type AnthropicChatCompletionResponse = {
  id: string;
  type: "message";
  role: "assistant";
  content: { type: "text"; text: string }[];
  model: string;
  stop_reason: string | null;
  stop_sequence: string | null;
  usage: { input_tokens: number; output_tokens: number };
};

/**
 * Given a list of OpenAI chat completion events, compiles them into a single
 * finalized Anthropic chat completion response so that non-streaming middleware
 * can operate on it as if it were a blocking response.
 */
export function mergeEventsForAnthropicChat(
  events: OpenAIChatCompletionStreamEvent[]
): AnthropicChatCompletionResponse {
  let merged: AnthropicChatCompletionResponse = {
    id: "",
    type: "message",
    role: "assistant",
    content: [],
    model: "",
    stop_reason: null,
    stop_sequence: null,
    usage: { input_tokens: 0, output_tokens: 0 },
  };
  merged = events.reduce((acc, event, i) => {
    // The first event will only contain role assignment and response metadata
    if (i === 0) {
      acc.id = event.id;
      acc.model = event.model;
      acc.content = [{ type: "text", text: "" }];
      return acc;
    }

    acc.stop_reason = event.choices[0].finish_reason ?? "";
    if (event.choices[0].delta.content) {
      acc.content[0].text += event.choices[0].delta.content;
    }

    return acc;
  }, merged);
  return merged;
}


code:
import { OpenAIChatCompletionStreamEvent } from "../index";

export type AnthropicTextCompletionResponse = {
  completion: string;
  stop_reason: string;
  truncated: boolean;
  stop: any;
  model: string;
  log_id: string;
  exception: null;
};

/**
 * Given a list of OpenAI chat completion events, compiles them into a single
 * finalized Anthropic completion response so that non-streaming middleware
 * can operate on it as if it were a blocking response.
 */
export function mergeEventsForAnthropicText(
  events: OpenAIChatCompletionStreamEvent[]
): AnthropicTextCompletionResponse {
  let merged: AnthropicTextCompletionResponse = {
    log_id: "",
    exception: null,
    model: "",
    completion: "",
    stop_reason: "",
    truncated: false,
    stop: null,
  };
  merged = events.reduce((acc, event, i) => {
    // The first event will only contain role assignment and response metadata
    if (i === 0) {
      acc.log_id = event.id;
      acc.model = event.model;
      acc.completion = "";
      acc.stop_reason = "";
      return acc;
    }

    acc.stop_reason = event.choices[0].finish_reason ?? "";
    if (event.choices[0].delta.content) {
      acc.completion += event.choices[0].delta.content;
    }

    return acc;
  }, merged);
  return merged;
}


code:
import { OpenAIChatCompletionStreamEvent } from "../index";

export type OpenAiChatCompletionResponse = {
  id: string;
  object: string;
  created: number;
  model: string;
  choices: {
    message: { role: string; content: string };
    finish_reason: string | null;
    index: number;
  }[];
};

/**
 * Given a list of OpenAI chat completion events, compiles them into a single
 * finalized OpenAI chat completion response so that non-streaming middleware
 * can operate on it as if it were a blocking response.
 */
export function mergeEventsForOpenAIChat(
  events: OpenAIChatCompletionStreamEvent[]
): OpenAiChatCompletionResponse {
  let merged: OpenAiChatCompletionResponse = {
    id: "",
    object: "",
    created: 0,
    model: "",
    choices: [],
  };
  merged = events.reduce((acc, event, i) => {
    // The first event will only contain role assignment and response metadata
    if (i === 0) {
      acc.id = event.id;
      acc.object = event.object;
      acc.created = event.created;
      acc.model = event.model;
      acc.choices = [
        {
          index: 0,
          message: {
            role: event.choices[0].delta.role ?? "assistant",
            content: "",
          },
          finish_reason: null,
        },
      ];
      return acc;
    }

    acc.choices[0].finish_reason = event.choices[0].finish_reason;
    if (event.choices[0].delta.content) {
      acc.choices[0].message.content += event.choices[0].delta.content;
    }

    return acc;
  }, merged);
  return merged;
}


code:
import { OpenAIChatCompletionStreamEvent } from "../index";

export type OpenAiTextCompletionResponse = {
  id: string;
  object: string;
  created: number;
  model: string;
  choices: {
    text: string;
    finish_reason: string | null;
    index: number;
    logprobs: null;
  }[];
};

/**
 * Given a list of OpenAI chat completion events, compiles them into a single
 * finalized OpenAI text completion response so that non-streaming middleware
 * can operate on it as if it were a blocking response.
 */
export function mergeEventsForOpenAIText(
  events: OpenAIChatCompletionStreamEvent[]
): OpenAiTextCompletionResponse {
  let merged: OpenAiTextCompletionResponse = {
    id: "",
    object: "",
    created: 0,
    model: "",
    choices: [],
  };
  merged = events.reduce((acc, event, i) => {
    // The first event will only contain role assignment and response metadata
    if (i === 0) {
      acc.id = event.id;
      acc.object = event.object;
      acc.created = event.created;
      acc.model = event.model;
      acc.choices = [
        {
          text: "",
          index: 0,
          finish_reason: null,
          logprobs: null,
        },
      ];
      return acc;
    }

    acc.choices[0].finish_reason = event.choices[0].finish_reason;
    if (event.choices[0].delta.content) {
      acc.choices[0].text += event.choices[0].delta.content;
    }

    return acc;
  }, merged);
  return merged;
}


code:
import pino from "pino";
import { Duplex, Readable } from "stream";
import { EventStreamMarshaller } from "@smithy/eventstream-serde-node";
import { fromUtf8, toUtf8 } from "@smithy/util-utf8";
import { Message } from "@smithy/eventstream-codec";

/**
 * Decodes a Readable stream, such as a proxied HTTP response, into a stream of
 * Message objects using the AWS SDK's EventStreamMarshaller. Error events in
 * the amazon eventstream protocol are decoded as Message objects and will not
 * emit an error event on the decoder stream.
 */
export function getAwsEventStreamDecoder(params: {
  input: Readable;
  logger: pino.Logger;
}): Duplex {
  const { input, logger } = params;
  const config = { utf8Encoder: toUtf8, utf8Decoder: fromUtf8 };
  const eventStream = new EventStreamMarshaller(config).deserialize(
    input,
    async (input: Record<string, Message>) => {
      const eventType = Object.keys(input)[0];
      let result;
      if (eventType === "chunk") {
        result = input[eventType];
      } else {
        // AWS unmarshaller treats non-chunk (errors and exceptions) oddly.
        result = { [eventType]: input[eventType] } as any;
      }
      return result;
    }
  );
  return new AWSEventStreamDecoder(eventStream, { logger });
}

class AWSEventStreamDecoder extends Duplex {
  private readonly asyncIterable: AsyncIterable<Message>;
  private iterator: AsyncIterator<Message>;
  private reading: boolean;
  private logger: pino.Logger;

  constructor(
    asyncIterable: AsyncIterable<Message>,
    options: { logger: pino.Logger }
  ) {
    super({ ...options, objectMode: true });
    this.asyncIterable = asyncIterable;
    this.iterator = this.asyncIterable[Symbol.asyncIterator]();
    this.reading = false;
    this.logger = options.logger.child({ module: "aws-eventstream-decoder" });
  }

  async _read(_size: number) {
    if (this.reading) return;
    this.reading = true;

    try {
      while (true) {
        const { value, done } = await this.iterator.next();
        if (done) {
          this.push(null);
          break;
        }
        if (!this.push(value)) break;
      }
    } catch (err) {
      // AWS SDK's EventStreamMarshaller emits errors in the stream itself as
      // whatever our deserializer returns, which will not be Error objects
      // because we want to pass the Message to the next stream for processing.
      // Any actual Error thrown here is some failure during deserialization.
      const isAwsError = !(err instanceof Error);

      if (isAwsError) {
        this.logger.warn({ err: err.headers }, "Received AWS error event");
        this.push(err);
        this.push(null);
      } else {
        this.logger.error(err, "Error during AWS stream deserialization");
        this.destroy(err);
      }
    } finally {
      this.reading = false;
    }
  }

  _write(_chunk: any, _encoding: string, callback: () => void) {
    callback();
  }

  _final(callback: () => void) {
    callback();
  }
}


code:
import { APIFormat } from "../../../../shared/key-management";
import { assertNever } from "../../../../shared/utils";
import {
  anthropicV2ToOpenAI,
  mergeEventsForAnthropicChat,
  mergeEventsForAnthropicText,
  mergeEventsForOpenAIChat,
  mergeEventsForOpenAIText,
  AnthropicV2StreamEvent,
  OpenAIChatCompletionStreamEvent,
} from "./index";

/**
 * Collects SSE events containing incremental chat completion responses and
 * compiles them into a single finalized response for downstream middleware.
 */
export class EventAggregator {
  private readonly format: APIFormat;
  private readonly events: OpenAIChatCompletionStreamEvent[];

  constructor({ format }: { format: APIFormat }) {
    this.events = [];
    this.format = format;
  }

  addEvent(event: OpenAIChatCompletionStreamEvent | AnthropicV2StreamEvent) {
    if (eventIsOpenAIEvent(event)) {
      this.events.push(event);
    } else {
      // horrible special case. previously all transformers' target format was
      // openai, so the event aggregator could conveniently assume all incoming
      // events were in openai format.
      // now we have added anthropic-chat-to-text, so aggregator needs to know
      // how to collapse events from two formats.
      // because that is annoying, we will simply transform anthropic events to
      // openai (even if the client didn't ask for openai) so we don't have to
      // write aggregation logic for anthropic chat (which is also a troublesome
      // stateful format).
      const openAIEvent = anthropicV2ToOpenAI({
        data: `event: completion\ndata: ${JSON.stringify(event)}\n\n`,
        lastPosition: -1,
        index: 0,
        fallbackId: event.log_id || "event-aggregator-fallback",
        fallbackModel: event.model || "claude-3-fallback",
      });
      if (openAIEvent.event) {
        this.events.push(openAIEvent.event);
      }
    }
  }

  getFinalResponse() {
    switch (this.format) {
      case "openai":
      case "google-ai":
      case "mistral-ai":
        return mergeEventsForOpenAIChat(this.events);
      case "openai-text":
        return mergeEventsForOpenAIText(this.events);
      case "anthropic-text":
        return mergeEventsForAnthropicText(this.events);
      case "anthropic-chat":
        return mergeEventsForAnthropicChat(this.events);
      case "openai-image":
        throw new Error(`SSE aggregation not supported for ${this.format}`);
      default:
        assertNever(this.format);
    }
  }
}

function eventIsOpenAIEvent(
  event: any
): event is OpenAIChatCompletionStreamEvent {
  return event?.object === "chat.completion.chunk";
}


code:
export type SSEResponseTransformArgs<S = Record<string, any>> = {
  data: string;
  lastPosition: number;
  index: number;
  fallbackId: string;
  fallbackModel: string;
  state?: S;
};

export type AnthropicV2StreamEvent = {
  log_id?: string;
  model?: string;
  completion: string;
  stop_reason: string | null;
};

export type OpenAIChatCompletionStreamEvent = {
  id: string;
  object: "chat.completion.chunk";
  created: number;
  model: string;
  choices: {
    index: number;
    delta: { role?: string; content?: string };
    finish_reason: string | null;
  }[];
};

export type StreamingCompletionTransformer<
  T = OpenAIChatCompletionStreamEvent,
  S = any,
> = (params: SSEResponseTransformArgs<S>) => {
  position: number;
  event?: T;
  state?: S;
};

export { openAITextToOpenAIChat } from "./transformers/openai-text-to-openai";
export { anthropicV1ToOpenAI } from "./transformers/anthropic-v1-to-openai";
export { anthropicV2ToOpenAI } from "./transformers/anthropic-v2-to-openai";
export { anthropicChatToAnthropicV2 } from "./transformers/anthropic-chat-to-anthropic-v2";
export { googleAIToOpenAI } from "./transformers/google-ai-to-openai";
export { passthroughToOpenAI } from "./transformers/passthrough-to-openai";
export { mergeEventsForOpenAIChat } from "./aggregators/openai-chat";
export { mergeEventsForOpenAIText } from "./aggregators/openai-text";
export { mergeEventsForAnthropicText } from "./aggregators/anthropic-text";
export { mergeEventsForAnthropicChat } from "./aggregators/anthropic-chat";


code:
export type ServerSentEvent = { id?: string; type?: string; data: string };

/** Given a string of SSE data, parse it into a `ServerSentEvent` object. */
export function parseEvent(event: string) {
  const buffer: ServerSentEvent = { data: "" };
  return event.split(/\r?\n/).reduce(parseLine, buffer);
}

function parseLine(event: ServerSentEvent, line: string) {
  const separator = line.indexOf(":");
  const field = separator === -1 ? line : line.slice(0, separator);
  const value = separator === -1 ? "" : line.slice(separator + 1);

  switch (field) {
    case "id":
      event.id = value.trim();
      break;
    case "event":
      event.type = value.trim();
      break;
    case "data":
      event.data += value.trimStart();
      break;
    default:
      break;
  }

  return event;
}


code:
import { Transform, TransformOptions } from "stream";
import { logger } from "../../../../logger";
import { APIFormat } from "../../../../shared/key-management";
import { assertNever } from "../../../../shared/utils";
import {
  anthropicChatToAnthropicV2,
  anthropicV1ToOpenAI,
  AnthropicV2StreamEvent,
  anthropicV2ToOpenAI,
  googleAIToOpenAI,
  OpenAIChatCompletionStreamEvent,
  openAITextToOpenAIChat,
  passthroughToOpenAI,
  StreamingCompletionTransformer,
} from "./index";

type SSEMessageTransformerOptions = TransformOptions & {
  requestedModel: string;
  requestId: string;
  inputFormat: APIFormat;
  inputApiVersion?: string;
  logger: typeof logger;
};

/**
 * Transforms SSE messages from one API format to OpenAI chat.completion.chunks.
 * Emits the original string SSE message as an "originalMessage" event.
 */
export class SSEMessageTransformer extends Transform {
  private lastPosition: number;
  private transformState: any;
  private msgCount: number;
  private readonly inputFormat: APIFormat;
  private readonly transformFn: StreamingCompletionTransformer<
    // TODO: Refactor transformers to not assume only OpenAI events as output
    OpenAIChatCompletionStreamEvent | AnthropicV2StreamEvent
  >;
  private readonly log;
  private readonly fallbackId: string;
  private readonly fallbackModel: string;

  constructor(options: SSEMessageTransformerOptions) {
    super({ ...options, readableObjectMode: true });
    this.log = options.logger?.child({ module: "sse-transformer" });
    this.lastPosition = 0;
    this.msgCount = 0;
    this.transformFn = getTransformer(
      options.inputFormat,
      options.inputApiVersion
    );
    this.inputFormat = options.inputFormat;
    this.fallbackId = options.requestId;
    this.fallbackModel = options.requestedModel;
    this.log.debug(
      {
        fn: this.transformFn.name,
        format: options.inputFormat,
        version: options.inputApiVersion,
      },
      "Selected SSE transformer"
    );
  }

  _transform(chunk: Buffer, _encoding: BufferEncoding, callback: Function) {
    try {
      const originalMessage = chunk.toString();
      const {
        event: transformedMessage,
        position: newPosition,
        state,
      } = this.transformFn({
        data: originalMessage,
        lastPosition: this.lastPosition,
        index: this.msgCount++,
        fallbackId: this.fallbackId,
        fallbackModel: this.fallbackModel,
        state: this.transformState,
      });
      this.lastPosition = newPosition;
      this.transformState = state;

      // Special case for Azure OpenAI, which is 99% the same as OpenAI but
      // sometimes emits an extra event at the beginning of the stream with the
      // content moderation system's response to the prompt. A lot of frontends
      // don't expect this and neither does our event aggregator so we drop it.
      if (this.inputFormat === "openai" && this.msgCount <= 1) {
        if (originalMessage.includes("prompt_filter_results")) {
          this.log.debug("Dropping Azure OpenAI content moderation SSE event");
          return callback();
        }
      }

      this.emit("originalMessage", originalMessage);

      // Some events may not be transformed, e.g. ping events
      if (!transformedMessage) return callback();

      if (this.msgCount === 1 && eventIsOpenAIEvent(transformedMessage)) {
        // TODO: does this need to be skipped for passthroughToOpenAI?
        this.push(createInitialMessage(transformedMessage));
      }
      this.push(transformedMessage);
      callback();
    } catch (err) {
      err.lastEvent = chunk?.toString();
      this.log.error(err, "Error transforming SSE message");
      callback(err);
    }
  }
}

function eventIsOpenAIEvent(
  event: any
): event is OpenAIChatCompletionStreamEvent {
  return event?.object === "chat.completion.chunk";
}

function getTransformer(
  responseApi: APIFormat,
  version?: string
): StreamingCompletionTransformer<
  OpenAIChatCompletionStreamEvent | AnthropicV2StreamEvent
> {
  switch (responseApi) {
    case "openai":
    case "mistral-ai":
      return passthroughToOpenAI;
    case "openai-text":
      return openAITextToOpenAIChat;
    case "anthropic-text":
      return version === "2023-01-01"
        ? anthropicV1ToOpenAI
        : anthropicV2ToOpenAI;
    case "anthropic-chat":
      return anthropicChatToAnthropicV2;
    case "google-ai":
      return googleAIToOpenAI;
    case "openai-image":
      throw new Error(`SSE transformation not supported for ${responseApi}`);
    default:
      assertNever(responseApi);
  }
}

/**
 * OpenAI streaming chat completions start with an event that contains only the
 * metadata and role (always 'assistant') for the response.  To simulate this
 * for APIs where the first event contains actual content, we create a fake
 * initial event with no content but correct metadata.
 */
function createInitialMessage(
  event: OpenAIChatCompletionStreamEvent
): OpenAIChatCompletionStreamEvent {
  return {
    ...event,
    choices: event.choices.map((choice) => ({
      ...choice,
      delta: { role: "assistant", content: "" },
    })),
  };
}


code:
import pino from "pino";
import { Transform, TransformOptions } from "stream";
import { Message } from "@smithy/eventstream-codec";
import { APIFormat } from "../../../../shared/key-management";
import { RetryableError } from "../index";
import { buildSpoofedSSE } from "../error-generator";

type SSEStreamAdapterOptions = TransformOptions & {
  contentType?: string;
  api: APIFormat;
  logger: pino.Logger;
};

/**
 * Receives a stream of events in a variety of formats and transforms them into
 * Server-Sent Events.
 *
 * This is an object-mode stream, so it expects to receive objects and will emit
 * strings.
 */
export class SSEStreamAdapter extends Transform {
  private readonly isAwsStream;
  private readonly isGoogleStream;
  private api: APIFormat;
  private partialMessage = "";
  private textDecoder = new TextDecoder("utf8");
  private log: pino.Logger;

  constructor(options: SSEStreamAdapterOptions) {
    super({ ...options, objectMode: true });
    this.isAwsStream =
      options?.contentType === "application/vnd.amazon.eventstream";
    this.isGoogleStream = options?.api === "google-ai";
    this.api = options.api;
    this.log = options.logger.child({ module: "sse-stream-adapter" });
  }

  protected processAwsMessage(message: Message): string | null {
    // Per amazon, headers and body are always present. headers is an object,
    // body is a Uint8Array, potentially zero-length.
    const { headers, body } = message;
    const eventType = headers[":event-type"]?.value;
    const messageType = headers[":message-type"]?.value;
    const contentType = headers[":content-type"]?.value;
    const exceptionType = headers[":exception-type"]?.value;
    const errorCode = headers[":error-code"]?.value;
    const bodyStr = this.textDecoder.decode(body);

    switch (messageType) {
      case "event":
        if (contentType === "application/json" && eventType === "chunk") {
          const { bytes } = JSON.parse(bodyStr);
          const event = Buffer.from(bytes, "base64").toString("utf8");
          const eventObj = JSON.parse(event);

          if ("completion" in eventObj) {
            return ["event: completion", `data: ${event}`].join(`\n`);
          } else {
            return [`event: ${eventObj.type}`, `data: ${event}`].join(`\n`);
          }
        }
      // Intentional fallthrough, as non-JSON events may as well be errors
      // noinspection FallThroughInSwitchStatementJS
      case "exception":
      case "error":
        const type = String(
          exceptionType || errorCode || "UnknownError"
        ).toLowerCase();
        switch (type) {
          case "throttlingexception":
            this.log.warn(
              "AWS request throttled after streaming has already started; retrying"
            );
            throw new RetryableError("AWS request throttled mid-stream");
          default:
            this.log.error({ message, type }, "Received bad AWS stream event");
            const error: any = new Error(`Got mysterious error chunk: ${type}`);
            error.lastEvent = message;
            this.emit("error", error);
            return null;
        }
      default:
        // Amazon says this can't ever happen...
        this.log.error({ message }, "Received very bad AWS stream event");
        return null;
    }
  }

  /** Processes an incoming array element from the Google AI JSON stream. */
  protected processGoogleObject(data: any): string | null {
    // Sometimes data has fields key and value, sometimes it's just the
    // candidates array.
    const candidates = data.value?.candidates ?? data.candidates ?? [{}];
    try {
      const hasParts = candidates[0].content?.parts?.length > 0;
      if (hasParts) {
        return `data: ${JSON.stringify(data)}`;
      } else {
        this.log.error({ event: data }, "Received bad Google AI event");
        return `data: ${buildSpoofedSSE({
          format: "google-ai",
          title: "Proxy stream error",
          message:
            "The proxy received malformed or unexpected data from Google AI while streaming.",
          obj: data,
          reqId: "proxy-sse-adapter-message",
          model: "",
        })}`;
      }
    } catch (error) {
      error.lastEvent = data;
      this.emit("error", error);
    }
    return null;
  }

  _transform(data: any, _enc: string, callback: (err?: Error | null) => void) {
    try {
      if (this.isAwsStream) {
        // `data` is a Message object
        const message = this.processAwsMessage(data);
        if (message) this.push(message + "\n\n");
      } else if (this.isGoogleStream) {
        // `data` is an element from the Google AI JSON stream
        const message = this.processGoogleObject(data);
        if (message) this.push(message + "\n\n");
      } else {
        // `data` is a string, but possibly only a partial message
        const fullMessages = (this.partialMessage + data).split(
          /\r\r|\n\n|\r\n\r\n/
        );
        this.partialMessage = fullMessages.pop() || "";

        for (const message of fullMessages) {
          // Mixing line endings will break some clients and our request queue
          // will have already sent \n for heartbeats, so we need to normalize
          // to \n.
          this.push(message.replace(/\r\n?/g, "\n") + "\n\n");
        }
      }
      callback();
    } catch (error) {
      error.lastEvent = data?.toString() ?? "[SSEStreamAdapter] no data";
      callback(error);
    }
  }

  _flush(callback: (err?: Error | null) => void) {
    callback();
  }
}


code:
import {
  AnthropicV2StreamEvent,
  StreamingCompletionTransformer,
} from "../index";
import { parseEvent, ServerSentEvent } from "../parse-sse";
import { logger } from "../../../../../logger";

const log = logger.child({
  module: "sse-transformer",
  transformer: "anthropic-chat-to-anthropic-v2",
});

export type AnthropicChatEventType =
  | "message_start"
  | "content_block_start"
  | "content_block_delta"
  | "content_block_stop"
  | "message_delta"
  | "message_stop";

type AnthropicChatStartEvent = {
  type: "message_start";
  message: {
    id: string;
    type: "message";
    role: "assistant";
    content: [];
    model: string;
    stop_reason: null;
    stop_sequence: null;
    usage: { input_tokens: number; output_tokens: number };
  };
};

type AnthropicChatContentBlockStartEvent = {
  type: "content_block_start";
  index: number;
  content_block: { type: "text"; text: string };
};

export type AnthropicChatContentBlockDeltaEvent = {
  type: "content_block_delta";
  index: number;
  delta: { type: "text_delta"; text: string };
};

type AnthropicChatContentBlockStopEvent = {
  type: "content_block_stop";
  index: number;
};

type AnthropicChatMessageDeltaEvent = {
  type: "message_delta";
  delta: {
    stop_reason: string;
    stop_sequence: null;
    usage: { output_tokens: number };
  };
};

type AnthropicChatMessageStopEvent = {
  type: "message_stop";
};

type AnthropicChatTransformerState = { content: string };

/**
 * Transforms an incoming Anthropic Chat SSE to an equivalent Anthropic V2
 * Text SSE.
 * For now we assume there is only one content block and message delta. In the
 * future Anthropic may add multi-turn responses or multiple content blocks
 * (probably for multimodal responses, image generation, etc) but as far as I
 * can tell this is not yet implemented.
 */
export const anthropicChatToAnthropicV2: StreamingCompletionTransformer<
  AnthropicV2StreamEvent,
  AnthropicChatTransformerState
> = (params) => {
  const { data } = params;

  const rawEvent = parseEvent(data);
  if (!rawEvent.data || !rawEvent.type) {
    return { position: -1 };
  }

  const deltaEvent = asAnthropicChatDelta(rawEvent);
  if (!deltaEvent) {
    return { position: -1 };
  }

  const newEvent = {
    log_id: params.fallbackId,
    model: params.fallbackModel,
    completion: deltaEvent.delta.text,
    stop_reason: null,
  };

  return { position: -1, event: newEvent };
};

export function asAnthropicChatDelta(
  event: ServerSentEvent
): AnthropicChatContentBlockDeltaEvent | null {
  if (
    !event.type ||
    !["content_block_start", "content_block_delta"].includes(event.type)
  ) {
    return null;
  }

  try {
    const parsed = JSON.parse(event.data);
    if (parsed.type === "content_block_delta") {
      return parsed;
    } else if (parsed.type === "content_block_start") {
      return {
        type: "content_block_delta",
        index: parsed.index,
        delta: { type: "text_delta", text: parsed.content_block?.text ?? "" },
      };
    } else {
      // noinspection ExceptionCaughtLocallyJS
      throw new Error("Invalid event type");
    }
  } catch (error) {
    log.warn({ error: error.stack, event }, "Received invalid event");
  }
  return null;
}


code:
import { StreamingCompletionTransformer } from "../index";
import { parseEvent } from "../parse-sse";
import { logger } from "../../../../../logger";
import { asAnthropicChatDelta } from "./anthropic-chat-to-anthropic-v2";

const log = logger.child({
  module: "sse-transformer",
  transformer: "anthropic-chat-to-openai",
});

/**
 * Transforms an incoming Anthropic Chat SSE to an equivalent OpenAI
 * chat.completion.chunks SSE.
 */
export const anthropicChatToOpenAI: StreamingCompletionTransformer = (
  params
) => {
  const { data } = params;

  const rawEvent = parseEvent(data);
  if (!rawEvent.data || !rawEvent.type) {
    return { position: -1 };
  }

  const deltaEvent = asAnthropicChatDelta(rawEvent);
  if (!deltaEvent) {
    return { position: -1 };
  }

  const newEvent = {
    id: params.fallbackId,
    object: "chat.completion.chunk" as const,
    created: Date.now(),
    model: params.fallbackModel,
    choices: [
      {
        index: params.index,
        delta: { content: deltaEvent.delta.text },
        finish_reason: null,
      },
    ],
  };

  return { position: -1, event: newEvent };
};


code:
import { StreamingCompletionTransformer } from "../index";
import { parseEvent, ServerSentEvent } from "../parse-sse";
import { logger } from "../../../../../logger";

const log = logger.child({
  module: "sse-transformer",
  transformer: "anthropic-v1-to-openai",
});

type AnthropicV1StreamEvent = {
  log_id?: string;
  model?: string;
  completion: string;
  stop_reason: string;
};

/**
 * Transforms an incoming Anthropic SSE (2023-01-01 API) to an equivalent
 * OpenAI chat.completion.chunk SSE.
 */
export const anthropicV1ToOpenAI: StreamingCompletionTransformer = (params) => {
  const { data, lastPosition } = params;

  const rawEvent = parseEvent(data);
  if (!rawEvent.data || rawEvent.data === "[DONE]") {
    return { position: lastPosition };
  }

  const completionEvent = asCompletion(rawEvent);
  if (!completionEvent) {
    return { position: lastPosition };
  }

  // Anthropic sends the full completion so far with each event whereas OpenAI
  // only sends the delta. To make the SSE events compatible, we remove
  // everything before `lastPosition` from the completion.
  const newEvent = {
    id: "ant-" + (completionEvent.log_id ?? params.fallbackId),
    object: "chat.completion.chunk" as const,
    created: Date.now(),
    model: completionEvent.model ?? params.fallbackModel,
    choices: [
      {
        index: 0,
        delta: { content: completionEvent.completion?.slice(lastPosition) },
        finish_reason: completionEvent.stop_reason,
      },
    ],
  };

  return { position: completionEvent.completion.length, event: newEvent };
};

function asCompletion(event: ServerSentEvent): AnthropicV1StreamEvent | null {
  try {
    const parsed = JSON.parse(event.data);
    if (parsed.completion !== undefined && parsed.stop_reason !== undefined) {
      return parsed;
    } else {
      // noinspection ExceptionCaughtLocallyJS
      throw new Error("Missing required fields");
    }
  } catch (error) {
    log.warn({ error: error.stack, event }, "Received invalid event");
  }
  return null;
}


code:
import {
  AnthropicV2StreamEvent,
  StreamingCompletionTransformer,
} from "../index";
import { parseEvent, ServerSentEvent } from "../parse-sse";
import { logger } from "../../../../../logger";

const log = logger.child({
  module: "sse-transformer",
  transformer: "anthropic-v2-to-openai",
});

/**
 * Transforms an incoming Anthropic SSE (2023-06-01 API) to an equivalent
 * OpenAI chat.completion.chunk SSE.
 */
export const anthropicV2ToOpenAI: StreamingCompletionTransformer = (params) => {
  const { data } = params;

  const rawEvent = parseEvent(data);
  if (!rawEvent.data || rawEvent.data === "[DONE]") {
    return { position: -1 };
  }

  const completionEvent = asCompletion(rawEvent);
  if (!completionEvent) {
    return { position: -1 };
  }

  const newEvent = {
    id: "ant-" + (completionEvent.log_id ?? params.fallbackId),
    object: "chat.completion.chunk" as const,
    created: Date.now(),
    model: completionEvent.model ?? params.fallbackModel,
    choices: [
      {
        index: 0,
        delta: { content: completionEvent.completion },
        finish_reason: completionEvent.stop_reason,
      },
    ],
  };

  return { position: completionEvent.completion.length, event: newEvent };
};

function asCompletion(event: ServerSentEvent): AnthropicV2StreamEvent | null {
  if (event.type === "ping") return null;

  try {
    const parsed = JSON.parse(event.data);
    if (parsed.completion !== undefined && parsed.stop_reason !== undefined) {
      return parsed;
    } else {
      // noinspection ExceptionCaughtLocallyJS
      throw new Error("Missing required fields");
    }
  } catch (error) {
    log.warn({ error: error.stack, event }, "Received invalid event");
  }
  return null;
}


code:
import { StreamingCompletionTransformer } from "../index";
import { parseEvent, ServerSentEvent } from "../parse-sse";
import { logger } from "../../../../../logger";

const log = logger.child({
  module: "sse-transformer",
  transformer: "google-ai-to-openai",
});

type GoogleAIStreamEvent = {
  candidates: {
    content: { parts: { text: string }[]; role: string };
    finishReason?: "STOP" | "MAX_TOKENS" | "SAFETY" | "RECITATION" | "OTHER";
    index: number;
    tokenCount?: number;
    safetyRatings: { category: string; probability: string }[];
  }[];
};

/**
 * Transforms an incoming Google AI SSE to an equivalent OpenAI
 * chat.completion.chunk SSE.
 */
export const googleAIToOpenAI: StreamingCompletionTransformer = (params) => {
  const { data, index } = params;

  const rawEvent = parseEvent(data);
  if (!rawEvent.data || rawEvent.data === "[DONE]") {
    return { position: -1 };
  }

  const completionEvent = asCompletion(rawEvent);
  if (!completionEvent) {
    return { position: -1 };
  }

  const parts = completionEvent.candidates[0].content.parts;
  let content = parts[0]?.text ?? "";

  // If this is the first chunk, try stripping speaker names from the response
  // e.g. "John: Hello" -> "Hello"
  if (index === 0) {
    content = content.replace(/^(.*?): /, "").trim();
  }

  const newEvent = {
    id: "goo-" + params.fallbackId,
    object: "chat.completion.chunk" as const,
    created: Date.now(),
    model: params.fallbackModel,
    choices: [
      {
        index: 0,
        delta: { content },
        finish_reason: completionEvent.candidates[0].finishReason ?? null,
      },
    ],
  };

  return { position: -1, event: newEvent };
};

function asCompletion(event: ServerSentEvent): GoogleAIStreamEvent | null {
  try {
    const parsed = JSON.parse(event.data) as GoogleAIStreamEvent;
    if (parsed.candidates?.length > 0) {
      return parsed;
    } else {
      // noinspection ExceptionCaughtLocallyJS
      throw new Error("Missing required fields");
    }
  } catch (error) {
    log.warn({ error: error.stack, event }, "Received invalid event");
  }
  return null;
}


code:
import { SSEResponseTransformArgs } from "../index";
import { parseEvent, ServerSentEvent } from "../parse-sse";
import { logger } from "../../../../../logger";

const log = logger.child({
  module: "sse-transformer",
  transformer: "openai-text-to-openai",
});

type OpenAITextCompletionStreamEvent = {
  id: string;
  object: "text_completion";
  created: number;
  choices: {
    text: string;
    index: number;
    logprobs: null;
    finish_reason: string | null;
  }[];
  model: string;
};

export const openAITextToOpenAIChat = (params: SSEResponseTransformArgs) => {
  const { data } = params;

  const rawEvent = parseEvent(data);
  if (!rawEvent.data || rawEvent.data === "[DONE]") {
    return { position: -1 };
  }

  const completionEvent = asCompletion(rawEvent);
  if (!completionEvent) {
    return { position: -1 };
  }

  const newEvent = {
    id: completionEvent.id,
    object: "chat.completion.chunk" as const,
    created: completionEvent.created,
    model: completionEvent.model,
    choices: [
      {
        index: completionEvent.choices[0].index,
        delta: { content: completionEvent.choices[0].text },
        finish_reason: completionEvent.choices[0].finish_reason,
      },
    ],
  };

  return { position: -1, event: newEvent };
};

function asCompletion(
  event: ServerSentEvent
): OpenAITextCompletionStreamEvent | null {
  try {
    const parsed = JSON.parse(event.data);
    if (Array.isArray(parsed.choices) && parsed.choices[0].text !== undefined) {
      return parsed;
    } else {
      // noinspection ExceptionCaughtLocallyJS
      throw new Error("Missing required fields");
    }
  } catch (error) {
    log.warn({ error: error.stack, event }, "Received invalid data event");
  }
  return null;
}


code:
import {
  OpenAIChatCompletionStreamEvent,
  SSEResponseTransformArgs,
} from "../index";
import { parseEvent, ServerSentEvent } from "../parse-sse";
import { logger } from "../../../../../logger";

const log = logger.child({
  module: "sse-transformer",
  transformer: "openai-to-openai",
});

export const passthroughToOpenAI = (params: SSEResponseTransformArgs) => {
  const { data } = params;

  const rawEvent = parseEvent(data);
  if (!rawEvent.data || rawEvent.data === "[DONE]") {
    return { position: -1 };
  }

  const completionEvent = asCompletion(rawEvent);
  if (!completionEvent) {
    return { position: -1 };
  }

  return { position: -1, event: completionEvent };
};

function asCompletion(
  event: ServerSentEvent
): OpenAIChatCompletionStreamEvent | null {
  try {
    return JSON.parse(event.data);
  } catch (error) {
    log.warn({ error: error.stack, event }, "Received invalid event");
  }
  return null;
}


code:
import { RequestHandler, Router } from "express";
import { createProxyMiddleware } from "http-proxy-middleware";
import { config } from "../config";
import { keyPool } from "../shared/key-management";
import {
  getMistralAIModelFamily,
  MistralAIModelFamily,
  ModelFamily,
} from "../shared/models";
import { logger } from "../logger";
import { createQueueMiddleware } from "./queue";
import { ipLimiter } from "./rate-limit";
import { handleProxyError } from "./middleware/common";
import {
  addKey,
  createOnProxyReqHandler,
  createPreprocessorMiddleware,
  finalizeBody,
} from "./middleware/request";
import {
  createOnProxyResHandler,
  ProxyResHandlerWithBody,
} from "./middleware/response";

// https://docs.mistral.ai/platform/endpoints
export const KNOWN_MISTRAL_AI_MODELS = [
  // Mistral 7b (open weight, legacy)
  "open-mistral-7b",
  "mistral-tiny-2312",
  // Mixtral 8x7b (open weight, legacy)
  "open-mixtral-8x7b",
  "mistral-small-2312",
  // Mixtral Small (newer 8x7b, closed weight)
  "mistral-small-latest",
  "mistral-small-2402",
  // Mistral Medium
  "mistral-medium-latest",
  "mistral-medium-2312",
  // Mistral Large
  "mistral-large-latest",
  "mistral-large-2402",
  // Deprecated identifiers (2024-05-01)
  "mistral-tiny",
  "mistral-small",
  "mistral-medium",
];

let modelsCache: any = null;
let modelsCacheTime = 0;

export function generateModelList(models = KNOWN_MISTRAL_AI_MODELS) {
  let available = new Set<MistralAIModelFamily>();
  for (const key of keyPool.list()) {
    if (key.isDisabled || key.service !== "mistral-ai") continue;
    key.modelFamilies.forEach((family) =>
      available.add(family as MistralAIModelFamily)
    );
  }
  const allowed = new Set<ModelFamily>(config.allowedModelFamilies);
  available = new Set([...available].filter((x) => allowed.has(x)));

  return models
    .map((id) => ({
      id,
      object: "model",
      created: new Date().getTime(),
      owned_by: "mistral-ai",
    }))
    .filter((model) => available.has(getMistralAIModelFamily(model.id)));
}

const handleModelRequest: RequestHandler = (_req, res) => {
  if (new Date().getTime() - modelsCacheTime < 1000 * 60){
    return res.status(200).json(modelsCache);
  }
  const result = generateModelList();
  modelsCache = { object: "list", data: result };
  modelsCacheTime = new Date().getTime();
  res.status(200).json(modelsCache);
};

const mistralAIResponseHandler: ProxyResHandlerWithBody = async (
  _proxyRes,
  req,
  res,
  body
) => {
  if (typeof body !== "object") {
    throw new Error("Expected body to be an object");
  }

  if (config.promptLogging) {
    const host = req.get("host");
    body.proxy_note = `Prompts are logged on this proxy instance. See ${host} for more information.`;
  }

  if (req.tokenizerInfo) {
    body.proxy_tokenizer = req.tokenizerInfo;
  }

  res.status(200).json(body);
};

const mistralAIProxy = createQueueMiddleware({
  proxyMiddleware: createProxyMiddleware({
    target: "https://api.mistral.ai",
    changeOrigin: true,
    selfHandleResponse: true,
    logger,
    on: {
      proxyReq: createOnProxyReqHandler({
        pipeline: [addKey, finalizeBody],
      }),
      proxyRes: createOnProxyResHandler([mistralAIResponseHandler]),
      error: handleProxyError,
    },
  }),
});

const mistralAIRouter = Router();
mistralAIRouter.get("/v1/models", handleModelRequest);
// General chat completion endpoint.
mistralAIRouter.post(
  "/v1/chat/completions",
  ipLimiter,
  createPreprocessorMiddleware({
    inApi: "mistral-ai",
    outApi: "mistral-ai",
    service: "mistral-ai",
  }),
  mistralAIProxy
);

export const mistralAI = mistralAIRouter;


code:
import { RequestHandler, Router, Request } from "express";
import { createProxyMiddleware } from "http-proxy-middleware";
import { config } from "../config";
import { logger } from "../logger";
import { createQueueMiddleware } from "./queue";
import { ipLimiter } from "./rate-limit";
import { handleProxyError } from "./middleware/common";
import {
  addKey,
  createPreprocessorMiddleware,
  finalizeBody,
  createOnProxyReqHandler,
} from "./middleware/request";
import {
  createOnProxyResHandler,
  ProxyResHandlerWithBody,
} from "./middleware/response";
import { generateModelList } from "./openai";
import {
  OpenAIImageGenerationResult,
} from "../shared/file-storage/mirror-generated-image";

const KNOWN_MODELS = ["dall-e-2", "dall-e-3"];

let modelListCache: any = null;
let modelListValid = 0;
const handleModelRequest: RequestHandler = (_req, res) => {
  if (new Date().getTime() - modelListValid < 1000 * 60) {
    return res.status(200).json(modelListCache);
  }
  const result = generateModelList(KNOWN_MODELS);
  modelListCache = { object: "list", data: result };
  modelListValid = new Date().getTime();
  res.status(200).json(modelListCache);
};

const openaiImagesResponseHandler: ProxyResHandlerWithBody = async (
  _proxyRes,
  req,
  res,
  body
) => {
  if (typeof body !== "object") {
    throw new Error("Expected body to be an object");
  }

  if (config.promptLogging) {
    const host = req.get("host");
    body.proxy_note = `Prompts are logged on this proxy instance. See ${host} for more information.`;
  }

  if (req.inboundApi === "openai") {
    req.log.info("Transforming OpenAI image response to OpenAI chat format");
    body = transformResponseForChat(body as OpenAIImageGenerationResult, req);
  }

  if (req.tokenizerInfo) {
    body.proxy_tokenizer = req.tokenizerInfo;
  }

  res.status(200).json(body);
};

/**
 * Transforms a DALL-E image generation response into a chat response, simply
 * embedding the image URL into the chat message as a Markdown image.
 */
function transformResponseForChat(
  imageBody: OpenAIImageGenerationResult,
  req: Request
): Record<string, any> {
  const prompt = imageBody.data[0].revised_prompt ?? req.body.prompt;
  const content = imageBody.data
    .map((item) => {
      const { url, b64_json } = item;
      if (b64_json) {
        return `![${prompt}](data:image/png;base64,${b64_json})`;
      } else {
        return `![${prompt}](${url})`;
      }
    })
    .join("\n\n");

  return {
    id: "dalle-" + req.id,
    object: "chat.completion",
    created: Date.now(),
    model: req.body.model,
    usage: {
      prompt_tokens: 0,
      completion_tokens: req.outputTokens,
      total_tokens: req.outputTokens,
    },
    choices: [
      {
        message: { role: "assistant", content },
        finish_reason: "stop",
        index: 0,
      },
    ],
  };
}

const openaiImagesProxy = createQueueMiddleware({
  proxyMiddleware: createProxyMiddleware({
    target: "https://api.openai.com",
    changeOrigin: true,
    selfHandleResponse: true,
    logger,
    pathRewrite: {
      "^/v1/chat/completions": "/v1/images/generations",
    },
    on: {
      proxyReq: createOnProxyReqHandler({ pipeline: [addKey, finalizeBody] }),
      proxyRes: createOnProxyResHandler([openaiImagesResponseHandler]),
      error: handleProxyError,
    },
  }),
});

const openaiImagesRouter = Router();
openaiImagesRouter.get("/v1/models", handleModelRequest);
openaiImagesRouter.post(
  "/v1/images/generations",
  ipLimiter,
  createPreprocessorMiddleware({
    inApi: "openai-image",
    outApi: "openai-image",
    service: "openai",
  }),
  openaiImagesProxy
);
openaiImagesRouter.post(
  "/v1/chat/completions",
  ipLimiter,
  createPreprocessorMiddleware({
    inApi: "openai",
    outApi: "openai-image",
    service: "openai",
  }),
  openaiImagesProxy
);
export const openaiImage = openaiImagesRouter;


code:
import { RequestHandler, Router } from "express";
import { createProxyMiddleware } from "http-proxy-middleware";
import { config } from "../config";
import { keyPool } from "../shared/key-management";
import {
  getOpenAIModelFamily,
  ModelFamily,
  OpenAIModelFamily,
} from "../shared/models";
import { logger } from "../logger";
import { createQueueMiddleware } from "./queue";
import { ipLimiter } from "./rate-limit";
import { handleProxyError } from "./middleware/common";
import {
  addKey,
  addKeyForEmbeddingsRequest,
  createEmbeddingsPreprocessorMiddleware,
  createOnProxyReqHandler,
  createPreprocessorMiddleware,
  finalizeBody,
  forceModel,
  RequestPreprocessor,
} from "./middleware/request";
import {
  createOnProxyResHandler,
  ProxyResHandlerWithBody,
} from "./middleware/response";

// https://platform.openai.com/docs/models/overview
export const KNOWN_OPENAI_MODELS = [
  "gpt-4-turbo-preview",
  "gpt-4-0125-preview",
  "gpt-4-1106-preview",
  "gpt-4-vision-preview",
  "gpt-4",
  "gpt-4-0613",
  "gpt-4-0314", // EOL 2024-06-13
  "gpt-4-32k",
  "gpt-4-32k-0613",
  "gpt-4-32k-0314",
  "gpt-3.5-turbo",
  "gpt-3.5-turbo-0301", // EOL 2024-06-13
  "gpt-3.5-turbo-0613",
  "gpt-3.5-turbo-16k",
  "gpt-3.5-turbo-16k-0613",
  "gpt-3.5-turbo-instruct",
  "gpt-3.5-turbo-instruct-0914",
  "text-embedding-ada-002",
];

let modelsCache: any = null;
let modelsCacheTime = 0;

export function generateModelList(models = KNOWN_OPENAI_MODELS) {
  let available = new Set<OpenAIModelFamily>();
  for (const key of keyPool.list()) {
    if (key.isDisabled || key.service !== "openai") continue;
    key.modelFamilies.forEach((family) =>
      available.add(family as OpenAIModelFamily)
    );
  }
  const allowed = new Set<ModelFamily>(config.allowedModelFamilies);
  available = new Set([...available].filter((x) => allowed.has(x)));

  return models
    .map((id) => ({
      id,
      object: "model",
      created: new Date().getTime(),
      owned_by: "openai",
      permission: [
        {
          id: "modelperm-" + id,
          object: "model_permission",
          created: new Date().getTime(),
          organization: "*",
          group: null,
          is_blocking: false,
        },
      ],
      root: id,
      parent: null,
    }))
    .filter((model) => available.has(getOpenAIModelFamily(model.id)));
}

const handleModelRequest: RequestHandler = (_req, res) => {
  if (new Date().getTime() - modelsCacheTime < 1000 * 60) {
    return res.status(200).json(modelsCache);
  }
  const result = generateModelList();
  modelsCache = { object: "list", data: result };
  modelsCacheTime = new Date().getTime();
  res.status(200).json(modelsCache);
};

/** Handles some turbo-instruct special cases. */
const rewriteForTurboInstruct: RequestPreprocessor = (req) => {
  // /v1/turbo-instruct/v1/chat/completions accepts either prompt or messages.
  // Depending on whichever is provided, we need to set the inbound format so
  // it is transformed correctly later.
  if (req.body.prompt && !req.body.messages) {
    req.inboundApi = "openai-text";
  } else if (req.body.messages && !req.body.prompt) {
    req.inboundApi = "openai";
    // Set model for user since they're using a client which is not aware of
    // turbo-instruct.
    req.body.model = "gpt-3.5-turbo-instruct";
  } else {
    throw new Error("`prompt` OR `messages` must be provided");
  }

  req.url = "/v1/completions";
};

const openaiResponseHandler: ProxyResHandlerWithBody = async (
  _proxyRes,
  req,
  res,
  body
) => {
  if (typeof body !== "object") {
    throw new Error("Expected body to be an object");
  }

  if (config.promptLogging) {
    const host = req.get("host");
    body.proxy_note = `Prompts are logged on this proxy instance. See ${host} for more information.`;
  }

  if (req.outboundApi === "openai-text" && req.inboundApi === "openai") {
    req.log.info("Transforming Turbo-Instruct response to Chat format");
    body = transformTurboInstructResponse(body);
  }

  if (req.tokenizerInfo) {
    body.proxy_tokenizer = req.tokenizerInfo;
  }

  res.status(200).json(body);
};

/** Only used for non-streaming responses. */
function transformTurboInstructResponse(
  turboInstructBody: Record<string, any>
): Record<string, any> {
  const transformed = { ...turboInstructBody };
  transformed.choices = [
    {
      ...turboInstructBody.choices[0],
      message: {
        role: "assistant",
        content: turboInstructBody.choices[0].text.trim(),
      },
    },
  ];
  delete transformed.choices[0].text;
  return transformed;
}

const openaiProxy = createQueueMiddleware({
  proxyMiddleware: createProxyMiddleware({
    target: "https://api.openai.com",
    changeOrigin: true,
    selfHandleResponse: true,
    logger,
    on: {
      proxyReq: createOnProxyReqHandler({ pipeline: [addKey, finalizeBody], }),
      proxyRes: createOnProxyResHandler([openaiResponseHandler]),
      error: handleProxyError,
    },
  }),
});

const openaiEmbeddingsProxy = createProxyMiddleware({
  target: "https://api.openai.com",
  changeOrigin: true,
  selfHandleResponse: false,
  logger,
  on: {
    proxyReq: createOnProxyReqHandler({
      pipeline: [addKeyForEmbeddingsRequest, finalizeBody],
    }),
    error: handleProxyError,
  },
});

const openaiRouter = Router();
openaiRouter.get("/v1/models", handleModelRequest);
// Native text completion endpoint, only for turbo-instruct.
openaiRouter.post(
  "/v1/completions",
  ipLimiter,
  createPreprocessorMiddleware({
    inApi: "openai-text",
    outApi: "openai-text",
    service: "openai",
  }),
  openaiProxy
);
// turbo-instruct compatibility endpoint, accepts either prompt or messages
openaiRouter.post(
  /\/v1\/turbo-instruct\/(v1\/)?chat\/completions/,
  ipLimiter,
  createPreprocessorMiddleware(
    { inApi: "openai", outApi: "openai-text", service: "openai" },
    {
      beforeTransform: [rewriteForTurboInstruct],
      afterTransform: [forceModel("gpt-3.5-turbo-instruct")],
    }
  ),
  openaiProxy
);
// General chat completion endpoint. Turbo-instruct is not supported here.
openaiRouter.post(
  "/v1/chat/completions",
  ipLimiter,
  createPreprocessorMiddleware({
    inApi: "openai",
    outApi: "openai",
    service: "openai",
  }),
  openaiProxy
);
// Embeddings endpoint.
openaiRouter.post(
  "/v1/embeddings",
  ipLimiter,
  createEmbeddingsPreprocessorMiddleware(),
  openaiEmbeddingsProxy
);

export const openai = openaiRouter;