- feature: added support for ollama in shinkai-visor (#151)

dcSpark · Feb 28, 2024 · 1458032 · 1458032
1 parent 53b12d0
commit 1458032
Show file tree

Hide file tree

Showing 6 changed files with 48 additions and 4 deletions.
diff --git a/apps/shinkai-visor/src/components/add-agent/add-agent.tsx b/apps/shinkai-visor/src/components/add-agent/add-agent.tsx
@@ -36,7 +36,7 @@ const formSchema = z
         'It just accepts alphanumeric characters and underscores',
       ),
     externalUrl: z.string().url(),
-    apiKey: z.string().min(4),
+    apiKey: z.string(),
     model: z.nativeEnum(Models),
     modelType: z.string(),
     isCustomModel: z.boolean().default(false).optional(),
@@ -45,7 +45,7 @@ const formSchema = z
   })
   .superRefine(
     (
-      { isCustomModel, model, modelType, modelCustom, modelTypeCustom },
+      { isCustomModel, model, modelType, modelCustom, modelTypeCustom, apiKey },
       ctx,
     ) => {
       if (isCustomModel) {
@@ -78,6 +78,13 @@ const formSchema = z
             message: 'Model Type is required',
           });
         }
+        if (!apiKey && model !== Models.Ollama) {
+          ctx.addIssue({
+            path: ['apiKey'],
+            code: z.ZodIssueCode.custom,
+            message: 'Api Key is required',
+          });
+        }
       }
     },
   );
@@ -126,6 +133,10 @@ export const AddAgent = () => {
       value: Models.TogetherComputer,
       label: intl.formatMessage({ id: 'togethercomputer' }),
     },
+    {
+      value: Models.Ollama,
+      label: intl.formatMessage({ id: 'ollama' }),
+    },
   ];
   const getModelObject = (model: Models, modelType: string) => {
     if (isCustomModelMode) {
@@ -137,6 +148,8 @@ export const AddAgent = () => {
         return { OpenAI: { model_type: modelType } };
       case Models.TogetherComputer:
         return { GenericAPI: { model_type: modelType } };
+      case Models.Ollama:
+        return { Ollama: { model_type: modelType } };
       default:
         throw new Error('unknown model');
     }

diff --git a/apps/shinkai-visor/src/components/add-agent/models.ts b/apps/shinkai-visor/src/components/add-agent/models.ts
@@ -1,6 +1,7 @@
 export enum Models {
   OpenAI = 'open-ai',
   TogetherComputer = 'togethercomputer',
+  Ollama = 'ollama',
 }
 
 export const modelsConfig = {
@@ -42,4 +43,25 @@ export const modelsConfig = {
       },
     ],
   },
+  [Models.Ollama]: {
+    apiUrl: 'http://localhost:11434',
+    modelTypes: [
+      {
+        name: 'Llama 2',
+        value: 'llama2',
+      },
+      {
+        name: 'Mistral',
+        value: 'mistral',
+      },
+      {
+        name: 'Mixtral',
+        value: 'mixtral',
+      },
+      {
+        name: 'LLaVA',
+        value: 'llava',
+      },
+    ],
+  },
 };
diff --git a/apps/shinkai-visor/src/lang/en.json b/apps/shinkai-visor/src/lang/en.json
@@ -109,5 +109,6 @@
   "shortcut-key": "Shortcut Key",
   "shortcut-key-description": "Modify the shortcut key to launch sidebar",
   "limit-reached": "Limit Reached",
-  "limit-reached-description": "You've used all of your queries for the month on this model/agent. Please start a new chat with another agent."
+  "limit-reached-description": "You've used all of your queries for the month on this model/agent. Please start a new chat with another agent.",
+  "ollama": "Ollama"
 }
diff --git a/apps/shinkai-visor/src/lang/es.json b/apps/shinkai-visor/src/lang/es.json
@@ -107,5 +107,6 @@
   "shortcut-key": "Tecla de acceso rápido",
   "shortcut-key-description": "Presiona la tecla de acceso rápido para abrir Shinkai Visor",
   "limit-reached": "Límite alcanzado",
-  "limit-reached-description": "Has utilizado todas tus consultas para el mes en este modelo/agente. Por favor, inicia un nuevo chat con otro agente."
+  "limit-reached-description": "Has utilizado todas tus consultas para el mes en este modelo/agente. Por favor, inicia un nuevo chat con otro agente.",
+  "ollama": "Ollama"
 }
diff --git a/libs/shinkai-message-ts/src/models/SchemaTypes.ts b/libs/shinkai-message-ts/src/models/SchemaTypes.ts
@@ -77,6 +77,11 @@ export interface SerializedAgent {
 export interface AgentAPIModel {
   OpenAI?: OpenAI;
   GenericAPI?: GenericAPI;
+  Ollama?: Ollama;
+}
+
+export interface Ollama {
+  model_type: string;
 }
 
 export interface OpenAI {

diff --git a/libs/shinkai-message-ts/src/wasm/SerializedAgentWrapper.ts b/libs/shinkai-message-ts/src/wasm/SerializedAgentWrapper.ts
@@ -32,6 +32,8 @@ export class SerializedAgentWrapper {
       agent.model.GenericAPI.model_type
     ) {
       modelStr = 'genericapi:' + agent.model.GenericAPI.model_type;
+    } else if (agent?.model?.Ollama) {
+      modelStr = 'ollama:' + agent.model.Ollama.model_type;
     } else {
       throw new Error('Invalid model: ' + JSON.stringify(agent.model));
     }