From 506d9cf68be6c14468a07cd7f9eb598f7992dc5d Mon Sep 17 00:00:00 2001 From: Thomas Kosiewski Date: Wed, 1 Oct 2025 15:44:03 +0200 Subject: [PATCH] feat: enable per-workspace model selection Implement workspace model persistence and command handling. Load metadata on mount, watch for updates, and save model changes. Route model updates through new IPC channel and AI service logic. Provide provider adapters, model catalog helpers, and UI feedback. Extend slash command parsing, suggestions, and tests for /model flow. --- src/components/AIView.tsx | 55 +++++++++++++++- src/components/ChatInput.tsx | 86 +++++++++++++++++++++++++ src/constants/ipc-constants.ts | 1 + src/constants/models.ts | 20 ++++++ src/debug/costs.ts | 3 +- src/main.ts | 23 +++++++ src/preload.ts | 2 + src/services/aiService.ts | 93 +++++++++++++++------------ src/services/providerAdapters.ts | 70 +++++++++++++++++++++ src/types/ipc.ts | 1 + src/types/workspace.ts | 4 ++ src/utils/commandParser.test.ts | 43 +++++++++++++ src/utils/commandParser.ts | 65 +++++++++++++++++++ src/utils/modelCatalog.ts | 105 +++++++++++++++++++++++++++++++ src/utils/slashCommands.test.ts | 30 +++++++++ src/utils/slashCommands.ts | 85 +++++++++++++++++++++++++ 16 files changed, 645 insertions(+), 41 deletions(-) create mode 100644 src/constants/models.ts create mode 100644 src/services/providerAdapters.ts create mode 100644 src/utils/modelCatalog.ts diff --git a/src/components/AIView.tsx b/src/components/AIView.tsx index e6837bbaab..7a05b66ec4 100644 --- a/src/components/AIView.tsx +++ b/src/components/AIView.tsx @@ -22,6 +22,7 @@ import { isReasoningDelta, isReasoningEnd, } from "../types/ipc"; +import { DEFAULT_MODEL } from "../constants/models"; // StreamingMessageAggregator is now imported from utils @@ -127,7 +128,7 @@ const AIViewInner: React.FC = ({ workspaceId, projectName, branch, message: string; details?: string; } | null>(null); - const [currentModel, setCurrentModel] = useState("anthropic:claude-opus-4-1"); + const [currentModel, setCurrentModel] = useState(DEFAULT_MODEL); const [editingMessage, setEditingMessage] = useState<{ id: string; content: string } | undefined>( undefined ); @@ -175,6 +176,46 @@ const AIViewInner: React.FC = ({ workspaceId, projectName, branch, const [loading, setLoading] = useState(false); + useEffect(() => { + let isMounted = true; + + const loadMetadata = async () => { + try { + const metadata = await window.api.workspace.getInfo(workspaceId); + if (isMounted) { + setCurrentModel(metadata?.model ?? DEFAULT_MODEL); + } + } catch (error) { + console.error("Failed to load workspace metadata:", error); + if (isMounted) { + setCurrentModel(DEFAULT_MODEL); + } + } + }; + + if (workspaceId) { + void loadMetadata(); + } + + return () => { + isMounted = false; + }; + }, [workspaceId]); + + useEffect(() => { + const unsubscribe = window.api.workspace.onMetadata(({ workspaceId: updatedId, metadata }) => { + if (updatedId === workspaceId) { + setCurrentModel(metadata?.model ?? DEFAULT_MODEL); + } + }); + + return () => { + if (typeof unsubscribe === "function") { + unsubscribe(); + } + }; + }, [workspaceId]); + // Handlers for editing messages const handleEditUserMessage = useCallback((messageId: string, content: string) => { setEditingMessage({ id: messageId, content }); @@ -375,6 +416,17 @@ const AIViewInner: React.FC = ({ workspaceId, projectName, branch, await window.api.workspace.clearHistory(workspaceId); }, [workspaceId, getAggregator]); + const handleSetModel = useCallback( + async (model: string) => { + const result = await window.api.workspace.setModel(workspaceId, model); + if (!result.success) { + throw new Error(result.error); + } + setCurrentModel(model); + }, + [workspaceId] + ); + const handleProviderConfig = useCallback( async (provider: string, keyPath: string[], value: string) => { const result = await window.api.providers.setProviderConfig(provider, keyPath, value); @@ -506,6 +558,7 @@ const AIViewInner: React.FC = ({ workspaceId, projectName, branch, onMessageSent={handleMessageSent} onClearHistory={handleClearHistory} onProviderConfig={handleProviderConfig} + onSetModel={handleSetModel} debugMode={debugMode} onDebugModeChange={setDebugMode} disabled={!projectName || !branch} diff --git a/src/components/ChatInput.tsx b/src/components/ChatInput.tsx index 309fe87aa8..77d24f9503 100644 --- a/src/components/ChatInput.tsx +++ b/src/components/ChatInput.tsx @@ -8,6 +8,7 @@ import { usePersistedState } from "../hooks/usePersistedState"; import { ThinkingSliderComponent } from "./ThinkingSlider"; import { useThinkingLevel } from "../hooks/useThinkingLevel"; import { getSlashCommandSuggestions, type SlashSuggestion } from "../utils/slashCommands"; +import { getModelAliasEntries } from "../constants/models"; const InputSection = styled.div` position: relative; @@ -19,6 +20,21 @@ const InputSection = styled.div` gap: 8px; `; +const MODEL_ALIAS_ENTRIES = getModelAliasEntries(); + +const renderAliasList = () => ( + <> + {MODEL_ALIAS_ENTRIES.map(({ alias, model }, index) => ( + + {alias} + {" -> "} + {model} + {index < MODEL_ALIAS_ENTRIES.length - 1 ?
: null} + + ))} + +); + const InputControls = styled.div` display: flex; gap: 10px; @@ -93,6 +109,7 @@ export interface ChatInputProps { onMessageSent?: () => void; // Optional callback after successful send onClearHistory: () => Promise; onProviderConfig?: (provider: string, keyPath: string[], value: string) => Promise; + onSetModel?: (model: string) => Promise; debugMode: boolean; onDebugModeChange: (enabled: boolean) => void; disabled?: boolean; @@ -162,6 +179,48 @@ const createCommandToast = (parsed: ParsedCommand): Toast | null => { ), }; + case "model-help": + return { + id: Date.now().toString(), + type: "error", + title: "Model Command", + message: "Select the AI model for this workspace.", + solution: ( + <> + Usage: + /model <provider> <model_name> +
+ /model <alias> +
+
+ Aliases: + {renderAliasList()} + + ), + }; + + case "model-invalid-input": + return { + id: Date.now().toString(), + type: "error", + title: "Invalid Model Input", + message: parsed.input + ? `Could not interpret '${parsed.input}' as a model.` + : "Model command requires a provider and model or a known alias.", + solution: ( + <> + Usage: + /model <provider> <model_name> +
+ /model <alias> +
+
+ Known Aliases: + {renderAliasList()} + + ), + }; + case "unknown-command": { const cmd = "/" + parsed.command + (parsed.subcommand ? " " + parsed.subcommand : ""); return { @@ -237,6 +296,7 @@ export const ChatInput: React.FC = ({ onMessageSent, onClearHistory, onProviderConfig, + onSetModel, debugMode, onDebugModeChange, disabled = false, @@ -359,6 +419,32 @@ export const ChatInput: React.FC = ({ return; } + if (parsed.type === "model-set" && onSetModel) { + setIsSending(true); + setInput(""); + + try { + await onSetModel(parsed.model); + setToast({ + id: Date.now().toString(), + type: "success", + title: "Model Updated", + message: `Using ${parsed.model}`, + }); + } catch (error) { + console.error("Failed to set model:", error); + setToast({ + id: Date.now().toString(), + type: "error", + message: error instanceof Error ? error.message : "Failed to set model", + }); + setInput(messageText); + } finally { + setIsSending(false); + } + return; + } + // Handle all other commands - show display toast const commandToast = createCommandToast(parsed); if (commandToast) { diff --git a/src/constants/ipc-constants.ts b/src/constants/ipc-constants.ts index 4bcee245cf..f8c4689f12 100644 --- a/src/constants/ipc-constants.ts +++ b/src/constants/ipc-constants.ts @@ -24,6 +24,7 @@ export const IPC_CHANNELS = { WORKSPACE_CLEAR_HISTORY: "workspace:clearHistory", WORKSPACE_STREAM_HISTORY: "workspace:streamHistory", WORKSPACE_GET_INFO: "workspace:getInfo", + WORKSPACE_SET_MODEL: "workspace:setModel", // Dynamic channel prefixes WORKSPACE_CHAT_PREFIX: "workspace:chat:", diff --git a/src/constants/models.ts b/src/constants/models.ts new file mode 100644 index 0000000000..72f213a183 --- /dev/null +++ b/src/constants/models.ts @@ -0,0 +1,20 @@ +export const DEFAULT_MODEL = "anthropic:claude-opus-4-1" as const; + +export const MODEL_ALIAS_MAP = { + opus: DEFAULT_MODEL, + sonnet: "anthropic:claude-sonnet-4", +} as const; + +export type ModelAlias = keyof typeof MODEL_ALIAS_MAP; + +export function resolveModelAlias(input: string): string | undefined { + const normalized = input.trim().toLowerCase(); + if (!normalized) { + return undefined; + } + return MODEL_ALIAS_MAP[normalized as ModelAlias]; +} + +export function getModelAliasEntries(): Array<{ alias: string; model: string }> { + return Object.entries(MODEL_ALIAS_MAP).map(([alias, model]) => ({ alias, model })); +} diff --git a/src/debug/costs.ts b/src/debug/costs.ts index 4728a5fee0..9eb638690e 100644 --- a/src/debug/costs.ts +++ b/src/debug/costs.ts @@ -3,6 +3,7 @@ import * as path from "path"; import { getSessionDir } from "../config"; import { CmuxMessage } from "../types/message"; import { calculateTokenStats } from "../utils/tokenStatsCalculator"; +import { DEFAULT_MODEL } from "../constants/models"; /** * Debug command to display cost/token statistics for a workspace @@ -34,7 +35,7 @@ export async function costsCommand(workspaceId: string) { // Detect model from first assistant message const firstAssistantMessage = messages.find((msg) => msg.role === "assistant"); - const model = firstAssistantMessage?.metadata?.model || "anthropic:claude-opus-4-1"; + const model = firstAssistantMessage?.metadata?.model || DEFAULT_MODEL; // Calculate stats using shared logic const stats = await calculateTokenStats(messages, model); diff --git a/src/main.ts b/src/main.ts index e517e50638..f2df817953 100644 --- a/src/main.ts +++ b/src/main.ts @@ -111,6 +111,7 @@ ipcMain.handle( id: workspaceId, projectName, workspacePath: result.path, + model: aiService.getDefaultModel(), }; await aiService.saveWorkspaceMetadata(workspaceId, metadata); @@ -205,6 +206,28 @@ ipcMain.handle(IPC_CHANNELS.WORKSPACE_GET_INFO, async (_event, workspaceId: stri return result.success ? result.data : null; }); +ipcMain.handle( + IPC_CHANNELS.WORKSPACE_SET_MODEL, + async (_event, workspaceId: string, model: string) => { + const result = await aiService.setWorkspaceModel(workspaceId, model); + if (!result.success) { + return { success: false, error: result.error }; + } + + const metadataResult = await aiService.getWorkspaceMetadata(workspaceId); + if (metadataResult.success) { + mainWindow?.webContents.send(IPC_CHANNELS.WORKSPACE_METADATA, { + workspaceId, + metadata: metadataResult.data, + }); + } else { + log.error("Failed to emit updated metadata:", metadataResult.error); + } + + return { success: true, data: undefined }; + } +); + ipcMain.handle( IPC_CHANNELS.WORKSPACE_SEND_MESSAGE, async ( diff --git a/src/preload.ts b/src/preload.ts index d88c6fef0c..a4914c425c 100644 --- a/src/preload.ts +++ b/src/preload.ts @@ -53,6 +53,8 @@ const api: IPCApi = { clearHistory: (workspaceId) => ipcRenderer.invoke(IPC_CHANNELS.WORKSPACE_CLEAR_HISTORY, workspaceId), getInfo: (workspaceId) => ipcRenderer.invoke(IPC_CHANNELS.WORKSPACE_GET_INFO, workspaceId), + setModel: (workspaceId, model) => + ipcRenderer.invoke(IPC_CHANNELS.WORKSPACE_SET_MODEL, workspaceId, model), onChat: (workspaceId, callback) => { const channel = getChatChannel(workspaceId); diff --git a/src/services/aiService.ts b/src/services/aiService.ts index 458669baf5..0f298542a1 100644 --- a/src/services/aiService.ts +++ b/src/services/aiService.ts @@ -2,7 +2,8 @@ import * as fs from "fs/promises"; import * as path from "path"; import { EventEmitter } from "events"; import { convertToModelMessages, type LanguageModel } from "ai"; -import { createAnthropic } from "@ai-sdk/anthropic"; +import { createOpenAI } from "@ai-sdk/openai"; +import { createGoogleGenerativeAI } from "@ai-sdk/google"; import { Result, Ok, Err } from "../types/result"; import { WorkspaceMetadata, WorkspaceMetadataSchema } from "../types/workspace"; import { CmuxMessage, createCmuxMessage } from "../types/message"; @@ -22,13 +23,15 @@ import { buildSystemMessage } from "./systemMessage"; import { getTokenizerForModel } from "../utils/tokenizer"; import { buildProviderOptions } from "../utils/providerOptions"; import { ThinkingLevel } from "../types/thinking"; +import { DEFAULT_MODEL } from "../constants/models"; +import { getProviderAdapter } from "./providerAdapters"; // Export a standalone version of getToolsForModel for use in backend export class AIService extends EventEmitter { private readonly METADATA_FILE = "metadata.json"; private streamManager = new StreamManager(); - private defaultModel = "anthropic:claude-opus-4-1"; // Default model string + private readonly defaultModel = DEFAULT_MODEL; // Default model string private historyService: HistoryService; constructor(historyService: HistoryService) { @@ -38,6 +41,10 @@ export class AIService extends EventEmitter { this.setupStreamEventForwarding(); } + getDefaultModel(): string { + return this.defaultModel; + } + /** * Forward all stream events from StreamManager to AIService consumers */ @@ -106,6 +113,25 @@ export class AIService extends EventEmitter { } } + async setWorkspaceModel(workspaceId: string, model: string): Promise> { + const metadataResult = await this.getWorkspaceMetadata(workspaceId); + if (!metadataResult.success) { + return Err(metadataResult.error); + } + + const normalizedModel = model.trim(); + if (!normalizedModel) { + return Err("Model cannot be empty"); + } + + const updatedMetadata: WorkspaceMetadata = { + ...metadataResult.data, + model: normalizedModel, + }; + + return this.saveWorkspaceMetadata(workspaceId, updatedMetadata); + } + /** * Split assistant messages that have text after tool calls with results. @@ -134,37 +160,23 @@ export class AIService extends EventEmitter { // Load providers configuration - the ONLY source of truth const providersConfig = loadProvidersConfig(); - const providerConfig = providersConfig?.[providerName]; + const adapter = getProviderAdapter(providerName); - if (!providerConfig) { + if (!adapter) { return Err({ type: "provider_not_configured", provider: providerName, }); } - // Handle Anthropic provider - if (providerName === "anthropic") { - // Check for API key in config - if (!providerConfig.apiKey) { - return Err({ - type: "api_key_not_found", - provider: providerName, - }); - } - - // Pass configuration verbatim to the provider, ensuring parity with Vercel AI SDK - const provider = createAnthropic(providerConfig); - return Ok(provider(modelId)); + const validation = adapter.validate(providerName, providersConfig?.[providerName]); + if (!validation.success) { + return Err(validation.error); } - // Add support for other providers here in the future - // if (providerName === "openai") { ... } - - return Err({ - type: "provider_not_configured", - provider: providerName, - }); + const providerConfig = validation.data; + const provider = adapter.instantiate(providerConfig, modelId); + return Ok(provider); } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error); return Err({ type: "unknown", raw: `Failed to create model: ${errorMessage}` }); @@ -186,8 +198,17 @@ export class AIService extends EventEmitter { abortSignal?: AbortSignal ): Promise> { try { + const metadataResult = await this.getWorkspaceMetadata(workspaceId); + if (!metadataResult.success) { + return Err({ type: "unknown", raw: metadataResult.error }); + } + + const workspaceMetadata = metadataResult.data; + const configuredModel = workspaceMetadata.model?.trim(); + const modelString = configuredModel && configuredModel.length > 0 ? configuredModel : this.defaultModel; + // Create model instance with early API key validation - const modelResult = await this.createModel(this.defaultModel); + const modelResult = await this.createModel(modelString); if (!modelResult.success) { return Err(modelResult.error); } @@ -204,7 +225,7 @@ export class AIService extends EventEmitter { const transformedMessages = transformModelMessages(modelMessages); // Apply cache control for Anthropic models AFTER transformation - const finalMessages = applyCacheControl(transformedMessages, this.defaultModel); + const finalMessages = applyCacheControl(transformedMessages, modelString); log.debug_obj(`${workspaceId}/3_final_messages.json`, finalMessages); @@ -215,29 +236,23 @@ export class AIService extends EventEmitter { // Continue anyway, as the API might be more lenient } - // Get workspace metadata to retrieve workspace path - const metadataResult = await this.getWorkspaceMetadata(workspaceId); - if (!metadataResult.success) { - return Err({ type: "unknown", raw: metadataResult.error }); - } - // Build system message from workspace metadata - const systemMessage = await buildSystemMessage(metadataResult.data); + const systemMessage = await buildSystemMessage(workspaceMetadata); // Count system message tokens for cost tracking - const tokenizer = getTokenizerForModel(this.defaultModel); + const tokenizer = getTokenizerForModel(modelString); const systemMessageTokens = await tokenizer.countTokens(systemMessage); - const workspacePath = metadataResult.data.workspacePath; + const workspacePath = workspaceMetadata.workspacePath; // Get model-specific tools with workspace path configuration - const tools = getToolsForModel(this.defaultModel, { cwd: workspacePath }); + const tools = getToolsForModel(modelString, { cwd: workspacePath }); // Create assistant message placeholder with historySequence from backend const assistantMessageId = `assistant-${Date.now()}-${Math.random().toString(36).substring(2, 11)}`; const assistantMessage = createCmuxMessage(assistantMessageId, "assistant", "", { timestamp: Date.now(), - model: this.defaultModel, + model: modelString, systemMessageTokens, }); @@ -251,14 +266,14 @@ export class AIService extends EventEmitter { const historySequence = assistantMessage.metadata?.historySequence ?? 0; // Build provider options based on thinking level - const providerOptions = buildProviderOptions(this.defaultModel, thinkingLevel || "off"); + const providerOptions = buildProviderOptions(modelString, thinkingLevel || "off"); // Delegate to StreamManager with model instance, system message, tools, historySequence, and initial metadata const streamResult = await this.streamManager.startStream( workspaceId, finalMessages, modelResult.data, - this.defaultModel, + modelString, historySequence, systemMessage, abortSignal, diff --git a/src/services/providerAdapters.ts b/src/services/providerAdapters.ts new file mode 100644 index 0000000000..43267b1b40 --- /dev/null +++ b/src/services/providerAdapters.ts @@ -0,0 +1,70 @@ +import type { LanguageModel } from "ai"; +import { createAnthropic } from "@ai-sdk/anthropic"; +import { createOpenAI } from "@ai-sdk/openai"; +import { createGoogleGenerativeAI } from "@ai-sdk/google"; +import type { ProviderConfig } from "../config"; +import { Err, Ok, type Result } from "../types/result"; +import type { SendMessageError } from "../types/errors"; + +interface ProviderAdapter { + validate: ( + providerName: string, + config: ProviderConfig | undefined + ) => Result; + instantiate: (config: ProviderConfig, modelId: string) => LanguageModel; +} + +interface ProviderAdapterOptions { + requiredKeys?: string[]; +} + +type ProviderFactory = (config: ProviderConfig) => (modelId: string) => LanguageModel; + +function createProviderAdapter( + factory: ProviderFactory, + options: ProviderAdapterOptions = {} +): ProviderAdapter { + const { requiredKeys = [] } = options; + + return { + validate(providerName, config) { + if (!config) { + return Err({ type: "provider_not_configured", provider: providerName }); + } + + for (const key of requiredKeys) { + const value = config[key]; + if (value === undefined || value === null || value === "") { + if (key === "apiKey") { + return Err({ type: "api_key_not_found", provider: providerName }); + } + + return Err({ + type: "unknown", + raw: `Missing required configuration '${key}' for provider '${providerName}'`, + }); + } + } + + return Ok(config); + }, + instantiate(config, modelId) { + const provider = factory(config); + return provider(modelId); + }, + }; +} + +const PROVIDER_ADAPTERS: Record = { + anthropic: createProviderAdapter(createAnthropic, { requiredKeys: ["apiKey"] }), + openai: createProviderAdapter(createOpenAI, { requiredKeys: ["apiKey"] }), + google: createProviderAdapter(createGoogleGenerativeAI, { requiredKeys: ["apiKey"] }), +}; + +export function getProviderAdapter(providerName: string): ProviderAdapter | undefined { + return PROVIDER_ADAPTERS[providerName]; +} + +export function registerProviderAdapter(providerName: string, adapter: ProviderAdapter): void { + PROVIDER_ADAPTERS[providerName] = adapter; +} diff --git a/src/types/ipc.ts b/src/types/ipc.ts index f93471db20..3195da8c84 100644 --- a/src/types/ipc.ts +++ b/src/types/ipc.ts @@ -132,6 +132,7 @@ export interface IPCApi { ): Promise>; clearHistory(workspaceId: string): Promise>; getInfo(workspaceId: string): Promise; + setModel(workspaceId: string, model: string): Promise>; // Event subscriptions (renderer-only) // These methods are designed to send current state immediately upon subscription, diff --git a/src/types/workspace.ts b/src/types/workspace.ts index ae82d960c3..41f08cc1ed 100644 --- a/src/types/workspace.ts +++ b/src/types/workspace.ts @@ -7,6 +7,7 @@ export const WorkspaceMetadataSchema = z.object({ id: z.string().min(1, "Workspace ID is required"), projectName: z.string().min(1, "Project name is required"), workspacePath: z.string().min(1, "Workspace path is required"), + model: z.string().min(1).optional(), }); /** @@ -22,6 +23,9 @@ export interface WorkspaceMetadata { /** Absolute path to the workspace worktree directory */ workspacePath: string; + + /** Selected model string in provider:model format */ + model?: string; } /** diff --git a/src/utils/commandParser.test.ts b/src/utils/commandParser.test.ts index e0b8742300..d5556ed63a 100644 --- a/src/utils/commandParser.test.ts +++ b/src/utils/commandParser.test.ts @@ -91,6 +91,49 @@ describe("commandParser", () => { }); }); + it("should show help for /model with no arguments", () => { + const result = parseCommand("/model"); + expect(result).toEqual({ type: "model-help" }); + }); + + it("should resolve model aliases", () => { + const result = parseCommand("/model opus"); + expect(result).toEqual({ + type: "model-set", + model: "anthropic:claude-opus-4-1", + requested: "opus", + source: "alias", + }); + }); + + it("should parse explicit provider and model arguments", () => { + const result = parseCommand("/model anthropic claude-sonnet-4"); + expect(result).toEqual({ + type: "model-set", + model: "anthropic:claude-sonnet-4", + requested: "anthropic claude-sonnet-4", + source: "provider-model", + }); + }); + + it("should parse provider:model syntax", () => { + const result = parseCommand("/model anthropic:claude-sonnet-4"); + expect(result).toEqual({ + type: "model-set", + model: "anthropic:claude-sonnet-4", + requested: "anthropic:claude-sonnet-4", + source: "explicit", + }); + }); + + it("should handle invalid model input", () => { + const result = parseCommand("/model anthropic"); + expect(result).toEqual({ + type: "model-invalid-input", + input: "anthropic", + }); + }); + it("should parse unknown commands", () => { expect(parseCommand("/foo")).toEqual({ type: "unknown-command", diff --git a/src/utils/commandParser.ts b/src/utils/commandParser.ts index 56af8b5da4..ce437f9329 100644 --- a/src/utils/commandParser.ts +++ b/src/utils/commandParser.ts @@ -1,3 +1,5 @@ +import { resolveModelAlias } from "../constants/models"; + /** * Command parser for parsing chat commands like /providers */ @@ -7,6 +9,9 @@ export type ParsedCommand = | { type: "providers-help" } | { type: "providers-invalid-subcommand"; subcommand: string } | { type: "providers-missing-args"; subcommand: string; argCount: number } + | { type: "model-set"; model: string; requested: string; source: "alias" | "explicit" | "provider-model" } + | { type: "model-help" } + | { type: "model-invalid-input"; input?: string } | { type: "clear" } | { type: "unknown-command"; command: string; subcommand?: string } | null; @@ -85,9 +90,69 @@ const providersCommandDefinition: SlashCommandDefinition = { }, children: [providersSetCommandDefinition], }; + +const modelCommandDefinition: SlashCommandDefinition = { + key: "model", + description: "Select the AI model", + handler: ({ cleanRemainingTokens }) => { + if (cleanRemainingTokens.length === 0) { + return { type: "model-help" }; + } + + const [firstToken, ...restTokens] = cleanRemainingTokens; + + if (restTokens.length === 0) { + const aliasModel = resolveModelAlias(firstToken); + if (aliasModel) { + return { + type: "model-set", + model: aliasModel, + requested: firstToken, + source: "alias", + }; + } + + if (firstToken.includes(":")) { + const [provider, ...modelParts] = firstToken.split(":"); + const modelName = modelParts.join(":"); + if (provider && modelName) { + return { + type: "model-set", + model: `${provider}:${modelName}`, + requested: firstToken, + source: "explicit", + }; + } + } + + return { + type: "model-invalid-input", + input: firstToken, + }; + } + + const provider = firstToken; + const modelName = restTokens.join(" ").trim(); + + if (!modelName) { + return { + type: "model-invalid-input", + input: provider, + }; + } + + return { + type: "model-set", + model: `${provider}:${modelName}`, + requested: `${provider} ${modelName}`, + source: "provider-model", + }; + }, +}; const SLASH_COMMAND_DEFINITIONS: ReadonlyArray = [ clearCommandDefinition, providersCommandDefinition, + modelCommandDefinition, ]; const SLASH_COMMAND_DEFINITION_MAP = new Map( diff --git a/src/utils/modelCatalog.ts b/src/utils/modelCatalog.ts new file mode 100644 index 0000000000..e23a0c037d --- /dev/null +++ b/src/utils/modelCatalog.ts @@ -0,0 +1,105 @@ +import modelsData from "./models.json"; + +interface RawModelData { + litellm_provider?: string; + mode?: string; + [key: string]: unknown; +} + +const PROVIDER_ALIASES: Record = { + anthropic: ["anthropic"], + openai: ["openai"], + google: [ + "google", + "google_ai_studio", + "vertex_ai-chat-models", + "vertex_ai-language-models", + "vertex_ai-code-chat-models", + "vertex_ai-code-text-models", + "vertex_ai-text-models", + ], +}; + +const TEXTUAL_MODE_KEYWORDS = ["chat", "completion", "text", "responses"]; +const EXCLUDED_MODE_KEYWORDS = ["audio", "image", "video", "embedding", "moderation", "rerank"]; + +function isTextualMode(mode: unknown): boolean { + if (typeof mode !== "string" || mode.length === 0) { + return true; + } + + const normalized = mode.toLowerCase(); + + if (EXCLUDED_MODE_KEYWORDS.some((keyword) => normalized.includes(keyword))) { + return false; + } + + return TEXTUAL_MODE_KEYWORDS.some((keyword) => normalized.includes(keyword)); +} + +function shouldIncludeModelName(name: string): boolean { + if (!name || typeof name !== "string") { + return false; + } + + if (name.includes("/")) { + return false; + } + + if (name.includes("@")) { + return false; + } + + if (name.startsWith("ft:")) { + return false; + } + + return true; +} + +function normalizeProvider(provider: string): string { + return provider.trim().toLowerCase(); +} + +const providerModelCache = new Map(); + +export function getModelsForProvider(provider: string): string[] { + const normalizedProvider = normalizeProvider(provider); + if (!normalizedProvider) { + return []; + } + + const cached = providerModelCache.get(normalizedProvider); + if (cached) { + return cached; + } + + const aliases = PROVIDER_ALIASES[normalizedProvider] ?? [normalizedProvider]; + const aliasSet = new Set(aliases.map((alias) => alias.toLowerCase())); + + const results = new Set(); + const data = modelsData as Record; + + for (const [name, info] of Object.entries(data)) { + if (!shouldIncludeModelName(name)) { + continue; + } + + const providerId = info.litellm_provider?.toLowerCase(); + if (!providerId || !aliasSet.has(providerId)) { + continue; + } + + if (!isTextualMode(info.mode)) { + continue; + } + + // Remove regional prefixes like "us." or "eu." when they mirror the same model name + const normalizedName = name.startsWith("us.") || name.startsWith("eu.") ? name.split(".").slice(1).join(".") : name; + results.add(normalizedName); + } + + const sorted = Array.from(results).sort((a, b) => a.localeCompare(b)); + providerModelCache.set(normalizedProvider, sorted); + return sorted; +} diff --git a/src/utils/slashCommands.test.ts b/src/utils/slashCommands.test.ts index 8614290712..4dc81bef09 100644 --- a/src/utils/slashCommands.test.ts +++ b/src/utils/slashCommands.test.ts @@ -13,6 +13,7 @@ describe("getSlashCommandSuggestions", () => { expect(labels).toContain("/clear"); expect(labels).toContain("/providers"); + expect(labels).toContain("/model"); }); it("filters top level commands by partial input", () => { @@ -51,4 +52,33 @@ describe("getSlashCommandSuggestions", () => { expect(suggestions).toHaveLength(1); expect(suggestions[0].display).toBe("apiKey"); }); + + it("suggests model aliases and providers after /model", () => { + const suggestions = getSlashCommandSuggestions("/model ", { + providerNames: ["anthropic"], + }); + + const labels = suggestions.map((s) => s.display); + expect(labels).toContain("opus"); + expect(labels).toContain("sonnet"); + expect(labels).toContain("anthropic"); + }); + + it("filters model alias suggestions by partial input", () => { + const suggestions = getSlashCommandSuggestions("/model op"); + expect(suggestions.map((s) => s.display)).toContain("opus"); + expect(suggestions.map((s) => s.display)).not.toContain("sonnet"); + }); + + it("suggests provider models after selecting a provider", () => { + const suggestions = getSlashCommandSuggestions("/model anthropic "); + const labels = suggestions.map((s) => s.display); + expect(labels).toContain("claude-opus-4-1"); + }); + + it("filters provider model suggestions by partial input", () => { + const suggestions = getSlashCommandSuggestions("/model openai gpt-4"); + const labels = suggestions.map((s) => s.display); + expect(labels.some((label) => label.startsWith("gpt-4"))).toBe(true); + }); }); diff --git a/src/utils/slashCommands.ts b/src/utils/slashCommands.ts index eb469adccf..a685bd29ce 100644 --- a/src/utils/slashCommands.ts +++ b/src/utils/slashCommands.ts @@ -1,4 +1,6 @@ import { getSlashCommandDefinitions, type SlashCommandDefinition } from "./commandParser"; +import { getModelAliasEntries } from "../constants/models"; +import { getModelsForProvider } from "./modelCatalog"; export interface SlashSuggestion { id: string; @@ -188,6 +190,85 @@ function buildProviderKeySuggestions( })); } +function buildModelFirstArgSuggestions( + partial: string, + providerNames: string[] | undefined +): SlashSuggestion[] { + const normalizedPartial = partial.trim().toLowerCase(); + + const aliasSuggestions = getModelAliasEntries() + .filter(({ alias }) => + normalizedPartial ? alias.toLowerCase().startsWith(normalizedPartial) : true + ) + .map(({ alias, model }) => ({ + id: `command:model:alias:${alias}`, + display: alias, + description: `Alias for ${model}`, + replacement: `/model ${alias}`, + })); + + const providerDefinitions = dedupeDefinitions([ + ...(providerNames ?? []).map((name) => ({ + key: name, + description: `${name} provider configuration`, + })), + ...DEFAULT_PROVIDER_NAMES, + ]); + + const providerSuggestions = filterAndMapSuggestions( + providerDefinitions, + partial, + (definition) => ({ + id: `command:model:provider:${definition.key}`, + display: definition.key, + description: definition.description, + replacement: `/model ${definition.key} `, + }) + ); + + return [...aliasSuggestions, ...providerSuggestions]; +} + +function buildModelCommandSuggestions( + stage: number, + partialToken: string, + context: SlashSuggestionContext, + completedTokens: string[], + tokens: string[] +): SlashSuggestion[] { + if (stage === 1) { + return buildModelFirstArgSuggestions(partialToken, context.providerNames); + } + + if (stage === 2) { + const providerName = completedTokens[1] ?? tokens[1]; + if (!providerName) { + return []; + } + + const models = getModelsForProvider(providerName); + if (models.length === 0) { + return []; + } + + const normalizedPartial = partialToken.trim().toLowerCase(); + + return models + .filter((model) => + normalizedPartial ? model.toLowerCase().includes(normalizedPartial) : true + ) + .slice(0, 25) + .map((model) => ({ + id: `command:model:${providerName}:${model}`, + display: model, + description: `${providerName}:${model}`, + replacement: `/model ${providerName} ${model}`, + })); + } + + return []; +} + export function getSlashCommandSuggestions( input: string, context: SlashSuggestionContext = {} @@ -248,6 +329,10 @@ export function getSlashCommandSuggestions( } } + if (definitionPath[0]?.key === "model") { + return buildModelCommandSuggestions(stage, partialToken, context, completedTokens, tokens); + } + if (definitionPath[0]?.key !== "providers") { return []; }