Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 8 additions & 1 deletion api/rag/rag.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,12 @@ def _get_rag_semaphore() -> asyncio.Semaphore:
return _RAG_PREPARE_SEMAPHORE


def _with_default_tag(model_name: str) -> str:
"""Ollama lists an untagged model as '<name>:latest', so 'nomic-embed-text' and
'nomic-embed-text:latest' are the same model. A ':' before the last '/' is a registry port."""
return model_name if ":" in model_name.rsplit("/", 1)[-1] else f"{model_name}:latest"


def check_ollama_model_exists(model_name: str, ollama_host: str | None = None) -> bool:
"""
Check if an Ollama model exists before attempting to use it.
Expand All @@ -54,7 +60,8 @@ def check_ollama_model_exists(model_name: str, ollama_host: str | None = None) -
# Remove /api prefix if present and add it back
ollama_host = ollama_host.removesuffix("/api")
ret: ollama.ListResponse = ollama.Client(host=ollama_host, timeout=5).list()
is_available = any(model_name == model.model for model in ret.models)
wanted = _with_default_tag(model_name)
is_available = any(wanted == _with_default_tag(model.model) for model in ret.models)
if is_available:
logger.info("Ollama model '%s' is available", model_name)
else:
Expand Down
30 changes: 30 additions & 0 deletions tests/unit/test_ollama_model_check.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,30 @@
from types import SimpleNamespace

import ollama
import pytest

from api.rag.rag import check_ollama_model_exists


def _fake_list(monkeypatch, names):
response = SimpleNamespace(models=[SimpleNamespace(model=name) for name in names])
monkeypatch.setattr(ollama.Client, "list", lambda self: response)


@pytest.mark.parametrize(
"configured, installed, expected",
[
# Ollama reports untagged pulls as ':latest' (issue #579).
("nomic-embed-text", ["nomic-embed-text:latest"], True),
("nomic-embed-text:latest", ["nomic-embed-text:latest"], True),
("qwen3:8b", ["qwen3:8b"], True),
("qwen3:8b", ["qwen3:1.7b"], False),
("nomic-embed-text", ["nomic-embed-text:v1.5"], False),
("nomic-embed-text:v1.5", ["nomic-embed-text:latest"], False),
("localhost:5000/embed", ["localhost:5000/embed:latest"], True),
("nomic-embed-text", [], False),
],
)
def test_check_ollama_model_exists(monkeypatch, configured, installed, expected):
_fake_list(monkeypatch, installed)
assert check_ollama_model_exists(configured, "http://localhost:11434") is expected