ollama call with schema-constrained output

sort_ym/analyze.py

@@ -2,7 +2,9 @@ from __future__ import annotations
 
 import hashlib
 import json
+import urllib.request
 from collections import Counter
+from dataclasses import dataclass
 from pathlib import Path
 
 ANALYSIS_CACHE_FILE = "lyrics_analysis.json"
@@ -126,3 +128,69 @@ def load_analysis(cache_dir: Path) -> dict[str, dict]:
     if cache_file.exists():
         return json.loads(cache_file.read_text(encoding="utf-8"))
     return {}
+
+
+@dataclass(frozen=True)
+class OllamaSettings:
+    host: str
+    model: str
+    prompt_version: int
+    timeout: float
+    keep_alive: str
+
+
+PROMPT_TEMPLATE = """Ты разбираешь текст песни для профиля музыкального вкуса одного слушателя.
+
+Трек: «{title}» — {artists}
+
+Текст:
+\"\"\"
+{lyrics}
+\"\"\"
+{themes_hint}
+Заполни поля схемы. Требования:
+- summary: 2-4 предложения по-русски о том, что в песне происходит и как это подано.
+- resonance: 1-2 предложения по-русски о том, чем этот текст может цеплять слушателя.
+- key_line: одна строка ДОСЛОВНО из текста выше, без изменений и без пересказа.
+- themes: 2-5 тем на английском в snake_case (1-3 слова, единственное число).
+- Остальные поля выбирай строго из допустимых значений схемы.
+"""
+
+THEMES_HINT_TEMPLATE = """
+Уже использованные темы в этой библиотеке (переиспользуй подходящую, иначе придумай новую):
+{themes}
+"""
+
+
+def build_prompt(title: str, artists: list[str], lyrics: str, known_themes: list[str]) -> str:
+    themes_hint = THEMES_HINT_TEMPLATE.format(themes=", ".join(known_themes)) if known_themes else ""
+    return PROMPT_TEMPLATE.format(
+        title=title,
+        artists=", ".join(artists),
+        lyrics=lyrics,
+        themes_hint=themes_hint,
+    )
+
+
+def call_ollama(settings: OllamaSettings, prompt: str) -> dict:
+    """think=True уводит рассуждение в отдельное поле message.thinking, а format (полная
+    JSON Schema, не строка "json") констрейнит только финальный ответ - модель думает
+    свободно, результат при этом гарантированно парсится. num_predict=-1: капать вывод
+    нельзя, обрыв на середине даст невалидный JSON вместо просто короткого ответа."""
+    payload = {
+        "model": settings.model,
+        "messages": [{"role": "user", "content": prompt}],
+        "stream": False,
+        "think": True,
+        "format": RESPONSE_SCHEMA,
+        "keep_alive": settings.keep_alive,
+        "options": {"num_predict": -1, "temperature": 0.65},
+    }
+    request = urllib.request.Request(
+        f"{settings.host.rstrip('/')}/api/chat",
+        data=json.dumps(payload, ensure_ascii=False).encode("utf-8"),
+        headers={"Content-Type": "application/json"},
+    )
+    with urllib.request.urlopen(request, timeout=settings.timeout) as response:
+        body = json.loads(response.read().decode("utf-8"))
+    return json.loads(body["message"]["content"])

tests/test_analyze.py

@@ -1,3 +1,4 @@
+import json
 import re
 from pathlib import Path
 
@@ -85,3 +86,86 @@ def test_top_themes_orders_by_frequency_and_limits():
 
 def test_load_analysis_returns_empty_dict_when_missing(tmp_path: Path):
     assert analyze.load_analysis(tmp_path) == {}
+
+
+class FakeResponse:
+    def __init__(self, payload):
+        self._payload = payload
+
+    def read(self):
+        return json.dumps(self._payload).encode("utf-8")
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *exc):
+        return False
+
+
+def _settings(**overrides):
+    base = dict(
+        host="http://localhost:11434",
+        model="qwen3.6-35b-a3b:latest",
+        prompt_version=1,
+        timeout=600,
+        keep_alive="30m",
+    )
+    base.update(overrides)
+    return analyze.OllamaSettings(**base)
+
+
+def test_call_ollama_sends_schema_constrained_thinking_request(monkeypatch):
+    captured = {}
+
+    def fake_urlopen(req, timeout=None):
+        captured["url"] = req.full_url
+        captured["timeout"] = timeout
+        captured["body"] = json.loads(req.data.decode("utf-8"))
+        return FakeResponse(
+            {"message": {"thinking": "рассуждение", "content": json.dumps({"mood": ["grief"]})}}
+        )
+
+    monkeypatch.setattr(analyze.urllib.request, "urlopen", fake_urlopen)
+
+    result = analyze.call_ollama(_settings(), "промпт")
+
+    assert captured["url"] == "http://localhost:11434/api/chat"
+    assert captured["timeout"] == 600
+    body = captured["body"]
+    assert body["model"] == "qwen3.6-35b-a3b:latest"
+    assert body["stream"] is False
+    assert body["think"] is True
+    assert body["format"] == analyze.RESPONSE_SCHEMA
+    assert body["keep_alive"] == "30m"
+    assert body["options"] == {"num_predict": -1, "temperature": 0.65}
+    assert body["messages"] == [{"role": "user", "content": "промпт"}]
+    assert result == {"mood": ["grief"]}, "в кэш идёт только content, не thinking"
+
+
+def test_call_ollama_strips_trailing_slash_from_host(monkeypatch):
+    captured = {}
+
+    def fake_urlopen(req, timeout=None):
+        captured["url"] = req.full_url
+        return FakeResponse({"message": {"content": "{}"}})
+
+    monkeypatch.setattr(analyze.urllib.request, "urlopen", fake_urlopen)
+
+    analyze.call_ollama(_settings(host="http://localhost:11434/"), "промпт")
+
+    assert captured["url"] == "http://localhost:11434/api/chat"
+
+
+def test_build_prompt_includes_lyrics_title_and_known_themes():
+    prompt = analyze.build_prompt("Тоска", ["Артист"], "Текст песни", ["loneliness", "city"])
+
+    assert "Тоска" in prompt
+    assert "Артист" in prompt
+    assert "Текст песни" in prompt
+    assert "loneliness, city" in prompt
+
+
+def test_build_prompt_omits_known_themes_section_when_empty():
+    prompt = analyze.build_prompt("Тоска", ["Артист"], "Текст песни", [])
+
+    assert "Уже использованные темы" not in prompt