oai-reverse-proxy/src/shared/api-support/kits/anthropic-chat/request-transformers.ts

import { AnthropicChatMessage, AnthropicV1MessagesSchema } from "./schema";
import { AnthropicV1TextSchema, APIRequestTransformer, OpenAIChatMessage } from "../../index";
import { BadRequestError } from "../../../errors";

import { OpenAIV1ChatCompletionSchema } from "../openai/schema";

/**
 * Represents the union of all content types without the `string` shorthand
 * for `text` content.
 */
type AnthropicChatMessageContentWithoutString = Exclude<
  AnthropicChatMessage["content"],
  string
>;
/** Represents a message with all shorthand `string` content expanded. */
type ConvertedAnthropicChatMessage = AnthropicChatMessage & {
  content: AnthropicChatMessageContentWithoutString;
};

export const transformOpenAIToAnthropicChat: APIRequestTransformer<
  typeof AnthropicV1MessagesSchema
> = async (req) => {
  const { body } = req;
  const result = OpenAIV1ChatCompletionSchema.safeParse(body);
  if (!result.success) {
    req.log.warn(
      { issues: result.error.issues, body },
      "Invalid OpenAI-to-Anthropic Chat request"
    );
    throw result.error;
  }

  req.headers["anthropic-version"] = "2023-06-01";

  const { messages, ...rest } = result.data;
  const { messages: newMessages, system } =
    openAIMessagesToClaudeChatPrompt(messages);

  return {
    system,
    messages: newMessages,
    model: rest.model,
    max_tokens: rest.max_tokens,
    stream: rest.stream,
    temperature: rest.temperature,
    top_p: rest.top_p,
    stop_sequences: typeof rest.stop === "string" ? [rest.stop] : rest.stop,
    ...(rest.user ? { metadata: { user_id: rest.user } } : {}),
    // Anthropic supports top_k, but OpenAI does not
    // OpenAI supports frequency_penalty, presence_penalty, logit_bias, n, seed,
    // and function calls, but Anthropic does not.
  };
};

/**
 * Converts an older Anthropic Text Completion prompt to the newer Messages API
 * by splitting the flat text into messages.
 */
export const transformAnthropicTextToAnthropicChat: APIRequestTransformer<
  typeof AnthropicV1MessagesSchema
> = async (req) => {
  const { body } = req;
  const result = AnthropicV1TextSchema.safeParse(body);
  if (!result.success) {
    req.log.warn(
      { issues: result.error.issues, body },
      "Invalid Anthropic Text-to-Anthropic Chat request"
    );
    throw result.error;
  }

  req.headers["anthropic-version"] = "2023-06-01";

  const { model, max_tokens_to_sample, prompt, ...rest } = result.data;
  validateAnthropicTextPrompt(prompt);

  // Iteratively slice the prompt into messages. Start from the beginning and
  // look for the next `\n\nHuman:` or `\n\nAssistant:`. Anything before the
  // first human message is a system message.
  let index = prompt.indexOf("\n\nHuman:");
  let remaining = prompt.slice(index);
  const system = prompt.slice(0, index);
  const messages: AnthropicChatMessage[] = [];
  while (remaining) {
    const isHuman = remaining.startsWith("\n\nHuman:");

    // Multiple messages from the same role are not permitted in Messages API.
    // We collect all messages until the next message from the opposite role.
    const thisRole = isHuman ? "\n\nHuman:" : "\n\nAssistant:";
    const nextRole = isHuman ? "\n\nAssistant:" : "\n\nHuman:";
    const nextIndex = remaining.indexOf(nextRole);

    // Collect text up to the next message, or the end of the prompt for the
    // Assistant prefill if present.
    const msg = remaining
      .slice(0, nextIndex === -1 ? undefined : nextIndex)
      .replace(thisRole, "")
      .trimStart();

    const role = isHuman ? "user" : "assistant";
    messages.push({ role, content: msg });
    remaining = remaining.slice(nextIndex);

    if (nextIndex === -1) break;
  }

  // fix "messages: final assistant content cannot end with trailing whitespace"
  const lastMessage = messages[messages.length - 1];
  if (
    lastMessage.role === "assistant" &&
    typeof lastMessage.content === "string"
  ) {
    messages[messages.length - 1].content = lastMessage.content.trimEnd();
  }

  return {
    model,
    system,
    messages,
    max_tokens: max_tokens_to_sample,
    ...rest,
  };
};

function validateAnthropicTextPrompt(prompt: string) {
  if (!prompt.includes("\n\nHuman:") || !prompt.includes("\n\nAssistant:")) {
    throw new BadRequestError(
      "Prompt must contain at least one human and one assistant message."
    );
  }
  // First human message must be before first assistant message
  const firstHuman = prompt.indexOf("\n\nHuman:");
  const firstAssistant = prompt.indexOf("\n\nAssistant:");
  if (firstAssistant < firstHuman) {
    throw new BadRequestError(
      "First Assistant message must come after the first Human message."
    );
  }
}

