From 9a83bde7e5d5745d3001542a04d8456088fc4986 Mon Sep 17 00:00:00 2001
From: "Timothy J. Baek" <timothyjrbeck@gmail.com>
Date: Fri, 8 Mar 2024 13:41:38 -0800
Subject: [PATCH] feat: max token option for litellm models

---
 src/lib/apis/litellm/index.ts                 |  4 +++-
 .../components/chat/Settings/Models.svelte    | 20 ++++++++++++++++++-
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/src/lib/apis/litellm/index.ts b/src/lib/apis/litellm/index.ts
index 6466ee35b..302e9c4a3 100644
--- a/src/lib/apis/litellm/index.ts
+++ b/src/lib/apis/litellm/index.ts
@@ -77,6 +77,7 @@ type AddLiteLLMModelForm = {
 	api_base: string;
 	api_key: string;
 	rpm: string;
+	max_tokens: string;
 };
 
 export const addLiteLLMModel = async (token: string = '', payload: AddLiteLLMModelForm) => {
@@ -95,7 +96,8 @@ export const addLiteLLMModel = async (token: string = '', payload: AddLiteLLMMod
 				model: payload.model,
 				...(payload.api_base === '' ? {} : { api_base: payload.api_base }),
 				...(payload.api_key === '' ? {} : { api_key: payload.api_key }),
-				...(isNaN(parseInt(payload.rpm)) ? {} : { rpm: parseInt(payload.rpm) })
+				...(isNaN(parseInt(payload.rpm)) ? {} : { rpm: parseInt(payload.rpm) }),
+				...(payload.max_tokens === '' ? {} : { max_tokens: payload.max_tokens })
 			}
 		})
 	})
diff --git a/src/lib/components/chat/Settings/Models.svelte b/src/lib/components/chat/Settings/Models.svelte
index bdd2cc77d..9a0799727 100644
--- a/src/lib/components/chat/Settings/Models.svelte
+++ b/src/lib/components/chat/Settings/Models.svelte
@@ -27,6 +27,7 @@
 	let liteLLMAPIBase = '';
 	let liteLLMAPIKey = '';
 	let liteLLMRPM = '';
+	let liteLLMMaxTokens = '';
 
 	let deleteLiteLLMModelId = '';
 
@@ -326,7 +327,8 @@
 				model: liteLLMModel,
 				api_base: liteLLMAPIBase,
 				api_key: liteLLMAPIKey,
-				rpm: liteLLMRPM
+				rpm: liteLLMRPM,
+				max_tokens: liteLLMMaxTokens
 			}).catch((error) => {
 				toast.error(error);
 				return null;
@@ -821,6 +823,22 @@
 									</div>
 								</div>
 							</div>
+
+							<div>
+								<div class="mb-1.5 text-sm font-medium">Max Tokens</div>
+								<div class="flex w-full">
+									<div class="flex-1">
+										<input
+											class="w-full rounded-lg py-2 px-4 text-sm dark:text-gray-300 dark:bg-gray-850 outline-none"
+											placeholder="Enter Max Tokens (litellm_params.max_tokens)"
+											bind:value={liteLLMMaxTokens}
+											type="number"
+											min="1"
+											autocomplete="off"
+										/>
+									</div>
+								</div>
+							</div>
 						{/if}
 					</div>