From 0c62b8bbfed85ee9284e59f3b7022f995894725a Mon Sep 17 00:00:00 2001
From: kwaroran <amasio1234@proton.me>
Date: Sun, 8 Dec 2024 15:46:54 +0900
Subject: [PATCH] Add pools

---
 src/lang/en.ts                           |  2 ++
 src/lib/Setting/Pages/BotSettings.svelte |  6 ++++++
 src/ts/globalApi.svelte.ts               |  1 +
 src/ts/model/modellist.ts                |  7 ++++---
 src/ts/process/request.ts                | 23 +++++++++++++++++++----
 src/ts/storage/database.svelte.ts        |  1 +
 6 files changed, 33 insertions(+), 7 deletions(-)

diff --git a/src/lang/en.ts b/src/lang/en.ts
index 85b25c67..3ec746a2 100644
--- a/src/lang/en.ts
+++ b/src/lang/en.ts
@@ -171,6 +171,7 @@ export const languageEnglish = {
         translatorPrompt: "The prompt that is used for translation. if it is blank, it will use the default prompt. you can also use ChatML formating with {{slot}} for the dest language, {{solt::content}} for the content, and {{slot::tnote}} for the translator note.",
         translateBeforeHTMLFormatting: "If enabled, it will translate the text before Regex scripts and HTML formatting. this could make the token lesser but could break the formatting.",
         autoTranslateCachedOnly: "If enabled, it will automatically translate only the text that the user has translated previously.",
+        APIPool: "If enabled, it will connect to RisuAI API Pool. Every user which API pool is enabled, the API key will be shared if it used for free, rate-limited models, making it user to make more request of rate-limited models by using other's API key that didn't used much." 
     },
     setup: {
         chooseProvider: "Choose AI Provider",
@@ -818,4 +819,5 @@ export const languageEnglish = {
     customFlags: "Custom Flags",
     enableCustomFlags: "Enable Custom Flags",
     googleCloudTokenization: "Google Cloud Tokenization",
+    APIPool: "API Pool"
 }
\ No newline at end of file
diff --git a/src/lib/Setting/Pages/BotSettings.svelte b/src/lib/Setting/Pages/BotSettings.svelte
index 7f04094a..d66c01a7 100644
--- a/src/lib/Setting/Pages/BotSettings.svelte
+++ b/src/lib/Setting/Pages/BotSettings.svelte
@@ -220,6 +220,12 @@
         {#if DBState.db.aiModel === 'reverse_proxy' || DBState.db.subModel === 'reverse_proxy'}
             <Check bind:check={DBState.db.reverseProxyOobaMode} name={`${language.reverseProxyOobaMode}`}/>
         {/if}
+        {#if modelInfo.flags.includes(LLMFlags.poolSupported) && DBState.db.useExperimental}
+            <Check bind:check={DBState.db.risuPool} name={language.APIPool}>
+                <Help key="APIPool" />
+                <Help key="experimental" />
+            </Check>
+        {/if}
         {#if modelInfo.provider === LLMProvider.NovelAI || subModelInfo.provider === LLMProvider.NovelAI}
             <Check bind:check={DBState.db.NAIadventure} name={language.textAdventureNAI}/>
 
diff --git a/src/ts/globalApi.svelte.ts b/src/ts/globalApi.svelte.ts
index 1e971a5e..cbf3c347 100644
--- a/src/ts/globalApi.svelte.ts
+++ b/src/ts/globalApi.svelte.ts
@@ -906,6 +906,7 @@ async function fetchWithProxy(url: string, arg: GlobalFetchArgs): Promise<Global
       "risu-url": encodeURIComponent(url),
       "Content-Type": arg.body instanceof URLSearchParams ? "application/x-www-form-urlencoded" : "application/json",
       ...(arg.useRisuToken && { "x-risu-tk": "use" }),
+      "risu-pool": DBState.db.risuPool ? 'true' : 'false'
     };
 
     const body = arg.body instanceof URLSearchParams ? arg.body.toString() : JSON.stringify(arg.body);
diff --git a/src/ts/model/modellist.ts b/src/ts/model/modellist.ts
index 888ec6e0..bbca5d7d 100644
--- a/src/ts/model/modellist.ts
+++ b/src/ts/model/modellist.ts
@@ -13,6 +13,7 @@ export enum LLMFlags{
     hasStreaming,
     requiresAlternateRole,
     mustStartWithUserInput,
+    poolSupported
 }
 
 export enum LLMProvider{
@@ -766,16 +767,16 @@ export const LLMModels: LLMModel[] = [
         id: 'gemini-exp-1121',
         provider: LLMProvider.GoogleCloud,
         format: LLMFormat.GoogleCloud,
-        flags: [LLMFlags.hasImageInput, LLMFlags.hasFirstSystemPrompt],
+        flags: [LLMFlags.hasImageInput, LLMFlags.hasFirstSystemPrompt, LLMFlags.poolSupported],
         parameters: ['temperature', 'top_k', 'top_p'],
-        tokenizer: LLMTokenizer.GoogleCloud
+        tokenizer: LLMTokenizer.GoogleCloud,
     },
     {
         name: "Gemini Exp 1206",
         id: 'gemini-exp-1206',
         provider: LLMProvider.GoogleCloud,
         format: LLMFormat.GoogleCloud,
-        flags: [LLMFlags.hasImageInput, LLMFlags.hasFirstSystemPrompt],
+        flags: [LLMFlags.hasImageInput, LLMFlags.hasFirstSystemPrompt, LLMFlags.poolSupported],
         recommended: true,
         parameters: ['temperature', 'top_k', 'top_p'],
         tokenizer: LLMTokenizer.GoogleCloud
diff --git a/src/ts/process/request.ts b/src/ts/process/request.ts
index 47c6a11e..8d88be23 100644
--- a/src/ts/process/request.ts
+++ b/src/ts/process/request.ts
@@ -1578,14 +1578,29 @@ async function requestGoogleCloudVertex(arg:RequestDataArgumentExtended):Promise
         }
     }
 
-    const url = arg.customURL ?? (arg.modelInfo.format === LLMFormat.VertexAIGemini ?
-        `https://${REGION}-aiplatform.googleapis.com/v1/projects/${PROJECT_ID}/locations/us-central1/publishers/google/models/${arg.modelInfo.internalID}:streamGenerateContent`
-        : `https://generativelanguage.googleapis.com/v1beta/models/${arg.modelInfo.internalID}:generateContent?key=${(arg.aiModel === 'reverse_proxy') ?  db.proxyKey : db.google.accessToken}`)
+    let url = ''
+    const pool = arg.modelInfo.flags.includes(LLMFlags.poolSupported) && db.risuPool && (!arg.customURL) && arg.modelInfo.format !== LLMFormat.VertexAIGemini
+    
+    if(arg.customURL){
+        const u = new URL(arg.customURL)
+        u.searchParams.set('key', db.proxyKey)
+        url = u.toString()
+    }
+    else if(arg.modelInfo.format === LLMFormat.VertexAIGemini){
+        url =`https://${REGION}-aiplatform.googleapis.com/v1/projects/${PROJECT_ID}/locations/us-central1/publishers/google/models/${arg.modelInfo.internalID}:streamGenerateContent`
+    }
+    else if(pool){
+        url = `https://sv.risuai.xyz/rapi/pool?model=${arg.modelInfo.internalID}&key=${db.google.accessToken}&type=google`
+    }
+    else{
+        url = `https://generativelanguage.googleapis.com/v1beta/models/${arg.modelInfo.internalID}:generateContent?key=${db.google.accessToken}`
+    }
     const res = await globalFetch(url, {
         headers: headers,
         body: body,
         chatId: arg.chatId,
-        abortSignal: arg.abortSignal
+        abortSignal: arg.abortSignal,
+        plainFetchForce: pool
     })
 
     if(!res.ok){
diff --git a/src/ts/storage/database.svelte.ts b/src/ts/storage/database.svelte.ts
index 03a1f62f..e2783885 100644
--- a/src/ts/storage/database.svelte.ts
+++ b/src/ts/storage/database.svelte.ts
@@ -854,6 +854,7 @@ export interface Database{
     customFlags: LLMFlags[]
     enableCustomFlags: boolean
     googleClaudeTokenizing: boolean
+    risuPool: boolean
 }
 
 interface SeparateParameters{