Manage models

n4ze3m · May 2, 2024 · dedf3c5 · dedf3c5
1 parent efc922d
commit dedf3c5
Show file tree

Hide file tree

Showing 30 changed files with 825 additions and 679 deletions.
diff --git a/app/ui/src/routes/settings/application.tsx b/app/ui/src/routes/settings/application.tsx
@@ -1,4 +1,4 @@
-import { Form, InputNumber, Switch, notification, Select } from "antd";
+import { Form, InputNumber, Switch, notification, Select, Input } from "antd";
 import React from "react";
 import api from "../../services/api";
 import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query";
@@ -25,6 +25,9 @@ export default function SettingsApplicationRoot() {
       defaultChunkOverlap: number;
       defaultChatModel: string;
       defaultEmbeddingModel: string;
+      hideDefaultModels: boolean;
+      dynamicallyFetchOllamaModels: boolean;
+      ollamaURL: string;
     };
   });
 
@@ -63,6 +66,22 @@ export default function SettingsApplicationRoot() {
     }
   );
 
+  const { mutateAsync: updateModelSettings, isLoading: isModelLoading } =
+    useMutation(onUpdateApplicatoon, {
+      onSuccess: (data) => {
+        queryClient.invalidateQueries(["fetchBotCreateConfig"]);
+        notification.success({
+          message: "Success",
+          description: data.message,
+        });
+      },
+      onError: (error: any) => {
+        notification.error({
+          message: "Error",
+          description: error?.response?.data?.message || "Something went wrong",
+        });
+      },
+    });
   const { mutateAsync: updateRagSettings, isLoading: isRagLoading } =
     useMutation(onRagApplicationUpdate, {
       onSuccess: (data) => {
@@ -135,6 +154,33 @@ export default function SettingsApplicationRoot() {
                   >
                     <Switch />
                   </Form.Item>
+                </div>
+                <div className="bg-gray-50 border-x border-b rounded-b-md rounded-x-md px-4 py-3 text-right sm:px-6 dark:bg-[#141414] dark:border-gray-600">
+                  <button
+                    disabled={isLoading}
+                    type="submit"
+                    className="inline-flex justify-center rounded-md border border-transparent bg-indigo-600 py-2 px-4 text-sm font-medium text-white  hover:bg-indigo-700 focus:outline-none focus:ring-2 focus:ring-indigo-500 focus:ring-offset-2"
+                  >
+                    {isLoading ? "Saving..." : "Save"}
+                  </button>
+                </div>
+              </div>
+            </Form>
+          </ApplicationCard>
+
+          <ApplicationCard
+            title="Models Settings"
+            description="Configure your models settings"
+          >
+            <Form
+              initialValues={{
+                ...data,
+              }}
+              layout="vertical"
+              onFinish={updateModelSettings}
+            >
+              <div className="sm:overflow-hidden ">
+                <div className="space-y-6 border-t border rounded-t-md  bg-white px-4 py-5 sm:p-6 dark:bg-[#171717] dark:border-gray-600">
                   <Form.Item
                     label="Default Chat Model"
                     name="defaultChatModel"
@@ -195,19 +241,49 @@ export default function SettingsApplicationRoot() {
                       loading={modeStatus === "loading"}
                     />
                   </Form.Item>
+                  <Form.Item
+                    label="Ollama URL"
+                    name="ollamaURL"
+                    rules={[
+                      {
+                        required: true,
+                        message: "Please input ollama url!",
+                      },
+                    ]}
+                  >
+                    <Input size="large" placeholder="Enter ollama url" />
+                  </Form.Item>
+
+                  <Form.Item
+                    label="Hide Default Models"
+                    name="hideDefaultModels"
+                    valuePropName="checked"
+                    help="This will hide all the default models and only show the models that are locally added or from ollama."
+                  >
+                    <Switch />
+                  </Form.Item>
+                  <Form.Item
+                    label="Dynamically Fetch Ollama Models"
+                    name="dynamicallyFetchOllamaModels"
+                    valuePropName="checked"
+                    help="This will dynamically fetch the models from ollama. You don't need to manually add the models."
+                  >
+                    <Switch />
+                  </Form.Item>
                 </div>
                 <div className="bg-gray-50 border-x border-b rounded-b-md rounded-x-md px-4 py-3 text-right sm:px-6 dark:bg-[#141414] dark:border-gray-600">
                   <button
-                    disabled={isLoading}
+                    disabled={isModelLoading}
                     type="submit"
                     className="inline-flex justify-center rounded-md border border-transparent bg-indigo-600 py-2 px-4 text-sm font-medium text-white  hover:bg-indigo-700 focus:outline-none focus:ring-2 focus:ring-indigo-500 focus:ring-offset-2"
                   >
-                    {isLoading ? "Saving..." : "Save"}
+                    {isModelLoading ? "Saving..." : "Save"}
                   </button>
                 </div>
               </div>
             </Form>
           </ApplicationCard>
+
           <ApplicationCard
             title="RAG Settings"
             description="Configure your RAG settings"

diff --git a/server/prisma/migrations/q_12_3/migration.sql b/server/prisma/migrations/q_12_3/migration.sql
@@ -0,0 +1,2 @@
+-- AlterTable
+ALTER TABLE "DialoqbaseSettings" ADD COLUMN     "ollamaURL" TEXT DEFAULT 'http://host.docker.internal:11434';
diff --git a/server/prisma/schema.prisma b/server/prisma/schema.prisma
@@ -96,6 +96,7 @@ model DialoqbaseSettings {
   hideDefaultModels            Boolean? @default(false)
   defaultChatModel             String   @default("gpt-3.5-turbo-dbase")
   defaultEmbeddingModel        String   @default("dialoqbase_eb_text-embedding-ada-002")
+  ollamaURL                   String? @default("http://host.docker.internal:11434")
 }
 
 model BotIntegration {

diff --git a/server/src/handlers/api/v1/admin/model.handler.ts b/server/src/handlers/api/v1/admin/model.handler.ts
@@ -9,6 +9,7 @@ import {
 } from "./type";
 import axios from "axios";
 import { removeTrailingSlash } from "../../../../utils/url";
+import { getSettings } from "../../../../utils/common";
 
 const _getModelFromUrl = async (url: string, apiKey?: string) => {
   try {
@@ -130,21 +131,29 @@ export const getAllModelsHandler = async (
   request: FastifyRequest,
   reply: FastifyReply
 ) => {
-  try {
-    const prisma = request.server.prisma;
-    const user = request.user;
+  const prisma = request.server.prisma;
+  const user = request.user;
 
-    if (!user.is_admin) {
-      return reply.status(403).send({
-        message: "Forbidden",
-      });
-    }
-    const allModels = await prisma.dialoqbaseModels.findMany({
-      where: {
-        deleted: false,
-      },
+  if (!user.is_admin) {
+    return reply.status(403).send({
+      message: "Forbidden",
     });
+  }
 
+  const settings = await getSettings(prisma);
+
+  const not_to_hide_providers = settings?.hideDefaultModels
+    ? [ "Local", "local", "ollama", "transformer", "Transformer"]
+    : undefined;
+  const allModels = await prisma.dialoqbaseModels.findMany({
+    where: {
+      deleted: false,
+      model_provider: {
+        in: not_to_hide_providers,
+      },
+    },
+  });
+  try {
     return {
       data: allModels.filter((model) => model.model_type !== "embedding"),
       embedding: allModels.filter((model) => model.model_type === "embedding"),
@@ -245,7 +254,7 @@ export const saveModelFromInputedUrlHandler = async (
         });
       }
 
-      let newModelId = model_id.trim() + `_custom_${new Date().getTime()}`;
+      let newModelId = model_id.trim() + `_dialoqbase_${new Date().getTime()}`;
       await prisma.dialoqbaseModels.create({
         data: {
           name: isModelExist.name,

diff --git a/server/src/handlers/api/v1/bot/bot/api.handler.ts b/server/src/handlers/api/v1/bot/bot/api.handler.ts
@@ -16,6 +16,7 @@ import {
   uniqueNamesGenerator,
 } from "unique-names-generator";
 import { validateDataSource } from "../../../../../utils/datasource-validation";
+import { getModelInfo } from "../../../../../utils/get-model-info";
 
 export const createBotAPIHandler = async (
   request: FastifyRequest<CreateBotAPIRequest>,
@@ -55,19 +56,11 @@ export const createBotAPIHandler = async (
       message: `Reach maximum limit of ${maxBotsAllowed} bots per user`,
     });
   }
-  const modelInfo = await prisma.dialoqbaseModels.findFirst({
-    where: {
-      hide: false,
-      deleted: false,
-      OR: [
-        {
-          model_id: model,
-        },
-        {
-          model_id: `${model}-dbase`,
-        },
-      ],
-    },
+
+  const modelInfo = await getModelInfo({
+    model,
+    prisma,
+    type: "chat",
   });
 
   if (!modelInfo) {
@@ -76,19 +69,10 @@ export const createBotAPIHandler = async (
     });
   }
 
-  const embeddingInfo = await prisma.dialoqbaseModels.findFirst({
-    where: {
-      OR: [
-        {
-          model_id: embedding,
-        },
-        {
-          model_id: `dialoqbase_eb_${embedding}`,
-        },
-      ],
-      hide: false,
-      deleted: false,
-    },
+  const embeddingInfo = await getModelInfo({
+    model: embedding,
+    prisma,
+    type: "embedding",
   });
 
   if (!embeddingInfo) {