From 4cc180ffff4fc2e0b4ef1d06c149f5df1789199e Mon Sep 17 00:00:00 2001
From: kwaroran <kwaworan@gmail.com>
Date: Thu, 7 Dec 2023 02:26:08 +0900
Subject: [PATCH] [feat] google palm2 unicorn

---
 src/lib/Setting/Pages/BotSettings.svelte |  7 +++
 src/lib/UI/ModelList.svelte              |  3 ++
 src/ts/process/request.ts                | 58 ++++++++++++++++++++++++
 src/ts/storage/database.ts               | 10 +++-
 4 files changed, 77 insertions(+), 1 deletion(-)

diff --git a/src/lib/Setting/Pages/BotSettings.svelte b/src/lib/Setting/Pages/BotSettings.svelte
index 30b85cc8..d3063081 100644
--- a/src/lib/Setting/Pages/BotSettings.svelte
+++ b/src/lib/Setting/Pages/BotSettings.svelte
@@ -119,6 +119,13 @@
     <span class="text-textcolor">Palm2 {language.apiKey}</span>
     <TextInput marginBottom={true} size={"sm"} placeholder="..." bind:value={$DataBase.palmAPI}/>
 {/if}
+{#if $DataBase.aiModel === 'palm2_unicorn' || $DataBase.subModel === 'palm2_unicorn'}
+    <span class="text-textcolor">Google Bearer Token</span>
+    <TextInput marginBottom={true} size={"sm"} placeholder="..." bind:value={$DataBase.google.accessToken}/>
+
+    <span class="text-textcolor">Google Project ID</span>
+    <TextInput marginBottom={true} size={"sm"} placeholder="..." bind:value={$DataBase.google.projectId}/>
+{/if}
 {#if $DataBase.aiModel.startsWith('novellist') || $DataBase.subModel.startsWith('novellist')}
     <span class="text-textcolor">NovelList {language.apiKey}</span>
     <TextInput marginBottom={true} size={"sm"} placeholder="..." bind:value={$DataBase.novellistAPI}/>
diff --git a/src/lib/UI/ModelList.svelte b/src/lib/UI/ModelList.svelte
index 4d18b9cf..ff705642 100644
--- a/src/lib/UI/ModelList.svelte
+++ b/src/lib/UI/ModelList.svelte
@@ -67,6 +67,8 @@
                 return "OpenRouter"
             case 'gptvi4_1106':
                 return "GPT-4 Turbo 1106 Vision"
+            case 'palm2_unicorn':
+                return "PaLM2 Unicorn"
             default:
                 if(name.startsWith("horde:::")){
                     return name.replace(":::", " ")
@@ -144,6 +146,7 @@
             <button class="hover:bg-selected px-6 py-2 text-lg" on:click={() => {changeModel('openrouter')}}>OpenRouter</button>
             {#if showUnrec}
                 <button class="hover:bg-selected px-6 py-2 text-lg" on:click={() => {changeModel('palm2')}}>Google PaLM2</button>
+                <button class="hover:bg-selected px-6 py-2 text-lg" on:click={() => {changeModel('palm2_unicorn')}}>Google PaLM2 Unicorn</button>
                 <button class="hover:bg-selected px-6 py-2 text-lg" on:click={() => {changeModel('kobold')}}>Kobold</button>
             {/if}
             <Arcodion name="Novellist">
diff --git a/src/ts/process/request.ts b/src/ts/process/request.ts
index 87e276fe..4e6ba009 100644
--- a/src/ts/process/request.ts
+++ b/src/ts/process/request.ts
@@ -894,6 +894,64 @@ export async function requestChatDataMain(arg:requestDataArgument, model:'model'
                 }
             }
         }
+        case 'palm2_unicorn':{
+            const bodyData = {
+                "instances": [
+                    {
+                        "content": stringlizeChat(formated, currentChar?.name ?? '', arg.continue)
+                    }
+                ],
+                "parameters": {
+                    "candidateCount": 1,
+                    "maxOutputTokens": maxTokens,
+                    "stopSequences": [
+                        "system:", "user:", "assistant:"
+                    ],
+                    "temperature": temperature,
+                }
+            };
+
+            const API_ENDPOINT="us-central1-aiplatform.googleapis.com"
+            const PROJECT_ID=db.google.projectId
+            const MODEL_ID="text-unicorn"
+            const LOCATION_ID="us-central1"
+
+            const url = `https://${API_ENDPOINT}/v1/projects/${PROJECT_ID}/locations/${LOCATION_ID}/publishers/google/models/${MODEL_ID}:predict`;
+            const res = await globalFetch(url, {
+                body: bodyData,
+                headers: {
+                    "Content-Type": "application/json",
+                    "Authorization": "Bearer " + db.google.accessToken
+                },
+                abortSignal
+            })
+            if(res.ok){
+                console.log(res.data)
+                if(res.data.predictions){
+                    let output:string = res.data.predictions[0].content
+                    const ind = output.search(/(system note)|(user)|(assistant):/gi)
+                    if(ind >= 0){
+                        output = output.substring(0, ind)
+                    }
+                    return {
+                        type: 'success',
+                        result: output
+                    }
+                }
+                else{
+                    return {
+                        type: 'fail',
+                        result: `${JSON.stringify(res.data)}`
+                    }
+                }
+            }
+            else{
+                return {
+                    type: 'fail',
+                    result: `${JSON.stringify(res.data)}`
+                }
+            }
+        }
         case "kobold":{
             const proompt = stringlizeChat(formated, currentChar?.name ?? '', arg.continue)
             const url = new URL(db.koboldURL)
diff --git a/src/ts/storage/database.ts b/src/ts/storage/database.ts
index 42020dba..d11261b7 100644
--- a/src/ts/storage/database.ts
+++ b/src/ts/storage/database.ts
@@ -352,6 +352,10 @@ export function setDatabase(data:Database){
         mode: 'instruct'
     }
     data.top_p ??= 1
+    //@ts-ignore
+    data.google ??= {}
+    data.google.accessToken ??= ''
+    data.google.projectId ??= ''
     changeLanguage(data.language)
     DataBase.set(data)
 }
@@ -547,6 +551,10 @@ export interface Database{
     allowAllExtentionFiles?:boolean
     translatorPrompt:string
     top_p: number,
+    google: {
+        accessToken: string
+        projectId: string
+    }
 }
 
 export interface customscript{
@@ -937,7 +945,7 @@ export const presetTemplate:botPreset = {
     reverseProxyOobaArgs: {
         mode: 'instruct'
     },
-    top_p: 1
+    top_p: 1,
 
 }