From 302bd11b33ca399307c811113ef0ff435209ad64 Mon Sep 17 00:00:00 2001
From: "Timothy J. Baek" <timothyjrbeck@gmail.com>
Date: Sun, 24 Mar 2024 23:11:55 -0700
Subject: [PATCH] fix: pull model from selector

---
 .../chat/ModelSelector/Selector.svelte        | 30 ++++++++++---------
 src/lib/utils/index.ts                        | 26 ++++++++++++++++
 2 files changed, 42 insertions(+), 14 deletions(-)

diff --git a/src/lib/components/chat/ModelSelector/Selector.svelte b/src/lib/components/chat/ModelSelector/Selector.svelte
index 0d9fa0ad3..c32201de3 100644
--- a/src/lib/components/chat/ModelSelector/Selector.svelte
+++ b/src/lib/components/chat/ModelSelector/Selector.svelte
@@ -10,9 +10,9 @@
 
 	import { cancelOllamaRequest, deleteModel, getOllamaVersion, pullModel } from '$lib/apis/ollama';
 
-	import { user, MODEL_DOWNLOAD_POOL } from '$lib/stores';
+	import { user, MODEL_DOWNLOAD_POOL, models } from '$lib/stores';
 	import { toast } from 'svelte-sonner';
-	import { splitStream } from '$lib/utils';
+	import { getModels, splitStream } from '$lib/utils';
 	import Tooltip from '$lib/components/common/Tooltip.svelte';
 
 	const i18n = getContext('i18n');
@@ -141,6 +141,8 @@
 						modelName: sanitizedModelTag
 					})
 				);
+
+				models.set(await getModels(localStorage.token));
 			} else {
 				toast.error('Download canceled');
 			}
@@ -233,6 +235,17 @@
 					</div>
 				{/each}
 
+				{#if !(searchValue.trim() in $MODEL_DOWNLOAD_POOL) && searchValue && ollamaVersion && $user.role === 'admin'}
+					<button
+						class="flex w-full font-medium line-clamp-1 select-none items-center rounded-button py-2 pl-3 pr-1.5 text-sm text-gray-700 dark:text-gray-100 outline-none transition-all duration-75 hover:bg-gray-100 dark:hover:bg-gray-850 rounded-lg cursor-pointer data-[highlighted]:bg-muted"
+						on:click={() => {
+							pullModelHandler();
+						}}
+					>
+						Pull "{searchValue}" from Ollama.com
+					</button>
+				{/if}
+
 				{#each Object.keys($MODEL_DOWNLOAD_POOL) as model}
 					<div
 						class="flex w-full justify-between font-medium select-none rounded-button py-2 pl-3 pr-1.5 text-sm text-gray-700 dark:text-gray-100 outline-none transition-all duration-75 rounded-lg cursor-pointer data-[highlighted]:bg-muted"
@@ -279,7 +292,7 @@
 							</div>
 						</div>
 
-						<div class="mr-3 translate-y-0.5">
+						<div class="mr-2 translate-y-0.5">
 							<Tooltip content="Cancel">
 								<button
 									class="text-gray-800 dark:text-gray-100"
@@ -309,17 +322,6 @@
 						</div>
 					</div>
 				{/each}
-
-				{#if !(searchValue.trim() in $MODEL_DOWNLOAD_POOL) && searchValue && ollamaVersion && $user.role === 'admin'}
-					<button
-						class="flex w-full font-medium line-clamp-1 select-none items-center rounded-button py-2 pl-3 pr-1.5 text-sm text-gray-700 dark:text-gray-100 outline-none transition-all duration-75 hover:bg-gray-100 dark:hover:bg-gray-850 rounded-lg cursor-pointer data-[highlighted]:bg-muted"
-						on:click={() => {
-							pullModelHandler();
-						}}
-					>
-						Pull "{searchValue}" from Ollama.com
-					</button>
-				{/if}
 			</div>
 		</slot>
 	</Select.Content>
diff --git a/src/lib/utils/index.ts b/src/lib/utils/index.ts
index a2c6cb29d..ad3e38932 100644
--- a/src/lib/utils/index.ts
+++ b/src/lib/utils/index.ts
@@ -1,5 +1,31 @@
 import { v4 as uuidv4 } from 'uuid';
 import sha256 from 'js-sha256';
+import { getOllamaModels } from '$lib/apis/ollama';
+import { getOpenAIModels } from '$lib/apis/openai';
+import { getLiteLLMModels } from '$lib/apis/litellm';
+
+export const getModels = async (token: string) => {
+	let models = await Promise.all([
+		await getOllamaModels(token).catch((error) => {
+			console.log(error);
+			return null;
+		}),
+		await getOpenAIModels(token).catch((error) => {
+			console.log(error);
+			return null;
+		}),
+		await getLiteLLMModels(token).catch((error) => {
+			console.log(error);
+			return null;
+		})
+	]);
+
+	models = models
+		.filter((models) => models)
+		.reduce((a, e, i, arr) => a.concat(e, ...(i < arr.length - 1 ? [{ name: 'hr' }] : [])), []);
+
+	return models;
+};
 
 //////////////////////////
 // Helper functions