Add Ability to configure default chat and embedding model

n4ze3m · Apr 22, 2024 · a05bbb0 · a05bbb0
1 parent 15b38c9
commit a05bbb0
Show file tree

Hide file tree

Showing 4 changed files with 82 additions and 6 deletions.
diff --git a/app/ui/src/components/Common/BotForm.tsx b/app/ui/src/components/Common/BotForm.tsx
@@ -480,9 +480,11 @@ export const BotForm = ({
           onFinish={createBot}
           form={form}
           className="space-y-6"
-          initialValues={{ 
-            embedding: "dialoqbase_eb_text-embedding-ada-002",
-            model: "gpt-3.5-turbo-dbase",
+          initialValues={{
+            embedding:
+              botConfig?.defaultEmbeddingModel ||
+              "dialoqbase_eb_text-embedding-ada-002",
+            model: botConfig?.defaultChatModel || "gpt-3.5-turbo-dbase",
             maxDepth: 2,
             maxLinks: 10,
             options: {
@@ -646,7 +648,7 @@ export const BotForm = ({
                   .toLowerCase()
                   .localeCompare((optionB?.label ?? "").toLowerCase())
               }
-              placeholder="Select a chat model"
+              placeholder="Select a Chat Model"
               options={botConfig.chatModel}
             />
           </Form.Item>
@@ -675,6 +677,7 @@ export const BotForm = ({
                   .localeCompare((optionB?.label ?? "").toLowerCase())
               }
               options={botConfig.embeddingModel}
+              placeholder="Select an Embedding Model"
             />
           </Form.Item>
 

diff --git a/app/ui/src/hooks/useCreateConfig.tsx b/app/ui/src/hooks/useCreateConfig.tsx
@@ -16,6 +16,8 @@ export const useCreateConfig = () => {
           label: string;
           value: string;
         }[];
+        defaultChatModel?: string;
+        defaultEmbeddingModel?: string;
       };
     }
   );

diff --git a/app/ui/src/routes/settings/application.tsx b/app/ui/src/routes/settings/application.tsx
@@ -1,16 +1,19 @@
-import { Form, InputNumber, Switch, notification } from "antd";
+import { Form, InputNumber, Switch, notification, Select } from "antd";
 import React from "react";
 import api from "../../services/api";
-import { useMutation, useQuery } from "@tanstack/react-query";
+import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query";
 
 import { SettingsLayout } from "../../Layout/SettingsLayout";
 import { SkeletonLoading } from "../../components/Common/SkeletonLoading";
 import { useNavigate } from "react-router-dom";
 import { ApplicationCard } from "../../components/Settings/Application/ApplicationCard";
+import { useCreateConfig } from "../../hooks/useCreateConfig";
 
 export default function SettingsApplicationRoot() {
   const [form] = Form.useForm();
   const navigate = useNavigate();
+  const { data: models, status: modeStatus } = useCreateConfig();
+  const queryClient = useQueryClient();
 
   const { data, status } = useQuery(["fetchApplicationSettings"], async () => {
     const response = await api.get("/admin/dialoqbase-settings");
@@ -20,6 +23,8 @@ export default function SettingsApplicationRoot() {
       allowUserToRegister: boolean;
       defaultChunkSize: number;
       defaultChunkOverlap: number;
+      defaultChatModel: string;
+      defaultEmbeddingModel: string;
     };
   });
 
@@ -43,6 +48,7 @@ export default function SettingsApplicationRoot() {
     onUpdateApplicatoon,
     {
       onSuccess: (data) => {
+        queryClient.invalidateQueries(["fetchBotCreateConfig"]);
         notification.success({
           message: "Success",
           description: data.message,
@@ -129,6 +135,66 @@ export default function SettingsApplicationRoot() {
                   >
                     <Switch />
                   </Form.Item>
+                  <Form.Item
+                    label="Default Chat Model"
+                    name="defaultChatModel"
+                    rules={[
+                      {
+                        required: true,
+                        message: "Please select default chat model!",
+                      },
+                    ]}
+                  >
+                    <Select
+                      showSearch
+                      filterOption={(input, option) =>
+                        (option?.label?.toLowerCase() ?? "").includes(
+                          input?.toLowerCase()
+                        ) ||
+                        (option?.value?.toLowerCase() ?? "").includes(
+                          input?.toLowerCase()
+                        )
+                      }
+                      filterSort={(optionA, optionB) =>
+                        (optionA?.label ?? "")
+                          .toLowerCase()
+                          .localeCompare((optionB?.label ?? "").toLowerCase())
+                      }
+                      placeholder="Select a Chat Model"
+                      options={models?.chatModel || []}
+                      loading={modeStatus === "loading"}
+                    />
+                  </Form.Item>
+                  <Form.Item
+                    label="Default Embedding Model"
+                    name="defaultEmbeddingModel"
+                    rules={[
+                      {
+                        required: true,
+                        message: "Please select default embedding model!",
+                      },
+                    ]}
+                  >
+                    <Select
+                      showSearch
+                      filterOption={(input, option) =>
+                        (option?.label?.toLowerCase() ?? "").includes(
+                          input?.toLowerCase()
+                        ) ||
+                        (option?.value?.toLowerCase() ?? "").includes(
+                          input?.toLowerCase()
+                        )
+                      }
+                      filterSort={(optionA, optionB) =>
+                        (optionA?.label ?? "")
+                          .toLowerCase()
+                          .localeCompare((optionB?.label ?? "").toLowerCase())
+                      }
+                      placeholder="Select a Embedding Model"
+                      options={models?.embeddingModel || []}
+                      loading={modeStatus === "loading"}
+                    />
+                  </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

diff --git a/server/src/handlers/api/v1/bot/bot/get.handler.ts b/server/src/handlers/api/v1/bot/bot/get.handler.ts
@@ -1,6 +1,7 @@
 import { FastifyReply, FastifyRequest } from "fastify";
 
 import { GetBotRequestById } from "./types";
+import { getSettings } from "../../../../../utils/common";
 
 export const getBotByIdEmbeddingsHandler = async (
   request: FastifyRequest<GetBotRequestById>,
@@ -153,10 +154,14 @@ export const getCreateBotConfigHandler = async (
         disabled: model.model_id === "dialoqbase_eb_dialoqbase-ollama",
       };
     });
+
+  const settings = await getSettings(prisma);
 
   return {
     chatModel,
     embeddingModel,
+    defaultChatModel: settings?.defaultChatModel,
+    defaultEmbeddingModel: settings?.defaultEmbeddingModel,
   };
 };