code:
import { createHash } from "crypto";
import { config } from "../../../config";
import { eventLogger } from "../../../shared/prompt-logging";
import { getModelFromBody, isTextGenerationRequest } from "../common";
import { ProxyResHandlerWithBody } from ".";
import {
  OpenAIChatMessage,
  AnthropicChatMessage,
} from "../../../shared/api-schemas";

/** If event logging is enabled, logs a chat completion event. */
export const logEvent: ProxyResHandlerWithBody = async (
  _proxyRes,
  req,
  _res,
  responseBody
) => {
  if (!config.eventLogging) {
    return;
  }
  if (typeof responseBody !== "object") {
    throw new Error("Expected body to be an object");
  }
  if (!["openai", "anthropic-chat"].includes(req.outboundApi)) {
    // only chat apis are supported
    return;
  }
  if (!req.user) {
    return;
  }

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

  const messages = req.body.messages as
    | OpenAIChatMessage[]
    | AnthropicChatMessage[];

  let hashes = [];
  hashes.push(hashMessages(messages));
  for (
    let i = 1;
    i <= Math.min(config.eventLoggingTrim!, messages.length);
    i++
  ) {
    hashes.push(hashMessages(messages.slice(0, -i)));
  }

  const model = getModelFromBody(req, responseBody);
  const userToken = req.user!.token;
  const family = req.modelFamily!;
  eventLogger.logEvent({
    ip: req.ip,
    type: "chat_completion",
    model,
    family,
    hashes,
    userToken,
    inputTokens: req.promptTokens ?? 0,
    outputTokens: req.outputTokens ?? 0,
  });
};

const hashMessages = (
  messages: OpenAIChatMessage[] | AnthropicChatMessage[]
): string => {
  let hasher = createHash("sha256");
  let messageTexts = [];
  for (const msg of messages) {
    if (!["system", "user", "assistant"].includes(msg.role)) continue;
    if (typeof msg.content === "string") {
      messageTexts.push(msg.content);
    } else if (Array.isArray(msg.content)) {
      if (msg.content[0].type === "text") {
        messageTexts.push(msg.content[0].text);
      }
    }
  }
  hasher.update(messageTexts.join("<|im_sep|>"));
  return hasher.digest("hex");
};


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,
  GoogleAIChatMessage,
  MistralAIChatMessage,
  OpenAIChatMessage,
} from "../../../shared/api-schemas";

/** 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);
  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[]
  | { contents: GoogleAIChatMessage[] }
  | { system: string; messages: 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":
      return req.body.messages;
    case "anthropic-chat":
      let system = req.body.system;
      if (Array.isArray(system)) {
        system = system
          .map((m: { type: string; text: string }) => m.text)
          .join("\n");
      }
      return { system, messages: req.body.messages };
    case "openai-text":
    case "anthropic-text":
    case "mistral-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 "google-ai":
      return { contents: req.body.contents };
    default:
      assertNever(req.outboundApi);
  }
};

const flattenMessages = (
  val:
    | string
    | OaiImageResult
    | OpenAIChatMessage[]
    | { contents: GoogleAIChatMessage[] }
    | { system: string; messages: AnthropicChatMessage[] }
    | MistralAIChatMessage[]
): string => {
  if (typeof val === "string") {
    return val.trim();
  }
  if (isAnthropicChatPrompt(val)) {
    const { system, messages } = val;
    return `System: ${system}\n\n${flattenAnthropicMessages(messages)}`;
  }
  if (isGoogleAIChatPrompt(val)) {
    return val.contents
      .map(({ parts, role }) => {
        const text = parts.map((p) => p.text).join("\n");
        return `${role}: ${text}`;
      })
      .join("\n");
  }
  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();
};

function isGoogleAIChatPrompt(
  val: unknown
): val is { contents: GoogleAIChatMessage[] } {
  return typeof val === "object" && val !== null && "contents" in val;
}

function isAnthropicChatPrompt(
  val: unknown
): val is { system: string; messages: AnthropicChatMessage[] } {
  return (
    typeof val === "object" &&
    val !== null &&
    "system" in val &&
    "messages" in val
  );
}

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 prompt = body.data[0].revised_prompt ?? req.body.prompt;
    const res = await mirrorGeneratedImage(
      req,
      prompt,
      body as OpenAIImageGenerationResult
    );
    req.log.info(
      { urls: res.data.map((item) => item.url) },
      "Saved generated image to user_content"
    );
  }
};


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 MistralChatCompletionResponse = {
  choices: {
    index: number;
    message: { role: string; content: string };
    finish_reason: string | null;
  }[];
};

/**
 * Given a list of OpenAI chat completion events, compiles them into a single
 * finalized Mistral chat completion response so that non-streaming middleware
 * can operate on it as if it were a blocking response.
 */
