training_coach/parser_ollama.py

import json
import logging
from time import perf_counter
from urllib import request

from training_coach.parser import (
    PARSER_SYSTEM_PROMPT,
    expected_response_format,
    log_parser_messages,
    log_parser_response_text,
    parse_model_response,
)
from training_coach.models import ParsedCheckIn


OLLAMA_URL = "http://127.0.0.1:11434/api/generate"
OLLAMA_MODEL = "qwen3:1.7B"
OLLAMA_TINY_CANDIDATE = "qwen3:4b"
logger = logging.getLogger(__name__)


def build_ollama_prompt(raw_text: str) -> str:
    schema = json.dumps(expected_response_format(), indent=2)
    return (
        "Parse this check-in into the expected response JSON schema.\n"
        "Return only the JSON object.\n\n"
        "Allowed top-level keys are: check_in, missing_fields, follow_up_items, "
        "follow_up_questions, context_signals, notes.\n\n"
        "The check_in object must use only these keys: raw_text, "
        "time_available_minutes, energy_level, sleep_quality, sleep_hours, "
        "soreness, pain_or_injury, pain_issues, mood_stress, notes.\n\n"
        f"Expected response JSON schema:\n{schema}\n\n"
        f"Check-in:\n{raw_text}"
    )


def generate_parser_response_ollama(
    raw_text: str,
    model_name: str = OLLAMA_MODEL,
    url: str = OLLAMA_URL,
) -> str:
    start_time = perf_counter()
    prompt = build_ollama_prompt(raw_text)
    logger.info(
        "event=parser_ollama_request model=%s prompt_chars=%s text_chars=%s",
        model_name,
        len(prompt),
        len(raw_text),
    )
    log_parser_messages(
        backend="ollama",
        model_name=model_name,
        messages=[
            {"role": "system", "content": PARSER_SYSTEM_PROMPT},
            {"role": "user", "content": prompt},
        ],
    )
    payload = {
        "model": model_name,
        "system": PARSER_SYSTEM_PROMPT,
        "prompt": prompt,
        "format": expected_response_format(),
        "stream": False,
        "think": False,
        "options": {
            "temperature": 0,
            "num_predict": 512,
        },
    }
    body = json.dumps(payload).encode("utf-8")
    http_request = request.Request(
        url,
        data=body,
        headers={"Content-Type": "application/json"},
        method="POST",
    )

    with request.urlopen(http_request, timeout=120) as response:
        data = json.loads(response.read().decode("utf-8"))

    response_text = data["response"].strip()
    logger.info(
        "event=parser_ollama_response model=%s response_chars=%s "
        "prompt_eval_count=%s eval_count=%s total_duration_ms=%s "
        "load_duration_ms=%s prompt_eval_duration_ms=%s eval_duration_ms=%s "
        "wall_elapsed_ms=%s",
        model_name,
        len(response_text),
        data.get("prompt_eval_count"),
        data.get("eval_count"),
        _ns_to_ms(data.get("total_duration")),
        _ns_to_ms(data.get("load_duration")),
        _ns_to_ms(data.get("prompt_eval_duration")),
        _ns_to_ms(data.get("eval_duration")),
        round((perf_counter() - start_time) * 1000),
    )
    log_parser_response_text(
        backend="ollama",
        model_name=model_name,
        response_text=response_text,
    )
    return response_text


def _ns_to_ms(value):
    if value is None:
        return None
    return round(value / 1_000_000)


def parse_check_in_with_ollama(
    raw_text: str,
    model_name: str = OLLAMA_MODEL,
) -> ParsedCheckIn:
    response_text = generate_parser_response_ollama(raw_text, model_name=model_name)
    return parse_model_response(response_text)