From 5c383c0634aca63866bae20cf48ec17ebc0718f8 Mon Sep 17 00:00:00 2001
From: duanfuxiang <duanfuxiang@haizhi.com>
Date: Wed, 18 Jun 2025 07:40:38 +0800
Subject: [PATCH] save user custom input model name

---
 .../components/ModelProviderSettings.tsx      | 97 ++++++++++++-------
 .../components/ProviderModelsPicker.tsx       | 59 ++++++++---
 src/types/settings.ts                         | 72 +++++++++-----
 3 files changed, 156 insertions(+), 72 deletions(-)

diff --git a/src/settings/components/ModelProviderSettings.tsx b/src/settings/components/ModelProviderSettings.tsx
index 1806139..751c536 100644
--- a/src/settings/components/ModelProviderSettings.tsx
+++ b/src/settings/components/ModelProviderSettings.tsx
@@ -44,7 +44,7 @@ const keyMap: Record<ApiProvider, ProviderSettingKey> = {
 	'OpenAICompatible': 'openaicompatibleProvider',
 };
 
-const getProviderSettingKey = (provider: ApiProvider): ProviderSettingKey => {
+export const getProviderSettingKey = (provider: ApiProvider): ProviderSettingKey => {
 	return keyMap[provider];
 };
 
@@ -70,7 +70,7 @@ const CustomProviderSettings: React.FC<CustomProviderSettingsProps> = ({ plugin,
 
 	// 获取已设置API Key的提供商列表
 	const getSettedProviders = (): ApiProvider[] => {
-		return providers.filter(provider => {			
+		return providers.filter(provider => {
 			const providerSetting = getProviderSetting(provider);
 			return providerSetting.apiKey && providerSetting.apiKey.trim() !== '';
 		});
@@ -79,7 +79,7 @@ const CustomProviderSettings: React.FC<CustomProviderSettingsProps> = ({ plugin,
 	// 一键配置模型
 	const handleOneClickConfig = () => {
 		const settedProviders = getSettedProviders();
-		
+
 		if (settedProviders.length === 0) {
 			// 提示用户未设置任何key
 			alert(t("settings.ModelProvider.noApiKeySet"));
@@ -88,7 +88,7 @@ const CustomProviderSettings: React.FC<CustomProviderSettingsProps> = ({ plugin,
 
 		// 选择chat和autocomplete的提供商（按providers排序选择最靠前的）
 		const selectedProvider = providers.find(provider => settedProviders.includes(provider));
-		
+
 		// 选择embedding的提供商（按embeddingProviders排序选择最靠前的）
 		const embeddingProvider = embeddingProviders.find(provider => settedProviders.includes(provider));
 
@@ -98,7 +98,7 @@ const CustomProviderSettings: React.FC<CustomProviderSettingsProps> = ({ plugin,
 
 		if (selectedProvider) {
 			const defaultModels = GetDefaultModelId(selectedProvider);
-			
+
 			// 设置chat和autocomplete模型
 			if (defaultModels.chat) {
 				newSettings.chatModelProvider = selectedProvider;
@@ -116,7 +116,7 @@ const CustomProviderSettings: React.FC<CustomProviderSettingsProps> = ({ plugin,
 
 		if (embeddingProvider) {
 			const embeddingDefaultModels = GetDefaultModelId(embeddingProvider);
-			
+
 			// 设置embedding模型
 			if (embeddingDefaultModels.embedding) {
 				newSettings.embeddingModelProvider = embeddingProvider;
@@ -171,37 +171,37 @@ const CustomProviderSettings: React.FC<CustomProviderSettingsProps> = ({ plugin,
 		});
 	};
 
-	const testApiConnection = async (provider: ApiProvider) => {
+	const testApiConnection = async (provider: ApiProvider, modelId?: string) => {
 		console.log(`Testing connection for ${provider}...`);
-		
+
 		try {
 			// 动态导入LLMManager以避免循环依赖
 			const { default: LLMManager } = await import('../../core/llm/manager');
 			const { GetDefaultModelId } = await import('../../utils/api');
-			
+
 			// 对于Ollama和OpenAICompatible，不支持测试API连接
 			if (provider === ApiProvider.Ollama || provider === ApiProvider.OpenAICompatible) {
 				throw new Error(t("settings.ModelProvider.testConnection.notSupported", { provider }));
 			}
-			
+
 			// 创建LLM管理器实例
 			const llmManager = new LLMManager(settings);
-			
+
 			// 获取提供商的默认聊天模型
 			const defaultModels = GetDefaultModelId(provider);
-			const testModelId = defaultModels.chat;
-			
+			const testModelId = modelId || defaultModels.chat;
+
 			// 对于没有默认模型的提供商，使用通用的测试模型
 			if (!testModelId) {
 				throw new Error(t("settings.ModelProvider.testConnection.noDefaultModel", { provider }));
 			}
-			
+
 			// 构造测试模型对象
 			const testModel = {
 				provider: provider,
 				modelId: testModelId
 			};
-			
+
 			// 构造简单的测试请求
 			const testRequest = {
 				messages: [
@@ -214,11 +214,11 @@ const CustomProviderSettings: React.FC<CustomProviderSettingsProps> = ({ plugin,
 				max_tokens: 10,
 				temperature: 0
 			};
-			
+
 			// 设置超时选项
 			const abortController = new AbortController();
 			const timeoutId = setTimeout(() => abortController.abort(), 10000); // 10秒超时
-			
+
 			try {
 				// 发起API调用测试
 				const response = await llmManager.generateResponse(
@@ -226,9 +226,9 @@ const CustomProviderSettings: React.FC<CustomProviderSettingsProps> = ({ plugin,
 					testRequest,
 					{ signal: abortController.signal }
 				);
-				
+
 				clearTimeout(timeoutId);
-				
+
 				// 检查响应是否有效
 				if (response && response.choices && response.choices.length > 0) {
 					console.log(`✅ ${provider} connection test successful:`, response.choices[0]?.message?.content);
@@ -241,13 +241,13 @@ const CustomProviderSettings: React.FC<CustomProviderSettingsProps> = ({ plugin,
 				clearTimeout(timeoutId);
 				throw apiError;
 			}
-			
+
 		} catch (error) {
 			console.error(`❌ ${provider} connection test failed:`, error);
-			
+
 			// 根据错误类型提供更具体的错误信息
 			let errorMessage = t("settings.ModelProvider.testConnection.connectionFailed");
-			
+
 			if (error.message?.includes('API key')) {
 				errorMessage = t("settings.ModelProvider.testConnection.invalidApiKey");
 			} else if (error.message?.includes('base URL') || error.message?.includes('baseURL')) {
@@ -285,27 +285,58 @@ const CustomProviderSettings: React.FC<CustomProviderSettingsProps> = ({ plugin,
 		return settings[providerKey] || {};
 	};
 
-	const updateChatModelId = (provider: ApiProvider, modelId: string) => {
+	const updateChatModelId = (
+		provider: ApiProvider,
+		modelId: string,
+		isCustom: boolean = false
+	) => {
+		console.log(`updateChatModelId: ${provider} -> ${modelId}, isCustom: ${isCustom}`)
+		const providerSettingKey = getProviderSettingKey(provider);
+		const providerSettings = settings[providerSettingKey] || {};
+		const currentModels = providerSettings.models || new Set<string>();
+		
 		handleSettingsUpdate({
 			...settings,
 			chatModelProvider: provider,
-			chatModelId: modelId
+			chatModelId: modelId,
+			[providerSettingKey]: {
+				...providerSettings,
+				models: isCustom ? new Set([...currentModels, modelId]) : currentModels
+			}
 		});
 	};
 
-	const updateApplyModelId = (provider: ApiProvider, modelId: string) => {
+	const updateApplyModelId = (provider: ApiProvider, modelId: string, isCustom: boolean = false) => {
+		console.log(`updateApplyModelId: ${provider} -> ${modelId}, isCustom: ${isCustom}`)
+		const providerSettingKey = getProviderSettingKey(provider);
+		const providerSettings = settings[providerSettingKey] || {};
+		const currentModels = providerSettings.models || new Set<string>();
+		
 		handleSettingsUpdate({
 			...settings,
 			applyModelProvider: provider,
-			applyModelId: modelId
+			applyModelId: modelId,
+			[providerSettingKey]: {
+				...providerSettings,
+				models: isCustom ? new Set([...currentModels, modelId]) : currentModels
+			}
 		});
 	};
 
-	const updateEmbeddingModelId = (provider: ApiProvider, modelId: string) => {
+	const updateEmbeddingModelId = (provider: ApiProvider, modelId: string, isCustom: boolean = false) => {
+		console.log(`updateEmbeddingModelId: ${provider} -> ${modelId}, isCustom: ${isCustom}`)
+		const providerSettingKey = getProviderSettingKey(provider);
+		const providerSettings = settings[providerSettingKey] || {};
+		const currentModels = providerSettings.models || new Set<string>();
+		
 		handleSettingsUpdate({
 			...settings,
 			embeddingModelProvider: provider,
-			embeddingModelId: modelId
+			embeddingModelId: modelId,
+			[providerSettingKey]: {
+				...providerSettings,
+				models: isCustom ? new Set([...currentModels, modelId]) : currentModels
+			}
 		});
 	};
 
@@ -313,7 +344,7 @@ const CustomProviderSettings: React.FC<CustomProviderSettingsProps> = ({ plugin,
 	const generateApiKeyDescription = (provider: ApiProvider): React.ReactNode => {
 		const apiUrl = getProviderApiUrl(provider);
 		const baseDescription = String(t("settings.ApiProvider.enterApiKeyDescription"));
-		
+
 		if (!apiUrl) {
 			// 如果没有URL，直接移除占位符
 			return baseDescription.replace('{provider_api_url}', '');
@@ -328,9 +359,9 @@ const CustomProviderSettings: React.FC<CustomProviderSettingsProps> = ({ plugin,
 		return (
 			<>
 				{parts[0]}
-				<a 
-					href={apiUrl} 
-					target="_blank" 
+				<a
+					href={apiUrl}
+					target="_blank"
 					rel="noopener noreferrer"
 					className="provider-api-link"
 				>
@@ -398,7 +429,7 @@ const CustomProviderSettings: React.FC<CustomProviderSettingsProps> = ({ plugin,
 			<div className="model-selection-section">
 				<div className="model-selection-header">
 					<h2 className="section-title">{t("settings.ModelProvider.modelSelection")}:</h2>
-					<button 
+					<button
 						className="one-click-config-btn"
 						onClick={handleOneClickConfig}
 						title={t("settings.ModelProvider.oneClickConfigTooltip")}
diff --git a/src/settings/components/ProviderModelsPicker.tsx b/src/settings/components/ProviderModelsPicker.tsx
index d5b8e8b..fa7e340 100644
--- a/src/settings/components/ProviderModelsPicker.tsx
+++ b/src/settings/components/ProviderModelsPicker.tsx
@@ -8,6 +8,8 @@ import { InfioSettings } from "../../types/settings";
 // import { PROVIDERS } from '../constants';
 import { GetAllProviders, GetEmbeddingProviderModelIds, GetEmbeddingProviders, GetProviderModelIds } from "../../utils/api";
 
+import { getProviderSettingKey } from "./ModelProviderSettings";
+
 type TextSegment = {
 	text: string;
 	isHighlighted: boolean;
@@ -21,6 +23,7 @@ type SearchableItem = {
 type HighlightedItem = {
 	id: string;
 	html: TextSegment[];
+	isCustom?: boolean;
 };
 
 // Type guard for Record<string, unknown>
@@ -154,7 +157,7 @@ export type ComboBoxComponentProps = {
 	settings?: InfioSettings | null;
 	isEmbedding?: boolean,
 	description?: string;
-	updateModel: (provider: ApiProvider, modelId: string) => void;
+	updateModel: (provider: ApiProvider, modelId: string, isCustom?: boolean) => void;
 };
 
 export const ComboBoxComponent: React.FC<ComboBoxComponentProps> = ({
@@ -178,24 +181,49 @@ export const ComboBoxComponent: React.FC<ComboBoxComponentProps> = ({
 
 	const [modelIds, setModelIds] = useState<string[]>([]);
 
+	// 统一处理模型选择和保存
+	const handleModelSelect = (provider: ApiProvider, modelId: string, isCustom?: boolean) => {
+		console.log(`handleModelSelect: ${provider} -> ${modelId}`)
+
+		// 检查是否是自定义模型（不在官方模型列表中）
+		// const isCustomModel = !modelIds.includes(modelId);
+
+		updateModel(provider, modelId, isCustom);
+	};
+
 	// Replace useMemo with useEffect for async fetching
 	useEffect(() => {
 		const fetchModelIds = async () => {
 			const ids = isEmbedding
 				? GetEmbeddingProviderModelIds(modelProvider)
 				: await GetProviderModelIds(modelProvider, settings);
+			console.log(`📝 Fetched ${ids.length} official models for ${modelProvider}:`, ids);
 			setModelIds(ids);
 		};
 
 		fetchModelIds();
 	}, [modelProvider, isEmbedding, settings]);
 
+	const combinedModelIds = useMemo(() => {
+		const providerKey = getProviderSettingKey(modelProvider);
+		const providerModels = settings?.[providerKey]?.models;
+		console.log(`🔍 Custom models in settings for ${modelProvider}:`, providerModels ? Array.from(providerModels) : 'none')
+		// Ensure providerModels is a Set of strings
+		if (!providerModels || !(providerModels instanceof Set)) {
+			console.log(`📋 Using only official models (${modelIds.length}):`, modelIds);
+			return modelIds;
+		}
+		const additionalModels = Array.from(providerModels).filter((model): model is string => typeof model === 'string');
+		console.log(`📋 Combined models: ${modelIds.length} official + ${additionalModels.length} custom`);
+		return [...modelIds, ...additionalModels];
+	}, [modelIds, settings, modelProvider]);
+
 	const searchableItems = useMemo(() => {
-		return modelIds.map((id) => ({
-			id,
-			html: id,
+		return combinedModelIds.map((id): SearchableItem => ({
+			id: String(id),
+			html: String(id),
 		}))
-	}, [modelIds])
+	}, [combinedModelIds])
 
 	// fuse, used for fuzzy search, simple configuration threshold can be adjusted as needed
 	const fuse: Fuse<SearchableItem> = useMemo(() => {
@@ -212,24 +240,25 @@ export const ComboBoxComponent: React.FC<ComboBoxComponentProps> = ({
 
 	// 根据 searchTerm 得到过滤后的数据列表
 	const filteredOptions = useMemo(() => {
-		let results: HighlightedItem[] = searchTerm
+		const results: HighlightedItem[] = searchTerm
 			? highlight(fuse.search(searchTerm))
 			: searchableItems.map(item => ({
 				...item,
 				html: typeof item.html === 'string' ? [{ text: item.html, isHighlighted: false }] : item.html
 			}))
-		
+
 		// 如果有搜索词，添加自定义选项（如果不存在完全匹配的话）
 		if (searchTerm && searchTerm.trim()) {
 			const exactMatch = searchableItems.some(item => item.id === searchTerm);
 			if (!exactMatch) {
 				results.unshift({
 					id: searchTerm,
-					html: [{ text: `${modelIds.length > 0 ? t("settings.ModelProvider.custom") : ''}${searchTerm}`, isHighlighted: false }]
+					html: [{ text: `${modelIds.length > 0 ? t("settings.ModelProvider.custom") : ''}${searchTerm}`, isHighlighted: false }],
+					isCustom: true
 				});
 			}
 		}
-		
+
 		return results
 	}, [searchableItems, searchTerm, fuse, modelIds.length])
 
@@ -251,10 +280,10 @@ export const ComboBoxComponent: React.FC<ComboBoxComponentProps> = ({
 		// Use proper type checking without type assertion
 		const availableProviders = providers;
 		const isValidProvider = (value: string): value is ApiProvider => {
-			// @ts-ignore
-			return (availableProviders as readonly string[]).includes(value);
+			// @ts-expect-error - checking if providers array includes the value
+			return availableProviders.includes(value);
 		};
-		
+
 		if (isValidProvider(newProvider)) {
 			setModelProvider(newProvider);
 		}
@@ -346,11 +375,11 @@ export const ComboBoxComponent: React.FC<ComboBoxComponentProps> = ({
 													if (filteredOptions.length > 0) {
 														const selectedOption = filteredOptions[selectedIndex];
 														if (selectedOption) {
-															updateModel(modelProvider, selectedOption.id);
+															handleModelSelect(modelProvider, selectedOption.id, selectedOption.isCustom);
 														}
 													} else if (searchTerm.trim()) {
 														// If no options but there is input content, use the input content directly
-														updateModel(modelProvider, searchTerm.trim());
+														handleModelSelect(modelProvider, searchTerm.trim(), true);
 													}
 													setSearchTerm("");
 													setIsOpen(false);
@@ -373,7 +402,7 @@ export const ComboBoxComponent: React.FC<ComboBoxComponentProps> = ({
 													ref={(el) => (itemRefs.current[index] = el)}
 													onMouseEnter={() => setSelectedIndex(index)}
 													onClick={() => {
-														updateModel(modelProvider, option.id);
+														handleModelSelect(modelProvider, option.id, option.isCustom);
 														setSearchTerm("");
 														setIsOpen(false);
 													}}
diff --git a/src/types/settings.ts b/src/types/settings.ts
index d79469a..aa589fd 100644
--- a/src/types/settings.ts
+++ b/src/types/settings.ts
@@ -19,144 +19,168 @@ const InfioProviderSchema = z.object({
 	name: z.literal('Infio'),
 	apiKey: z.string().catch(''),
 	baseUrl: z.string().catch(''),
-	useCustomUrl: z.boolean().catch(false)
+	useCustomUrl: z.boolean().catch(false),
+	models: z.set(z.string()).catch(new Set())
 }).catch({
 	name: 'Infio',
 	apiKey: '',
 	baseUrl: '',
-	useCustomUrl: false
+	useCustomUrl: false,
+	models: new Set()
 })
 
 const OpenRouterProviderSchema = z.object({
 	name: z.literal('OpenRouter'),
 	apiKey: z.string().catch(''),
 	baseUrl: z.string().catch(''),
-	useCustomUrl: z.boolean().catch(false)
+	useCustomUrl: z.boolean().catch(false),
+	models: z.set(z.string()).catch(new Set())
 }).catch({
 	name: 'OpenRouter',
 	apiKey: '',
 	baseUrl: '',
-	useCustomUrl: false
+	useCustomUrl: false,
+	models: new Set()
 })
 
 const SiliconFlowProviderSchema = z.object({
 	name: z.literal('SiliconFlow'),
 	apiKey: z.string().catch(''),
 	baseUrl: z.string().catch(''),
-	useCustomUrl: z.boolean().catch(false)
+	useCustomUrl: z.boolean().catch(false),
+	models: z.set(z.string()).catch(new Set())
 }).catch({
 	name: 'SiliconFlow',
 	apiKey: '',
 	baseUrl: '',
-	useCustomUrl: false
+	useCustomUrl: false,
+	models: new Set()
 })
 
 const AlibabaQwenProviderSchema = z.object({
 	name: z.literal('AlibabaQwen'),
 	apiKey: z.string().catch(''),
 	baseUrl: z.string().catch(''),
-	useCustomUrl: z.boolean().catch(false)
+	useCustomUrl: z.boolean().catch(false),
+	models: z.set(z.string()).catch(new Set())
 }).catch({
 	name: 'AlibabaQwen',
 	apiKey: '',
 	baseUrl: '',
-	useCustomUrl: false
+	useCustomUrl: false,
+	models: new Set()
 })
 
 const AnthropicProviderSchema = z.object({
 	name: z.literal('Anthropic'),
 	apiKey: z.string().catch(''),
 	baseUrl: z.string().optional(),
-	useCustomUrl: z.boolean().catch(false)
+	useCustomUrl: z.boolean().catch(false),
+	models: z.set(z.string()).catch(new Set())
 }).catch({
 	name: 'Anthropic',
 	apiKey: '',
 	baseUrl: '',
-	useCustomUrl: false
+	useCustomUrl: false,
+	models: new Set()
 })
 
 const DeepSeekProviderSchema = z.object({
 	name: z.literal('DeepSeek'),
 	apiKey: z.string().catch(''),
 	baseUrl: z.string().catch(''),
-	useCustomUrl: z.boolean().catch(false)
+	useCustomUrl: z.boolean().catch(false),
+	models: z.set(z.string()).catch(new Set())
 }).catch({
 	name: 'DeepSeek',
 	apiKey: '',
 	baseUrl: '',
-	useCustomUrl: false
+	useCustomUrl: false,
+	models: new Set()
 })
 
 const GoogleProviderSchema = z.object({
 	name: z.literal('Google'),
 	apiKey: z.string().catch(''),
 	baseUrl: z.string().catch(''),
-	useCustomUrl: z.boolean().catch(false)
+	useCustomUrl: z.boolean().catch(false),
+	models: z.set(z.string()).catch(new Set())
 }).catch({
 	name: 'Google',
 	apiKey: '',
 	baseUrl: '',
-	useCustomUrl: false
+	useCustomUrl: false,
+	models: new Set()
 })
 
 const OpenAIProviderSchema = z.object({
 	name: z.literal('OpenAI'),
 	apiKey: z.string().catch(''),
 	baseUrl: z.string().optional(),
-	useCustomUrl: z.boolean().catch(false)
+	useCustomUrl: z.boolean().catch(false),
+	models: z.set(z.string()).catch(new Set())
 }).catch({
 	name: 'OpenAI',
 	apiKey: '',
 	baseUrl: '',
-	useCustomUrl: false
+	useCustomUrl: false,
+	models: new Set()
 })
 
 const OpenAICompatibleProviderSchema = z.object({
 	name: z.literal('OpenAICompatible'),
 	apiKey: z.string().catch(''),
 	baseUrl: z.string().optional(),
-	useCustomUrl: z.boolean().catch(true)
+	useCustomUrl: z.boolean().catch(true),
+	models: z.set(z.string()).catch(new Set())
 }).catch({
 	name: 'OpenAICompatible',
 	apiKey: '',
 	baseUrl: '',
-	useCustomUrl: true
+	useCustomUrl: true,
+	models: new Set()
 })
 
 const OllamaProviderSchema = z.object({
 	name: z.literal('Ollama'),
 	apiKey: z.string().catch('ollama'),
 	baseUrl: z.string().catch(''),
-	useCustomUrl: z.boolean().catch(false)
+	useCustomUrl: z.boolean().catch(false),
+	models: z.set(z.string()).catch(new Set())
 }).catch({
 	name: 'Ollama',
 	apiKey: 'ollama',
 	baseUrl: '',
-	useCustomUrl: true
+	useCustomUrl: true,
+	models: new Set()
 })
 
 const GroqProviderSchema = z.object({
 	name: z.literal('Groq'),
 	apiKey: z.string().catch(''),
 	baseUrl: z.string().catch(''),
-	useCustomUrl: z.boolean().catch(false)
+	useCustomUrl: z.boolean().catch(false),
+	models: z.set(z.string()).catch(new Set())
 }).catch({
 	name: 'Groq',
 	apiKey: '',
 	baseUrl: '',
-	useCustomUrl: false
+	useCustomUrl: false,
+	models: new Set()
 })
 
 const GrokProviderSchema = z.object({
 	name: z.literal('Grok'),
 	apiKey: z.string().catch(''),
 	baseUrl: z.string().catch(''),
-	useCustomUrl: z.boolean().catch(false)
+	useCustomUrl: z.boolean().catch(false),
+	models: z.set(z.string()).catch(new Set())
 }).catch({
 	name: 'Grok',
 	apiKey: '',
 	baseUrl: '',
-	useCustomUrl: false
+	useCustomUrl: false,
+	models: new Set()
 })
 
 const ollamaModelSchema = z.object({