From 407da24c5b4530b0f993e3cb1f9660b222792300 Mon Sep 17 00:00:00 2001 From: Ari Webb Date: Thu, 18 Dec 2025 09:49:20 -0800 Subject: [PATCH] Byok support (#277) --- src/agent/modify.ts | 15 +- src/cli/App.tsx | 22 ++- src/cli/components/ModelSelector.tsx | 229 +++++++++++++++++---------- 3 files changed, 165 insertions(+), 101 deletions(-) diff --git a/src/agent/modify.ts b/src/agent/modify.ts index 1653999..5035b4b 100644 --- a/src/agent/modify.ts +++ b/src/agent/modify.ts @@ -97,7 +97,7 @@ function buildModelSettings( } settings = googleVertexSettings; } else { - // For unknown providers, return generic settings with parallel_tool_calls + // For BYOK/unknown providers, return generic settings with parallel_tool_calls settings = { parallel_tool_calls: true }; } @@ -130,14 +130,13 @@ export async function updateAgentLLMConfig( const modelSettings = buildModelSettings(modelHandle, updateArgs); const contextWindow = updateArgs?.context_window as number | undefined; + const hasModelSettings = Object.keys(modelSettings).length > 0; - if (modelSettings || contextWindow) { - await client.agents.update(agentId, { - model: modelHandle, - ...(modelSettings && { model_settings: modelSettings }), - ...(contextWindow && { context_window_limit: contextWindow }), - }); - } + await client.agents.update(agentId, { + model: modelHandle, + ...(hasModelSettings && { model_settings: modelSettings }), + ...(contextWindow && { context_window_limit: contextWindow }), + }); const finalAgent = await client.agents.retrieve(agentId); return finalAgent.llm_config; diff --git a/src/cli/App.tsx b/src/cli/App.tsx index 628a978..9fa5fc7 100644 --- a/src/cli/App.tsx +++ b/src/cli/App.tsx @@ -433,6 +433,7 @@ export default function App({ | null >(null); const [llmConfig, setLlmConfig] = useState(null); + const [currentModelId, setCurrentModelId] = useState(null); const [agentName, setAgentName] = useState(null); const [agentDescription, setAgentDescription] = useState(null); const [agentLastRunAt, setAgentLastRunAt] = useState(null); @@ -3514,7 +3515,18 @@ ${recentCommits} try { // Find the selected model from models.json first (for loading message) const { models } = await import("../agent/model"); - const selectedModel = models.find((m) => m.id === modelId); + let selectedModel = models.find((m) => m.id === modelId); + + // If not found in static list, it might be a BYOK model where id === handle + if (!selectedModel && modelId.includes("/")) { + // Treat it as a BYOK model - the modelId is actually the handle + selectedModel = { + id: modelId, + handle: modelId, + label: modelId.split("/").pop() ?? modelId, + description: "Custom model", + } as unknown as (typeof models)[number]; + } if (!selectedModel) { // Create a failed command in the transcript @@ -3553,6 +3565,7 @@ ${recentCommits} selectedModel.updateArgs, ); setLlmConfig(updatedConfig); + setCurrentModelId(modelId); // After switching models, only switch toolset if it actually changes const { isOpenAIModel, isGeminiModel } = await import( @@ -4319,12 +4332,7 @@ Plan file path: ${planFilePath}`; {/* Model Selector - conditionally mounted as overlay */} {activeOverlay === "model" && ( diff --git a/src/cli/components/ModelSelector.tsx b/src/cli/components/ModelSelector.tsx index 45ea782..c8467d5 100644 --- a/src/cli/components/ModelSelector.tsx +++ b/src/cli/components/ModelSelector.tsx @@ -1,6 +1,6 @@ // Import useInput from vendored Ink for bracketed paste support import { Box, Text, useInput } from "ink"; -import { useEffect, useMemo, useRef, useState } from "react"; +import { useCallback, useEffect, useMemo, useRef, useState } from "react"; import { clearAvailableModelsCache, getAvailableModelHandles, @@ -9,6 +9,11 @@ import { import { models } from "../../agent/model"; import { colors } from "./colors"; +const PAGE_SIZE = 10; + +type ModelCategory = "supported" | "all"; +const MODEL_CATEGORIES: ModelCategory[] = ["supported", "all"]; + type UiModel = { id: string; handle: string; @@ -20,27 +25,28 @@ type UiModel = { }; interface ModelSelectorProps { - currentModel?: string; - currentEnableReasoner?: boolean; + currentModelId?: string; onSelect: (modelId: string) => void; onCancel: () => void; } export function ModelSelector({ - currentModel, - currentEnableReasoner, + currentModelId, onSelect, onCancel, }: ModelSelectorProps) { const typedModels = models as UiModel[]; - const [showAll, setShowAll] = useState(false); + const [category, setCategory] = useState("supported"); + const [currentPage, setCurrentPage] = useState(0); const [selectedIndex, setSelectedIndex] = useState(0); + // undefined: not loaded yet (show spinner) // Set: loaded and filtered // null: error fallback (show all models + warning) - const [availableModels, setAvailableModels] = useState< + const [availableHandles, setAvailableHandles] = useState< Set | null | undefined >(undefined); + const [allApiHandles, setAllApiHandles] = useState([]); const [isLoading, setIsLoading] = useState(true); const [error, setError] = useState(null); const [isCached, setIsCached] = useState(false); @@ -70,7 +76,8 @@ export function ModelSelector({ if (!mountedRef.current) return; - setAvailableModels(result.handles); + setAvailableHandles(result.handles); + setAllApiHandles(Array.from(result.handles)); setIsCached(!forceRefresh && cacheInfoBefore.isFresh); setIsLoading(false); setRefreshing(false); @@ -80,7 +87,8 @@ export function ModelSelector({ setIsLoading(false); setRefreshing(false); // Fallback: show all models if API fails - setAvailableModels(null); + setAvailableHandles(null); + setAllApiHandles([]); } }); @@ -88,44 +96,79 @@ export function ModelSelector({ loadModels.current(false); }, []); - // Filter models based on availability - const filteredModels = useMemo(() => { - // Not loaded yet: render nothing (avoid briefly showing unfiltered models) - if (availableModels === undefined) return []; - // Error fallback: show all models with warning - if (availableModels === null) return typedModels; - // Loaded: filter to only show models the user has access to - return typedModels.filter((model) => availableModels.has(model.handle)); - }, [typedModels, availableModels]); + // Handles from models.json (for filtering "all" category) + const staticModelHandles = useMemo( + () => new Set(typedModels.map((m) => m.handle)), + [typedModels], + ); - const featuredModels = useMemo( - () => filteredModels.filter((model) => model.isFeatured), - [filteredModels], + // Supported models: models.json entries that are available + const supportedModels = useMemo(() => { + if (availableHandles === undefined) return []; + if (availableHandles === null) return typedModels; // fallback + return typedModels.filter((m) => availableHandles.has(m.handle)); + }, [typedModels, availableHandles]); + + // All other models: API handles not in models.json + const otherModelHandles = useMemo(() => { + return allApiHandles.filter((handle) => !staticModelHandles.has(handle)); + }, [allApiHandles, staticModelHandles]); + + // Get the list for current category + const currentList: UiModel[] = useMemo(() => { + if (category === "supported") { + return supportedModels; + } + // For "all" category, convert handles to simple UiModel objects + return otherModelHandles.map((handle) => ({ + id: handle, + handle, + label: handle, + description: "", + })); + }, [category, supportedModels, otherModelHandles]); + + // Pagination + const totalPages = useMemo( + () => Math.max(1, Math.ceil(currentList.length / PAGE_SIZE)), + [currentList.length], ); const visibleModels = useMemo(() => { - if (showAll) return filteredModels; - if (featuredModels.length > 0) return featuredModels; - return filteredModels.slice(0, 5); - }, [featuredModels, showAll, filteredModels]); + const start = currentPage * PAGE_SIZE; + return currentList.slice(start, start + PAGE_SIZE); + }, [currentList, currentPage]); + + // Reset page and selection when category changes + const cycleCategory = useCallback(() => { + setCategory((current) => { + const idx = MODEL_CATEGORIES.indexOf(current); + return MODEL_CATEGORIES[ + (idx + 1) % MODEL_CATEGORIES.length + ] as ModelCategory; + }); + setCurrentPage(0); + setSelectedIndex(0); + }, []); // Set initial selection to current model on mount const initializedRef = useRef(false); useEffect(() => { - if (!initializedRef.current) { - const index = visibleModels.findIndex((m) => m.handle === currentModel); + if (!initializedRef.current && visibleModels.length > 0) { + const index = visibleModels.findIndex((m) => m.id === currentModelId); if (index >= 0) { setSelectedIndex(index); } initializedRef.current = true; } - }, [visibleModels, currentModel]); + }, [visibleModels, currentModelId]); - const hasMoreModels = - !showAll && filteredModels.length > visibleModels.length; - const totalItems = hasMoreModels - ? visibleModels.length + 1 - : visibleModels.length; + // Clamp selectedIndex when list changes + useEffect(() => { + if (selectedIndex >= visibleModels.length && visibleModels.length > 0) { + setSelectedIndex(visibleModels.length - 1); + } + }, [selectedIndex, visibleModels.length]); useInput( (input, key) => { @@ -141,6 +184,11 @@ export function ModelSelector({ return; } + if (key.tab) { + cycleCategory(); + return; + } + // Disable other inputs while loading if (isLoading || refreshing || visibleModels.length === 0) { return; @@ -149,16 +197,31 @@ export function ModelSelector({ if (key.upArrow) { setSelectedIndex((prev) => Math.max(0, prev - 1)); } else if (key.downArrow) { - setSelectedIndex((prev) => Math.min(totalItems - 1, prev + 1)); - } else if (key.return) { - if (hasMoreModels && selectedIndex === visibleModels.length) { - setShowAll(true); + setSelectedIndex((prev) => + Math.min(visibleModels.length - 1, prev + 1), + ); + } else if (input === "j" || input === "J") { + // Previous page + if (currentPage > 0) { + setCurrentPage((prev) => prev - 1); setSelectedIndex(0); - } else { - const selectedModel = visibleModels[selectedIndex]; - if (selectedModel) { - onSelect(selectedModel.id); - } + } + } else if (input === "k" || input === "K") { + // Next page + if (currentPage < totalPages - 1) { + setCurrentPage((prev) => prev + 1); + setSelectedIndex(0); + } + } else if (key.leftArrow && currentPage > 0) { + setCurrentPage((prev) => prev - 1); + setSelectedIndex(0); + } else if (key.rightArrow && currentPage < totalPages - 1) { + setCurrentPage((prev) => prev + 1); + setSelectedIndex(0); + } else if (key.return) { + const selectedModel = visibleModels[selectedIndex]; + if (selectedModel) { + onSelect(selectedModel.id); } } }, @@ -166,17 +229,42 @@ export function ModelSelector({ { isActive: true }, ); + const getCategoryLabel = (cat: ModelCategory) => { + if (cat === "supported") return `Supported (${supportedModels.length})`; + return `All Available Models (${otherModelHandles.length})`; + }; + return ( - Select Model (↑↓ to navigate, Enter to select, ESC to cancel) + Select Model (↑↓ navigate, ←→/jk page, Enter select, ESC cancel) + {!isLoading && !refreshing && ( + + Category: + {MODEL_CATEGORIES.map((cat, i) => ( + + {i > 0 && · } + + {getCategoryLabel(cat)} + + + ))} + (Tab to switch) + + )} {!isLoading && !refreshing && ( - {isCached - ? "Cached models (press 'r' to refresh)" - : "Press 'r' to refresh"} + Page {currentPage + 1}/{totalPages} + {isCached ? " · cached" : ""} · 'r' to refresh )} @@ -201,10 +289,12 @@ export function ModelSelector({ )} - {!isLoading && visibleModels.length === 0 && ( + {!isLoading && !refreshing && visibleModels.length === 0 && ( - - No models available. Please check your Letta configuration. + + {category === "supported" + ? "No supported models available." + : "No additional models available."} )} @@ -212,26 +302,7 @@ export function ModelSelector({ {visibleModels.map((model, index) => { const isSelected = index === selectedIndex; - - // Check if this model is current by comparing handle and relevant settings - let isCurrent = model.handle === currentModel; - - // For models with the same handle, also check specific configuration settings - if (isCurrent && model.handle?.startsWith("anthropic/")) { - // For Anthropic models, check enable_reasoner setting - const modelEnableReasoner = model.updateArgs?.enable_reasoner; - - // If the model explicitly sets enable_reasoner, check if it matches current settings - if (modelEnableReasoner !== undefined) { - // Model has explicit enable_reasoner setting, compare with current - isCurrent = - isCurrent && modelEnableReasoner === currentEnableReasoner; - } else { - // If model doesn't explicitly set enable_reasoner, it defaults to enabled (or undefined) - // It's current if currentEnableReasoner is not explicitly false - isCurrent = isCurrent && currentEnableReasoner !== false; - } - } + const isCurrent = model.id === currentModelId; return ( @@ -252,27 +323,13 @@ export function ModelSelector({ (current) )} - {model.description} + {model.description && ( + {model.description} + )} ); })} - {!showAll && filteredModels.length > visibleModels.length && ( - - - {selectedIndex === visibleModels.length ? "›" : " "} - - - Show all models ({filteredModels.length} available) - - - )} );