From b9d529d94e65831cfa9c6290aceffb020d7a9d07 Mon Sep 17 00:00:00 2001 From: Richard Tang Date: Wed, 18 Mar 2026 11:19:44 -0700 Subject: [PATCH] feat: support separate worker llm setup --- core/framework/config.py | 131 +++++++ core/framework/server/session_manager.py | 6 +- quickstart.ps1 | 3 + quickstart.sh | 2 + scripts/setup_worker_model.ps1 | 284 ++++++++++++++ scripts/setup_worker_model.sh | 459 +++++++++++++++++++++++ 6 files changed, 884 insertions(+), 1 deletion(-) create mode 100644 scripts/setup_worker_model.ps1 create mode 100755 scripts/setup_worker_model.sh diff --git a/core/framework/config.py b/core/framework/config.py index 382cb173..27125168 100644 --- a/core/framework/config.py +++ b/core/framework/config.py @@ -61,6 +61,137 @@ def get_preferred_model() -> str: return "anthropic/claude-sonnet-4-20250514" +def get_preferred_worker_model() -> str | None: + """Return the user's preferred worker LLM model, or None if not configured. + + Reads from the ``worker_llm`` section of ~/.hive/configuration.json. + Returns None when no worker-specific model is set, so callers can + fall back to the default (queen) model via ``get_preferred_model()``. + """ + worker_llm = get_hive_config().get("worker_llm", {}) + if worker_llm.get("provider") and worker_llm.get("model"): + provider = str(worker_llm["provider"]) + model = str(worker_llm["model"]).strip() + if provider.lower() == "openrouter" and model.lower().startswith("openrouter/"): + model = model[len("openrouter/") :] + if model: + return f"{provider}/{model}" + return None + + +def get_worker_api_key() -> str | None: + """Return the API key for the worker LLM, falling back to the default key.""" + worker_llm = get_hive_config().get("worker_llm", {}) + if not worker_llm: + return get_api_key() + + # Worker-specific subscription / env var + if worker_llm.get("use_claude_code_subscription"): + try: + from framework.runner.runner import get_claude_code_token + + token = get_claude_code_token() + if token: + return token + except ImportError: + pass + + if worker_llm.get("use_codex_subscription"): + try: + from framework.runner.runner import get_codex_token + + token = get_codex_token() + if token: + return token + except ImportError: + pass + + if worker_llm.get("use_kimi_code_subscription"): + try: + from framework.runner.runner import get_kimi_code_token + + token = get_kimi_code_token() + if token: + return token + except ImportError: + pass + + api_key_env_var = worker_llm.get("api_key_env_var") + if api_key_env_var: + return os.environ.get(api_key_env_var) + + # Fall back to default key + return get_api_key() + + +def get_worker_api_base() -> str | None: + """Return the api_base for the worker LLM, falling back to the default.""" + worker_llm = get_hive_config().get("worker_llm", {}) + if not worker_llm: + return get_api_base() + + if worker_llm.get("use_codex_subscription"): + return "https://chatgpt.com/backend-api/codex" + if worker_llm.get("use_kimi_code_subscription"): + return "https://api.kimi.com/coding" + if worker_llm.get("api_base"): + return worker_llm["api_base"] + if str(worker_llm.get("provider", "")).lower() == "openrouter": + return OPENROUTER_API_BASE + return None + + +def get_worker_llm_extra_kwargs() -> dict[str, Any]: + """Return extra kwargs for the worker LLM provider.""" + worker_llm = get_hive_config().get("worker_llm", {}) + if not worker_llm: + return get_llm_extra_kwargs() + + if worker_llm.get("use_claude_code_subscription"): + api_key = get_worker_api_key() + if api_key: + return { + "extra_headers": {"authorization": f"Bearer {api_key}"}, + } + if worker_llm.get("use_codex_subscription"): + api_key = get_worker_api_key() + if api_key: + headers: dict[str, str] = { + "Authorization": f"Bearer {api_key}", + "User-Agent": "CodexBar", + } + try: + from framework.runner.runner import get_codex_account_id + + account_id = get_codex_account_id() + if account_id: + headers["ChatGPT-Account-Id"] = account_id + except ImportError: + pass + return { + "extra_headers": headers, + "store": False, + "allowed_openai_params": ["store"], + } + return {} + + +def get_worker_max_tokens() -> int: + """Return max_tokens for the worker LLM, falling back to default.""" + worker_llm = get_hive_config().get("worker_llm", {}) + if worker_llm and "max_tokens" in worker_llm: + return worker_llm["max_tokens"] + return get_max_tokens() + + +def get_worker_max_context_tokens() -> int: + """Return max_context_tokens for the worker LLM, falling back to default.""" + worker_llm = get_hive_config().get("worker_llm", {}) + if worker_llm and "max_context_tokens" in worker_llm: + return worker_llm["max_context_tokens"] + return get_max_context_tokens() + + def get_max_tokens() -> int: """Return the configured max_tokens, falling back to DEFAULT_MAX_TOKENS.""" return get_hive_config().get("llm", {}).get("max_tokens", DEFAULT_MAX_TOKENS) diff --git a/core/framework/server/session_manager.py b/core/framework/server/session_manager.py index d220b20f..0c560684 100644 --- a/core/framework/server/session_manager.py +++ b/core/framework/server/session_manager.py @@ -287,7 +287,11 @@ class SessionManager: try: # Blocking I/O — load in executor loop = asyncio.get_running_loop() - resolved_model = model or self._model + + # Prioritize: explicit model arg > worker-specific model > session default + from framework.config import get_preferred_worker_model + + resolved_model = model or get_preferred_worker_model() or self._model runner = await loop.run_in_executor( None, lambda: AgentRunner.load( diff --git a/quickstart.ps1 b/quickstart.ps1 index 92ce3b4f..52596547 100644 --- a/quickstart.ps1 +++ b/quickstart.ps1 @@ -1880,6 +1880,9 @@ if ($SelectedProviderId) { Write-Host " -> " -NoNewline Write-Color -Text $SelectedModel -Color DarkGray } + Write-Color -Text " To use a different model for worker agents, run:" -Color DarkGray + Write-Host " " -NoNewline + Write-Color -Text ".\scripts\setup_worker_model.ps1" -Color Cyan Write-Host "" } diff --git a/quickstart.sh b/quickstart.sh index d9dc19a6..91f9c5a0 100755 --- a/quickstart.sh +++ b/quickstart.sh @@ -1767,6 +1767,8 @@ if [ -n "$SELECTED_PROVIDER_ID" ]; then else echo -e " ${CYAN}$SELECTED_PROVIDER_ID${NC} → ${DIM}$SELECTED_MODEL${NC}" fi + echo -e " ${DIM}To use a different model for worker agents, run:${NC}" + echo -e " ${CYAN}./scripts/setup_worker_model.sh${NC}" echo "" fi diff --git a/scripts/setup_worker_model.ps1 b/scripts/setup_worker_model.ps1 new file mode 100644 index 00000000..5f185473 --- /dev/null +++ b/scripts/setup_worker_model.ps1 @@ -0,0 +1,284 @@ +#Requires -Version 5.1 +<# +.SYNOPSIS + setup_worker_model.ps1 - Configure a separate LLM model for worker agents + +.DESCRIPTION + Worker agents can use a different (e.g. cheaper/faster) model than the + queen agent. This script writes a "worker_llm" section to + ~/.hive/configuration.json. If no worker model is configured, workers + fall back to the default (queen) model. + +.NOTES + Run from the project root: .\scripts\setup_worker_model.ps1 +#> + +$ErrorActionPreference = "Continue" +$ScriptDir = Split-Path -Parent $MyInvocation.MyCommand.Definition +$ProjectDir = Split-Path -Parent $ScriptDir +$UvHelperPath = Join-Path $ScriptDir "uv-discovery.ps1" +$HiveConfigDir = Join-Path $env:USERPROFILE ".hive" +$HiveConfigFile = Join-Path $HiveConfigDir "configuration.json" + +. $UvHelperPath + +# ============================================================ +# Colors / helpers +# ============================================================ + +function Write-Color { + param( + [string]$Text, + [ConsoleColor]$Color = [ConsoleColor]::White, + [switch]$NoNewline + ) + $prev = $Host.UI.RawUI.ForegroundColor + $Host.UI.RawUI.ForegroundColor = $Color + if ($NoNewline) { Write-Host $Text -NoNewline } + else { Write-Host $Text } + $Host.UI.RawUI.ForegroundColor = $prev +} + +function Write-Ok { + param([string]$Text) + Write-Color -Text "$([char]0x2B22) $Text" -Color Green +} + +function Write-Warn { + param([string]$Text) + Write-Color -Text "$([char]0x2B22) $Text" -Color Yellow +} + +# ============================================================ +# Provider / model definitions +# ============================================================ + +$ProviderEnvVars = @{ + "ANTHROPIC_API_KEY" = @{ Name = "Anthropic (Claude)"; Id = "anthropic" } + "OPENAI_API_KEY" = @{ Name = "OpenAI (GPT)"; Id = "openai" } + "GEMINI_API_KEY" = @{ Name = "Google Gemini"; Id = "gemini" } + "GROQ_API_KEY" = @{ Name = "Groq"; Id = "groq" } + "CEREBRAS_API_KEY" = @{ Name = "Cerebras"; Id = "cerebras" } + "OPENROUTER_API_KEY" = @{ Name = "OpenRouter"; Id = "openrouter" } + "MISTRAL_API_KEY" = @{ Name = "Mistral"; Id = "mistral" } + "TOGETHER_API_KEY" = @{ Name = "Together AI"; Id = "together" } + "DEEPSEEK_API_KEY" = @{ Name = "DeepSeek"; Id = "deepseek" } +} + +$DefaultModels = @{ + "anthropic" = "claude-haiku-4-5-20251001" + "openai" = "gpt-5-mini" + "gemini" = "gemini-3-flash-preview" + "groq" = "moonshotai/kimi-k2-instruct-0905" + "cerebras" = "zai-glm-4.7" + "mistral" = "mistral-large-latest" + "together_ai" = "meta-llama/Llama-3.3-70B-Instruct-Turbo" + "deepseek" = "deepseek-chat" +} + +$ModelChoices = @{ + "anthropic" = @( + @{ Id = "claude-haiku-4-5-20251001"; Label = "Haiku 4.5 - Fast + cheap (recommended for workers)"; MaxTokens = 8192; MaxContext = 180000 } + @{ Id = "claude-sonnet-4-20250514"; Label = "Sonnet 4 - Fast + capable"; MaxTokens = 8192; MaxContext = 180000 } + @{ Id = "claude-sonnet-4-5-20250929"; Label = "Sonnet 4.5 - Best balance"; MaxTokens = 16384; MaxContext = 180000 } + @{ Id = "claude-opus-4-6"; Label = "Opus 4.6 - Most capable"; MaxTokens = 32768; MaxContext = 180000 } + ) + "openai" = @( + @{ Id = "gpt-5-mini"; Label = "GPT-5 Mini - Fast + cheap (recommended for workers)"; MaxTokens = 16384; MaxContext = 120000 } + @{ Id = "gpt-5.2"; Label = "GPT-5.2 - Most capable"; MaxTokens = 16384; MaxContext = 120000 } + ) + "gemini" = @( + @{ Id = "gemini-3-flash-preview"; Label = "Gemini 3 Flash - Fast (recommended for workers)"; MaxTokens = 8192; MaxContext = 900000 } + @{ Id = "gemini-3.1-pro-preview"; Label = "Gemini 3.1 Pro - Best quality"; MaxTokens = 8192; MaxContext = 900000 } + ) + "groq" = @( + @{ Id = "moonshotai/kimi-k2-instruct-0905"; Label = "Kimi K2 - Best quality (recommended)"; MaxTokens = 8192; MaxContext = 120000 } + @{ Id = "openai/gpt-oss-120b"; Label = "GPT-OSS 120B - Fast reasoning"; MaxTokens = 8192; MaxContext = 120000 } + ) + "cerebras" = @( + @{ Id = "zai-glm-4.7"; Label = "ZAI-GLM 4.7 - Best quality (recommended)"; MaxTokens = 8192; MaxContext = 120000 } + @{ Id = "qwen3-235b-a22b-instruct-2507"; Label = "Qwen3 235B - Frontier reasoning"; MaxTokens = 8192; MaxContext = 120000 } + ) +} + +# ============================================================ +# Main +# ============================================================ + +$uvInfo = Find-Uv +if (-not $uvInfo) { + Write-Color -Text "uv not found. Run quickstart.ps1 first." -Color Red + exit 1 +} +$UvCmd = $uvInfo.Path + +Write-Host "" +Write-Color -Text "$([char]0x2B22) Worker Model Setup" -Color Yellow +Write-Host "" +Write-Color -Text "Configure a separate LLM model for worker agents." -Color DarkGray +Write-Color -Text "Worker agents will use this model instead of the default queen model." -Color DarkGray +Write-Host "" + +# Show current configuration +if (Test-Path $HiveConfigFile) { + try { + Push-Location $ProjectDir + $currentConfig = & $UvCmd run python -c " +from framework.config import get_preferred_model, get_preferred_worker_model +print(f'Queen: {get_preferred_model()}') +wm = get_preferred_worker_model() +print(f'Worker: {wm if wm else chr(34) + \"(same as queen)\" + chr(34)}') +" 2>$null + Pop-Location + if ($currentConfig) { + Write-Color -Text "Current configuration:" -Color White + foreach ($line in $currentConfig) { + Write-Color -Text " $line" -Color DarkGray + } + Write-Host "" + } + } catch { + Pop-Location + } +} + +# Detect available providers +$AvailableProviders = @() +foreach ($envVar in $ProviderEnvVars.Keys) { + $val = [System.Environment]::GetEnvironmentVariable($envVar, "User") + if (-not $val) { $val = [System.Environment]::GetEnvironmentVariable($envVar) } + if ($val) { + $AvailableProviders += @{ + EnvVar = $envVar + Name = $ProviderEnvVars[$envVar].Name + Id = $ProviderEnvVars[$envVar].Id + } + } +} + +if ($AvailableProviders.Count -eq 0) { + Write-Color -Text "No API keys found." -Color Red + Write-Host "Run .\quickstart.ps1 first to set up your LLM provider." + exit 1 +} + +# Pick provider +$SelectedProvider = $null +if ($AvailableProviders.Count -eq 1) { + $SelectedProvider = $AvailableProviders[0] + Write-Ok "Provider: $($SelectedProvider.Name)" +} else { + Write-Color -Text "Select provider for worker agents:" -Color White + Write-Host "" + for ($i = 0; $i -lt $AvailableProviders.Count; $i++) { + Write-Host " " -NoNewline + Write-Color -Text "$($i+1))" -Color Cyan -NoNewline + Write-Host " $($AvailableProviders[$i].Name)" + } + Write-Host "" + while ($true) { + $choice = Read-Host "Enter choice (1-$($AvailableProviders.Count))" + $idx = [int]$choice - 1 + if ($idx -ge 0 -and $idx -lt $AvailableProviders.Count) { + $SelectedProvider = $AvailableProviders[$idx] + Write-Host "" + Write-Ok "Provider: $($SelectedProvider.Name)" + break + } + Write-Color -Text "Invalid choice." -Color Red + } +} + +$SelectedProviderId = $SelectedProvider.Id +$SelectedEnvVar = $SelectedProvider.EnvVar +$SelectedModel = $DefaultModels[$SelectedProviderId] +$SelectedMaxTokens = 8192 +$SelectedMaxContextTokens = 120000 +$SelectedApiBase = "" + +if ($SelectedProviderId -eq "openrouter") { + $SelectedApiBase = "https://openrouter.ai/api/v1" +} + +# Select model +$choices = $ModelChoices[$SelectedProviderId] +if ($choices -and $choices.Count -gt 0) { + Write-Host "" + Write-Color -Text "Select worker model:" -Color White + for ($i = 0; $i -lt $choices.Count; $i++) { + Write-Host " " -NoNewline + Write-Color -Text "$($i+1))" -Color Cyan -NoNewline + Write-Host " $($choices[$i].Label)" + } + Write-Host "" + while ($true) { + $choice = Read-Host "Enter choice (1-$($choices.Count)) [1]" + if (-not $choice) { $choice = "1" } + $idx = [int]$choice - 1 + if ($idx -ge 0 -and $idx -lt $choices.Count) { + $SelectedModel = $choices[$idx].Id + $SelectedMaxTokens = $choices[$idx].MaxTokens + $SelectedMaxContextTokens = $choices[$idx].MaxContext + Write-Host "" + Write-Ok "Worker model: $SelectedModel" + break + } + Write-Color -Text "Invalid choice." -Color Red + } +} else { + Write-Host "" + Write-Ok "Worker model: $SelectedModel" +} + +# Confirm and save +Write-Host "" +$confirm = Read-Host "Save this worker model configuration? [Y/n]" +if ($confirm -and $confirm -notmatch "^[Yy]") { + Write-Host "" + Write-Host "Cancelled. Worker agents will continue using the default model." + exit 0 +} + +Write-Host "" +Write-Host " Saving worker model configuration... " -NoNewline + +# Read existing config, add worker_llm section +if (-not (Test-Path $HiveConfigDir)) { + New-Item -ItemType Directory -Path $HiveConfigDir -Force | Out-Null +} + +try { + if (Test-Path $HiveConfigFile) { + $config = Get-Content -Path $HiveConfigFile -Raw | ConvertFrom-Json + } else { + $config = @{} + } +} catch { + $config = @{} +} + +$workerLlm = @{ + provider = $SelectedProviderId + model = $SelectedModel + max_tokens = $SelectedMaxTokens + max_context_tokens = $SelectedMaxContextTokens +} + +if ($SelectedEnvVar) { + $workerLlm["api_key_env_var"] = $SelectedEnvVar +} +if ($SelectedApiBase) { + $workerLlm["api_base"] = $SelectedApiBase +} + +$config | Add-Member -NotePropertyName "worker_llm" -NotePropertyValue $workerLlm -Force +$config | ConvertTo-Json -Depth 4 | Set-Content -Path $HiveConfigFile -Encoding UTF8 +Write-Ok "done" +Write-Color -Text " ~/.hive/configuration.json (worker_llm section)" -Color DarkGray + +Write-Host "" +Write-Ok "Worker model configured successfully." +Write-Color -Text " Worker agents will now use: $SelectedProviderId/$SelectedModel" -Color DarkGray +Write-Color -Text " Run this script again to change, or remove the worker_llm section" -Color DarkGray +Write-Color -Text " from ~/.hive/configuration.json to revert to the default." -Color DarkGray +Write-Host "" diff --git a/scripts/setup_worker_model.sh b/scripts/setup_worker_model.sh new file mode 100755 index 00000000..00b23ed7 --- /dev/null +++ b/scripts/setup_worker_model.sh @@ -0,0 +1,459 @@ +#!/bin/bash +# +# setup_worker_model.sh - Configure a separate LLM model for worker agents +# +# Worker agents can use a different (e.g. cheaper/faster) model than the +# queen agent. This script writes a "worker_llm" section to +# ~/.hive/configuration.json. If no worker model is configured, workers +# fall back to the default (queen) model. +# + +set -e + +# Detect Bash version for compatibility +BASH_MAJOR_VERSION="${BASH_VERSINFO[0]}" +USE_ASSOC_ARRAYS=false +if [ "$BASH_MAJOR_VERSION" -ge 4 ]; then + USE_ASSOC_ARRAYS=true +fi + +# Colors for output +RED='\033[0;31m' +GREEN='\033[0;32m' +YELLOW='\033[1;33m' +BLUE='\033[0;34m' +CYAN='\033[0;36m' +BOLD='\033[1m' +DIM='\033[2m' +NC='\033[0m' + +# Get the directory where this script is located, then the project root +SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )" +PROJECT_DIR="$( cd "$SCRIPT_DIR/.." && pwd )" +HIVE_CONFIG_FILE="$HOME/.hive/configuration.json" + +# Helper function for prompts +prompt_yes_no() { + local prompt="$1" + local default="${2:-y}" + local response + + if [ "$default" = "y" ]; then + prompt="$prompt [Y/n] " + else + prompt="$prompt [y/N] " + fi + read -r -p "$prompt" response + response="${response:-$default}" + [[ "$response" =~ ^[Yy] ]] +} + +# ── Provider / model definitions (same as quickstart) ──────────────── + +if [ "$USE_ASSOC_ARRAYS" = true ]; then + declare -A PROVIDER_NAMES=( + ["ANTHROPIC_API_KEY"]="Anthropic (Claude)" + ["OPENAI_API_KEY"]="OpenAI (GPT)" + ["MINIMAX_API_KEY"]="MiniMax" + ["GEMINI_API_KEY"]="Google Gemini" + ["GOOGLE_API_KEY"]="Google AI" + ["GROQ_API_KEY"]="Groq" + ["CEREBRAS_API_KEY"]="Cerebras" + ["OPENROUTER_API_KEY"]="OpenRouter" + ["MISTRAL_API_KEY"]="Mistral" + ["TOGETHER_API_KEY"]="Together AI" + ["DEEPSEEK_API_KEY"]="DeepSeek" + ) + + declare -A PROVIDER_IDS=( + ["ANTHROPIC_API_KEY"]="anthropic" + ["OPENAI_API_KEY"]="openai" + ["MINIMAX_API_KEY"]="minimax" + ["GEMINI_API_KEY"]="gemini" + ["GOOGLE_API_KEY"]="google" + ["GROQ_API_KEY"]="groq" + ["CEREBRAS_API_KEY"]="cerebras" + ["OPENROUTER_API_KEY"]="openrouter" + ["MISTRAL_API_KEY"]="mistral" + ["TOGETHER_API_KEY"]="together" + ["DEEPSEEK_API_KEY"]="deepseek" + ) + + declare -A DEFAULT_MODELS=( + ["anthropic"]="claude-haiku-4-5-20251001" + ["openai"]="gpt-5-mini" + ["minimax"]="MiniMax-M2.5" + ["gemini"]="gemini-3-flash-preview" + ["groq"]="moonshotai/kimi-k2-instruct-0905" + ["cerebras"]="zai-glm-4.7" + ["mistral"]="mistral-large-latest" + ["together_ai"]="meta-llama/Llama-3.3-70B-Instruct-Turbo" + ["deepseek"]="deepseek-chat" + ) + + declare -A MODEL_CHOICES_ID=( + ["anthropic:0"]="claude-haiku-4-5-20251001" + ["anthropic:1"]="claude-sonnet-4-20250514" + ["anthropic:2"]="claude-sonnet-4-5-20250929" + ["anthropic:3"]="claude-opus-4-6" + ["openai:0"]="gpt-5-mini" + ["openai:1"]="gpt-5.2" + ["gemini:0"]="gemini-3-flash-preview" + ["gemini:1"]="gemini-3.1-pro-preview" + ["groq:0"]="moonshotai/kimi-k2-instruct-0905" + ["groq:1"]="openai/gpt-oss-120b" + ["cerebras:0"]="zai-glm-4.7" + ["cerebras:1"]="qwen3-235b-a22b-instruct-2507" + ) + + declare -A MODEL_CHOICES_LABEL=( + ["anthropic:0"]="Haiku 4.5 - Fast + cheap (recommended for workers)" + ["anthropic:1"]="Sonnet 4 - Fast + capable" + ["anthropic:2"]="Sonnet 4.5 - Best balance" + ["anthropic:3"]="Opus 4.6 - Most capable" + ["openai:0"]="GPT-5 Mini - Fast + cheap (recommended for workers)" + ["openai:1"]="GPT-5.2 - Most capable" + ["gemini:0"]="Gemini 3 Flash - Fast (recommended for workers)" + ["gemini:1"]="Gemini 3.1 Pro - Best quality" + ["groq:0"]="Kimi K2 - Best quality (recommended)" + ["groq:1"]="GPT-OSS 120B - Fast reasoning" + ["cerebras:0"]="ZAI-GLM 4.7 - Best quality (recommended)" + ["cerebras:1"]="Qwen3 235B - Frontier reasoning" + ) + + declare -A MODEL_CHOICES_MAXTOKENS=( + ["anthropic:0"]=8192 + ["anthropic:1"]=8192 + ["anthropic:2"]=16384 + ["anthropic:3"]=32768 + ["openai:0"]=16384 + ["openai:1"]=16384 + ["gemini:0"]=8192 + ["gemini:1"]=8192 + ["groq:0"]=8192 + ["groq:1"]=8192 + ["cerebras:0"]=8192 + ["cerebras:1"]=8192 + ) + + declare -A MODEL_CHOICES_MAXCONTEXTTOKENS=( + ["anthropic:0"]=180000 + ["anthropic:1"]=180000 + ["anthropic:2"]=180000 + ["anthropic:3"]=180000 + ["openai:0"]=120000 + ["openai:1"]=120000 + ["gemini:0"]=900000 + ["gemini:1"]=900000 + ["groq:0"]=120000 + ["groq:1"]=120000 + ["cerebras:0"]=120000 + ["cerebras:1"]=120000 + ) + + declare -A MODEL_CHOICES_COUNT=( + ["anthropic"]=4 + ["openai"]=2 + ["gemini"]=2 + ["groq"]=2 + ["cerebras"]=2 + ) + + get_provider_name() { echo "${PROVIDER_NAMES[$1]}"; } + get_provider_id() { echo "${PROVIDER_IDS[$1]}"; } + get_default_model() { echo "${DEFAULT_MODELS[$1]}"; } + get_model_choice_count() { echo "${MODEL_CHOICES_COUNT[$1]:-0}"; } + get_model_choice_id() { echo "${MODEL_CHOICES_ID[$1:$2]}"; } + get_model_choice_label() { echo "${MODEL_CHOICES_LABEL[$1:$2]}"; } + get_model_choice_maxtokens() { echo "${MODEL_CHOICES_MAXTOKENS[$1:$2]}"; } + get_model_choice_maxcontexttokens() { echo "${MODEL_CHOICES_MAXCONTEXTTOKENS[$1:$2]}"; } +else + # Bash 3.2 fallback + PROVIDER_ENV_VARS=(ANTHROPIC_API_KEY OPENAI_API_KEY MINIMAX_API_KEY GEMINI_API_KEY GOOGLE_API_KEY GROQ_API_KEY CEREBRAS_API_KEY OPENROUTER_API_KEY MISTRAL_API_KEY TOGETHER_API_KEY DEEPSEEK_API_KEY) + PROVIDER_DISPLAY_NAMES=("Anthropic (Claude)" "OpenAI (GPT)" "MiniMax" "Google Gemini" "Google AI" "Groq" "Cerebras" "OpenRouter" "Mistral" "Together AI" "DeepSeek") + PROVIDER_ID_LIST=(anthropic openai minimax gemini google groq cerebras openrouter mistral together deepseek) + + MODEL_PROVIDER_IDS=(anthropic openai minimax gemini groq cerebras mistral together_ai deepseek) + MODEL_DEFAULTS=("claude-haiku-4-5-20251001" "gpt-5-mini" "MiniMax-M2.5" "gemini-3-flash-preview" "moonshotai/kimi-k2-instruct-0905" "zai-glm-4.7" "mistral-large-latest" "meta-llama/Llama-3.3-70B-Instruct-Turbo" "deepseek-chat") + + get_provider_name() { + local env_var="$1"; local i=0 + while [ $i -lt ${#PROVIDER_ENV_VARS[@]} ]; do + if [ "${PROVIDER_ENV_VARS[$i]}" = "$env_var" ]; then echo "${PROVIDER_DISPLAY_NAMES[$i]}"; return; fi + i=$((i + 1)) + done + } + get_provider_id() { + local env_var="$1"; local i=0 + while [ $i -lt ${#PROVIDER_ENV_VARS[@]} ]; do + if [ "${PROVIDER_ENV_VARS[$i]}" = "$env_var" ]; then echo "${PROVIDER_ID_LIST[$i]}"; return; fi + i=$((i + 1)) + done + } + get_default_model() { + local provider_id="$1"; local i=0 + while [ $i -lt ${#MODEL_PROVIDER_IDS[@]} ]; do + if [ "${MODEL_PROVIDER_IDS[$i]}" = "$provider_id" ]; then echo "${MODEL_DEFAULTS[$i]}"; return; fi + i=$((i + 1)) + done + } + + MC_PROVIDERS=(anthropic anthropic anthropic anthropic openai openai gemini gemini groq groq cerebras cerebras) + MC_IDS=("claude-haiku-4-5-20251001" "claude-sonnet-4-20250514" "claude-sonnet-4-5-20250929" "claude-opus-4-6" "gpt-5-mini" "gpt-5.2" "gemini-3-flash-preview" "gemini-3.1-pro-preview" "moonshotai/kimi-k2-instruct-0905" "openai/gpt-oss-120b" "zai-glm-4.7" "qwen3-235b-a22b-instruct-2507") + MC_LABELS=("Haiku 4.5 - Fast + cheap (recommended for workers)" "Sonnet 4 - Fast + capable" "Sonnet 4.5 - Best balance" "Opus 4.6 - Most capable" "GPT-5 Mini - Fast + cheap (recommended for workers)" "GPT-5.2 - Most capable" "Gemini 3 Flash - Fast (recommended for workers)" "Gemini 3.1 Pro - Best quality" "Kimi K2 - Best quality (recommended)" "GPT-OSS 120B - Fast reasoning" "ZAI-GLM 4.7 - Best quality (recommended)" "Qwen3 235B - Frontier reasoning") + MC_MAXTOKENS=(8192 8192 16384 32768 16384 16384 8192 8192 8192 8192 8192 8192) + MC_MAXCONTEXTTOKENS=(180000 180000 180000 180000 120000 120000 900000 900000 120000 120000 120000 120000) + + get_model_choice_count() { + local p="$1"; local cnt=0; local i=0 + while [ $i -lt ${#MC_PROVIDERS[@]} ]; do + if [ "${MC_PROVIDERS[$i]}" = "$p" ]; then cnt=$((cnt + 1)); fi + i=$((i + 1)) + done + echo "$cnt" + } + _mc_nth() { + local p="$1"; local n="$2"; local cnt=0; local i=0 + while [ $i -lt ${#MC_PROVIDERS[@]} ]; do + if [ "${MC_PROVIDERS[$i]}" = "$p" ]; then + if [ "$cnt" -eq "$n" ]; then echo "$i"; return; fi + cnt=$((cnt + 1)) + fi + i=$((i + 1)) + done + } + get_model_choice_id() { local idx=$(_mc_nth "$1" "$2"); echo "${MC_IDS[$idx]}"; } + get_model_choice_label() { local idx=$(_mc_nth "$1" "$2"); echo "${MC_LABELS[$idx]}"; } + get_model_choice_maxtokens() { local idx=$(_mc_nth "$1" "$2"); echo "${MC_MAXTOKENS[$idx]}"; } + get_model_choice_maxcontexttokens() { local idx=$(_mc_nth "$1" "$2"); echo "${MC_MAXCONTEXTTOKENS[$idx]}"; } +fi + +# ── Model selection prompt ─────────────────────────────────────────── + +select_model() { + local provider_id="$1" + local count + count=$(get_model_choice_count "$provider_id") + + if [ "$count" -eq 0 ]; then + SELECTED_MODEL="$(get_default_model "$provider_id")" + SELECTED_MAX_TOKENS=8192 + SELECTED_MAX_CONTEXT_TOKENS=120000 + echo "" + echo -e "${GREEN}⬢${NC} Worker model: ${DIM}$SELECTED_MODEL${NC}" + return + fi + + echo "" + echo -e "${BOLD}Select worker model:${NC}" + local i=0 + while [ "$i" -lt "$count" ]; do + echo -e " ${CYAN}$((i+1)))${NC} $(get_model_choice_label "$provider_id" "$i")" + i=$((i + 1)) + done + echo "" + + while true; do + local default_idx=1 + read -r -p "Enter choice (1-$count) [$default_idx]: " choice || true + choice="${choice:-$default_idx}" + if [[ "$choice" =~ ^[0-9]+$ ]] && [ "$choice" -ge 1 ] && [ "$choice" -le "$count" ]; then + local idx=$((choice - 1)) + SELECTED_MODEL="$(get_model_choice_id "$provider_id" "$idx")" + SELECTED_MAX_TOKENS="$(get_model_choice_maxtokens "$provider_id" "$idx")" + SELECTED_MAX_CONTEXT_TOKENS="$(get_model_choice_maxcontexttokens "$provider_id" "$idx")" + echo "" + echo -e "${GREEN}⬢${NC} Worker model: ${DIM}$SELECTED_MODEL${NC}" + return + fi + echo -e "${RED}Invalid choice. Please enter 1-$count${NC}" + done +} + +# ── Save worker_llm section to configuration.json ──────────────────── + +save_worker_configuration() { + local provider_id="$1" + local env_var="$2" + local model="$3" + local max_tokens="$4" + local max_context_tokens="$5" + local api_base="${6:-}" + + if [ -z "$model" ]; then + model="$(get_default_model "$provider_id")" + fi + if [ -z "$max_tokens" ]; then max_tokens=8192; fi + if [ -z "$max_context_tokens" ]; then max_context_tokens=120000; fi + + cd "$PROJECT_DIR" + uv run python - \ + "$provider_id" \ + "$env_var" \ + "$model" \ + "$max_tokens" \ + "$max_context_tokens" \ + "$api_base" 2>/dev/null <<'PY' +import json +import sys +from pathlib import Path + +( + provider_id, + env_var, + model, + max_tokens, + max_context_tokens, + api_base, +) = sys.argv[1:7] + +cfg_path = Path.home() / ".hive" / "configuration.json" +cfg_path.parent.mkdir(parents=True, exist_ok=True) + +try: + with open(cfg_path, encoding="utf-8-sig") as f: + config = json.load(f) +except (OSError, json.JSONDecodeError): + config = {} + +config["worker_llm"] = { + "provider": provider_id, + "model": model, + "max_tokens": int(max_tokens), + "max_context_tokens": int(max_context_tokens), +} +if env_var: + config["worker_llm"]["api_key_env_var"] = env_var +if api_base: + config["worker_llm"]["api_base"] = api_base + +tmp_path = cfg_path.with_name(cfg_path.name + ".tmp") +with open(tmp_path, "w", encoding="utf-8") as f: + json.dump(config, f, indent=2) +tmp_path.replace(cfg_path) +print(json.dumps(config.get("worker_llm", {}), indent=2)) +PY +} + +# ── Main ───────────────────────────────────────────────────────────── + +echo "" +echo -e "${YELLOW}⬢${NC}${DIM}⬡${NC}${YELLOW}⬢${NC}${DIM}⬡${NC}${YELLOW}⬢${NC} ${BOLD}Worker Model Setup${NC} ${YELLOW}⬢${NC}${DIM}⬡${NC}${YELLOW}⬢${NC}${DIM}⬡${NC}${YELLOW}⬢${NC}" +echo "" +echo -e "${DIM}Configure a separate LLM model for worker agents.${NC}" +echo -e "${DIM}Worker agents will use this model instead of the default queen model.${NC}" +echo "" + +# Show current configuration +if [ -f "$HIVE_CONFIG_FILE" ]; then + CURRENT_QUEEN=$(cd "$PROJECT_DIR" && uv run python -c " +from framework.config import get_preferred_model, get_preferred_worker_model +print(f'Queen: {get_preferred_model()}') +wm = get_preferred_worker_model() +print(f'Worker: {wm if wm else \"(same as queen)\"}') +" 2>/dev/null) || true + if [ -n "$CURRENT_QUEEN" ]; then + echo -e "${BOLD}Current configuration:${NC}" + echo -e " ${DIM}$CURRENT_QUEEN${NC}" | head -1 + echo -e " ${DIM}$(echo "$CURRENT_QUEEN" | tail -1)${NC}" + echo "" + fi +fi + +# Source shell rc to pick up env vars +SHELL_RC_FILE="$HOME/.bashrc" +if [ -n "$ZSH_VERSION" ] || [ "$SHELL" = "/bin/zsh" ]; then + SHELL_RC_FILE="$HOME/.zshrc" +fi +set +e +if [ -f "$SHELL_RC_FILE" ]; then + eval "$(grep -E '^export [A-Z_]+=' "$SHELL_RC_FILE" 2>/dev/null)" +fi +set -e + +# Detect available providers +AVAIL_PROVIDERS=() +AVAIL_ENV_VARS=() + +ENV_VARS_TO_CHECK=(ANTHROPIC_API_KEY OPENAI_API_KEY GEMINI_API_KEY GOOGLE_API_KEY GROQ_API_KEY CEREBRAS_API_KEY OPENROUTER_API_KEY MISTRAL_API_KEY TOGETHER_API_KEY DEEPSEEK_API_KEY) + +for ev in "${ENV_VARS_TO_CHECK[@]}"; do + if [ -n "${!ev:-}" ]; then + AVAIL_PROVIDERS+=("$(get_provider_name "$ev")") + AVAIL_ENV_VARS+=("$ev") + fi +done + +if [ ${#AVAIL_PROVIDERS[@]} -eq 0 ]; then + echo -e "${RED}No API keys found.${NC}" + echo -e "Run ${CYAN}./quickstart.sh${NC} first to set up your LLM provider." + exit 1 +fi + +# Pick provider +SELECTED_PROVIDER_ID="" +SELECTED_ENV_VAR="" +SELECTED_MODEL="" +SELECTED_MAX_TOKENS=8192 +SELECTED_MAX_CONTEXT_TOKENS=120000 +SELECTED_API_BASE="" + +if [ ${#AVAIL_PROVIDERS[@]} -eq 1 ]; then + SELECTED_ENV_VAR="${AVAIL_ENV_VARS[0]}" + SELECTED_PROVIDER_ID="$(get_provider_id "$SELECTED_ENV_VAR")" + echo -e "${GREEN}⬢${NC} Provider: $(get_provider_name "$SELECTED_ENV_VAR")" +else + echo -e "${BOLD}Select provider for worker agents:${NC}" + echo "" + local_i=1 + for prov in "${AVAIL_PROVIDERS[@]}"; do + echo -e " ${CYAN}${local_i})${NC} $prov" + local_i=$((local_i + 1)) + done + echo "" + while true; do + read -r -p "Enter choice (1-${#AVAIL_PROVIDERS[@]}): " choice || true + if [[ "$choice" =~ ^[0-9]+$ ]] && [ "$choice" -ge 1 ] && [ "$choice" -le "${#AVAIL_PROVIDERS[@]}" ]; then + idx=$((choice - 1)) + SELECTED_ENV_VAR="${AVAIL_ENV_VARS[$idx]}" + SELECTED_PROVIDER_ID="$(get_provider_id "$SELECTED_ENV_VAR")" + echo "" + echo -e "${GREEN}⬢${NC} Provider: ${AVAIL_PROVIDERS[$idx]}" + break + fi + echo -e "${RED}Invalid choice.${NC}" + done +fi + +# OpenRouter: custom api_base +if [ "$SELECTED_PROVIDER_ID" = "openrouter" ]; then + SELECTED_API_BASE="https://openrouter.ai/api/v1" +fi + +# Select model +select_model "$SELECTED_PROVIDER_ID" + +# Option to clear worker model +echo "" +if prompt_yes_no "Save this worker model configuration?"; then + echo "" + echo -n " Saving worker model configuration... " + if save_worker_configuration "$SELECTED_PROVIDER_ID" "$SELECTED_ENV_VAR" "$SELECTED_MODEL" "$SELECTED_MAX_TOKENS" "$SELECTED_MAX_CONTEXT_TOKENS" "$SELECTED_API_BASE" > /dev/null; then + echo -e "${GREEN}done${NC}" + echo -e " ${DIM}~/.hive/configuration.json (worker_llm section)${NC}" + else + echo -e "${RED}failed${NC}" + exit 1 + fi +else + echo "" + echo "Cancelled. Worker agents will continue using the default model." + exit 0 +fi + +echo "" +echo -e "${GREEN}⬢${NC} Worker model configured successfully." +echo -e " ${DIM}Worker agents will now use: ${SELECTED_PROVIDER_ID}/${SELECTED_MODEL}${NC}" +echo -e " ${DIM}Run this script again to change, or remove the worker_llm section${NC}" +echo -e " ${DIM}from ~/.hive/configuration.json to revert to the default.${NC}" +echo ""