From 13a87275528940331efb9b976bdf81add6209d59 Mon Sep 17 00:00:00 2001
From: Faisal Amir <urmauur@gmail.com>
Date: Tue, 2 Sep 2025 22:16:43 +0700
Subject: [PATCH] feat: enable reasoning configuration

---
 .../browser/extensions/engines/AIEngine.ts    |   3 +-
 extensions/llamacpp-extension/src/index.ts    |  11 +-
 web-app/src/containers/Capabilities.tsx       |   4 +-
 web-app/src/containers/ChatInput.tsx          | 138 +++++++++++++++---
 .../src/containers/DropdownModelProvider.tsx  |   2 +-
 web-app/src/hooks/useChat.ts                  |   4 +
 .../src/services/__tests__/providers.test.ts  |   2 +-
 web-app/src/services/models.ts                |  28 ++--
 web-app/src/services/providers.ts             |  85 ++++++-----
 web-app/src/types/modelProviders.d.ts         |   2 +
 web-app/src/types/models.ts                   |   1 +
 11 files changed, 210 insertions(+), 70 deletions(-)

diff --git a/core/src/browser/extensions/engines/AIEngine.ts b/core/src/browser/extensions/engines/AIEngine.ts
index 7a223e468..29daf837b 100644
--- a/core/src/browser/extensions/engines/AIEngine.ts
+++ b/core/src/browser/extensions/engines/AIEngine.ts
@@ -8,6 +8,7 @@ export interface chatCompletionRequestMessage {
   content: string | null | Content[] // Content can be a string OR an array of content parts
   reasoning?: string | null // Some models return reasoning in completed responses
   reasoning_content?: string | null // Some models return reasoning in completed responses
+  reasoning_effort?: string | null
   name?: string
   tool_calls?: any[] // Simplified tool_call_id?: string
 }
@@ -281,5 +282,5 @@ export abstract class AIEngine extends BaseExtension {
    * Check if a tool is supported by the model
    * @param modelId
    */
-  abstract isToolSupported(modelId: string): Promise<boolean>
+  abstract isModelCapabilitySupported(modelId: string, capability: string): Promise<boolean>
 }
diff --git a/extensions/llamacpp-extension/src/index.ts b/extensions/llamacpp-extension/src/index.ts
index d584b3d08..5c2298f8e 100644
--- a/extensions/llamacpp-extension/src/index.ts
+++ b/extensions/llamacpp-extension/src/index.ts
@@ -71,6 +71,7 @@ type LlamacppConfig = {
   rope_scale: number
   rope_freq_base: number
   rope_freq_scale: number
+  reasoning_budget: boolean
   ctx_shift: boolean
 }
 
