feat: Add lead-worker model selection and real-time model display in GUI (#2964)

Co-authored-by: jack <jack@deck.local>
2025-06-18 05:40:20 +02:00
parent d8b6e6011b
commit 657718d8c0
16 changed files with 475 additions and 9 deletions
@@ -1,4 +1,4 @@
-import React, { useEffect, useRef, useState, useMemo, useCallback } from 'react';
+import React, { useEffect, useRef, useState, useMemo, useCallback, createContext, useContext } from 'react';
 import { getApiUrl } from '../config';
 import FlappyGoose from './FlappyGoose';
 import GooseMessage from './GooseMessage';
@@ -37,6 +37,10 @@ import {
  TextContent,
 } from '../types/message';

+// Context for sharing current model info
+const CurrentModelContext = createContext<{ model: string; mode: string } | null>(null);
+export const useCurrentModelInfo = () => useContext(CurrentModelContext);
+
 export interface ChatType {
  id: string;
  title: string;
@@ -144,6 +148,7 @@ function ChatContent({
    handleSubmit: _submitMessage,
    updateMessageStreamBody,
    notifications,
+    currentModelInfo,
  } = useMessageStream({
    api: getApiUrl('/reply'),
    initialMessages: chat.messages,
@@ -504,7 +509,8 @@ function ChatContent({
  }, new Map());

  return (
-    <div className="flex flex-col w-full h-screen items-center justify-center">
+    <CurrentModelContext.Provider value={currentModelInfo}>
+      <div className="flex flex-col w-full h-screen items-center justify-center">
      {/* Loader when generating recipe */}
      {isGeneratingRecipe && <LayingEggLoader />}
      <MoreMenuLayout
@@ -647,5 +653,6 @@ function ChatContent({
        summaryContent={summaryContent}
      />
    </div>
+    </CurrentModelContext.Provider>
  );
 }
@@ -2,8 +2,12 @@ import { Sliders } from 'lucide-react';
 import React, { useEffect, useState, useRef } from 'react';
 import { useModelAndProvider } from '../../../ModelAndProviderContext';
 import { AddModelModal } from '../subcomponents/AddModelModal';
+import { LeadWorkerSettings } from '../subcomponents/LeadWorkerSettings';
 import { View } from '../../../../App';
 import { Tooltip, TooltipTrigger, TooltipContent, TooltipProvider } from '../../../ui/Tooltip';
+import Modal from '../../../Modal';
+import { useCurrentModelInfo } from '../../../ChatView';
+import { useConfig } from '../../../ConfigContext';

 interface ModelsBottomBarProps {
  dropdownRef: React.RefObject<HTMLDivElement>;
@@ -12,15 +16,36 @@ interface ModelsBottomBarProps {
 export default function ModelsBottomBar({ dropdownRef, setView }: ModelsBottomBarProps) {
  const { currentModel, currentProvider, getCurrentModelAndProviderForDisplay } =
    useModelAndProvider();
+  const currentModelInfo = useCurrentModelInfo();
+  const { read } = useConfig();
  const [isModelMenuOpen, setIsModelMenuOpen] = useState(false);
  const [displayProvider, setDisplayProvider] = useState<string | null>(null);
  const [isAddModelModalOpen, setIsAddModelModalOpen] = useState(false);
+  const [isLeadWorkerModalOpen, setIsLeadWorkerModalOpen] = useState(false);
+  const [isLeadWorkerActive, setIsLeadWorkerActive] = useState(false);
  const menuRef = useRef<HTMLDivElement>(null);
  const [isModelTruncated, setIsModelTruncated] = useState(false);
  // eslint-disable-next-line no-undef
  const modelRef = useRef<HTMLSpanElement>(null);
  const [isTooltipOpen, setIsTooltipOpen] = useState(false);

+  // Check if lead/worker mode is active
+  useEffect(() => {
+    const checkLeadWorker = async () => {
+      try {
+        const leadModel = await read('GOOSE_LEAD_MODEL', false);
+        setIsLeadWorkerActive(!!leadModel);
+      } catch (error) {
+        setIsLeadWorkerActive(false);
+      }
+    };
+    checkLeadWorker();
+  }, [read]);
+
+  // Determine which model to display - activeModel takes priority when lead/worker is active
+  const displayModel = (isLeadWorkerActive && currentModelInfo?.model) ? currentModelInfo.model : (currentModel || 'Select Model');
+  const modelMode = currentModelInfo?.mode;
+
  // Update display provider when current provider changes
  useEffect(() => {
    if (currentProvider) {
@@ -40,7 +65,7 @@ export default function ModelsBottomBar({ dropdownRef, setView }: ModelsBottomBa
    checkTruncation();
    window.addEventListener('resize', checkTruncation);
    return () => window.removeEventListener('resize', checkTruncation);
-  }, [currentModel]);
+  }, [displayModel]);

  useEffect(() => {
    setIsTooltipOpen(false);
@@ -79,12 +104,22 @@ export default function ModelsBottomBar({ dropdownRef, setView }: ModelsBottomBa
                  ref={modelRef}
                  className="truncate max-w-[130px] md:max-w-[200px] lg:max-w-[360px] min-w-0 block"
                >
-                  {currentModel || 'Select Model'}
+                  {displayModel}
+                  {isLeadWorkerActive && modelMode && (
+                    <span className="ml-1 text-[10px] opacity-60">
+                      ({modelMode})
+                    </span>
+                  )}
                </span>
              </TooltipTrigger>
              {isModelTruncated && (
                <TooltipContent className="max-w-96 overflow-auto scrollbar-thin" side="top">
-                  {currentModel || 'Select Model'}
+                  {displayModel}
+                  {isLeadWorkerActive && modelMode && (
+                    <span className="ml-1 text-[10px] opacity-60">
+                      ({modelMode})
+                    </span>
+                  )}
                </TooltipContent>
              )}
            </Tooltip>
@@ -110,6 +145,17 @@ export default function ModelsBottomBar({ dropdownRef, setView }: ModelsBottomBa
                <span className="text-sm">Change Model</span>
                <Sliders className="w-4 h-4 ml-2 rotate-90" />
              </div>
+              <div
+                className="flex items-center justify-between text-textStandard p-2 cursor-pointer transition-colors hover:bg-bgStandard
+                    border-t border-borderSubtle"
+                onClick={() => {
+                  setIsModelMenuOpen(false);
+                  setIsLeadWorkerModalOpen(true);
+                }}
+              >
+                <span className="text-sm">Lead/Worker Settings</span>
+                <Sliders className="w-4 h-4 ml-2" />
+              </div>
            </div>
          </div>
        )}
@@ -118,6 +164,12 @@ export default function ModelsBottomBar({ dropdownRef, setView }: ModelsBottomBa
      {isAddModelModalOpen ? (
        <AddModelModal setView={setView} onClose={() => setIsAddModelModalOpen(false)} />
      ) : null}
+      
+      {isLeadWorkerModalOpen ? (
+        <Modal onClose={() => setIsLeadWorkerModalOpen(false)}>
+          <LeadWorkerSettings onClose={() => setIsLeadWorkerModalOpen(false)} />
+        </Modal>
+      ) : null}
    </div>
  );
 }
@@ -0,0 +1,262 @@
+import { useState, useEffect } from 'react';
+import { useConfig } from '../../../ConfigContext';
+import { useModelAndProvider } from '../../../ModelAndProviderContext';
+import { Button } from '../../../ui/button';
+import { Select } from '../../../ui/Select';
+import { Input } from '../../../ui/input';
+import { Info } from 'lucide-react';
+
+interface LeadWorkerSettingsProps {
+  onClose: () => void;
+}
+
+export function LeadWorkerSettings({ onClose }: LeadWorkerSettingsProps) {
+  const { read, upsert, getProviders, remove } = useConfig();
+  const { currentModel } = useModelAndProvider();
+  const [leadModel, setLeadModel] = useState<string>('');
+  const [workerModel, setWorkerModel] = useState<string>('');
+  const [leadProvider, setLeadProvider] = useState<string>('');
+  const [workerProvider, setWorkerProvider] = useState<string>('');
+  const [leadTurns, setLeadTurns] = useState<number>(3);
+  const [failureThreshold, setFailureThreshold] = useState<number>(2);
+  const [fallbackTurns, setFallbackTurns] = useState<number>(2);
+  const [isEnabled, setIsEnabled] = useState(false);
+  const [modelOptions, setModelOptions] = useState<{ value: string; label: string; provider: string }[]>([]);
+  const [isLoading, setIsLoading] = useState(true);
+
+  // Load current configuration
+  useEffect(() => {
+    const loadConfig = async () => {
+      try {
+        setIsLoading(true);
+        const [
+          leadModelConfig,
+          leadProviderConfig,
+          leadTurnsConfig,
+          failureThresholdConfig,
+          fallbackTurnsConfig,
+        ] = await Promise.all([
+          read('GOOSE_LEAD_MODEL', false),
+          read('GOOSE_LEAD_PROVIDER', false),
+          read('GOOSE_LEAD_TURNS', false),
+          read('GOOSE_LEAD_FAILURE_THRESHOLD', false),
+          read('GOOSE_LEAD_FALLBACK_TURNS', false),
+        ]);
+
+        if (leadModelConfig) {
+          setLeadModel(leadModelConfig as string);
+          setIsEnabled(true);
+        }
+        if (leadProviderConfig) setLeadProvider(leadProviderConfig as string);
+        if (leadTurnsConfig) setLeadTurns(Number(leadTurnsConfig));
+        if (failureThresholdConfig) setFailureThreshold(Number(failureThresholdConfig));
+        if (fallbackTurnsConfig) setFallbackTurns(Number(fallbackTurnsConfig));
+        
+        // Set worker model to current model or from config
+        const workerModelConfig = await read('GOOSE_MODEL', false);
+        if (workerModelConfig) {
+          setWorkerModel(workerModelConfig as string);
+        } else if (currentModel) {
+          setWorkerModel(currentModel as string);
+        }
+        
+        const workerProviderConfig = await read('GOOSE_PROVIDER', false);
+        if (workerProviderConfig) {
+          setWorkerProvider(workerProviderConfig as string);
+        }
+
+        // Load available models
+        const providers = await getProviders(false);
+        const activeProviders = providers.filter((p) => p.is_configured);
+        const options: { value: string; label: string; provider: string }[] = [];
+        
+        activeProviders.forEach(({ metadata, name }) => {
+          if (metadata.known_models) {
+            metadata.known_models.forEach((model) => {
+              options.push({
+                value: model.name,
+                label: `${model.name} (${metadata.display_name})`,
+                provider: name,
+              });
+            });
+          }
+        });
+        
+        setModelOptions(options);
+      } catch (error) {
+        console.error('Error loading configuration:', error);
+      } finally {
+        setIsLoading(false);
+      }
+    };
+
+    loadConfig();
+  }, [read, getProviders, currentModel]);
+
+  const handleSave = async () => {
+    try {
+      if (isEnabled && leadModel && workerModel) {
+        // Save lead/worker configuration
+        await Promise.all([
+          upsert('GOOSE_LEAD_MODEL', leadModel, false),
+          leadProvider && upsert('GOOSE_LEAD_PROVIDER', leadProvider, false),
+          upsert('GOOSE_MODEL', workerModel, false),
+          workerProvider && upsert('GOOSE_PROVIDER', workerProvider, false),
+          upsert('GOOSE_LEAD_TURNS', leadTurns, false),
+          upsert('GOOSE_LEAD_FAILURE_THRESHOLD', failureThreshold, false),
+          upsert('GOOSE_LEAD_FALLBACK_TURNS', fallbackTurns, false),
+        ]);
+      } else {
+        // Remove lead/worker configuration
+        await Promise.all([
+          remove('GOOSE_LEAD_MODEL', false),
+          remove('GOOSE_LEAD_PROVIDER', false),
+          remove('GOOSE_LEAD_TURNS', false),
+          remove('GOOSE_LEAD_FAILURE_THRESHOLD', false),
+          remove('GOOSE_LEAD_FALLBACK_TURNS', false),
+        ]);
+      }
+      onClose();
+    } catch (error) {
+      console.error('Error saving configuration:', error);
+    }
+  };
+
+  if (isLoading) {
+    return <div className="p-4">Loading...</div>;
+  }
+
+  return (
+    <div className="p-4 space-y-4">
+      <div className="space-y-2">
+        <h3 className="text-lg font-medium text-textProminent">Lead/Worker Mode</h3>
+        <p className="text-sm text-textSubtle">
+          Configure a lead model for planning and a worker model for execution
+        </p>
+      </div>
+
+      <div className="flex items-center space-x-2">
+        <input
+          type="checkbox"
+          id="enable-lead-worker"
+          checked={isEnabled}
+          onChange={(e) => setIsEnabled(e.target.checked)}
+          className="rounded border-borderStandard"
+        />
+        <label htmlFor="enable-lead-worker" className="text-sm text-textStandard">
+          Enable lead/worker mode
+        </label>
+      </div>
+
+      <div className="space-y-4">
+        <div className="space-y-2">
+          <label className="text-sm text-textSubtle">Lead Model</label>
+          <Select
+            options={modelOptions}
+            value={modelOptions.find((opt) => opt.value === leadModel) || null}
+            onChange={(newValue: unknown) => {
+              const option = newValue as { value: string; provider: string } | null;
+              if (option) {
+                setLeadModel(option.value);
+                setLeadProvider(option.provider);
+              }
+            }}
+            placeholder="Select lead model..."
+            isDisabled={!isEnabled}
+          />
+          <p className="text-xs text-textSubtle">
+            Strong model for initial planning and fallback recovery
+          </p>
+        </div>
+
+        <div className="space-y-2">
+          <label className="text-sm text-textSubtle">Worker Model</label>
+          <Select
+            options={modelOptions}
+            value={modelOptions.find((opt) => opt.value === workerModel) || null}
+            onChange={(newValue: unknown) => {
+              const option = newValue as { value: string; provider: string } | null;
+              if (option) {
+                setWorkerModel(option.value);
+                setWorkerProvider(option.provider);
+              }
+            }}
+            placeholder="Select worker model..."
+            isDisabled={!isEnabled}
+          />
+          <p className="text-xs text-textSubtle">
+            Fast model for routine execution tasks
+          </p>
+        </div>
+
+        <div className="space-y-4 pt-4 border-t border-borderSubtle">
+          <div className="space-y-2">
+            <label className="text-sm text-textSubtle flex items-center gap-1">
+              Initial Lead Turns
+              <Info size={14} className="text-textSubtle" />
+            </label>
+            <Input
+              type="number"
+              min={1}
+              max={10}
+              value={leadTurns}
+              onChange={(e) => setLeadTurns(Number(e.target.value))}
+              className="w-20"
+              disabled={!isEnabled}
+            />
+            <p className="text-xs text-textSubtle">
+              Number of turns to use the lead model at the start
+            </p>
+          </div>
+
+          <div className="space-y-2">
+            <label className="text-sm text-textSubtle flex items-center gap-1">
+              Failure Threshold
+              <Info size={14} className="text-textSubtle" />
+            </label>
+            <Input
+              type="number"
+              min={1}
+              max={5}
+              value={failureThreshold}
+              onChange={(e) => setFailureThreshold(Number(e.target.value))}
+              className="w-20"
+              disabled={!isEnabled}
+            />
+            <p className="text-xs text-textSubtle">
+              Consecutive failures before switching back to lead
+            </p>
+          </div>
+
+          <div className="space-y-2">
+            <label className="text-sm text-textSubtle flex items-center gap-1">
+              Fallback Turns
+              <Info size={14} className="text-textSubtle" />
+            </label>
+            <Input
+              type="number"
+              min={1}
+              max={5}
+              value={fallbackTurns}
+              onChange={(e) => setFallbackTurns(Number(e.target.value))}
+              className="w-20"
+              disabled={!isEnabled}
+            />
+            <p className="text-xs text-textSubtle">
+              Turns to use lead model during fallback
+            </p>
+          </div>
+        </div>
+      </div>
+
+      <div className="flex justify-end space-x-2 pt-4 border-t border-borderSubtle">
+        <Button variant="ghost" onClick={onClose}>
+          Cancel
+        </Button>
+        <Button onClick={handleSave} disabled={isEnabled && (!leadModel || !workerModel)}>
+          Save Settings
+        </Button>
+      </div>
+    </div>
+  );
+}
@@ -0,0 +1,10 @@
+import { useCurrentModelInfo } from '../components/ChatView';
+
+export function useCurrentModel() {
+  const modelInfo = useCurrentModelInfo();
+  
+  return { 
+    currentModel: modelInfo?.model || null,
+    isLoading: false 
+  };
+}
@@ -24,6 +24,7 @@ type MessageEvent =
  | { type: 'Message'; message: Message }
  | { type: 'Error'; error: string }
  | { type: 'Finish'; reason: string }
+  | { type: 'ModelChange'; model: string; mode: string }
  | NotificationEvent;

 export interface UseMessageStreamOptions {
@@ -140,6 +141,9 @@ export interface UseMessageStreamHelpers {
  updateMessageStreamBody?: (newBody: object) => void;

  notifications: NotificationEvent[];
+  
+  /** Current model info from the backend */
+  currentModelInfo: { model: string; mode: string } | null;
 }

 /**
@@ -168,6 +172,7 @@ export function useMessageStream({
  });

  const [notifications, setNotifications] = useState<NotificationEvent[]>([]);
+  const [currentModelInfo, setCurrentModelInfo] = useState<{ model: string; mode: string } | null>(null);

  // expose a way to update the body so we can update the session id when CLE occurs
  const updateMessageStreamBody = useCallback((newBody: object) => {
@@ -273,6 +278,16 @@ export function useMessageStream({
                    break;
                  }

+                  case 'ModelChange': {
+                    // Update the current model in the frontend
+                    const modelInfo = {
+                      model: parsedEvent.model,
+                      mode: parsedEvent.mode,
+                    };
+                    setCurrentModelInfo(modelInfo);
+                    break;
+                  }
+
                  case 'Error':
                    throw new Error(parsedEvent.error);

@@ -543,5 +558,6 @@ export function useMessageStream({
    addToolResult,
    updateMessageStreamBody,
    notifications,
+    currentModelInfo,
  };
 }