From eea41c04d9e7cc13e87b8ee6c53bd45eb357da0f Mon Sep 17 00:00:00 2001
From: Lucas Oliveira <lucas.edu.oli@hotmail.com>
Date: Thu, 18 Sep 2025 17:39:42 -0300
Subject: [PATCH] Refactored to have models inside ollama

---
 frontend/src/app/onboarding/advanced.tsx      | 99 +++++++++++--------
 .../src/app/onboarding/ibm-onboarding.tsx     | 44 ++++++++-
 .../src/app/onboarding/model-selector.tsx     |  1 +
 .../src/app/onboarding/ollama-onboarding.tsx  | 44 +++++++--
 .../src/app/onboarding/openai-onboarding.tsx  | 37 ++++++-
 frontend/src/app/onboarding/page.tsx          | 15 ++-
 6 files changed, 188 insertions(+), 52 deletions(-)

diff --git a/frontend/src/app/onboarding/advanced.tsx b/frontend/src/app/onboarding/advanced.tsx
index ae734444..32071ba1 100644
--- a/frontend/src/app/onboarding/advanced.tsx
+++ b/frontend/src/app/onboarding/advanced.tsx
@@ -1,4 +1,3 @@
-import { useState } from "react";
 import { LabelWrapper } from "@/components/label-wrapper";
 import OllamaLogo from "@/components/logo/ollama-logo";
 import {
@@ -12,59 +11,77 @@ import { Switch } from "@/components/ui/switch";
 import { ModelSelector } from "./model-selector";
 
 export function AdvancedOnboarding({
-  modelProvider,
+  icon,
+  languageModels,
+  embeddingModels,
+  languageModel,
+  embeddingModel,
+  setLanguageModel,
+  setEmbeddingModel,
+  sampleDataset,
+  setSampleDataset,
 }: {
-  modelProvider: string;
+  icon?: React.ReactNode;
+  languageModels?: { value: string; label: string }[];
+  embeddingModels?: { value: string; label: string }[];
+  languageModel?: string;
+  embeddingModel?: string;
+  setLanguageModel?: (model: string) => void;
+  setEmbeddingModel?: (model: string) => void;
+  sampleDataset: boolean;
+  setSampleDataset: (dataset: boolean) => void;
 }) {
-  const [model, setModel] = useState("");
-  const options = [
-    { value: "gpt-4", label: "GPT-4" },
-    { value: "gpt-4-turbo", label: "GPT-4 Turbo" },
-    { value: "gpt-3.5-turbo", label: "GPT-3.5 Turbo" },
-    { value: "claude-3-opus", label: "Claude 3 Opus" },
-    { value: "claude-3-sonnet", label: "Claude 3 Sonnet" },
-  ];
+  const hasEmbeddingModels =
+    embeddingModels && embeddingModel && setEmbeddingModel;
+  const hasLanguageModels = languageModels && languageModel && setLanguageModel;
   return (
     <Accordion type="single" collapsible>
       <AccordionItem value="item-1">
         <AccordionTrigger>Advanced settings</AccordionTrigger>
         <AccordionContent className="space-y-6">
-          <LabelWrapper
-            label="Embedding model"
-            description="It’s recommended that you use XYZ, ABC, or DEF models for best performance."
-            helperText="The embedding model for your Ollama server."
-            id="embedding-model"
-            required={true}
-          >
-            <ModelSelector
-              options={options}
-              icon={<OllamaLogo className="w-4 h-4" />}
-              value={model}
-              onValueChange={setModel}
-            />
-          </LabelWrapper>
-          <LabelWrapper
-            label="Language model"
-            description="It’s recommended that you use XYZ, ABC, or DEF models for best performance."
-            helperText="The embedding model for your Ollama server."
-            id="embedding-model"
-            required={true}
-          >
-            <ModelSelector
-              options={options}
-              icon={<OllamaLogo className="w-4 h-4" />}
-              value={model}
-              onValueChange={setModel}
-            />
-          </LabelWrapper>
-          <Separator />
+          {hasEmbeddingModels && (
+            <LabelWrapper
+              label="Embedding model"
+              description="It’s recommended that you use XYZ, ABC, or DEF models for best performance."
+              helperText="The embedding model for your Ollama server."
+              id="embedding-model"
+              required={true}
+            >
+              <ModelSelector
+                options={embeddingModels}
+                icon={icon}
+                value={embeddingModel}
+                onValueChange={setEmbeddingModel}
+              />
+            </LabelWrapper>
+          )}
+          {hasLanguageModels && (
+            <LabelWrapper
+              label="Language model"
+              description="It’s recommended that you use XYZ, ABC, or DEF models for best performance."
+              helperText="The embedding model for your Ollama server."
+              id="embedding-model"
+              required={true}
+            >
+              <ModelSelector
+                options={languageModels}
+                icon={icon}
+                value={languageModel}
+                onValueChange={setLanguageModel}
+              />
+            </LabelWrapper>
+          )}
+          {(hasLanguageModels || hasEmbeddingModels) && <Separator />}
           <LabelWrapper
             label="Sample dataset"
             description="Ingest two small PDFs"
             id="sample-dataset"
             flex
           >
-            <Switch />
+            <Switch
+              checked={sampleDataset}
+              onCheckedChange={setSampleDataset}
+            />
           </LabelWrapper>
         </AccordionContent>
       </AccordionItem>
diff --git a/frontend/src/app/onboarding/ibm-onboarding.tsx b/frontend/src/app/onboarding/ibm-onboarding.tsx
index fccbc732..f77e2d99 100644
--- a/frontend/src/app/onboarding/ibm-onboarding.tsx
+++ b/frontend/src/app/onboarding/ibm-onboarding.tsx
@@ -1,14 +1,46 @@
+import { useState } from "react";
 import { LabelInput } from "@/components/label-input";
+import IBMLogo from "@/components/logo/ibm-logo";
 import type { Settings } from "../api/queries/useGetSettingsQuery";
 import { AdvancedOnboarding } from "./advanced";
 
 export function IBMOnboarding({
   settings,
   setSettings,
+  sampleDataset,
+  setSampleDataset,
 }: {
   settings: Settings;
   setSettings: (settings: Settings) => void;
+  sampleDataset: boolean;
+  setSampleDataset: (dataset: boolean) => void;
 }) {
+  const languageModels = [
+    { value: "gpt-oss", label: "gpt-oss" },
+    { value: "llama3.1", label: "llama3.1" },
+    { value: "llama3.2", label: "llama3.2" },
+    { value: "llama3.3", label: "llama3.3" },
+    { value: "llama3.4", label: "llama3.4" },
+    { value: "llama3.5", label: "llama3.5" },
+  ];
+  const embeddingModels = [
+    { value: "text-embedding-3-small", label: "text-embedding-3-small" },
+  ];
+  const [languageModel, setLanguageModel] = useState("gpt-oss");
+  const [embeddingModel, setEmbeddingModel] = useState(
+    "text-embedding-3-small",
+  );
+  const handleLanguageModelChange = (model: string) => {
+    setLanguageModel(model);
+  };
+
+  const handleEmbeddingModelChange = (model: string) => {
+    setEmbeddingModel(model);
+  };
+
+  const handleSampleDatasetChange = (dataset: boolean) => {
+    setSampleDataset(dataset);
+  };
   return (
     <>
       <LabelInput
@@ -32,7 +64,17 @@ export function IBMOnboarding({
         required
         placeholder="..."
       />
-      <AdvancedOnboarding modelProvider="watsonx" />
+      <AdvancedOnboarding
+        icon={<IBMLogo className="w-4 h-4" />}
+        languageModels={languageModels}
+        embeddingModels={embeddingModels}
+        languageModel={languageModel}
+        embeddingModel={embeddingModel}
+        sampleDataset={sampleDataset}
+        setLanguageModel={handleLanguageModelChange}
+        setEmbeddingModel={handleEmbeddingModelChange}
+        setSampleDataset={handleSampleDatasetChange}
+      />
     </>
   );
 }
diff --git a/frontend/src/app/onboarding/model-selector.tsx b/frontend/src/app/onboarding/model-selector.tsx
index 5314cb0e..c6cdebab 100644
--- a/frontend/src/app/onboarding/model-selector.tsx
+++ b/frontend/src/app/onboarding/model-selector.tsx
@@ -35,6 +35,7 @@ export function ModelSelector({
   return (
     <Popover open={open} onOpenChange={setOpen}>
       <PopoverTrigger asChild>
+        {/** biome-ignore lint/a11y/useSemanticElements: has to be a Button */}
         <Button
           variant="outline"
           role="combobox"
diff --git a/frontend/src/app/onboarding/ollama-onboarding.tsx b/frontend/src/app/onboarding/ollama-onboarding.tsx
index 3dea6611..886144f4 100644
--- a/frontend/src/app/onboarding/ollama-onboarding.tsx
+++ b/frontend/src/app/onboarding/ollama-onboarding.tsx
@@ -9,13 +9,21 @@ import { ModelSelector } from "./model-selector";
 export function OllamaOnboarding({
   settings,
   setSettings,
+  sampleDataset,
+  setSampleDataset,
 }: {
   settings: Settings;
   setSettings: (settings: Settings) => void;
+  sampleDataset: boolean;
+  setSampleDataset: (dataset: boolean) => void;
 }) {
   const [open, setOpen] = useState(false);
   const [value, setValue] = useState("");
-  const frameworks = [
+  const [languageModel, setLanguageModel] = useState("gpt-oss");
+  const [embeddingModel, setEmbeddingModel] = useState(
+    "text-embedding-3-small",
+  );
+  const languageModels = [
     { value: "gpt-oss", label: "gpt-oss", default: true },
     { value: "llama3.1", label: "llama3.1" },
     { value: "llama3.2", label: "llama3.2" },
@@ -23,6 +31,13 @@ export function OllamaOnboarding({
     { value: "llama3.4", label: "llama3.4" },
     { value: "llama3.5", label: "llama3.5" },
   ];
+  const embeddingModels = [
+    { value: "text-embedding-3-small", label: "text-embedding-3-small" },
+  ];
+
+  const handleSampleDatasetChange = (dataset: boolean) => {
+    setSampleDataset(dataset);
+  };
   return (
     <>
       <LabelInput
@@ -32,22 +47,37 @@ export function OllamaOnboarding({
         required
         placeholder="http://..."
       />
-
       <LabelWrapper
-        label="Embedding Model"
+        label="Embedding model"
         helperText="The embedding model for your Ollama server."
         id="embedding-model"
         required={true}
       >
         <ModelSelector
-          options={frameworks}
+          options={embeddingModels}
           icon={<OllamaLogo className="w-4 h-4" />}
-          value={value}
-          onValueChange={setValue}
+          value={embeddingModel}
+          onValueChange={setEmbeddingModel}
+        />
+      </LabelWrapper>
+      <LabelWrapper
+        label="Language model"
+        helperText="The embedding model for your Ollama server."
+        id="embedding-model"
+        required={true}
+      >
+        <ModelSelector
+          options={languageModels}
+          icon={<OllamaLogo className="w-4 h-4" />}
+          value={languageModel}
+          onValueChange={setLanguageModel}
         />
       </LabelWrapper>
 
-      <AdvancedOnboarding modelProvider="ollama" />
+      <AdvancedOnboarding
+        sampleDataset={sampleDataset}
+        setSampleDataset={handleSampleDatasetChange}
+      />
     </>
   );
 }
diff --git a/frontend/src/app/onboarding/openai-onboarding.tsx b/frontend/src/app/onboarding/openai-onboarding.tsx
index 9ef81019..2efbe360 100644
--- a/frontend/src/app/onboarding/openai-onboarding.tsx
+++ b/frontend/src/app/onboarding/openai-onboarding.tsx
@@ -1,14 +1,39 @@
+import { useState } from "react";
 import { LabelInput } from "@/components/label-input";
+import OpenAILogo from "@/components/logo/openai-logo";
 import type { Settings } from "../api/queries/useGetSettingsQuery";
 import { AdvancedOnboarding } from "./advanced";
 
 export function OpenAIOnboarding({
   settings,
   setSettings,
+  sampleDataset,
+  setSampleDataset,
 }: {
   settings: Settings;
   setSettings: (settings: Settings) => void;
+  sampleDataset: boolean;
+  setSampleDataset: (dataset: boolean) => void;
 }) {
+  const [languageModel, setLanguageModel] = useState("gpt-4o-mini");
+  const [embeddingModel, setEmbeddingModel] = useState(
+    "text-embedding-3-small",
+  );
+  const languageModels = [{ value: "gpt-4o-mini", label: "gpt-4o-mini" }];
+  const embeddingModels = [
+    { value: "text-embedding-3-small", label: "text-embedding-3-small" },
+  ];
+  const handleLanguageModelChange = (model: string) => {
+    setLanguageModel(model);
+  };
+
+  const handleEmbeddingModelChange = (model: string) => {
+    setEmbeddingModel(model);
+  };
+
+  const handleSampleDatasetChange = (dataset: boolean) => {
+    setSampleDataset(dataset);
+  };
   return (
     <>
       <LabelInput
@@ -18,7 +43,17 @@ export function OpenAIOnboarding({
         required
         placeholder="sk-..."
       />
-      <AdvancedOnboarding modelProvider="openai" />
+      <AdvancedOnboarding
+        icon={<OpenAILogo className="w-4 h-4" />}
+        languageModels={languageModels}
+        embeddingModels={embeddingModels}
+        languageModel={languageModel}
+        embeddingModel={embeddingModel}
+        sampleDataset={sampleDataset}
+        setLanguageModel={handleLanguageModelChange}
+        setSampleDataset={handleSampleDatasetChange}
+        setEmbeddingModel={handleEmbeddingModelChange}
+      />
     </>
   );
 }
diff --git a/frontend/src/app/onboarding/page.tsx b/frontend/src/app/onboarding/page.tsx
index 67b31ae2..84ab5d60 100644
--- a/frontend/src/app/onboarding/page.tsx
+++ b/frontend/src/app/onboarding/page.tsx
@@ -28,6 +28,8 @@ function OnboardingPage() {
   const { isAuthenticated } = useAuth();
 
   const [modelProvider, setModelProvider] = useState<string>("openai");
+
+  const [sampleDataset, setSampleDataset] = useState<boolean>(false);
   // Fetch settings using React Query
   const { data: settingsDb = {} } = useGetSettingsQuery({
     enabled: isAuthenticated,
@@ -64,7 +66,7 @@ function OnboardingPage() {
         backgroundPosition: "center",
       }}
     >
-      <div className="flex flex-col items-center gap-5 min-h-[480px] w-full">
+      <div className="flex flex-col items-center gap-5 min-h-[550px] w-full">
         <div className="flex flex-col items-center justify-center gap-4">
           <h1 className="text-2xl font-medium font-chivo">
             Configure your models
@@ -94,15 +96,24 @@ function OnboardingPage() {
                 <OpenAIOnboarding
                   settings={settings}
                   setSettings={setSettings}
+                  sampleDataset={sampleDataset}
+                  setSampleDataset={setSampleDataset}
                 />
               </TabsContent>
               <TabsContent value="watsonx">
-                <IBMOnboarding settings={settings} setSettings={setSettings} />
+                <IBMOnboarding
+                  settings={settings}
+                  setSettings={setSettings}
+                  sampleDataset={sampleDataset}
+                  setSampleDataset={setSampleDataset}
+                />
               </TabsContent>
               <TabsContent value="ollama">
                 <OllamaOnboarding
                   settings={settings}
                   setSettings={setSettings}
+                  sampleDataset={sampleDataset}
+                  setSampleDataset={setSampleDataset}
                 />
               </TabsContent>
             </CardContent>