feat: cancel an Ollama model download and uninstall models

app/api/routes/llm.py

@@ -6,13 +6,16 @@
 from fastapi import APIRouter
 from fastapi import Depends
 from fastapi.responses import StreamingResponse
+from starlette.background import BackgroundTask
 from sqlmodel import Session
 
 from app.api.deps import get_db
 from app.api.schemas.enums import LLMProvider
 from app.api.schemas.system import (
     DownloadableModel,
     LLMModelChoice,
+    LLMModelDeleteRequest,
+    LLMModelDeleteResult,
     LLMModels,
     LLMProviderModels,
     LLMPullRequest,
@@ -21,7 +24,11 @@
     ModelInfo,
 )
 from app.services import llm
-from app.services.llm.errors import LLMNotAvailableError
+from app.services.llm.errors import (
+    LLMModelInUseError,
+    LLMModelNotInstalledError,
+    LLMNotAvailableError,
+)
 from app.services.llm.ollama import OllamaLLM
 
 router = APIRouter(prefix="/system/llm", tags=["system"])
@@ -96,4 +103,50 @@ def pull_llm_model(req: LLMPullRequest, db: Session = Depends(get_db)):
         req.model,
         on_success=lambda: registry.provider_catalog(LLMProvider.ollama, refresh=True),
     )
-    return StreamingResponse(stream, media_type="application/x-ndjson")
+    return StreamingResponse(
+        stream,
+        media_type="application/x-ndjson",
+        background=BackgroundTask(stream.close),
+    )
+
+
+@router.delete(
+    "/models",
+    response_model=LLMModelDeleteResult,
+    responses={
+        404: {"description": "Ollama does not have the model installed"},
+        409: {"description": "The model is a default or currently being pulled"},
+        422: {"description": "Invalid request or model is not allowed"},
+        502: {"description": "Ollama returned an unexpected response"},
+        503: {"description": "Ollama service unavailable"},
+    },
+    summary="Delete an installed Ollama model",
+)
+def delete_llm_model(req: LLMModelDeleteRequest, db: Session = Depends(get_db)):
+    registry = llm.effective_registry(db)
+    registry.validate_pull(req.provider, req.model)
+
+    client = registry.client(LLMProvider.ollama)
+    if not isinstance(client, OllamaLLM):
+        raise LLMNotAvailableError("Model deletion is only supported for Ollama.")
+
+    catalog = registry.provider_catalog(LLMProvider.ollama, refresh=True)
+    if catalog.error:
+        raise type(catalog.error)(catalog.error.message)
+    if not any(model.name == req.model and model.installed is True for model in catalog.models):
+        raise LLMModelNotInstalledError(f"Ollama does not have model {req.model} installed.")
+
+    ollama_default = registry.settings.providers[LLMProvider.ollama].default_model
+    if (
+        registry.default.provider is LLMProvider.ollama
+        and registry.default.model == req.model
+    ):
+        raise LLMModelInUseError(f"Cannot delete {req.model}: it is the global default model.")
+    if ollama_default == req.model:
+        raise LLMModelInUseError(f"Cannot delete {req.model}: it is the Ollama provider default model.")
+    if client.is_model_pulling(req.model):
+        raise LLMModelInUseError(f"Cannot delete {req.model}: a download is currently active.")
+
+    client.delete_model(req.model)
+    registry.provider_catalog(LLMProvider.ollama, refresh=True)
+    return LLMModelDeleteResult(model=req.model, deleted=True)

app/api/schemas/system.py

@@ -103,6 +103,21 @@ def _validate_model_name(cls, value: str) -> str:
         return validate_model_name(value)
 
 
+class LLMModelDeleteRequest(BaseModel):
+    model: str
+    provider: str | None = "ollama"
+
+    @field_validator("model")
+    @classmethod
+    def _validate_model_name(cls, value: str) -> str:
+        return validate_model_name(value)
+
+
+class LLMModelDeleteResult(BaseModel):
+    model: str
+    deleted: bool
+
+
 class LLMTestRequest(BaseModel):
     provider: LLMProvider | None = None
     model: str | None = None

