Leonid PershinandCursor
d934fc9366
Write preferred chat tag into Assistent/ollama-roles.json so UI and warmup align with manifest default: true. Co-authored-by: Cursor <cursoragent@cursor.com>
57 lines
1.7 KiB
Python
57 lines
1.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
|