diff --git a/admin/app/controllers/ollama_controller.ts b/admin/app/controllers/ollama_controller.ts index cac3994..fe50783 100644 --- a/admin/app/controllers/ollama_controller.ts +++ b/admin/app/controllers/ollama_controller.ts @@ -5,7 +5,7 @@ import { RagService } from '#services/rag_service' import Service from '#models/service' import KVStore from '#models/kv_store' import { modelNameSchema } from '#validators/download' -import { chatSchema, getAvailableModelsSchema } from '#validators/ollama' +import { chatSchema, getAvailableModelsSchema, unloadChatModelsSchema } from '#validators/ollama' import { inject } from '@adonisjs/core' import type { HttpContext } from '@adonisjs/core/http' import { RAG_CONTEXT_LIMITS, SYSTEM_PROMPTS } from '../../constants/ollama.js' @@ -33,6 +33,19 @@ export default class OllamaController { }) } + /** + * Send Ollama `keep_alive: 0` hints to every currently-loaded chat model + * except the embedding model and (optionally) a target model to preserve. + * Used by the chat UI to enforce the "one chat model at a time" invariant + * on model-switch, session-switch, and page-load. Best-effort: a failure + * here should not block the calling flow. + */ + async unloadChatModels({ request, response }: HttpContext) { + const { targetModel } = await request.validateUsing(unloadChatModelsSchema) + const unloaded = await this.ollamaService.unloadAllChatModelsExcept(targetModel ?? null) + return response.status(200).json({ unloaded }) + } + async chat({ request, response }: HttpContext) { const reqData = await request.validateUsing(chatSchema) diff --git a/admin/app/services/ollama_service.ts b/admin/app/services/ollama_service.ts index 075fda4..8f805be 100644 --- a/admin/app/services/ollama_service.ts +++ b/admin/app/services/ollama_service.ts @@ -594,6 +594,80 @@ export class OllamaService { } } + /** + * Embedding model name exempt from the chat-model unload sweep. Hardcoded + * to avoid a circular import with `RagService.EMBEDDING_MODEL`; keep in + * sync with that constant if it ever changes. + */ + private static readonly EMBEDDING_MODEL_EXEMPT = 'nomic-embed-text:v1.5' + + /** + * Enforces the "at most one chat model resident in VRAM" invariant by firing + * `keep_alive: 0` against every currently-loaded model except (a) the + * embedding model (always exempt) and (b) `targetModel` (the one we want + * loaded next — leaving it alone preserves a hot model when the target is + * already loaded). + * + * Best-effort: queries `/api/ps` and POSTs unload hints in parallel. Network + * or Ollama errors are swallowed and logged — neither chat nor page-load + * should fail just because the unload housekeeping didn't go through. + * + * Returns the list of model names that were sent the unload hint, so the + * caller (and tests) can confirm what actually happened. + * + * Pass `targetModel: null` to unload every chat model (used for the future + * "free up VRAM" path; not exposed yet but the helper supports it). + * + * Note that `keep_alive: 0` is a post-completion hint, not a force-kill — + * Ollama defers eviction until the runner is idle, so in-flight inference + * on the same model is never interrupted. See the design doc for the race + * analysis behind this. + */ + public async unloadAllChatModelsExcept(targetModel: string | null): Promise { + await this._ensureDependencies() + if (!this.baseUrl) return [] + + let loadedModels: string[] = [] + try { + const response = await axios.get(`${this.baseUrl}/api/ps`, { timeout: 5000 }) + loadedModels = (response.data?.models ?? []) + .map((m: { name?: string }) => m.name) + .filter((name: unknown): name is string => typeof name === 'string') + } catch (err: any) { + logger.warn( + `[OllamaService] unloadAllChatModelsExcept: /api/ps unreachable, skipping unload sweep: ${err?.message ?? err}` + ) + return [] + } + + const toUnload = loadedModels.filter( + (name) => name !== OllamaService.EMBEDDING_MODEL_EXEMPT && name !== targetModel + ) + + await Promise.all( + toUnload.map(async (modelName) => { + try { + await axios.post( + `${this.baseUrl}/api/generate`, + { model: modelName, prompt: '', keep_alive: 0 }, + { timeout: 10000 } + ) + } catch (err: any) { + logger.warn( + `[OllamaService] Failed to send unload hint for ${modelName}: ${err?.message ?? err}` + ) + } + }) + ) + + if (toUnload.length > 0) { + logger.info( + `[OllamaService] Sent unload hint for ${toUnload.length} chat model(s): ${toUnload.join(', ')}` + ) + } + return toUnload + } + public async getModels(includeEmbeddings = false): Promise { await this._ensureDependencies() if (!this.baseUrl) { diff --git a/admin/app/validators/ollama.ts b/admin/app/validators/ollama.ts index d83ed3b..1784ada 100644 --- a/admin/app/validators/ollama.ts +++ b/admin/app/validators/ollama.ts @@ -14,6 +14,12 @@ export const chatSchema = vine.compile( }) ) +export const unloadChatModelsSchema = vine.compile( + vine.object({ + targetModel: vine.string().trim().minLength(1).nullable().optional(), + }) +) + export const getAvailableModelsSchema = vine.compile( vine.object({ sort: vine.enum(['pulls', 'name'] as const).optional(), diff --git a/admin/inertia/components/chat/index.tsx b/admin/inertia/components/chat/index.tsx index 04f7c15..9be60d4 100644 --- a/admin/inertia/components/chat/index.tsx +++ b/admin/inertia/components/chat/index.tsx @@ -33,6 +33,8 @@ export default function Chat({ const [activeSessionId, setActiveSessionId] = useState(null) const [messages, setMessages] = useState([]) const [selectedModel, setSelectedModel] = useState('') + const [pendingModelSwitch, setPendingModelSwitch] = useState(null) + const pageLoadNormalizedRef = useRef(false) const [isStreamingResponse, setIsStreamingResponse] = useState(false) const streamAbortRef = useRef(null) @@ -151,6 +153,62 @@ export default function Chat({ } }, [selectedModel]) + // Page-load normalization: enforce the "one chat model at a time" invariant + // when the chat page first mounts. Anything stacked from a prior session + // gets `keep_alive: 0` so it can be evicted; the embedding model is exempt + // server-side. We wait for `selectedModel` to be populated by the + // first-installed / lastModel effect so the request has a target to preserve. + useEffect(() => { + if (!enabled) return + if (!selectedModel) return + if (pageLoadNormalizedRef.current) return + pageLoadNormalizedRef.current = true + api.unloadChatModels(selectedModel).catch((err) => { + console.warn('Failed to normalize loaded models on chat-page mount:', err) + }) + }, [enabled, selectedModel]) + + const handleUserSelectedModel = useCallback( + (newModel: string) => { + if (newModel === selectedModel) return + // No active chat session yet → no conversation to lose, no popup needed. + // Just update the dropdown silently; the next "New Chat" will use it. + if (!activeSessionId) { + setSelectedModel(newModel) + return + } + // Active session: defer the actual model swap until the user confirms. + // Setting `pendingModelSwitch` drives the dropdown's effective value + // *and* opens the confirm modal — clearing it on cancel reverts the + // visible selection without us having to touch `selectedModel`. + setPendingModelSwitch(newModel) + }, + [selectedModel, activeSessionId] + ) + + const handleConfirmModelSwitch = useCallback(async () => { + const newModel = pendingModelSwitch + if (!newModel) return + // Best-effort unload of the previously-active chat model. Fire-and-forget: + // Ollama queues the eviction until the runner is idle, so an in-flight + // request on the old model finishes cleanly. We don't await this before + // clearing the session — UI responsiveness wins over housekeeping. + api.unloadChatModels(newModel).catch((err) => { + console.warn('Failed to unload previous chat model:', err) + }) + setSelectedModel(newModel) + setPendingModelSwitch(null) + // Clear the active session and messages — the next user message will + // lazily create a new session via the existing handleSendMessage path, + // which already calls api.createChatSession with `selectedModel`. + setActiveSessionId(null) + setMessages([]) + }, [pendingModelSwitch]) + + const handleCancelModelSwitch = useCallback(() => { + setPendingModelSwitch(null) + }, []) + const handleNewChat = useCallback(() => { // Just clear the active session and messages - don't create a session yet setActiveSessionId(null) @@ -202,8 +260,19 @@ export default function Chat({ if (sessionData?.model) { setSelectedModel(sessionData.model) } + + // Enforce the one-chat-model-at-a-time invariant: ask the backend to + // unload anything that isn't the target session's model. Fire-and-forget; + // this is housekeeping. Note we pass the *session's* model here rather + // than reading `selectedModel`, because setSelectedModel above is async + // and the effect-driven page-load normalize wouldn't catch a sidebar + // click after the first render. + const targetModel = sessionData?.model ?? selectedModel ?? null + api.unloadChatModels(targetModel).catch((err) => { + console.warn('Failed to unload non-target chat models on session switch:', err) + }) }, - [installedModels, queryClient] + [installedModels, queryClient, selectedModel] ) const handleSendMessage = useCallback( @@ -352,6 +421,23 @@ export default function Chat({ ) return ( + <> + {pendingModelSwitch && ( + +

+ Switching to {pendingModelSwitch} will start a new chat. Your current + conversation stays available in the sidebar. +

+ + )}
setSelectedModel(e.target.value)} + value={pendingModelSwitch ?? selectedModel} + onChange={(e) => handleUserSelectedModel(e.target.value)} className="px-3 py-1.5 border border-border-default rounded-lg text-sm focus:outline-none focus:ring-2 focus:ring-desert-green focus:border-transparent bg-surface-primary" > {installedModels.map((model) => ( @@ -433,5 +519,6 @@ export default function Chat({ />
+ ) } diff --git a/admin/inertia/lib/api.ts b/admin/inertia/lib/api.ts index cd7eba8..b6372e9 100644 --- a/admin/inertia/lib/api.ts +++ b/admin/inertia/lib/api.ts @@ -272,6 +272,24 @@ class API { })() } + /** + * Ask the backend to send Ollama `keep_alive: 0` to every currently-loaded + * chat model except `targetModel` (and the embedding model, which is always + * exempt server-side). Fire-and-forget — the chat UI doesn't await this + * before creating a new session, since unload is housekeeping. + * + * Pass `null` to unload every chat model. + */ + async unloadChatModels(targetModel: string | null) { + return catchInternal(async () => { + const response = await this.client.post<{ unloaded: string[] }>( + '/ollama/unload-chat-models', + { targetModel } + ) + return response.data + })() + } + async getAvailableModels(params: { query?: string; recommendedOnly?: boolean; limit?: number; force?: boolean }) { return catchInternal(async () => { const response = await this.client.get<{ diff --git a/admin/start/routes.ts b/admin/start/routes.ts index 1868b0e..8435778 100644 --- a/admin/start/routes.ts +++ b/admin/start/routes.ts @@ -118,6 +118,7 @@ router router.post('/models', [OllamaController, 'dispatchModelDownload']) router.delete('/models', [OllamaController, 'deleteModel']) router.get('/installed-models', [OllamaController, 'installedModels']) + router.post('/unload-chat-models', [OllamaController, 'unloadChatModels']) router.post('/configure-remote', [OllamaController, 'configureRemote']) router.get('/remote-status', [OllamaController, 'remoteStatus']) })