返回 presentation-ai
useLocalModels.ts
根目录 / src / hooks / presentation / useLocalModels.ts
1 "use client";
2
3 import { createLogger } from "@/lib/observability/logger";
4 import { useQuery } from "@tanstack/react-query";
5 import { useEffect, useState } from "react";
6
7 interface ModelInfo {
8 id: string;
9 name: string;
10 provider: "ollama" | "lmstudio";
11 }
12
13 const localModelLogger = createLogger("client:local-models");
14 const LOCAL_MODELS_API_URL = "/api/presentation/local-models";
15
16 function mergeModels(...groups: ModelInfo[][]): ModelInfo[] {
17 const merged: ModelInfo[] = [];
18 const seen = new Set<string>();
19
20 for (const group of groups) {
21 for (const model of group) {
22 if (seen.has(model.id)) {
23 continue;
24 }
25
26 seen.add(model.id);
27 merged.push(model);
28 }
29 }
30
31 return merged;
32 }
33
34 function getSavedLocalModel(): ModelInfo | null {
35 const selectedModel = getSelectedModel();
36 if (!selectedModel?.modelId) {
37 return null;
38 }
39
40 if (
41 selectedModel.modelProvider !== "ollama" &&
42 selectedModel.modelProvider !== "lmstudio"
43 ) {
44 return null;
45 }
46
47 return {
48 id: `${selectedModel.modelProvider}-${selectedModel.modelId}`,
49 name: selectedModel.modelId,
50 provider: selectedModel.modelProvider,
51 };
52 }
53
54 interface LocalModelsApiResponse {
55 models?: ModelInfo[];
56 }
57
58 async function fetchLocalModels(): Promise<ModelInfo[]> {
59 try {
60 const response = await fetch(LOCAL_MODELS_API_URL, {
61 cache: "no-store",
62 });
63
64 if (!response.ok) {
65 throw new Error(`Local models API responded with ${response.status}`);
66 }
67
68 const data = (await response.json()) as LocalModelsApiResponse;
69 const models = Array.isArray(data.models) ? data.models : [];
70
71 localModelLogger.info("Local model discovery completed", {
72 total: models.length,
73 });
74
75 if (models.length === 0) {
76 localModelLogger.warn(
77 "No live local models detected; falling back to downloadable Ollama suggestions",
78 {
79 downloadableModels: downloadableModels.map((model) => model.name),
80 },
81 );
82 }
83
84 return models;
85 } catch (error) {
86 localModelLogger.warn("Failed to refresh local models from the server", {
87 error: error instanceof Error ? error.message : String(error),
88 });
89 return [];
90 }
91 }
92
93 export const downloadableModels: ModelInfo[] = [
94 {
95 id: "ollama-llama3.1:8b",
96 name: "llama3.1:8b",
97 provider: "ollama",
98 },
99 {
100 id: "ollama-llama3.1:70b",
101 name: "llama3.1:70b",
102 provider: "ollama",
103 },
104 {
105 id: "ollama-llama3.2:3b",
106 name: "llama3.2:3b",
107 provider: "ollama",
108 },
109 {
110 id: "ollama-llama3.2:8b",
111 name: "llama3.2:8b",
112 provider: "ollama",
113 },
114 {
115 id: "ollama-mistral:7b",
116 name: "mistral:7b",
117 provider: "ollama",
118 },
119 {
120 id: "ollama-codellama:7b",
121 name: "codellama:7b",
122 provider: "ollama",
123 },
124 {
125 id: "ollama-qwen2.5:7b",
126 name: "qwen2.5:7b",
127 provider: "ollama",
128 },
129 {
130 id: "ollama-gemma2:9b",
131 name: "gemma2:9b",
132 provider: "ollama",
133 },
134 {
135 id: "ollama-phi3:3.8b",
136 name: "phi3:3.8b",
137 provider: "ollama",
138 },
139 {
140 id: "ollama-neural-chat:7b",
141 name: "neural-chat:7b",
142 provider: "ollama",
143 },
144 ];
145
146 export const fallbackModels: ModelInfo[] = downloadableModels;
147
148 const MODELS_CACHE_KEY = "presentation-models-cache";
149 const SELECTED_MODEL_KEY = "presentation-selected-model";
150 const CACHE_EXPIRY_KEY = "presentation-models-cache-expiry";
151 const CACHE_DURATION = 5 * 60 * 1000;
152
153 function getCachedModels(): ModelInfo[] | null {
154 if (typeof window === "undefined") {
155 return null;
156 }
157
158 try {
159 const cached = localStorage.getItem(MODELS_CACHE_KEY);
160 const expiry = localStorage.getItem(CACHE_EXPIRY_KEY);
161
162 if (cached && expiry && Date.now() < parseInt(expiry, 10)) {
163 return JSON.parse(cached) as ModelInfo[];
164 }
165
166 return null;
167 } catch {
168 return null;
169 }
170 }
171
172 function setCachedModels(models: ModelInfo[]): void {
173 if (typeof window === "undefined") {
174 return;
175 }
176
177 try {
178 localStorage.setItem(MODELS_CACHE_KEY, JSON.stringify(models));
179 localStorage.setItem(
180 CACHE_EXPIRY_KEY,
181 (Date.now() + CACHE_DURATION).toString(),
182 );
183 } catch {
184 // Ignore localStorage errors.
185 }
186 }
187
188 export function getSelectedModel(): {
189 modelProvider: string;
190 modelId: string;
191 } | null {
192 if (typeof window === "undefined") {
193 return null;
194 }
195
196 try {
197 const selected = localStorage.getItem(SELECTED_MODEL_KEY);
198 return selected
199 ? (JSON.parse(selected) as { modelProvider: string; modelId: string })
200 : null;
201 } catch (error) {
202 localModelLogger.error(
203 "Failed to read selected model from localStorage",
204 error,
205 );
206 return null;
207 }
208 }
209
210 export function setSelectedModel(modelProvider: string, modelId: string): void {
211 if (typeof window === "undefined") {
212 return;
213 }
214
215 try {
216 localStorage.setItem(
217 SELECTED_MODEL_KEY,
218 JSON.stringify({ modelProvider, modelId }),
219 );
220 } catch (error) {
221 localModelLogger.error("Failed to save selected model to localStorage", error);
222 }
223 }
224
225 export function useLocalModels() {
226 const [isInitialLoad, setIsInitialLoad] = useState(true);
227 const cachedModels = getCachedModels();
228
229 const query = useQuery({
230 queryKey: ["local-models"],
231 queryFn: async () => {
232 localModelLogger.info("Refreshing local model list");
233 const freshModels = await fetchLocalModels();
234
235 if (freshModels.length > 0) {
236 setCachedModels(freshModels);
237 return freshModels;
238 }
239
240 if (cachedModels && cachedModels.length > 0) {
241 localModelLogger.info(
242 "Local discovery returned no live models; keeping cached models in place",
243 {
244 cachedCount: cachedModels.length,
245 },
246 );
247 return cachedModels;
248 }
249
250 setCachedModels([]);
251 return [];
252 },
253 staleTime: 5 * 60 * 1000,
254 retry: 1,
255 retryDelay: 1000,
256 initialData: cachedModels || undefined,
257 select: (data) => {
258 const savedLocalModel = getSavedLocalModel();
259 const localModels = mergeModels(
260 data,
261 savedLocalModel ? [savedLocalModel] : [],
262 );
263
264 return {
265 localModels,
266 downloadableModels,
267 showDownloadable: downloadableModels.length > 0,
268 };
269 },
270 });
271
272 useEffect(() => {
273 if (isInitialLoad) {
274 setIsInitialLoad(false);
275 }
276 }, [isInitialLoad]);
277
278 return {
279 ...query,
280 isInitialLoad,
281 };
282 }
283
283 lines TYPESCRIPT