From 3978a464483a7fb2edf0fe4361f4df66841e3f56 Mon Sep 17 00:00:00 2001
From: Alfredo Gallardo <agallardol@outlook.com>
Date: Wed, 7 Aug 2024 18:36:31 -0400
Subject: [PATCH] - fix: added missing logic about installed models

---
 .../components/ollama-models-repository.tsx   | 21 ++++++++++++++++---
 1 file changed, 18 insertions(+), 3 deletions(-)

diff --git a/apps/shinkai-desktop/src/components/shinkai-node-manager/components/ollama-models-repository.tsx b/apps/shinkai-desktop/src/components/shinkai-node-manager/components/ollama-models-repository.tsx
index 43a39a296..d21c6ef87 100644
--- a/apps/shinkai-desktop/src/components/shinkai-node-manager/components/ollama-models-repository.tsx
+++ b/apps/shinkai-desktop/src/components/shinkai-node-manager/components/ollama-models-repository.tsx
@@ -40,8 +40,14 @@ import {
   FILTERED_OLLAMA_MODELS_REPOSITORY,
   OllamaModelDefinition,
 } from '../../..//lib/shinkai-node-manager/ollama-models';
-import { useOllamaPullingQuery } from '../../../lib/shinkai-node-manager/ollama-client';
-import { useShinkaiNodeGetDefaultModel } from '../../../lib/shinkai-node-manager/shinkai-node-manager-client';
+import {
+  useOllamaListQuery,
+  useOllamaPullingQuery,
+} from '../../../lib/shinkai-node-manager/ollama-client';
+import {
+  useShinkaiNodeGetDefaultModel,
+  useShinkaiNodeGetOllamaApiUrlQuery,
+} from '../../../lib/shinkai-node-manager/shinkai-node-manager-client';
 import { OllamaModelInstallButton } from './ollama-model-install-button';
 
 export const OllamaModelsRepository = ({
@@ -51,6 +57,10 @@ export const OllamaModelsRepository = ({
   const { t } = useTranslation();
   const { data: defaultModel } = useShinkaiNodeGetDefaultModel();
   const { data: pullingModelsMap } = useOllamaPullingQuery();
+  const { data: ollamaApiUrl } = useShinkaiNodeGetOllamaApiUrlQuery();
+
+  const ollamaConfig = { host: ollamaApiUrl || 'http://127.0.0.1:11435' };
+  const { data: installedOllamaModels } = useOllamaListQuery(ollamaConfig);
   const installedOllamaModelsMap = useMap<string, ModelResponse>();
   const selectedTagMap = useMap<string, string>();
   const tableContainerRef = useRef<HTMLDivElement>(null);
@@ -220,7 +230,12 @@ export const OllamaModelsRepository = ({
       selectedTagMap.set(model.name, defaultTag);
     });
   }, [installedOllamaModelsMap, pullingModelsMap, selectedTagMap]);
-
+  useEffect(() => {
+    installedOllamaModels?.models &&
+      installedOllamaModels.models.forEach((modelResponse) => {
+        installedOllamaModelsMap.set(modelResponse.name, modelResponse);
+      });
+  }, [installedOllamaModels?.models, installedOllamaModelsMap]);
   const getFullName = (model: string, tag: string): string => {
     return `${model}:${tag}`;
   };