export function mergeEventsForMistralChat(
  events: OpenAIChatCompletionStreamEvent[]
): MistralChatCompletionResponse {
  let merged: MistralChatCompletionResponse = {
    choices: [
      { index: 0, message: { role: "", content: "" }, finish_reason: "" },
    ],
  };
  merged = events.reduce((acc, event, i) => {
    // The first event will only contain role assignment and response metadata
    if (i === 0) {
      acc.choices[0].message.role = event.choices[0].delta.role ?? "assistant";
      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 MistralTextCompletionResponse = {
  outputs: {
    text: string;
    stop_reason: string | null;
  }[];
};

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

    acc.outputs[0].text += event.choices[0].delta.content ?? "";
    acc.outputs[0].stop_reason = event.choices[0].finish_reason ?? "";

    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 events (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 express from "express";
import { APIFormat } from "../../../../shared/key-management";
import { assertNever } from "../../../../shared/utils";
import {
  anthropicV2ToOpenAI,
  mergeEventsForAnthropicChat,
  mergeEventsForAnthropicText,
  mergeEventsForOpenAIChat,
  mergeEventsForOpenAIText,
  mergeEventsForMistralChat,
  mergeEventsForMistralText,
  AnthropicV2StreamEvent,
  OpenAIChatCompletionStreamEvent,
  mistralAIToOpenAI,
  MistralAIStreamEvent,
  MistralChatCompletionEvent,
} 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 model: string;
  private readonly requestFormat: APIFormat;
  private readonly responseFormat: APIFormat;
  private readonly events: OpenAIChatCompletionStreamEvent[];

  constructor({ body, inboundApi, outboundApi }: express.Request) {
    this.events = [];
    this.requestFormat = inboundApi;
    this.responseFormat = outboundApi;
    this.model = body.model;
  }

  addEvent(
    event:
      | OpenAIChatCompletionStreamEvent
      | AnthropicV2StreamEvent
      | MistralAIStreamEvent
  ) {
    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 some transformers that convert between non-openai
      // formats, so aggregator needs to know how to collapse for more than
      // just openai.
      // because writing aggregation logic for every possible output format is
      // annoying, we will just transform any non-openai output events to openai
      // format (even if the client did not request openai at all) so that we
      // still only need to write aggregators for openai SSEs.
      let openAIEvent: OpenAIChatCompletionStreamEvent | undefined;
      switch (this.requestFormat) {
        case "anthropic-text":
          assertIsAnthropicV2Event(event);
          openAIEvent = anthropicV2ToOpenAI({
            data: `event: completion\ndata: ${JSON.stringify(event)}\n\n`,
            lastPosition: -1,
            index: 0,
            fallbackId: event.log_id || "fallback-" + Date.now(),
            fallbackModel: event.model || this.model || "fallback-claude-3",
          })?.event;
          break;
        case "mistral-ai":
          assertIsMistralChatEvent(event);
          openAIEvent = mistralAIToOpenAI({
            data: `data: ${JSON.stringify(event)}\n\n`,
            lastPosition: -1,
            index: 0,
            fallbackId: "fallback-" + Date.now(),
            fallbackModel: this.model || "fallback-mistral",
          })?.event;
          break;
      }
      if (openAIEvent) {
        this.events.push(openAIEvent);
      }
    }
  }

  getFinalResponse() {
    switch (this.responseFormat) {
      case "openai":
      case "google-ai": // TODO: this is probably wrong now that we support native Google Makersuite prompts
        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 "mistral-ai":
        return mergeEventsForMistralChat(this.events);
      case "mistral-text":
        return mergeEventsForMistralText(this.events);
      case "openai-image":
        throw new Error(
          `SSE aggregation not supported for ${this.responseFormat}`
        );
      default:
        assertNever(this.responseFormat);
    }
  }

  hasEvents() {
    return this.events.length > 0;
  }
}

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

function assertIsAnthropicV2Event(event: any): asserts event is AnthropicV2StreamEvent {
  if (!event?.completion) {
    throw new Error(`Bad event for Anthropic V2 SSE aggregation`);
  }
}

function assertIsMistralChatEvent(
  event: any
): asserts event is MistralChatCompletionEvent {
  if (!event?.choices) {
    throw new Error(`Bad event for Mistral SSE aggregation`);
  }
}


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

export type MistralChatCompletionEvent = {
  choices: {
    index: number;
    message: { role: string; content: string };
    stop_reason: string | null;
  }[];
};
export type MistralTextCompletionEvent = {
  outputs: { text: string; stop_reason: string | null }[];
};
export type MistralAIStreamEvent = {
  "amazon-bedrock-invocationMetrics"?: {
    inputTokenCount: number;
    outputTokenCount: number;
    invocationLatency: number;
    firstByteLatency: number;
  };
} & (MistralChatCompletionEvent | MistralTextCompletionEvent);

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 { anthropicChatToOpenAI } from "./transformers/anthropic-chat-to-openai";
export { googleAIToOpenAI } from "./transformers/google-ai-to-openai";
export { mistralAIToOpenAI } from "./transformers/mistral-ai-to-openai";
export { mistralTextToMistralChat } from "./transformers/mistral-text-to-mistral-chat";
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";
export { mergeEventsForMistralChat } from "./aggregators/mistral-chat";
export { mergeEventsForMistralText } from "./aggregators/mistral-text";


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 {
  anthropicChatToOpenAI,
  anthropicChatToAnthropicV2,
  anthropicV1ToOpenAI,
  AnthropicV2StreamEvent,
  anthropicV2ToOpenAI,
  googleAIToOpenAI,
  OpenAIChatCompletionStreamEvent,
  openAITextToOpenAIChat,
  mistralAIToOpenAI,
  mistralTextToMistralChat,
  passthroughToOpenAI,
  StreamingCompletionTransformer,
  MistralChatCompletionEvent,
} from "./index";

type SSEMessageTransformerOptions = TransformOptions & {
  requestedModel: string;
  requestId: string;
  inputFormat: APIFormat;
  inputApiVersion?: string;
  outputFormat?: APIFormat;
  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
    | MistralChatCompletionEvent
  >;
  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,
      options.outputFormat
    );
    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,
  // In most cases, we are transforming back to OpenAI. Some responses can be
  // translated between two non-OpenAI formats, eg Anthropic Chat -> Anthropic
  // Text, or Mistral Text -> Mistral Chat.
  requestApi: APIFormat = "openai"
): StreamingCompletionTransformer<
  | OpenAIChatCompletionStreamEvent
  | AnthropicV2StreamEvent
  | MistralChatCompletionEvent
> {
  switch (responseApi) {
    case "openai":
      return passthroughToOpenAI;
    case "openai-text":
      return openAITextToOpenAIChat;
    case "anthropic-text":
      return version === "2023-01-01"
        ? anthropicV1ToOpenAI
        : anthropicV2ToOpenAI;
    case "anthropic-chat":
      return requestApi === "anthropic-text"
        ? anthropicChatToAnthropicV2 // User's legacy text prompt was converted to chat, and response must be converted back to text
        : anthropicChatToOpenAI;
    case "google-ai":
      return googleAIToOpenAI;
    case "mistral-ai":
      return mistralAIToOpenAI;
    case "mistral-text":
      return requestApi === "mistral-ai"
        ? mistralTextToMistralChat // User's chat request was converted to text, and response must be converted back to chat
        : mistralAIToOpenAI;
    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 { BadRequestError, RetryableError } from "../../../../shared/errors";

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 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.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 if (eventObj.type) {
            return [`event: ${eventObj.type}`, `data: ${event}`].join(`\n`);
          } else {
            return `data: ${event}`;
          }
        }
      // noinspection FallThroughInSwitchStatementJS -- non-JSON data is unexpected
      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");
          case "validationexception":
            try {
              const { message } = JSON.parse(bodyStr);
              this.log.error({ message }, "Received AWS validation error");
              this.emit(
                "error",
                new BadRequestError(`AWS validation error: ${message}`)
              );
              return null;
            } catch (error) {
              this.log.error(
                { body: bodyStr, error },
                "Could not parse AWS validation error"
              );
            }
          // noinspection FallThroughInSwitchStatementJS -- who knows what this is
          default:
            let text;
            try {
              text = JSON.parse(bodyStr).message;
            } catch (error) {
              text = bodyStr;
            }
            const error: any = new Error(
              `Got mysterious error chunk: [${type}] ${text}`
            );
            error.lastEvent = text;
            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;
    }
  }

  _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 {
        // `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: 0,
        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 (isSafetyStop(completionEvent)) {
    content = `[Proxy Warning] Gemini safety filter triggered: ${JSON.stringify(
      completionEvent.candidates[0].safetyRatings
    )}`;
  }

  // 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 isSafetyStop(completion: GoogleAIStreamEvent) {
  const isSafetyStop = ["SAFETY", "OTHER"].includes(
    completion.candidates[0].finishReason ?? ""
  );
  const hasNoContent = completion.candidates[0].content?.parts?.length === 0;
  return isSafetyStop && hasNoContent;
}

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 { logger } from "../../../../../logger";
import { MistralAIStreamEvent, SSEResponseTransformArgs } from "../index";
import { parseEvent, ServerSentEvent } from "../parse-sse";

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

export const mistralAIToOpenAI = (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 };
  }

  if ("choices" in completionEvent) {
    const newChatEvent = {
      id: params.fallbackId,
      object: "chat.completion.chunk" as const,
      created: Date.now(),
      model: params.fallbackModel,
      choices: [
        {
          index: completionEvent.choices[0].index,
          delta: { content: completionEvent.choices[0].message.content },
          finish_reason: completionEvent.choices[0].stop_reason,
        },
      ],
    };
    return { position: -1, event: newChatEvent };
  } else if ("outputs" in completionEvent) {
    const newTextEvent = {
      id: params.fallbackId,
      object: "chat.completion.chunk" as const,
      created: Date.now(),
      model: params.fallbackModel,
      choices: [
        {
          index: 0,
          delta: { content: completionEvent.outputs[0].text },
          finish_reason: completionEvent.outputs[0].stop_reason,
        },
      ],
    };
    return { position: -1, event: newTextEvent };
  }

  // should never happen
  return { position: -1 };
};

