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 []; }