110 lines
3.7 KiB
Python
110 lines
3.7 KiB
Python
"""Ollama manifest use: chat|memory parsing."""
|
|
|
|
from pathlib import Path
|
|
|
|
from gpu_rent.llm_runtime import (
|
|
MEMORY_EMBED_MODEL,
|
|
ensure_memory_model_entries,
|
|
ollama_roles_payload,
|
|
parse_ollama_models,
|
|
preferred_ollama_model,
|
|
write_ollama_models_preset,
|
|
)
|
|
|
|
|
|
def test_parse_use_chat_and_memory(tmp_path: Path):
|
|
path = tmp_path / "ollama-models.yaml"
|
|
path.write_text(
|
|
"""
|
|
models:
|
|
- name: chat-model:7b
|
|
use: chat
|
|
default: true
|
|
- name: nomic-embed-text
|
|
use: memory
|
|
""",
|
|
encoding="utf-8",
|
|
)
|
|
entries = parse_ollama_models(path)
|
|
assert len(entries) == 2
|
|
assert entries[0].use == "chat" and entries[0].default
|
|
assert entries[1].use == "memory" and not entries[1].default
|
|
assert preferred_ollama_model(path) == "chat-model:7b"
|
|
roles = ollama_roles_payload(entries)
|
|
assert roles["chat"] == ["chat-model:7b"]
|
|
assert roles["memory"] == ["nomic-embed-text"]
|
|
assert roles["default_chat"] == "chat-model:7b"
|
|
|
|
|
|
def test_ensure_memory_appended():
|
|
from gpu_rent.llm_runtime import OllamaModelEntry
|
|
|
|
entries = [OllamaModelEntry(name="vl:7b", default=True, use="chat")]
|
|
out = ensure_memory_model_entries(entries)
|
|
assert any(e.use == "memory" and e.name == MEMORY_EMBED_MODEL for e in out)
|
|
|
|
|
|
def test_write_preset_includes_memory(tmp_path: Path):
|
|
path = tmp_path / "m.yaml"
|
|
write_ollama_models_preset(path, "recommended")
|
|
text = path.read_text(encoding="utf-8")
|
|
assert "use: chat" in text
|
|
assert "use: memory" in text
|
|
assert MEMORY_EMBED_MODEL in text
|
|
entries = parse_ollama_models(path)
|
|
assert preferred_ollama_model(path) is not None
|
|
assert preferred_ollama_model(path) != MEMORY_EMBED_MODEL
|
|
|
|
|
|
def test_cpu_ollama_tag_does_not_stack():
|
|
from gpu_rent.llm_runtime import (
|
|
cpu_ollama_tag,
|
|
finalize_memory_role_tags,
|
|
is_stacked_cpu_ollama_tag,
|
|
ollama_model_root,
|
|
pick_ollama_from_tag,
|
|
)
|
|
|
|
assert ollama_model_root("nomic-embed-text-cpu-cpu-cpu") == "nomic-embed-text"
|
|
assert cpu_ollama_tag("nomic-embed-text") == "nomic-embed-text-cpu"
|
|
assert cpu_ollama_tag("nomic-embed-text-cpu") == "nomic-embed-text-cpu"
|
|
assert cpu_ollama_tag("nomic-embed-text-cpu-cpu:latest") == "nomic-embed-text-cpu"
|
|
assert is_stacked_cpu_ollama_tag("nomic-embed-text-cpu-cpu")
|
|
assert not is_stacked_cpu_ollama_tag("nomic-embed-text-cpu")
|
|
|
|
have = {
|
|
"nomic-embed-text",
|
|
"nomic-embed-text-cpu",
|
|
"nomic-embed-text-cpu-cpu",
|
|
"nomic-embed-text-cpu-cpu-cpu",
|
|
}
|
|
assert pick_ollama_from_tag(have, "nomic-embed-text") == "nomic-embed-text"
|
|
assert pick_ollama_from_tag(have, "nomic-embed-text-cpu") == "nomic-embed-text"
|
|
# Re-provision must not FROM a stacked tag
|
|
assert pick_ollama_from_tag(
|
|
{"nomic-embed-text-cpu-cpu", "nomic-embed-text-cpu"},
|
|
"nomic-embed-text-cpu",
|
|
) == "nomic-embed-text-cpu"
|
|
|
|
assert finalize_memory_role_tags(
|
|
["nomic-embed-text", "nomic-embed-text-cpu-cpu-cpu"],
|
|
prefer_cpu=True,
|
|
) == ["nomic-embed-text-cpu"]
|
|
assert finalize_memory_role_tags(
|
|
["nomic-embed-text-cpu-cpu", "nomic-embed-text"],
|
|
prefer_cpu=True,
|
|
) == ["nomic-embed-text-cpu"]
|
|
|
|
|
|
def test_roles_payload_collapses_duplicate_memory_roots():
|
|
from gpu_rent.llm_runtime import OllamaModelEntry, ollama_roles_payload
|
|
|
|
roles = ollama_roles_payload(
|
|
[
|
|
OllamaModelEntry(name="chat:7b", default=True, use="chat"),
|
|
OllamaModelEntry(name="nomic-embed-text", use="memory"),
|
|
OllamaModelEntry(name="nomic-embed-text-cpu-cpu", use="memory"),
|
|
]
|
|
)
|
|
assert roles["memory"] == ["nomic-embed-text"]
|