function asCompletion(event: ServerSentEvent): MistralAIStreamEvent | null {
  try {
    const parsed = JSON.parse(event.data);
    if (
      (Array.isArray(parsed.choices) &&
        parsed.choices[0].message !== undefined) ||
      (Array.isArray(parsed.outputs) && parsed.outputs[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 {
  MistralChatCompletionEvent,
  MistralTextCompletionEvent,
  StreamingCompletionTransformer,
} from "../index";
import { parseEvent, ServerSentEvent } from "../parse-sse";
import { logger } from "../../../../../logger";

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

/**
 * Transforms an incoming Mistral Text SSE to an equivalent Mistral Chat SSE.
 * This is generally used when a client sends a Mistral Chat prompt, but we
 * convert it to Mistral Text before sending it to the API to work around
 * some bugs in Mistral/AWS prompt templating. In these cases we need to convert
 * the response back to Mistral Chat.
 */
export const mistralTextToMistralChat: StreamingCompletionTransformer<
  MistralChatCompletionEvent
> = (params) => {
  const { data } = params;

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

  const textCompletion = asTextCompletion(rawEvent);
  if (!textCompletion) {
    return { position: -1 };
  }

  const chatEvent: MistralChatCompletionEvent = {
    choices: [
      {
        index: 0,
        message: { role: "assistant", content: textCompletion.outputs[0].text },
        stop_reason: textCompletion.outputs[0].stop_reason,
      },
    ],
  };
  return { position: -1, event: chatEvent };
};

function asTextCompletion(
  event: ServerSentEvent
): MistralTextCompletionEvent | null {
  try {
    const parsed = JSON.parse(event.data);
    if (Array.isArray(parsed.outputs) && parsed.outputs[0].text !== undefined) {
      return parsed;
    } else {
      // noinspection ExceptionCaughtLocallyJS
      throw new Error("Missing required fields");
    }
  } catch (error: any) {
    log.warn({ error: error.stack, event }, "Received invalid data 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 { Request, RequestHandler, Router } from "express";
import { BadRequestError } from "../shared/errors";
import { keyPool } from "../shared/key-management";
import {
  getMistralAIModelFamily,
  MistralAIModelFamily,
  ModelFamily,
} from "../shared/models";
import { config } from "../config";
import { ipLimiter } from "./rate-limit";
import {
  addKey,
  createPreprocessorMiddleware,
  finalizeBody,
} from "./middleware/request";
import { ProxyResHandlerWithBody } from "./middleware/response";
import { createQueuedProxyMiddleware } from "./middleware/request/proxy-middleware-factory";

// Mistral can't settle on a single naming scheme and deprecates models within
// months of releasing them so this list is hard to keep up to date. 2024-07-28
// https://docs.mistral.ai/platform/endpoints
export const KNOWN_MISTRAL_AI_MODELS = [
  /*
  Mistral Nemo
  "A 12B model built with the partnership with Nvidia.  It is easy to use and a
  drop-in replacement in any system using Mistral 7B that it supersedes."
  */
  "open-mistral-nemo",
  "open-mistral-nemo-2407",
  /*
  Mistral Large
  "Our flagship model with state-of-the-art reasoning, knowledge, and coding
  capabilities."
  */
  "mistral-large-latest",
  "mistral-large-2407",
  "mistral-large-2402", // deprecated
  /*
  Codestral
  "A cutting-edge generative model that has been specifically designed and
  optimized for code generation tasks, including fill-in-the-middle and code
  completion."
  note: this uses a separate bidi completion endpoint that is not implemented
  */
  "codestral-latest",
  "codestral-2405",
  /* So-called "Research Models" */
  "open-mistral-7b",
  "open-mixtral-8x7b",
  "open-mistral-8x22b",
  "open-codestral-mamba",
  /* Deprecated production models */
  "mistral-small-latest",
  "mistral-small-2402",
  "mistral-medium-latest",
  "mistral-medium-2312",
  "mistral-tiny",
  "mistral-tiny-2312",
];

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");
  }

  let newBody = body;
  if (req.inboundApi === "mistral-text" && req.outboundApi === "mistral-ai") {
    newBody = transformMistralTextToMistralChat(body);
  }

  res.status(200).json({ ...newBody, proxy: body.proxy });
};

export function transformMistralTextToMistralChat(textBody: any) {
  return {
    ...textBody,
    choices: [
      { message: { content: textBody.outputs[0].text, role: "assistant" } },
    ],
    outputs: undefined,
  };
}

const mistralAIProxy = createQueuedProxyMiddleware({
  target: "https://api.mistral.ai",
  mutations: [addKey, finalizeBody],
  blockingResponseHandler: mistralAIResponseHandler,
});

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",
    },
    { beforeTransform: [detectMistralInputApi] }
  ),
  mistralAIProxy
);

/**
 * We can't determine if a request is Mistral text or chat just from the path
 * because they both use the same endpoint. We need to check the request body
 * for either `messages` or `prompt`.
 * @param req
 */
export function detectMistralInputApi(req: Request) {
  const { messages, prompt } = req.body;
  if (messages) {
    req.inboundApi = "mistral-ai";
    req.outboundApi = "mistral-ai";
  } else if (prompt && req.service === "mistral-ai") {
    // Mistral La Plateforme doesn't expose a text completions endpoint.
    throw new BadRequestError(
      "Mistral (via La Plateforme API) does not support text completions. This format is only supported on Mistral via the AWS API."
    );
  } else if (prompt && req.service === "aws") {
    req.inboundApi = "mistral-text";
    req.outboundApi = "mistral-text";
  }
}

export const mistralAI = mistralAIRouter;


code:
import { Request, RequestHandler, Router } from "express";
import { OpenAIImageGenerationResult } from "../shared/file-storage/mirror-generated-image";
import { generateModelList } from "./openai";
import { ipLimiter } from "./rate-limit";
import {
  addKey,
  createPreprocessorMiddleware,
  finalizeBody,
} from "./middleware/request";
import { ProxyResHandlerWithBody } from "./middleware/response";
import { ProxyReqManager } from "./middleware/request/proxy-req-manager";
import { createQueuedProxyMiddleware } from "./middleware/request/proxy-middleware-factory";

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("openai").filter((m: { id: string }) =>
    KNOWN_MODELS.includes(m.id)
  );
  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");
  }

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

  res.status(200).json({ ...newBody, proxy: body.proxy });
};

/**
 * 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,
      },
    ],
  };
}

function replacePath(manager: ProxyReqManager) {
  const req = manager.request;
  const pathname = req.url.split("?")[0];
  req.log.debug({ pathname }, "OpenAI image path filter");
  if (req.path.startsWith("/v1/chat/completions")) {
    manager.setPath("/v1/images/generations");
  }
}

const openaiImagesProxy = createQueuedProxyMiddleware({
  target: "https://api.openai.com",
  mutations: [replacePath, addKey, finalizeBody],
  blockingResponseHandler: openaiImagesResponseHandler,
});

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 { Request, RequestHandler, Router } from "express";
import { config } from "../config";
import { AzureOpenAIKey, keyPool, OpenAIKey } from "../shared/key-management";
import { getOpenAIModelFamily } from "../shared/models";
import { ipLimiter } from "./rate-limit";
import {
  addKey,
  addKeyForEmbeddingsRequest,
  createEmbeddingsPreprocessorMiddleware,
  createPreprocessorMiddleware,
  finalizeBody,
  RequestPreprocessor,
} from "./middleware/request";
import { ProxyResHandlerWithBody } from "./middleware/response";
import { createQueuedProxyMiddleware } from "./middleware/request/proxy-middleware-factory";

// https://platform.openai.com/docs/models/overview
let modelsCache: any = null;
let modelsCacheTime = 0;

export function generateModelList(service: "openai" | "azure") {
  const keys = keyPool
    .list()
    .filter((k) => k.service === service && !k.isDisabled) as
    | OpenAIKey[]
    | AzureOpenAIKey[];
  if (keys.length === 0) return [];

  const allowedModelFamilies = new Set(config.allowedModelFamilies);
  const modelFamilies = new Set(
    keys
      .flatMap((k) => k.modelFamilies)
      .filter((f) => allowedModelFamilies.has(f))
  );

  const modelIds = new Set(
    keys
      .flatMap((k) => k.modelIds)
      .filter((id) => {
        const allowed = modelFamilies.has(getOpenAIModelFamily(id));
        const known = ["gpt", "o1", "dall-e", "chatgpt", "text-embedding"].some(
          (prefix) => id.startsWith(prefix)
        );
        const isFinetune = id.includes("ft");
        return allowed && known && !isFinetune;
      })
  );

  return Array.from(modelIds).map((id) => ({
    id,
    object: "model",
    created: new Date().getTime(),
    owned_by: service,
    permission: [
      {
        id: "modelperm-" + id,
        object: "model_permission",
        created: new Date().getTime(),
        organization: "*",
        group: null,
        is_blocking: false,
      },
    ],
    root: id,
    parent: null,
  }));
}

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

  if (!config.openaiKey) return { object: "list", data: [] };

  const result = generateModelList("openai");

  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");
  }

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

  res.status(200).json({ ...newBody, proxy: body.proxy });
};

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 = createQueuedProxyMiddleware({
  mutations: [addKey, finalizeBody],
  target: "https://api.openai.com",
  blockingResponseHandler: openaiResponseHandler,
});

const openaiEmbeddingsProxy = createQueuedProxyMiddleware({
  mutations: [addKeyForEmbeddingsRequest, finalizeBody],
  target: "https://api.openai.com",
});

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" },
    { afterTransform: [fixupMaxTokens] }
  ),
  openaiProxy
);
// Embeddings endpoint.
openaiRouter.post(
  "/v1/embeddings",
  ipLimiter,
  createEmbeddingsPreprocessorMiddleware(),
  openaiEmbeddingsProxy
);

function forceModel(model: string): RequestPreprocessor {
  return (req: Request) => void (req.body.model = model);
}

function fixupMaxTokens(req: Request) {
  if (!req.body.max_completion_tokens) {
    req.body.max_completion_tokens = req.body.max_tokens;
  }
  delete req.body.max_tokens;
}

export const openai = openaiRouter;


code:
/**
 * Very scuffed request queue. OpenAI's GPT-4 keys have a very strict rate limit
 * of 40000 generated tokens per minute. We don't actually know how many tokens
 * a given key has generated, so our queue will simply retry requests that fail
 * with a non-billing related 429 over and over again until they succeed.
 *
 * When a request to a proxied endpoint is received, we create a closure around
 * the call to http-proxy-middleware and attach it to the request. This allows
 * us to pause the request until we have a key available. Further, if the
 * proxied request encounters a retryable error, we can simply put the request
 * back in the queue and it will be retried later using the same closure.
 */

import crypto from "crypto";
import { Handler, Request } from "express";
import { config } from "../config";
import { BadRequestError, TooManyRequestsError } from "../shared/errors";
import { keyPool } from "../shared/key-management";
import {
  getModelFamilyForRequest,
  MODEL_FAMILIES,
  ModelFamily,
} from "../shared/models";
import { initializeSseStream } from "../shared/streaming";
import { logger } from "../logger";
import { getUniqueIps } from "./rate-limit";
import { ProxyReqMutator, RequestPreprocessor } from "./middleware/request";
import { sendErrorToClient } from "./middleware/response/error-generator";
import { ProxyReqManager } from "./middleware/request/proxy-req-manager";
import { classifyErrorAndSend } from "./middleware/common";

const queue: Request[] = [];
const log = logger.child({ module: "request-queue" });

/** Maximum number of queue slots for individual users. */
const USER_CONCURRENCY_LIMIT = parseInt(
  process.env.USER_CONCURRENCY_LIMIT ?? "1"
);
const MIN_HEARTBEAT_SIZE = parseInt(process.env.MIN_HEARTBEAT_SIZE_B ?? "512");
const MAX_HEARTBEAT_SIZE =
  1024 * parseInt(process.env.MAX_HEARTBEAT_SIZE_KB ?? "1024");
const HEARTBEAT_INTERVAL =
  1000 * parseInt(process.env.HEARTBEAT_INTERVAL_SEC ?? "5");
const LOAD_THRESHOLD = parseFloat(process.env.LOAD_THRESHOLD ?? "150");
const PAYLOAD_SCALE_FACTOR = parseFloat(
  process.env.PAYLOAD_SCALE_FACTOR ?? "6"
);
const QUEUE_JOIN_TIMEOUT = 5000;

/**
 * Returns an identifier for a request. This is used to determine if a
 * request is already in the queue.
 *
 * This can be (in order of preference):
 * - user token assigned by the proxy operator
 * - x-risu-tk header, if the request is from RisuAI.xyz
 * - 'shared-ip' if the request is from a shared IP address like Agnai.chat
 * - IP address
 */
function getIdentifier(req: Request) {
  if (req.user) return req.user.token;
  if (req.risuToken) return req.risuToken;
  // if (isFromSharedIp(req)) return "shared-ip";
  return req.ip;
}

const sharesIdentifierWith = (incoming: Request) => (queued: Request) =>
  getIdentifier(queued) === getIdentifier(incoming);

async function enqueue(req: Request) {
  if (req.socket.destroyed || req.res?.writableEnded) {
    // In rare cases, a request can be disconnected after it is dequeued for a
    // retry, but before it is re-enqueued. In this case we may miss the abort
    // and the request will loop in the queue forever.
    req.log.warn("Attempt to enqueue aborted request.");
    throw new Error("Attempt to enqueue aborted request.");
  }

  const enqueuedRequestCount = queue.filter(sharesIdentifierWith(req)).length;

  if (enqueuedRequestCount >= USER_CONCURRENCY_LIMIT) {
    throw new TooManyRequestsError(
      "Your IP or user token already has another request in the queue."
    );
  }

  // shitty hack to remove hpm's event listeners on retried requests
  removeProxyMiddlewareEventListeners(req);

  // If the request opted into streaming, we need to register a heartbeat
  // handler to keep the connection alive while it waits in the queue. We
  // deregister the handler when the request is dequeued.
  const { stream } = req.body;
  if (stream === "true" || stream === true || req.isStreaming) {
    const res = req.res!;
    if (!res.headersSent) {
      await initStreaming(req);
    }
    registerHeartbeat(req);
  } else if (getProxyLoad() > LOAD_THRESHOLD) {
    throw new BadRequestError(
      "Due to heavy traffic on this proxy, you must enable streaming in your chat client to use this endpoint."
    );
  }

  queue.push(req);
  req.queueOutTime = 0;

  const removeFromQueue = () => {
    req.log.info(`Removing aborted request from queue.`);
    const index = queue.indexOf(req);
    if (index !== -1) {
      queue.splice(index, 1);
    }
    if (req.heartbeatInterval) clearInterval(req.heartbeatInterval);
    if (req.monitorInterval) clearInterval(req.monitorInterval);
  };
  req.onAborted = removeFromQueue;
  req.res!.once("close", removeFromQueue);

  if (req.retryCount ?? 0 > 0) {
    req.log.info({ retries: req.retryCount }, `Enqueued request for retry.`);
  } else {
    const size = req.socket.bytesRead;
    const endpoint = req.url?.split("?")[0];
    req.log.info({ size, endpoint }, `Enqueued new request.`);
  }
}

export 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);
}

