From 5d73bf8b4eac4a9cb7b6d271d0edcee3c2754a5a Mon Sep 17 00:00:00 2001
From: Archeb <archebasic@hotmail.com>
Date: Sat, 22 Jun 2024 15:07:27 +0800
Subject: [PATCH] supports model selection

---
 index.mjs | 18 +++++++++++++++++-
 usage.md  |  2 +-
 2 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/index.mjs b/index.mjs
index 52f9ee9..c9d86f1 100644
--- a/index.mjs
+++ b/index.mjs
@@ -8,6 +8,15 @@ const app = express();
 const port = process.env.PORT || 8080;
 const validApiKey = process.env.PASSWORD;
 
+const modelMappping = {
+	"claude-3-5-sonnet-20240620": "claude_3_5_sonnet",
+	"claude-3-20240229": "claude_3_opus",
+	"claude-3-sonnet-20240229": "claude_3_sonnet",
+	"claude-3-haiku-20240307": "claude_3_haiku",
+	"claude-2.1": "claude_2",
+	"claude-2.0": "claude_2",
+};
+
 // import config.js
 try {
 	var { config } = await import("./config.mjs");
@@ -54,7 +63,14 @@ app.post("/v1/messages", apiKeyAuth, (req, res) => {
 			var session = sessions[randomSession];
 			console.log("Using session " + randomSession);
 
-			var proxyModel = process.env.AI_MODEL || "claude_3_opus";
+			// decide which model to use
+			if (jsonBody.model && modelMappping[jsonBody.model]) {
+				var proxyModel = modelMappping[jsonBody.model];
+			}else{
+				var proxyModel = process.env.AI_MODEL || "claude_3_opus";
+			}
+			console.log("Using model " + proxyModel);
+			
 			// 检查该session是否已经创建对应模型的对应user chat mode
 			if (process.env.USE_CUSTOM_MODE == "true") {
 				if (config.sessions[session.configIndex].user_chat_mode_id && config.sessions[session.configIndex].user_chat_mode_id[proxyModel]) {
diff --git a/usage.md b/usage.md
index 1ac0918..15a8687 100644
--- a/usage.md
+++ b/usage.md
@@ -56,7 +56,7 @@ export const config = {
 
 # 修改模型
 
-修改`AI_MODEL`环境变量可以切换模型，支持的模型名字如下 (请参考官网获取最新模型):
+Claude系列模型直接在酒馆中选择即可使用，修改`AI_MODEL`环境变量可以切换Claude以外的模型，支持的模型名字如下 (请参考官网获取最新模型):
 
 ```
 gpt_4o