@@ -1389,6 +1390,9 @@ export default class llamacpp_extension extends AIEngine {
     // This is an expert level settings and should only be used by people
     // who knows what they are doing.
     // Takes a regex with matching tensor name as input
+    if (!cfg.reasoning_budget) {
+      args.push('--reasoning-budget', '0')
+    }
     if (cfg.override_tensor_buffer_t)
       args.push('--override-tensor', cfg.override_tensor_buffer_t)
     // offload multimodal projector model to the GPU by default. if there is not enough memory
@@ -1827,7 +1831,10 @@ export default class llamacpp_extension extends AIEngine {
    * @param modelId
    * @returns
    */
-  async isToolSupported(modelId: string): Promise<boolean> {
+  async isModelCapabilitySupported(
+    modelId: string,
+    capability: string
+  ): Promise<boolean> {
     const janDataFolderPath = await getJanDataFolderPath()
     const modelConfigPath = await joinPath([
       this.providerPath,
@@ -1846,7 +1853,7 @@ export default class llamacpp_extension extends AIEngine {
     ])
     return (await readGgufMetadata(modelPath)).metadata?.[
       'tokenizer.chat_template'
-    ]?.includes('tools')
+    ]?.includes(capability)
   }
 
   /**
diff --git a/web-app/src/containers/Capabilities.tsx b/web-app/src/containers/Capabilities.tsx
index e2e09030a..fa8bcf5b5 100644
--- a/web-app/src/containers/Capabilities.tsx
+++ b/web-app/src/containers/Capabilities.tsx
@@ -7,7 +7,7 @@ import {
 import {
   IconEye,
   IconTool,
-  IconAtom,
+  IconBrain,
   IconWorld,
   IconCodeCircle2,
 } from '@tabler/icons-react'
@@ -30,7 +30,7 @@ const Capabilities = ({ capabilities }: CapabilitiesProps) => {
         } else if (capability === 'tools') {
           icon = <IconTool className="size-3.5" />
         } else if (capability === 'reasoning') {
-          icon = <IconAtom className="size-3.5" />
+          icon = <IconBrain className="size-3.5" />
         } else if (capability === 'embeddings') {
           icon = <IconCodeCircle2 className="size-3.5" />
         } else if (capability === 'web_search') {
diff --git a/web-app/src/containers/ChatInput.tsx b/web-app/src/containers/ChatInput.tsx
index f799f6b50..b8f61e2b4 100644
--- a/web-app/src/containers/ChatInput.tsx
+++ b/web-app/src/containers/ChatInput.tsx
@@ -12,15 +12,22 @@ import {
   TooltipProvider,
   TooltipTrigger,
 } from '@/components/ui/tooltip'
+import {
+  Popover,
+  PopoverContent,
+  PopoverTrigger,
+} from '@/components/ui/popover'
 import { ArrowRight } from 'lucide-react'
 import {
   IconPhoto,
   IconWorld,
-  IconAtom,
+  IconBrain,
   IconTool,
   IconCodeCircle2,
   IconPlayerStopFilled,
   IconX,
+  IconChevronUp,
+  IconChevronDown,
 } from '@tabler/icons-react'
 import { useTranslation } from '@/i18n/react-i18next-compat'
 import { useGeneralSetting } from '@/hooks/useGeneralSetting'
@@ -33,7 +40,7 @@ import DropdownModelProvider from '@/containers/DropdownModelProvider'
 import { ModelLoader } from '@/containers/loaders/ModelLoader'
 import DropdownToolsAvailable from '@/containers/DropdownToolsAvailable'
 import { getConnectedServers } from '@/services/mcp'
-import { checkMmprojExists } from '@/services/models'
+import { checkMmprojExists, stopModel } from '@/services/models'
 
 type ChatInputProps = {
   className?: string
@@ -61,6 +68,7 @@ const ChatInput = ({ model, className, initialMessage }: ChatInputProps) => {
   const maxRows = 10
 
   const { selectedModel, selectedProvider } = useModelProvider()
+
   const { sendMessage } = useChat()
   const [message, setMessage] = useState('')
   const [dropdownToolsAvailable, setDropdownToolsAvailable] = useState(false)
@@ -77,6 +85,7 @@ const ChatInput = ({ model, className, initialMessage }: ChatInputProps) => {
   const [connectedServers, setConnectedServers] = useState<string[]>([])
   const [isDragOver, setIsDragOver] = useState(false)
   const [hasMmproj, setHasMmproj] = useState(false)
+  const [reasoningEffortOpen, setReasoningEffortOpen] = useState(false)
 
   // Check for connected MCP servers
   useEffect(() => {
@@ -654,6 +663,114 @@ const ChatInput = ({ model, className, initialMessage }: ChatInputProps) => {
                   </TooltipProvider>
                 )}
 
+                {selectedModel?.capabilities?.includes('reasoning') && (
+                  <div className="flex gap-0.5 items-center">
+                    <TooltipProvider>
+                      <Tooltip>
+                        <TooltipTrigger asChild>
+                          <div
+                            className="h-7 p-1 flex items-center justify-center rounded-sm hover:bg-main-view-fg/10 transition-all duration-200 ease-in-out gap-1 cursor-pointer"
+                            onClick={async () => {
+                              if (
+                                selectedModel?.reasoning &&
+                                selectedProvider
+                              ) {
+                                // Toggle reasoning budget
+                                selectedModel.reasoning.reasoning_budget =
+                                  !selectedModel.reasoning.reasoning_budget
+
+                                // If model is loaded, restart it with new settings
+                                try {
+                                  await stopModel(selectedModel.id)
+                                } catch (error) {
+                                  console.error(
+                                    'Error restarting model with new reasoning budget:',
+                                    error
+                                  )
+                                }
+                              }
+                            }}
+                          >
+                            <IconBrain
+                              size={18}
+                              className={
+                                selectedModel?.reasoning?.reasoning_budget
+                                  ? 'text-main-view-fg/50'
+                                  : 'text-main-view-fg/20'
+                              }
+                            />
+                          </div>
+                        </TooltipTrigger>
+                        <TooltipContent>
+                          <p>
+                            {t('reasoning')}:{' '}
+                            {selectedModel?.reasoning?.reasoning_budget
+                              ? 'On'
+                              : 'Off'}
+                          </p>
+                        </TooltipContent>
+                      </Tooltip>
+                    </TooltipProvider>
+                    {selectedModel?.reasoning?.reasoning_budget &&
+                      selectedModel?.reasoning?.reasoning_effort && (
+                        <Popover
+                          open={reasoningEffortOpen}
+                          onOpenChange={setReasoningEffortOpen}
+                        >
+                          <PopoverTrigger asChild>
+                            <div className="h-7 mr-1 text-sm rounded-sm text-main-view-fg/70 cursor-pointer flex items-center gap-1">
+                              <span className="capitalize">
+                                {selectedModel?.reasoning?.reasoning_effort ||
+                                  'auto'}
+                              </span>
+                              {reasoningEffortOpen ? (
+                                <IconChevronUp
+                                  size={12}
+                                  className="text-main-view-fg/50"
+                                />
+                              ) : (
+                                <IconChevronDown
+                                  size={12}
+                                  className="text-main-view-fg/50"
+                                />
+                              )}
+                            </div>
+                          </PopoverTrigger>
+                          <PopoverContent className="w-24 p-0" align="start">
+                            <div className="py-1">
+                              {['auto', 'low', 'medium', 'high'].map(
+                                (effort) => (
+                                  <div
+                                    key={effort}
+                                    className="px-2 py-1.5 text-sm cursor-pointer hover:bg-main-view-fg/20 capitalize"
+                                    onClick={async () => {
+                                      if (selectedModel?.reasoning) {
+                                        selectedModel.reasoning.reasoning_effort =
+                                          effort
+                                        setReasoningEffortOpen(false)
+                                        // Restart model with new reasoning effort
+                                        try {
+                                          await stopModel(selectedModel.id)
+                                        } catch (error) {
+                                          console.error(
+                                            'Error restarting model with new reasoning effort:',
+                                            error
+                                          )
+                                        }
+                                      }
+                                    }}
+                                  >
+                                    {effort}
+                                  </div>
+                                )
+                              )}
+                            </div>
+                          </PopoverContent>
+                        </Popover>
+                      )}
+                  </div>
+                )}
+
                 {selectedModel?.capabilities?.includes('tools') &&
                   hasActiveMCPServers && (
                     <TooltipProvider>
@@ -728,23 +845,6 @@ const ChatInput = ({ model, className, initialMessage }: ChatInputProps) => {
                     </Tooltip>
                   </TooltipProvider>
                 )}
-                {selectedModel?.capabilities?.includes('reasoning') && (
-                  <TooltipProvider>
-                    <Tooltip>
-                      <TooltipTrigger asChild>
-                        <div className="h-7 p-1 flex items-center justify-center rounded-sm hover:bg-main-view-fg/10 transition-all duration-200 ease-in-out gap-1">
-                          <IconAtom
-                            size={18}
-                            className="text-main-view-fg/50"
-                          />
-                        </div>
-                      </TooltipTrigger>
-                      <TooltipContent>
-                        <p>{t('reasoning')}</p>
-                      </TooltipContent>
-                    </Tooltip>
-                  </TooltipProvider>
-                )}
               </div>
             </div>
 
diff --git a/web-app/src/containers/DropdownModelProvider.tsx b/web-app/src/containers/DropdownModelProvider.tsx
index e47c31503..1973fc9c5 100644
--- a/web-app/src/containers/DropdownModelProvider.tsx
+++ b/web-app/src/containers/DropdownModelProvider.tsx
@@ -393,7 +393,7 @@ const DropdownModelProvider = ({
 
   return (
     <Popover open={open} onOpenChange={onOpenChange}>
-      <div className="bg-main-view-fg/5 hover:bg-main-view-fg/8 px-2 py-1 flex items-center gap-1.5 rounded-sm max-h-[32px] mr-0.5">
+      <div className="bg-main-view-fg/5 hover:bg-main-view-fg/8 px-2 py-1 flex items-center gap-1.5 rounded-sm max-h-[32px] mr-1">
         <PopoverTrigger asChild>
           <button
             title={displayModel}
diff --git a/web-app/src/hooks/useChat.ts b/web-app/src/hooks/useChat.ts
index 134dc1ae1..3bbbd362f 100644
--- a/web-app/src/hooks/useChat.ts
+++ b/web-app/src/hooks/useChat.ts
@@ -289,6 +289,10 @@ export const useChat = () => {
             {
               ...modelSettings,
               ...currentAssistant.parameters,
+              ...(selectedModel?.reasoning?.reasoning_budget &&
+                selectedModel?.reasoning?.reasoning_budget !== 'auto' && {
+                  reasoning_effort: selectedModel?.reasoning?.reasoning_effort,
+                }),
             } as unknown as Record<string, object>
           )
 
diff --git a/web-app/src/services/__tests__/providers.test.ts b/web-app/src/services/__tests__/providers.test.ts
index 6660ffa30..6355d8763 100644
--- a/web-app/src/services/__tests__/providers.test.ts
+++ b/web-app/src/services/__tests__/providers.test.ts
@@ -69,7 +69,7 @@ vi.mock('../models', () => ({
       { id: 'llama-2-7b', name: 'Llama 2 7B', description: 'Llama model' },
     ])
   ),
-  isToolSupported: vi.fn(() => Promise.resolve(false)),
+  isModelCapabilitySupported: vi.fn(() => Promise.resolve(false)),
 }))
 
 vi.mock('@/lib/extension', () => ({
diff --git a/web-app/src/services/models.ts b/web-app/src/services/models.ts
index f85134998..4ff9b25f0 100644
--- a/web-app/src/services/models.ts
+++ b/web-app/src/services/models.ts
@@ -415,14 +415,19 @@ export const startModel = async (
     return keyMappings[key] || key
   }
 
-  const settings = modelConfig?.settings
-    ? Object.fromEntries(
-        Object.entries(modelConfig.settings).map(([key, value]) => [
-          mapSettingKey(key),
-          value.controller_props?.value,
-        ])
-      )
-    : undefined
+  const settings = {
+    ...(modelConfig?.settings
+      ? Object.fromEntries(
+          Object.entries(modelConfig.settings).map(([key, value]) => [
+            mapSettingKey(key),
+            value.controller_props?.value,
+          ])
+        )
+      : {}),
+    ...(modelConfig?.reasoning?.reasoning_budget !== undefined && {
+      reasoning_budget: modelConfig.reasoning.reasoning_budget,
+    }),
+  }
 
   return engine.load(model, settings).catch((error) => {
     console.error(
@@ -439,11 +444,14 @@ export const startModel = async (
  * @param modelId
  * @returns
  */
-export const isToolSupported = async (modelId: string): Promise<boolean> => {
+export const isModelCapabilitySupported = async (
+  modelId: string,
+  capability: string
+): Promise<boolean> => {
   const engine = getEngine()
   if (!engine) return false
 
-  return engine.isToolSupported(modelId)
+  return engine.isModelCapabilitySupported(modelId, capability)
 }
 
 /**
diff --git a/web-app/src/services/providers.ts b/web-app/src/services/providers.ts
index e9f05fd09..c7529d5a5 100644
--- a/web-app/src/services/providers.ts
+++ b/web-app/src/services/providers.ts
@@ -3,7 +3,7 @@ import { predefinedProviders } from '@/consts/providers'
 import { EngineManager, SettingComponentProps } from '@janhq/core'
 import { ModelCapabilities } from '@/types/models'
 import { modelSettings } from '@/lib/predefined'
-import { fetchModels, isToolSupported } from './models'
+import { fetchModels, isModelCapabilitySupported } from './models'
 import { ExtensionManager } from '@/lib/extension'
 import { fetch as fetchTauri } from '@tauri-apps/plugin-http'
 
@@ -63,39 +63,56 @@ export const getProviders = async (): Promise<ModelProvider[]> => {
         }
       }) as ProviderSetting[],
       models: await Promise.all(
-        models.map(
-          async (model) =>
-            ({
-              id: model.id,
-              model: model.id,
-              name: model.name,
-              description: model.description,
-              capabilities:
-                'capabilities' in model
-                  ? (model.capabilities as string[])
-                  : (await isToolSupported(model.id))
-                    ? [ModelCapabilities.TOOLS]
-                    : [],
-              provider: providerName,
-              settings: Object.values(modelSettings).reduce(
-                (acc, setting) => {
-                  let value = setting.controller_props.value
-                  if (setting.key === 'ctx_len') {
-                    value = 8192 // Default context length for Llama.cpp models
-                  }
-                  acc[setting.key] = {
-                    ...setting,
-                    controller_props: {
-                      ...setting.controller_props,
-                      value: value,
-                    },
-                  }
-                  return acc
-                },
-                {} as Record<string, ProviderSetting>
-              ),
-            }) as Model
-        )
+        models.map(async (model) => {
+          return {
+            id: model.id,
+            model: model.id,
+            name: model.name,
+            description: model.description,
+            capabilities: [
+              ...(model.capabilities?.includes(ModelCapabilities.TOOLS) ||
+              (await isModelCapabilitySupported(model.id, 'tools'))
+                ? [ModelCapabilities.TOOLS]
+                : []),
+              ...(model.capabilities?.includes(ModelCapabilities.REASONING) ||
+              (await isModelCapabilitySupported(model.id, 'reasoning'))
+                ? [ModelCapabilities.REASONING]
+                : []),
+            ],
+            provider: providerName,
+            settings: Object.values(modelSettings).reduce(
+              (acc, setting) => {
+                let value = setting.controller_props.value
+                if (setting.key === 'ctx_len') {
+                  value = 8192 // Default context length for Llama.cpp models
+                }
+                acc[setting.key] = {
+                  ...setting,
+                  controller_props: {
+                    ...setting.controller_props,
+                    value: value,
+                  },
+                }
+                return acc
+              },
+              {} as Record<string, ProviderSetting>
+            ),
+            ...(model.capabilities?.includes(ModelCapabilities.REASONING) ||
+            (await isModelCapabilitySupported(model.id, 'reasoning'))
+              ? {
+                  reasoning: {
+                    reasoning_budget: true,
+                    ...((await isModelCapabilitySupported(
+                      model.id,
+                      'reasoning_effort'
+                    ))
+                      ? { reasoning_effort: 'auto' }
+                      : {}),
+                  },
+                }
+              : {}),
+          } as Model
+        })
       ),
     }
     runtimeProviders.push(provider)
diff --git a/web-app/src/types/modelProviders.d.ts b/web-app/src/types/modelProviders.d.ts
index eb035e471..fe006a359 100644
--- a/web-app/src/types/modelProviders.d.ts
+++ b/web-app/src/types/modelProviders.d.ts
@@ -33,6 +33,8 @@ type Model = {
   format?: string
   capabilities?: string[]
   settings?: Record<string, ProviderSetting>
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  reasoning?: Record<string, any>
 }
 
 /**
diff --git a/web-app/src/types/models.ts b/web-app/src/types/models.ts
index f88541bb1..db2f48854 100644
--- a/web-app/src/types/models.ts
+++ b/web-app/src/types/models.ts
@@ -7,6 +7,7 @@ export enum ModelCapabilities {
   COMPLETION = 'completion',
   TOOLS = 'tools',
   EMBEDDINGS = 'embeddings',
+  REASONING = 'reasoning',
   IMAGE_GENERATION = 'image_generation',
   AUDIO_GENERATION = 'audio_generation',
   TEXT_TO_IMAGE = 'text_to_image',