function getQueueForPartition(partition: ModelFamily): Request[] {
  return queue.filter((req) => getModelFamilyForRequest(req) === partition);
}

export function dequeue(partition: ModelFamily): Request | undefined {
  const modelQueue = getQueueForPartition(partition);

  if (modelQueue.length === 0) {
    return undefined;
  }

  const req = modelQueue.reduce((prev, curr) =>
    prev.startTime +
      config.tokensPunishmentFactor *
        ((prev.promptTokens ?? 0) + (prev.outputTokens ?? 0)) <
    curr.startTime +
      config.tokensPunishmentFactor *
        ((curr.promptTokens ?? 0) + (curr.outputTokens ?? 0))
      ? prev
      : curr
  );
  queue.splice(queue.indexOf(req), 1);

  if (req.onAborted) {
    req.res!.off("close", req.onAborted);
    req.onAborted = undefined;
  }

  if (req.heartbeatInterval) clearInterval(req.heartbeatInterval);
  if (req.monitorInterval) clearInterval(req.monitorInterval);

  // Track the time leaving the queue now, but don't add it to the wait times
  // yet because we don't know if the request will succeed or fail. We track
  // the time now and not after the request succeeds because we don't want to
  // include the model processing time.
  req.queueOutTime = Date.now();
  return req;
}