function openAIMessagesToClaudeChatPrompt(messages: OpenAIChatMessage[]): {
  messages: AnthropicChatMessage[];
  system: string;
} {
  // Similar formats, but Claude doesn't use `name` property and doesn't have
  // a `system` role.  Also, Claude does not allow consecutive messages from
  // the same role, so we need to merge them.
  // 1. Collect all system messages up to the first non-system message and set
  // that as the `system` prompt.
  // 2. Iterate through messages and:
  //   - If the message is from system, reassign it to assistant with System:
  //     prefix.
  //   - If message is from same role as previous, append it to the previous
  //     message rather than creating a new one.
  //   - Otherwise, create a new message and prefix with `name` if present.

  // TODO: When a Claude message has multiple `text` contents, does the internal
  // message flattening insert newlines between them?  If not, we may need to
  // do that here...

  let firstNonSystem = -1;
  const result: { messages: ConvertedAnthropicChatMessage[]; system: string } =
    { messages: [], system: "" };
  for (let i = 0; i < messages.length; i++) {
    const msg = messages[i];
    const isSystem = isSystemOpenAIRole(msg.role);

    if (firstNonSystem === -1 && isSystem) {
      // Still merging initial system messages into the system prompt
      result.system += getFirstTextContent(msg.content) + "\n";
      continue;
    }

    if (firstNonSystem === -1 && !isSystem) {
      // Encountered the first non-system message
      firstNonSystem = i;

      if (msg.role === "assistant") {
        // There is an annoying rule that the first message must be from the user.
        // This is commonly not the case with roleplay prompts that start with a
        // block of system messages followed by an assistant message. We will try
        // to reconcile this by splicing the last line of the system prompt into
        // a beginning user message -- this is *commonly* ST's [Start a new chat]
        // nudge, which works okay as a user message.

        // Find the last non-empty line in the system prompt
        const execResult = /(?:[^\r\n]*\r?\n)*([^\r\n]+)(?:\r?\n)*/d.exec(
          result.system
        );

        let text = "";
        if (execResult) {
          text = execResult[1];
          // Remove last line from system so it doesn't get duplicated
          const [_, [lastLineStart]] = execResult.indices || [];
          result.system = result.system.slice(0, lastLineStart);
        } else {
          // This is a bad prompt; there's no system content to move to user and
          // it starts with assistant. We don't have any good options.
          text = "[ Joining chat... ]";
        }

        result.messages.push({
          role: "user",
          content: [{ type: "text", text }],
        });
      }
    }

    const last = result.messages[result.messages.length - 1];
    // I have to handle tools as system messages to be exhaustive here but the
    // experience will be bad.
    const role = isSystemOpenAIRole(msg.role) ? "assistant" : msg.role;

    // Here we will lose the original name if it was a system message, but that
    // is generally okay because the system message is usually a prompt and not
    // a character in the chat.
    const name = msg.role === "system" ? "System" : msg.name?.trim();
    const content = convertOpenAIContent(msg.content);

    // Prepend the display name to the first text content in the current message
    // if it exists. We don't need to add the name to every content block.
    if (name?.length) {
      const firstTextContent = content.find((c) => c.type === "text");
      if (firstTextContent && "text" in firstTextContent) {
        // This mutates the element in `content`.
        firstTextContent.text = `${name}: ${firstTextContent.text}`;
      }
    }

    // Merge messages if necessary. If two assistant roles are consecutive but
    // had different names, the final converted assistant message will have
    // multiple characters in it, but the name prefixes should assist the model
    // in differentiating between speakers.
    if (last && last.role === role) {
      last.content.push(...content);
    } else {
      result.messages.push({ role, content });
    }
  }

  result.system = result.system.trimEnd();
  return result;
}

function isSystemOpenAIRole(
  role: OpenAIChatMessage["role"]
): role is "system" | "function" | "tool" {
  return ["system", "function", "tool"].includes(role);
}

function getFirstTextContent(content: OpenAIChatMessage["content"]) {
  if (typeof content === "string") return content;
  for (const c of content) {
    if ("text" in c) return c.text;
  }
  return "[ No text content in this message ]";
}

function convertOpenAIContent(
  content: OpenAIChatMessage["content"]
): AnthropicChatMessageContentWithoutString {
  if (typeof content === "string") {
    return [{ type: "text", text: content.trimEnd() }];
  }

  return content.map((c) => {
    if ("text" in c) {
      return { type: "text", text: c.text.trimEnd() };
    } else if ("image_url" in c) {
      const url = c.image_url.url;
      try {
        const mimeType = url.split(";")[0].split(":")[1];
        const data = url.split(",")[1];
        return {
          type: "image",
          source: { type: "base64", media_type: mimeType, data },
        };
      } catch (e) {
        return {
          type: "text",
          text: `[ Unsupported image URL: ${url.slice(0, 200)} ]`,
        };
      }
    } else {
      const type = String((c as any)?.type);
      return { type: "text", text: `[ Unsupported content type: ${type} ]` };
    }
  });
}