989a4f3e98
Empirical A/B/C measurement against the live RTX 5050 Ollama stack (qwen2.5:3b + nomic-embed-text) showed keep_alive=0 unloads the embed model ~2s after every call, so each turn after a brief idle gap pays a cold reload. VRAM is not the constraint (~4.4-4.7 GB free with both models resident) and keep_alive=0 never evicted the chat model, so CPU embedding (num_gpu=0) gave no benefit. A short positive keep_alive is the fastest of the three: it keeps the ~0.3 GB embed model resident across consecutive turns at negligible VRAM cost. Add tests/test_embeddings.py covering the warm-across-turns behaviour. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
53 lines
2.2 KiB
Python
53 lines
2.2 KiB
Python
"""Tests for the Ollama embedding client.
|
|
|
|
Behaviour under test: the embedding request keeps the embed model warm across
|
|
consecutive conversation turns. With ``keep_alive=0`` Ollama unloads the embed
|
|
model ~2s after every call, so each turn after a short idle gap pays a cold
|
|
reload. A short positive ``keep_alive`` keeps it resident between turns at a
|
|
negligible VRAM cost (nomic-embed-text is ~0.3 GB).
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from jarvis.memory.embeddings import get_embedding
|
|
|
|
|
|
def _mock_response(vec):
|
|
resp = MagicMock()
|
|
resp.raise_for_status.return_value = None
|
|
resp.json.return_value = {"embedding": vec}
|
|
return resp
|
|
|
|
|
|
def test_get_embedding_posts_to_embeddings_endpoint():
|
|
with patch("jarvis.memory.embeddings.requests") as mock_requests:
|
|
mock_requests.post.return_value = _mock_response([0.1, 0.2, 0.3])
|
|
vec = get_embedding("hello", "http://localhost:11434", "nomic-embed-text")
|
|
|
|
assert vec == [0.1, 0.2, 0.3]
|
|
args, kwargs = mock_requests.post.call_args
|
|
assert args[0].endswith("/api/embeddings")
|
|
assert kwargs["json"]["model"] == "nomic-embed-text"
|
|
assert kwargs["json"]["prompt"] == "hello"
|
|
|
|
|
|
def test_get_embedding_keeps_model_warm_between_turns():
|
|
"""The request must not unload the model after each call (keep_alive > 0)."""
|
|
with patch("jarvis.memory.embeddings.requests") as mock_requests:
|
|
mock_requests.post.return_value = _mock_response([0.0])
|
|
get_embedding("warm me", "http://localhost:11434", "nomic-embed-text")
|
|
|
|
_, kwargs = mock_requests.post.call_args
|
|
keep_alive = kwargs["json"].get("keep_alive")
|
|
# A falsy/zero keep_alive evicts the model immediately, forcing a cold
|
|
# reload on the next turn. Anything truthy positive keeps it resident.
|
|
assert keep_alive, f"embedding keep_alive should be a positive duration, got {keep_alive!r}"
|
|
assert keep_alive != 0
|
|
|
|
|
|
def test_get_embedding_returns_none_on_error():
|
|
with patch("jarvis.memory.embeddings.requests") as mock_requests:
|
|
mock_requests.post.side_effect = RuntimeError("boom")
|
|
assert get_embedding("x", "http://localhost:11434", "nomic-embed-text") is None
|