app/services/llm/__init__.py

@@ -20,6 +20,8 @@
     LLMConfigError,
     LLMError,
     LLMExternalConsentError,
+    LLMModelInUseError,
+    LLMModelNotInstalledError,
     LLMModelNotFoundError,
     LLMNotAvailableError,
     LLMRateLimitError,
@@ -48,6 +50,8 @@ def get_registry() -> LLMRegistry:
     "LLMError",
     "LLMExternalConsentError",
     "LLMHealth",
+    "LLMModelInUseError",
+    "LLMModelNotInstalledError",
     "LLMModelNotFoundError",
     "LLMNotAvailableError",
     "LLMProvider",

app/services/llm/errors.py

@@ -68,6 +68,20 @@ class LLMModelNotFoundError(LLMError):
     error_code = "LLM_MODEL_NOT_FOUND"
 
 
+class LLMModelNotInstalledError(LLMError):
+    """Ollama does not have the requested model installed."""
+
+    status_code = 404
+    error_code = "LLM_MODEL_NOT_INSTALLED"
+
+
+class LLMModelInUseError(LLMError):
+    """The model cannot be deleted while it is a default or a pull is active."""
+
+    status_code = 409
+    error_code = "LLM_MODEL_IN_USE"
+
+
 class LLMResponseError(LLMError):
     """The server answered, but not with a usable JSON object."""
 

app/services/llm/ollama.py

@@ -6,17 +6,94 @@
 
 import json
 import logging
+import threading
 import time
 from collections.abc import Callable, Iterator
 
 import requests
 
 from app.api.schemas.enums import HealthState, LLMProvider
 from app.services.llm.base import LLM, PROBE_TIMEOUT_SECONDS, LLMHealth, ModelInfo
-from app.services.llm.errors import LLMError, LLMModelNotFoundError, LLMResponseError, LLMUnavailableError
+from app.services.llm.errors import (
+    LLMError,
+    LLMModelNotFoundError,
+    LLMModelNotInstalledError,
+    LLMResponseError,
+    LLMUnavailableError,
+)
 
 logger = logging.getLogger(__name__)
 
+_ACTIVE_PULL_MODELS: set[str] = set()
+_ACTIVE_PULL_COUNTS: dict[str, int] = {}
+_ACTIVE_PULL_LOCK = threading.Lock()
+
+
+def _register_active_pull(model: str) -> None:
+    with _ACTIVE_PULL_LOCK:
+        _ACTIVE_PULL_COUNTS[model] = _ACTIVE_PULL_COUNTS.get(model, 0) + 1
+        _ACTIVE_PULL_MODELS.add(model)
+
+
+def _unregister_active_pull(model: str) -> None:
+    with _ACTIVE_PULL_LOCK:
+        count = _ACTIVE_PULL_COUNTS.get(model, 0)
+        if count <= 1:
+            _ACTIVE_PULL_COUNTS.pop(model, None)
+            _ACTIVE_PULL_MODELS.discard(model)
+        else:
+            _ACTIVE_PULL_COUNTS[model] = count - 1
+
+
+class _PullStream(Iterator[str]):
+    def __init__(self, response: requests.Response, iterator: Iterator[str], model: str):
+        self._response = response
+        self._iterator = iterator
+        self._model = model
+        self._started = False
+        self._finished = False
+        self._closed = False
+
+    def __iter__(self) -> "_PullStream":
+        return self
+
+    def __next__(self) -> str:
+        if self._closed:
+            raise StopIteration
+        self._started = True
+        try:
+            return next(self._iterator)
+        except StopIteration:
+            self._finished = True
+            self.close()
+            raise
+        except BaseException:
+            self.close()
+            raise
+
+    def close(self) -> None:
+        if self._closed:
+            return
+        self._closed = True
+        try:
+            close = getattr(self._iterator, "close", None)
+            if close:
+                close()
+        finally:
+            try:
+                self._response.close()
+            finally:
+                _unregister_active_pull(self._model)
+                if not self._started and not self._finished:
+                    logger.info("pull of %s cancelled by the client", self._model)
+
+    def __del__(self):
+        try:
+            self.close()
+        except Exception:
+            logger.exception("Failed to close Ollama pull stream for %s", self._model)
+
+
 OPTIONS = {
     "num_ctx": 32768,
     # Backstop for runaway generation; the largest real output (ICS 209) is ~8k tokens
@@ -95,14 +172,58 @@ def _to_models(tags: dict, running: list[str]) -> list[ModelInfo]:
             if m.get("name")
         ]
 