/**
 * Naive way to keep the queue moving by continuously dequeuing requests. Not
 * ideal because it limits throughput but we probably won't have enough traffic
 * or keys for this to be a problem.  If it does we can dequeue multiple
 * per tick.
 **/
function processQueue() {
  // This isn't completely correct, because a key can service multiple models.
  // Currently if a key is locked out on one model it will also stop servicing
  // the others, because we only track rate limits for the key as a whole.

  const reqs: (Request | undefined)[] = [];
  MODEL_FAMILIES.forEach((modelFamily) => {
    const lockout = keyPool.getLockoutPeriod(modelFamily);
    if (lockout === 0) {
      reqs.push(dequeue(modelFamily));
    }
  });

  reqs.filter(Boolean).forEach((req) => {
    if (req?.proceed) {
      const modelFamily = getModelFamilyForRequest(req!);
      req.log.info(
        { retries: req.retryCount, partition: modelFamily },
        `Dequeuing request.`
      );
      req.proceed();
    }
  });
  setTimeout(processQueue, 50);
}

/**
 * Kill stalled requests after 5 minutes, and remove tracked wait times after 2
 * minutes.
 **/
function cleanQueue() {
  const now = Date.now();
  const oldRequests = queue.filter(
    (req) => now - (req.startTime ?? now) > 5 * 60 * 1000
  );
  oldRequests.forEach((req) => {
    req.log.info(`Removing request from queue after 5 minutes.`);
    killQueuedRequest(req);
  });

  const index = waitTimes.findIndex(
    (waitTime) => now - waitTime.end > 300 * 1000
  );
  const removed = waitTimes.splice(0, index + 1);
  log.trace(
    { stalledRequests: oldRequests.length, prunedWaitTimes: removed.length },
    `Cleaning up request queue.`
  );
  setTimeout(cleanQueue, 20 * 1000);
}

