From 12276a1f59a91d9fba1dc097da6f070c31923428 Mon Sep 17 00:00:00 2001
From: random-username-423 <54107-handle635243@users.noreply.gitgud.io>
Date: Wed, 6 Dec 2023 03:21:27 +0000
Subject: [PATCH] Fix AWS Claude Model Reassigning
 (khanon/oai-reverse-proxy!55)

---
 docs/aws-configuration.md |  7 +++--
 src/proxy/aws.ts          | 58 +++++++++++++++++++++++++++++++++------
 2 files changed, 53 insertions(+), 12 deletions(-)

diff --git a/docs/aws-configuration.md b/docs/aws-configuration.md
index 7039332..576a3e1 100644
--- a/docs/aws-configuration.md
+++ b/docs/aws-configuration.md
@@ -45,10 +45,11 @@ You can also request Claude Instant, but support for this isn't fully implemente
 ### Supported model IDs
 Users can send these model IDs to the proxy to invoke the corresponding models.
 - **Claude**
-  - `anthropic.claude-v1` (~18k context)
-  - `anthropic.claude-v2` (~100k context)
+  - `anthropic.claude-v1` (~18k context, claude 1.3)
+  - `anthropic.claude-v2` (~100k context, claude 2.0)
+  - `anthropic.claude-v2:1` (~200k context, claude 2.1)
 - **Claude Instant**
-  - `anthropic.claude-instant-v1`
+  - `anthropic.claude-instant-v1` (~100k context, claude instant 1.2)
 
 ## Note regarding logging
 
diff --git a/src/proxy/aws.ts b/src/proxy/aws.ts
index 9c8d6dd..e30d02f 100644
--- a/src/proxy/aws.ts
+++ b/src/proxy/aws.ts
@@ -178,18 +178,58 @@ awsRouter.post(
  * - frontends sending OpenAI model names because they expect the proxy to
  *   translate them
  */
+const LATEST_AWS_V2_MINOR_VERSION = '1';
+
 function maybeReassignModel(req: Request) {
   const model = req.body.model;
-  // User's client sent an AWS model already
-  if (model.includes("anthropic.claude")) return;
-  // User's client is sending Anthropic-style model names, check for v1
-  if (model.match(/^claude-v?1/)) {
-    req.body.model = "anthropic.claude-v1";
-  } else {
-    // User's client requested v2 or possibly some OpenAI model, default to v2
-    req.body.model = "anthropic.claude-v2:1";
+
+  // If the string already includes "anthropic.claude", return it unmodified
+  if (model.includes("anthropic.claude")) {
+    return;
   }
-  // TODO: Handle claude-instant
+
+  // Define a regular expression pattern to match the Claude version strings
+  const pattern = /^(claude-)?(instant-)?(v)?(\d+)(\.(\d+))?(-\d+k)?$/i;
+
+  // Execute the pattern on the model string
+  const match = model.match(pattern);
+
+  // If there's no match, return the latest v2 model
+  if (!match) {
+    req.body.model = `anthropic.claude-v2:${LATEST_AWS_V2_MINOR_VERSION}`;
+    return;
+  }
+
+  // Extract parts of the version string
+  const [, , instant, v, major, , minor] = match;
+
+  // If 'instant' is part of the version, return the fixed instant model string
+  if (instant) {
+    req.body.model = 'anthropic.claude-instant-v1';
+    return;
+  }
+
+  // If the major version is '1', return the fixed v1 model string
+  if (major === '1') {
+    req.body.model = 'anthropic.claude-v1';
+    return;
+  }
+
+  // If the major version is '2'
+  if (major === '2') {
+    // If the minor version is explicitly '0', return "anthropic.claude-v2" which is claude-2.0
+    if (minor === '0') {
+      req.body.model = 'anthropic.claude-v2';
+      return;
+    }
+    // Otherwise, return the v2 model string with the latest minor version
+    req.body.model = `anthropic.claude-v2:${LATEST_AWS_V2_MINOR_VERSION}`;
+    return;
+  }
+
+  // If none of the above conditions are met, return the latest v2 model by default
+  req.body.model = `anthropic.claude-v2:${LATEST_AWS_V2_MINOR_VERSION}`;
+  return;
 }
 
 export const aws = awsRouter;