-    def pull_model(self, model: str, on_success: Callable[[], None] | None = None) -> Iterator[str]:
+    def pull_model(self, model: str, on_success: Callable[[], None] | None = None) -> "_PullStream":
         """Stream pull progress from Ollama as newline-delimited JSON.
 
         Opens the connection before returning the generator so a down Ollama raises
         LLMUnavailableError (HTTP 503) instead of committing a 200 response.
         """
-        response = self._open_pull_stream(model)
-        return self._stream_pull(response, on_success=on_success)
+        _register_active_pull(model)
+        try:
+            response = self._open_pull_stream(model)
+        except BaseException:
+            _unregister_active_pull(model)
+            raise
+        try:
+            return self._stream_pull(response, on_success=on_success, model=model)
+        except BaseException:
+            try:
+                response.close()
+            finally:
+                _unregister_active_pull(model)
+            raise
+
+    @staticmethod
+    def is_model_pulling(model: str) -> bool:
+        with _ACTIVE_PULL_LOCK:
+            return model in _ACTIVE_PULL_MODELS
+
+    def delete_model(self, model: str) -> None:
+        """Delete an installed model from Ollama."""
+        try:
+            response = requests.delete(
+                self.host + "/api/delete",
+                json={"model": model},
+                timeout=PROBE_TIMEOUT_SECONDS,
+            )
+        except requests.Timeout as exc:
+            raise LLMUnavailableError(f"Ollama at {self.host} timed out deleting {model}") from exc
+        except requests.RequestException as exc:
+            raise LLMUnavailableError(f"Could not reach Ollama at {self.host}. Check OLLAMA_HOST: {exc}") from exc
+
+        try:
+            if response.status_code == 404:
+                raise LLMModelNotInstalledError(f"Ollama does not have model {model} installed.")
+            if response.status_code >= 500:
+                raise LLMUnavailableError(
+                    f"Ollama had a server error {response.status_code}: {_reason(response)}"
+                )
+            if not 200 <= response.status_code < 300:
+                raise LLMResponseError(
+                    f"Ollama refused to delete {model} with {response.status_code}: {_reason(response)}"
+                )
+        finally:
+            response.close()
 
     def _open_pull_stream(self, model: str) -> requests.Response:
         url = self.host + "/api/pull"
@@ -121,7 +242,15 @@ def _open_pull_stream(self, model: str) -> requests.Response:
             raise LLMUnavailableError(f"Could not reach Ollama at {self.host}. Check OLLAMA_HOST: {exc}") from exc
 
     def _stream_pull(
-        self, response: requests.Response, on_success: Callable[[], None] | None = None
+        self, response: requests.Response, on_success: Callable[[], None] | None = None, model: str = ""
+    ) -> _PullStream:
+        return _PullStream(response, self._iterate_pull(response, on_success, model), model)
+
+    def _iterate_pull(
+        self,
+        response: requests.Response,
+        on_success: Callable[[], None] | None = None,
+        model: str = "",
     ) -> Iterator[str]:
         saw_error = False
         success_line: str | None = None
@@ -164,6 +293,9 @@ def inspect_line(line: str) -> bool:
                 if buffer.strip():
                     if not inspect_line(buffer):
                         yield buffer.strip() + "\n"
+        except GeneratorExit:
+            logger.info("pull of %s cancelled by the client", model)
+            raise
         except requests.RequestException as exc:
             saw_error = True
             yield json.dumps({"error": str(exc)}) + "\n"