export function start() {
  MODEL_FAMILIES.forEach((modelFamily) => {
    historicalEmas.set(modelFamily, 0);
    currentEmas.set(modelFamily, 0);
    estimates.set(modelFamily, 0);
  });
  processQueue();
  cleanQueue();
  log.info(`Started request queue.`);
}

let waitTimes: {
  partition: ModelFamily;
  start: number;
  end: number;
}[] = [];

/** Adds a successful request to the list of wait times. */
export function trackWaitTime(req: Request) {
  waitTimes.push({
    partition: getModelFamilyForRequest(req),
    start: req.startTime!,
    end: req.queueOutTime ?? Date.now(),
  });
}

const WAIT_TIME_INTERVAL = 3000;
const ALPHA_HISTORICAL = 0.2;
const ALPHA_CURRENT = 0.3;
const historicalEmas: Map<ModelFamily, number> = new Map();
const currentEmas: Map<ModelFamily, number> = new Map();
const estimates: Map<ModelFamily, number> = new Map();

export function getEstimatedWaitTime(partition: ModelFamily) {
  return estimates.get(partition) ?? 0;
}

/**
 * Returns estimated wait time for the given queue partition in milliseconds.
 * Requests which are deprioritized are not included in the calculation as they
 * would skew the results due to their longer wait times.
 */
function calculateWaitTime(partition: ModelFamily) {
  const now = Date.now();
  const recentWaits = waitTimes
    .filter((wait) => {
      const isSamePartition = wait.partition === partition;
      const isRecent = now - wait.end < 300 * 1000;
      return isSamePartition && isRecent;
    })
    .map((wait) => wait.end - wait.start);
  const recentAverage = recentWaits.length
    ? recentWaits.reduce((sum, wait) => sum + wait, 0) / recentWaits.length
    : 0;

  const historicalEma = historicalEmas.get(partition) ?? 0;
  historicalEmas.set(
    partition,
    ALPHA_HISTORICAL * recentAverage + (1 - ALPHA_HISTORICAL) * historicalEma
  );

  const currentWaits = queue
    .filter((req) => getModelFamilyForRequest(req) === partition)
    .map((req) => now - req.startTime!);
  const longestCurrentWait = Math.max(...currentWaits, 0);

  const currentEma = currentEmas.get(partition) ?? 0;
  currentEmas.set(
    partition,
    ALPHA_CURRENT * longestCurrentWait + (1 - ALPHA_CURRENT) * currentEma
  );

  return (historicalEma + currentEma) / 2;
}

setInterval(() => {
  MODEL_FAMILIES.forEach((modelFamily) => {
    estimates.set(modelFamily, calculateWaitTime(modelFamily));
  });
}, WAIT_TIME_INTERVAL);

export function getQueueLength(partition: ModelFamily | "all" = "all") {
  if (partition === "all") {
    return queue.length;
  }
  const modelQueue = getQueueForPartition(partition);
  return modelQueue.length;
}

export function createQueueMiddleware({
  mutations = [],
  proxyMiddleware,
}: {
  mutations?: ProxyReqMutator[];
  proxyMiddleware: Handler;
}): Handler {
  return async (req, res, next) => {
    req.proceed = async () => {
      // canonicalize the stream field which is set in a few places not always
      // consistently
      req.isStreaming = req.isStreaming || String(req.body.stream) === "true";
      req.body.stream = req.isStreaming;

      try {
        // Just before executing the proxyMiddleware, we will create a
        // ProxyReqManager to track modifications to the request. This allows
        // us to revert those changes if the proxied request fails with a
        // retryable error. That happens in proxyMiddleware's onProxyRes
        // handler.
        const changeManager = new ProxyReqManager(req);
        req.changeManager = changeManager;
        for (const mutator of mutations) {
          await mutator(changeManager);
        }
      } catch (err) {
        // Failure during request preparation is a fatal error.
        return classifyErrorAndSend(err, req, res);
      }

      proxyMiddleware(req, res, next);
    };

    try {
      await enqueue(req);
    } catch (err: any) {
      const title =
        err.status === 429
          ? "Proxy queue error (too many concurrent requests)"
          : "Proxy queue error (streaming required)";
      sendErrorToClient({
        options: {
          title,
          message: err.message,
          format: req.inboundApi,
          reqId: req.id,
          model: req.body?.model,
        },
        req,
        res,
      });
    }
  };
}

function killQueuedRequest(req: Request) {
  if (!req.res || req.res.writableEnded) {
    req.log.warn(`Attempted to terminate request that has already ended.`);
    queue.splice(queue.indexOf(req), 1);
    return;
  }
  const res = req.res;
  try {
    const message = `Your request has been terminated by the proxy because it has been in the queue for more than 5 minutes.`;
    sendErrorToClient({
      options: {
        title: "Proxy queue error (request killed)",
        message,
        format: req.inboundApi,
        reqId: req.id,
        model: req.body?.model,
      },
      req,
      res,
    });
  } catch (e) {
    req.log.error(e, `Error killing stalled request.`);
  }
}

