- fix: installation flow (#405)

dcSpark · Aug 9, 2024 · dddd0af · dddd0af
1 parent ff7ed18
commit dddd0af
Show file tree

Hide file tree

Showing 3 changed files with 1,370 additions and 1,370 deletions.
diff --git a/...inkai-desktop/src/components/shinkai-node-manager/components/ollama-models-repository.tsx b/...inkai-desktop/src/components/shinkai-node-manager/components/ollama-models-repository.tsx
@@ -219,6 +219,10 @@ export const OllamaModelsRepository = ({
     getScrollElement: () => tableContainerRef.current,
   });
   useEffect(() => {
+    installedOllamaModels?.models &&
+    installedOllamaModels.models.forEach((modelResponse) => {
+      installedOllamaModelsMap.set(modelResponse.name, modelResponse);
+    });
     FILTERED_OLLAMA_MODELS_REPOSITORY.forEach((model) => {
       const defaultTag: string =
         model.tags?.find((tag) =>
@@ -231,13 +235,8 @@ export const OllamaModelsRepository = ({
         model.tags[0].name;
       selectedTagMap.set(model.name, defaultTag);
     });
-  }, [installedOllamaModelsMap, pullingModelsMap, selectedTagMap]);
-  useEffect(() => {
-    installedOllamaModels?.models &&
-      installedOllamaModels.models.forEach((modelResponse) => {
-        installedOllamaModelsMap.set(modelResponse.name, modelResponse);
-      });
-  }, [installedOllamaModels?.models, installedOllamaModelsMap]);
+  }, [installedOllamaModels?.models, installedOllamaModelsMap, pullingModelsMap, selectedTagMap]);
+
   const getFullName = (model: string, tag: string): string => {
     return `${model}:${tag}`;
   };

diff --git a/apps/shinkai-desktop/src/lib/shinkai-node-manager/ollama-client.ts b/apps/shinkai-desktop/src/lib/shinkai-node-manager/ollama-client.ts
@@ -110,7 +110,7 @@ export const useOllamaPullMutation = (
         queryKey: [FunctionKey.SCAN_OLLAMA_MODELS],
       });
       queryClient.invalidateQueries({
-        queryKey: ['shinkai_node_set_default_options', 'ollama_pulling_models'],
+        queryKey: ['shinkai_node_set_default_options', 'ollama_list', 'ollama_pulling_models'],
       });
       if (options?.onSuccess) {
         options.onSuccess(...onSuccessParameters);
@@ -131,12 +131,13 @@ export const useOllamaRemoveMutation = (
       const response = await ollamaClient.delete({
         model: input.model,
       });
+      pullingModelsMap.delete(input.model);
       return response;
     },
     ...options,
     onSuccess: (...onSuccessParameters) => {
       queryClient.invalidateQueries({
-        queryKey: ['ollama_list'],
+        queryKey: ['ollama_list', 'ollama_pulling_models'],
       });
       queryClient.invalidateQueries({
         queryKey: [FunctionKey.SCAN_OLLAMA_MODELS],