gramps
b3746949a9
model_pull.py: ensure_model() checks llama-server availability at startup, falls back to Ollama pull API if model not found. Integrated into app.py lifespan after assess_hardware(). 11 tests cover all paths: available/unreachable/pull success/fail.
143 lines
4.7 KiB
Python
143 lines
4.7 KiB
Python
import asyncio
|
|
|
|
import httpx
|
|
|
|
from model_pull import ensure_model, _model_available_on_llama, _model_available_on_ollama, _pull_via_ollama
|
|
|
|
|
|
class _MockAsyncResponse:
|
|
def __init__(self, status_code=200, json_data=None):
|
|
self.status_code = status_code
|
|
self._json_data = json_data or {}
|
|
|
|
def json(self):
|
|
return self._json_data
|
|
|
|
|
|
class _MockStreamResponse:
|
|
def __init__(self, status_code=200, lines=None):
|
|
self.status_code = status_code
|
|
self._lines = lines or []
|
|
|
|
async def __aenter__(self):
|
|
return self
|
|
|
|
async def __aexit__(self, *args):
|
|
pass
|
|
|
|
def aiter_lines(self):
|
|
class _AIter:
|
|
def __init__(self, lines):
|
|
self._lines = iter(lines)
|
|
def __aiter__(self):
|
|
return self
|
|
async def __anext__(self):
|
|
try:
|
|
return next(self._lines)
|
|
except StopIteration:
|
|
raise StopAsyncIteration
|
|
return _AIter(self._lines)
|
|
|
|
|
|
def _mock_models_on_llama(*models):
|
|
async def _get(*args, **kwargs):
|
|
return _MockAsyncResponse(json_data={
|
|
"data": [{"id": m} for m in models]
|
|
})
|
|
return _get
|
|
|
|
|
|
async def _mock_ollama_available(*args, **kwargs):
|
|
return _MockAsyncResponse(status_code=200, json_data={"name": "qwen2.5:latest"})
|
|
|
|
|
|
async def _mock_ollama_unavailable(*args, **kwargs):
|
|
raise httpx.ConnectError("refused")
|
|
|
|
|
|
def _mock_ollama_pull_ok(*args, **kwargs):
|
|
return _MockStreamResponse(200, [
|
|
'{"status": "pulling manifest"}',
|
|
'{"status": "success"}',
|
|
])
|
|
|
|
|
|
def _mock_ollama_pull_fail(*args, **kwargs):
|
|
return _MockStreamResponse(500, [])
|
|
|
|
|
|
def _mock_ollama_connect_error(*args, **kwargs):
|
|
raise httpx.ConnectError("refused")
|
|
|
|
|
|
def test_available_on_llama(monkeypatch):
|
|
monkeypatch.setattr(httpx.AsyncClient, "get", _mock_models_on_llama("qwen2.5-7b-instruct"))
|
|
assert asyncio.run(_model_available_on_llama("qwen2.5-7b-instruct")) is True
|
|
assert asyncio.run(_model_available_on_llama("nonexistent-model")) is False
|
|
|
|
|
|
def test_available_on_llama_unreachable(monkeypatch):
|
|
async def _connect_error(*args, **kwargs):
|
|
raise httpx.ConnectError("refused")
|
|
monkeypatch.setattr(httpx.AsyncClient, "get", _connect_error)
|
|
assert asyncio.run(_model_available_on_llama("qwen2.5-7b-instruct")) is False
|
|
|
|
|
|
def test_available_on_ollama(monkeypatch):
|
|
monkeypatch.setattr(httpx.AsyncClient, "post", _mock_ollama_available)
|
|
assert asyncio.run(_model_available_on_ollama("qwen2.5:latest")) is True
|
|
|
|
|
|
def test_available_on_ollama_unreachable(monkeypatch):
|
|
monkeypatch.setattr(httpx.AsyncClient, "post", _mock_ollama_unavailable)
|
|
assert asyncio.run(_model_available_on_ollama("qwen2.5:latest")) is False
|
|
|
|
|
|
def test_pull_via_ollama_success(monkeypatch):
|
|
monkeypatch.setattr(httpx.AsyncClient, "stream", _mock_ollama_pull_ok)
|
|
assert asyncio.run(_pull_via_ollama("qwen2.5:latest")) is True
|
|
|
|
|
|
def test_pull_via_ollama_fail(monkeypatch):
|
|
monkeypatch.setattr(httpx.AsyncClient, "stream", _mock_ollama_pull_fail)
|
|
assert asyncio.run(_pull_via_ollama("qwen2.5:latest")) is False
|
|
|
|
|
|
def test_pull_via_ollama_connect_error(monkeypatch):
|
|
monkeypatch.setattr(httpx.AsyncClient, "stream", _mock_ollama_connect_error)
|
|
assert asyncio.run(_pull_via_ollama("qwen2.5:latest")) is False
|
|
|
|
|
|
def test_ensure_model_already_on_llama(monkeypatch):
|
|
monkeypatch.setattr(httpx.AsyncClient, "get", _mock_models_on_llama("qwen2.5-7b-instruct"))
|
|
assert asyncio.run(ensure_model("qwen2.5-7b-instruct")) is True
|
|
|
|
|
|
def test_ensure_model_not_on_llama_but_on_ollama(monkeypatch):
|
|
async def _get(*args, **kwargs):
|
|
return _MockAsyncResponse(json_data={"data": []})
|
|
|
|
monkeypatch.setattr(httpx.AsyncClient, "get", _get)
|
|
monkeypatch.setattr(httpx.AsyncClient, "post", _mock_ollama_available)
|
|
assert asyncio.run(ensure_model("qwen2.5-7b-instruct")) is True
|
|
|
|
|
|
def test_ensure_model_needs_pull(monkeypatch):
|
|
async def _get(*args, **kwargs):
|
|
return _MockAsyncResponse(json_data={"data": []})
|
|
|
|
monkeypatch.setattr(httpx.AsyncClient, "get", _get)
|
|
monkeypatch.setattr(httpx.AsyncClient, "post", _mock_ollama_unavailable)
|
|
monkeypatch.setattr(httpx.AsyncClient, "stream", _mock_ollama_pull_ok)
|
|
assert asyncio.run(ensure_model("qwen2.5-7b-instruct")) is True
|
|
|
|
|
|
def test_ensure_model_pull_fails(monkeypatch):
|
|
async def _get(*args, **kwargs):
|
|
return _MockAsyncResponse(json_data={"data": []})
|
|
|
|
monkeypatch.setattr(httpx.AsyncClient, "get", _get)
|
|
monkeypatch.setattr(httpx.AsyncClient, "post", _mock_ollama_unavailable)
|
|
monkeypatch.setattr(httpx.AsyncClient, "stream", _mock_ollama_pull_fail)
|
|
assert asyncio.run(ensure_model("qwen2.5-7b-instruct")) is False
|