async function initStreaming(req: Request) {
  const res = req.res!;
  initializeSseStream(res);

  const joinMsg = `: joining queue at position ${
    queue.length
  }\n\n${getHeartbeatPayload()}`;

  let drainTimeout: NodeJS.Timeout;
  const welcome = new Promise<void>((resolve, reject) => {
    const onDrain = () => {
      clearTimeout(drainTimeout);
      req.log.debug(`Client finished consuming join message.`);
      res.off("drain", onDrain);
      resolve();
    };

    drainTimeout = setTimeout(() => {
      res.off("drain", onDrain);
      res.destroy();
      reject(new Error("Unreponsive streaming client; killing connection"));
    }, QUEUE_JOIN_TIMEOUT);

    if (!res.write(joinMsg)) {
      req.log.warn("Kernel buffer is full; holding client request.");
      res.once("drain", onDrain);
    } else {
      clearTimeout(drainTimeout);
      resolve();
    }
  });

  await welcome;
}

/**
 * http-proxy-middleware attaches a bunch of event listeners to the req and
 * res objects which causes problems with our approach to re-enqueuing failed
 * proxied requests. This function removes those event listeners.
 * We don't have references to the original event listeners, so we have to
 * look through the list and remove HPM's listeners by looking for particular
 * strings in the listener functions. This is an astoundingly shitty way to do
 * this, but it's the best I can come up with.
 */
function removeProxyMiddlewareEventListeners(req: Request) {
  // node_modules/http-proxy-middleware/dist/plugins/default/debug-proxy-errors-plugin.js:29
  // res.listeners('close')
  const RES_ONCLOSE = `Destroying proxyRes in proxyRes close event`;
  // node_modules/http-proxy-middleware/dist/plugins/default/debug-proxy-errors-plugin.js:19
  // res.listeners('error')
  const RES_ONERROR = `Socket error in proxyReq event`;
  // node_modules/http-proxy/lib/http-proxy/passes/web-incoming.js:146
  // req.listeners('aborted')
  const REQ_ONABORTED = `proxyReq.abort()`;
  // node_modules/http-proxy/lib/http-proxy/passes/web-incoming.js:156
  // req.listeners('error')
  const REQ_ONERROR = `if (req.socket.destroyed`;

  const res = req.res!;

  const resOnClose = res
    .listeners("close")
    .find((listener) => listener.toString().includes(RES_ONCLOSE));
  if (resOnClose) {
    res.removeListener("close", resOnClose as any);
  }

  const resOnError = res
    .listeners("error")
    .find((listener) => listener.toString().includes(RES_ONERROR));
  if (resOnError) {
    res.removeListener("error", resOnError as any);
  }

  const reqOnAborted = req
    .listeners("aborted")
    .find((listener) => listener.toString().includes(REQ_ONABORTED));
  if (reqOnAborted) {
    req.removeListener("aborted", reqOnAborted as any);
  }

  const reqOnError = req
    .listeners("error")
    .find((listener) => listener.toString().includes(REQ_ONERROR));
  if (reqOnError) {
    req.removeListener("error", reqOnError as any);
  }
}

export function registerHeartbeat(req: Request) {
  const res = req.res!;

  let isBufferFull = false;
  let bufferFullCount = 0;
  req.heartbeatInterval = setInterval(() => {
    if (isBufferFull) {
      bufferFullCount++;
      if (bufferFullCount >= 3) {
        req.log.error("Heartbeat skipped too many times; killing connection.");
        res.destroy();
      } else {
        req.log.warn({ bufferFullCount }, "Heartbeat skipped; buffer is full.");
      }
      return;
    }

    const data = getHeartbeatPayload();
    if (!res.write(data)) {
      isBufferFull = true;
      res.once("drain", () => (isBufferFull = false));
    }
  }, HEARTBEAT_INTERVAL);
  monitorHeartbeat(req);
}

function monitorHeartbeat(req: Request) {
  const res = req.res!;

  let lastBytesSent = 0;
  req.monitorInterval = setInterval(() => {
    const bytesSent = res.socket?.bytesWritten ?? 0;
    const bytesSinceLast = bytesSent - lastBytesSent;
    req.log.debug(
      {
        previousBytesSent: lastBytesSent,
        currentBytesSent: bytesSent,
      },
      "Heartbeat monitor check."
    );
    lastBytesSent = bytesSent;

    const minBytes = Math.floor(getHeartbeatSize() / 2);
    if (bytesSinceLast < minBytes) {
      req.log.warn(
        { minBytes, bytesSinceLast },
        "Queued request is not processing heartbeats enough data or server is overloaded; killing connection."
      );
      res.destroy();
    }
  }, HEARTBEAT_INTERVAL * 2);
}

/** Sends larger heartbeats when the queue is overloaded */
function getHeartbeatSize() {
  const load = getProxyLoad();

  if (load <= LOAD_THRESHOLD) {
    return MIN_HEARTBEAT_SIZE;
  } else {
    const excessLoad = load - LOAD_THRESHOLD;
    const size =
      MIN_HEARTBEAT_SIZE + Math.pow(excessLoad * PAYLOAD_SCALE_FACTOR, 2);
    if (size > MAX_HEARTBEAT_SIZE) return MAX_HEARTBEAT_SIZE;
    return size;
  }
}

function getHeartbeatPayload() {
  const size = getHeartbeatSize();
  const data =
    process.env.NODE_ENV === "production"
      ? crypto.randomBytes(size).toString("base64")
      : `payload size: ${size}`;

  return `: queue heartbeat ${data}\n\n`;
}

function getProxyLoad() {
  return Math.max(getUniqueIps(), queue.length);
}


code:
import { Request, Response, NextFunction } from "express";
import { config } from "../config";

const ONE_MINUTE_MS = 60 * 1000;

type Timestamp = number;
/** Tracks time of last attempts from each IP address or token. */
const lastAttempts = new Map<string, Timestamp[]>();
/** Tracks time of exempted attempts from shared IPs like Agnai.chat. */
const exemptedRequests: Timestamp[] = [];

