diff --git a/studio/frontend/src/components/assistant-ui/model-selector.tsx b/studio/frontend/src/components/assistant-ui/model-selector.tsx index fa6980824f..b97c34152f 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector.tsx @@ -21,7 +21,9 @@ import { Search01Icon, } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; -import { useMemo, useState } from "react"; +import { type KeyboardEvent, useMemo, useState } from "react"; +import { Input } from "../ui/input"; +import { HubModelPicker, LoraModelPicker } from "./model-selector/pickers"; import type { DeletedModelRef, ExternalModelOption, @@ -29,8 +31,6 @@ import type { ModelOption, ModelSelectorChangeMeta, } from "./model-selector/types"; -import { HubModelPicker, LoraModelPicker } from "./model-selector/pickers"; -import { Input } from "../ui/input"; const PROVIDER_LOGO_EXT: Record = { openai: "svg", @@ -241,11 +241,55 @@ function ModelSelectorContent({ return "hub"; }, [externalModels, loraModels, value]); + function focusActiveModelOption(root: HTMLElement): boolean { + const option = + root.querySelector( + '[role="tabpanel"]:not([hidden]) [data-model-picker-active-option="true"]', + ) ?? + root.querySelector( + '[data-model-picker-active-option="true"]', + ) ?? + root.querySelector( + '[role="tabpanel"]:not([hidden]) [data-model-picker-option]', + ) ?? + root.querySelector( + "[data-model-picker-option]", + ); + if (!option) { + return false; + } + option.focus(); + return true; + } + + function handlePickerEntryKeyDown(event: KeyboardEvent) { + if (event.key !== "ArrowDown") { + return; + } + + const target = event.target; + if (!(target instanceof HTMLElement)) { + return; + } + const isPickerSearchInput = target.matches( + "[data-model-picker-search-input]", + ); + const isTabTrigger = Boolean(target.closest('[role="tab"]')); + if (!isPickerSearchInput && !isTabTrigger) { + return; + } + + if (focusActiveModelOption(event.currentTarget)) { + event.preventDefault(); + } + } + return ( ( + "[data-model-picker-option]", + ); + if (!option) { + return false; + } + option.focus(); + return true; +} + +type ModelRowOptionProps = { + id: string; + tabIndex: number; + onFocus: () => void; + onKeyDown: (event: KeyboardEvent) => void; + "data-model-picker-option": true; + "data-model-picker-active-option"?: "true"; + "aria-current"?: "true"; +}; + +function useRovingModelList({ + label, + optionKeys, + selectedOptionKey, + onNavigatePastStart, + onNavigatePastEnd, +}: { + label: string; + optionKeys: string[]; + selectedOptionKey?: string; + onNavigatePastStart?: () => void; + onNavigatePastEnd?: () => void; +}) { + const rawListboxId = useId(); + const listboxId = `model-picker-${rawListboxId.replace(/:/g, "")}`; + const [rovingOptionKey, setRovingOptionKey] = useState(null); + + const preferredOptionKey = + selectedOptionKey && optionKeys.includes(selectedOptionKey) + ? selectedOptionKey + : (optionKeys[0] ?? null); + const activeOptionKey = + rovingOptionKey && optionKeys.includes(rovingOptionKey) + ? rovingOptionKey + : preferredOptionKey; + + const getOptionDomId = useCallback( + (optionKey: string) => { + const index = optionKeys.indexOf(optionKey); + return index === -1 ? undefined : `${listboxId}-option-${index}`; + }, + [listboxId, optionKeys], + ); + + const focusOption = useCallback( + (optionKey: string) => { + const id = getOptionDomId(optionKey); + if (!id) { + return; + } + document.getElementById(id)?.focus(); + }, + [getOptionDomId], + ); + + const moveFocus = useCallback( + ( + fromOptionKey: string, + direction: "next" | "previous" | "first" | "last", + ) => { + if (optionKeys.length === 0) { + return; + } + + const currentIndex = optionKeys.indexOf(fromOptionKey); + let nextIndex = currentIndex === -1 ? 0 : currentIndex; + if (direction === "next") { + if (currentIndex >= optionKeys.length - 1) { + onNavigatePastEnd?.(); + return; + } + nextIndex = Math.min(optionKeys.length - 1, nextIndex + 1); + } else if (direction === "previous") { + if (currentIndex <= 0) { + onNavigatePastStart?.(); + return; + } + nextIndex = Math.max(0, nextIndex - 1); + } else if (direction === "first") { + nextIndex = 0; + } else { + nextIndex = optionKeys.length - 1; + } + + const nextOptionKey = optionKeys[nextIndex]; + setRovingOptionKey(nextOptionKey); + focusOption(nextOptionKey); + }, + [focusOption, onNavigatePastEnd, onNavigatePastStart, optionKeys], + ); + + const getOptionProps = useCallback( + (optionKey: string, selected: boolean): ModelRowOptionProps => ({ + id: getOptionDomId(optionKey) ?? `${listboxId}-option-missing`, + tabIndex: 0, + onFocus: () => { + setRovingOptionKey(optionKey); + }, + onKeyDown: (event) => { + if (event.key === "ArrowDown") { + event.preventDefault(); + moveFocus(optionKey, "next"); + } else if (event.key === "ArrowUp") { + event.preventDefault(); + moveFocus(optionKey, "previous"); + } else if (event.key === "Home") { + event.preventDefault(); + moveFocus(optionKey, "first"); + } else if (event.key === "End") { + event.preventDefault(); + moveFocus(optionKey, "last"); + } + }, + "data-model-picker-option": true, + "data-model-picker-active-option": + optionKey === activeOptionKey ? "true" : undefined, + "aria-current": selected ? "true" : undefined, + }), + [activeOptionKey, getOptionDomId, listboxId, moveFocus], + ); + + return { + activeOptionKey, + focusOption, + getOptionProps, + moveFocus, + listboxProps: { + id: listboxId, + "data-model-picker-list": true, + "aria-label": label, + }, + }; } function ListLabel({ @@ -132,6 +288,8 @@ function ModelRow({ vramEst, gpuGb, tooltipText, + optionProps, + onArrowDownIntoChildren, }: { label: string; meta?: string | null; @@ -141,6 +299,8 @@ function ModelRow({ vramEst?: number; gpuGb?: number; tooltipText?: ReactNode; + optionProps?: ModelRowOptionProps; + onArrowDownIntoChildren?: () => boolean; }) { const exceeds = vramStatus === "exceeds"; const showVramTooltip = @@ -157,9 +317,17 @@ function ModelRow({ const content = (