From 2b715d2a75bc5580a4dafa11ddb98db9496a85d6 Mon Sep 17 00:00:00 2001 From: Andrew Pimenov Date: Sun, 16 Aug 2026 19:35:39 +0300 Subject: [PATCH] ollama call with schema-constrained output --- sort_ym/analyze.py | 68 +++++++++++++++++++++++++++++++++++ tests/test_analyze.py | 84 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 152 insertions(+) diff --git a/sort_ym/analyze.py b/sort_ym/analyze.py index 65ad920..5bee9a5 100644 --- a/sort_ym/analyze.py +++ b/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"]) diff --git a/tests/test_analyze.py b/tests/test_analyze.py index 117521e..731bb48 100644 --- a/tests/test_analyze.py +++ b/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