const isRecentAttempt = (now: Timestamp) => (attempt: Timestamp) =>
  attempt > now - ONE_MINUTE_MS;

/**
 * Returns duration in seconds to wait before retrying for Retry-After header.
 */
const getRetryAfter = (ip: string, type: "text" | "image") => {
  const now = Date.now();
  const attempts = lastAttempts.get(ip) || [];
  const validAttempts = attempts.filter(isRecentAttempt(now));

  const limit =
    type === "text" ? config.textModelRateLimit : config.imageModelRateLimit;

  if (validAttempts.length >= limit) {
    return (validAttempts[0] - now + ONE_MINUTE_MS) / 1000;
  } else {
    lastAttempts.set(ip, [...validAttempts, now]);
    return 0;
  }
};

const getStatus = (ip: string, type: "text" | "image") => {
  const now = Date.now();
  const attempts = lastAttempts.get(ip) || [];
  const validAttempts = attempts.filter(isRecentAttempt(now));

  const limit =
    type === "text" ? config.textModelRateLimit : config.imageModelRateLimit;

  return {
    remaining: Math.max(0, limit - validAttempts.length),
    reset: validAttempts.length > 0 ? validAttempts[0] + ONE_MINUTE_MS : now,
  };
};

/** Prunes attempts and IPs that are no longer relevant after one minute. */
const clearOldAttempts = () => {
  const now = Date.now();
  for (const [ip, attempts] of lastAttempts.entries()) {
    const validAttempts = attempts.filter(isRecentAttempt(now));
    if (validAttempts.length === 0) {
      lastAttempts.delete(ip);
    } else {
      lastAttempts.set(ip, validAttempts);
    }
  }
};
setInterval(clearOldAttempts, 10 * 1000);

/** Prunes exempted requests which are older than one minute. */
const clearOldExemptions = () => {
  const now = Date.now();
  const validExemptions = exemptedRequests.filter(isRecentAttempt(now));
  exemptedRequests.splice(0, exemptedRequests.length, ...validExemptions);
};
setInterval(clearOldExemptions, 10 * 1000);

export const getUniqueIps = () => lastAttempts.size;

/**
 * Can be used to manually remove the most recent attempt from an IP address,
 * ie. in case a prompt triggered OpenAI's content filter and therefore did not
 * result in a generation.
 */
export const refundLastAttempt = (req: Request) => {
  const key = req.user?.token || req.risuToken || req.ip;
  const attempts = lastAttempts.get(key) || [];
  attempts.pop();
};

export const ipLimiter = async (
  req: Request,
  res: Response,
  next: NextFunction
) => {
  const imageLimit = config.imageModelRateLimit;
  const textLimit = config.textModelRateLimit;

  if (!textLimit && !imageLimit) return next();
  if (req.user?.type === "special") return next();

  const path = req.baseUrl + req.path;
  const type =
    path.includes("openai-image") || path.includes("images/generations")
      ? "image"
      : "text";
  const limit = type === "image" ? imageLimit : textLimit;

  // If user is authenticated, key rate limiting by their token. Otherwise, key
  // rate limiting by their IP address. Mitigates key sharing.
  const rateLimitKey = req.user?.token || req.risuToken || req.ip;

  const { remaining, reset } = getStatus(rateLimitKey, type);
  res.set("X-RateLimit-Limit", limit.toString());
  res.set("X-RateLimit-Remaining", remaining.toString());
  res.set("X-RateLimit-Reset", reset.toString());

  const retryAfterTime = getRetryAfter(rateLimitKey, type);
  if (retryAfterTime > 0) {
    const waitSec = Math.ceil(retryAfterTime).toString();
    res.set("Retry-After", waitSec);
    res.status(429).json({
      error: {
        type: "proxy_rate_limited",
        message: `This model type is rate limited to ${limit} prompts per minute. Please try again in ${waitSec} seconds.`,
      },
    });
  } else {
    next();
  }
};


code:
import express from "express";
import { addV1 } from "./add-v1";
import { anthropic } from "./anthropic";
import { aws } from "./aws";
import { azure } from "./azure";
import { checkRisuToken } from "./check-risu-token";
import { gatekeeper } from "./gatekeeper";
import { gcp } from "./gcp";
import { googleAI } from "./google-ai";
import { mistralAI } from "./mistral-ai";
import { openai } from "./openai";
import { openaiImage } from "./openai-image";
import { sendErrorToClient } from "./middleware/response/error-generator";

const proxyRouter = express.Router();

// Remove `expect: 100-continue` header from requests due to incompatibility
// with node-http-proxy.
proxyRouter.use((req, _res, next) => {
  if (req.headers.expect) {
    delete req.headers.expect;
  }
  next();
});

// Apply body parsers.
proxyRouter.use(
  express.json({ limit: "100mb" }),
  express.urlencoded({ extended: true, limit: "100mb" })
);

// Apply auth/rate limits.
proxyRouter.use(gatekeeper);
proxyRouter.use(checkRisuToken);

// Initialize request queue metadata.
proxyRouter.use((req, _res, next) => {
  req.startTime = Date.now();
  req.retryCount = 0;
  next();
});

// Proxy endpoints.
proxyRouter.use("/openai", addV1, openai);
proxyRouter.use("/openai-image", addV1, openaiImage);
proxyRouter.use("/anthropic", addV1, anthropic);
proxyRouter.use("/google-ai", addV1, googleAI);
proxyRouter.use("/mistral-ai", addV1, mistralAI);
proxyRouter.use("/aws", aws);
proxyRouter.use("/gcp/claude", addV1, gcp);
proxyRouter.use("/azure/openai", addV1, azure);

// Redirect browser requests to the homepage.
proxyRouter.get("*", (req, res, next) => {
  const isBrowser = req.headers["user-agent"]?.includes("Mozilla");
  if (isBrowser) {
    res.redirect("/");
  } else {
    next();
  }
});

// Send a fake client error if user specifies an invalid proxy endpoint.
proxyRouter.use((req, res) => {
  sendErrorToClient({
    req,
    res,
    options: {
      title: "Proxy error (HTTP 404 Not Found)",
      message: "The requested proxy endpoint does not exist.",
      model: req.body?.model,
      reqId: req.id,
      format: "unknown",
      obj: {
        proxy_note:
          "Your chat client is using the wrong endpoint. Check the Service Info page for the list of available endpoints.",
        requested_url: req.originalUrl,
      },
    },
  });
});

export { proxyRouter as proxyRouter };