voice-agent/mac-worker/app/core/ollama_client.py

import json
import logging
import re
from typing import Any

import httpx

from . import profiles
from .config import settings

log = logging.getLogger("mac-worker.ollama")


def _strip_codefence(s: str) -> str:
    s = s.strip()
    if s.startswith("```"):
        s = re.sub(r"^```(?:json)?", "", s, count=1).strip()
        if s.endswith("```"):
            s = s[:-3].strip()
    return s


def _parse_json(text: str) -> dict:
    text = _strip_codefence(text)
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        m = re.search(r"\{.*\}", text, re.S)
        if m:
            return json.loads(m.group(0))
        raise


def _truncate(text: str, max_chars: int = 60000) -> str:
    if len(text) <= max_chars:
        return text
    head = text[: max_chars // 2]
    tail = text[-max_chars // 2 :]
    return head + "\n\n[... gekürzt ...]\n\n" + tail


def render(template: str, **vars: str) -> str:
    """Ersetzt ausschließlich die übergebenen `{name}`-Platzhalter.

    Bewusst KEIN str.format() — sonst würden literale `{...}` im Template
    (z. B. JSON-Schema-Beispiele wie `{"owner", "task"}` in den Profil-Prompts)
    als Platzhalter interpretiert und einen KeyError werfen.
    """
    out = template
    for k, v in vars.items():
        out = out.replace("{" + k + "}", str(v))
    return out


def _resolve_model(model: str | None) -> str:
    """Wählt das tatsächlich zu verwendende Ollama-Modell.

    Leerer oder fehlender `model`-Parameter → Worker-Default aus .env.
    """
    return (model or "").strip() or settings.ollama_model


async def _chat(prompt: str, model: str | None = None) -> str:
    """Ruft Ollama im Streaming-Modus auf.

    Streaming hat zwei wichtige Vorteile:
    1. Read-Timeout gilt pro Chunk statt für die gesamte Antwort —
       lange Generierungen laufen nicht mehr in den Gesamt-Timeout.
    2. Wir können in den Logs sehen, dass das Modell tatsächlich arbeitet
       (Anzahl Chunks).

    `keep_alive`: das Modell bleibt 30 min im RAM zwischen Calls,
    damit nachfolgende Anfragen nicht erneut die Ladezeit zahlen.
    """
    use_model = _resolve_model(model)
    url = f"{settings.ollama_url.rstrip('/')}/api/generate"
    payload = {
        "model": use_model,
        "prompt": prompt,
        "stream": True,
        "keep_alive": "30m",
        "options": {"temperature": 0.2},
    }
    log.info("Ollama generate model=%s prompt_len=%d (streaming)", use_model, len(prompt))
    timeout = httpx.Timeout(
        connect=15.0,
        read=settings.ollama_timeout,   # pro Chunk — großzügig wegen initial Modell-Load
        write=30.0,
        pool=10.0,
    )
    chunks: list[str] = []
    chunk_count = 0
    async with httpx.AsyncClient(timeout=timeout) as c:
        async with c.stream("POST", url, json=payload) as r:
            r.raise_for_status()
            async for line in r.aiter_lines():
                if not line.strip():
                    continue
                try:
                    obj = json.loads(line)
                except json.JSONDecodeError:
                    log.warning("Ollama: ungültige JSON-Zeile übersprungen: %r", line[:100])
                    continue
                chunks.append(obj.get("response", ""))
                chunk_count += 1
                if obj.get("done"):
                    break
    log.info("Ollama: %d Chunks empfangen, Gesamtlänge %d Zeichen", chunk_count, sum(map(len, chunks)))
    return "".join(chunks)


async def summarize(
    transcript: str, title: str = "", profile_name: str | None = None, model: str | None = None,
) -> dict[str, Any]:
    profile = profiles.get(profile_name)
    prompt = render(
        profile.summary_prompt,
        title=title or "(ohne Titel)",
        transcript=_truncate(transcript),
    )
    log.info("Summarize profile=%s model=%s", profile.name, _resolve_model(model))
    raw = await _chat(prompt, model=model)
    try:
        return _parse_json(raw)
    except Exception:  # noqa: BLE001
        log.warning("Failed to parse summary JSON — gebe Roh-Antwort als _raw zurück")
        return {"_raw": raw.strip(), "_parse_error": True}


async def make_protocol(
    transcript: str, summary: dict, title: str = "", profile_name: str | None = None,
    model: str | None = None,
) -> dict:
    profile = profiles.get(profile_name)
    prompt = render(
        profile.protocol_prompt,
        title=title or "(ohne Titel)",
        summary_json=json.dumps(summary, ensure_ascii=False),
        transcript=_truncate(transcript),
    )
    log.info("Protocol profile=%s model=%s", profile.name, _resolve_model(model))
    raw = await _chat(prompt, model=model)
    try:
        return _parse_json(raw)
    except Exception:  # noqa: BLE001
        log.warning("Failed to parse protocol JSON — Fallback auf Summary")
        fallback = dict(summary) if isinstance(summary, dict) else {}
        fallback.setdefault("title", title)
        return fallback


async def health() -> dict:
    url = f"{settings.ollama_url.rstrip('/')}/api/tags"
    async with httpx.AsyncClient(timeout=10) as c:
        r = await c.get(url)
        r.raise_for_status()
        return r.json()


async def list_models() -> list[dict]:
    """Liefert die lokal gepullten Ollama-Modelle inkl. Größe und Mod-Zeit.

    Antwort-Form pro Eintrag: {"name", "size", "modified_at"} (Größe in Bytes).
    """
    url = f"{settings.ollama_url.rstrip('/')}/api/tags"
    async with httpx.AsyncClient(timeout=10) as c:
        r = await c.get(url)
        r.raise_for_status()
        data = r.json()
    out: list[dict] = []
    for m in data.get("models", []):
        out.append({
            "name": m.get("name") or m.get("model") or "",
            "size": m.get("size", 0),
            "modified_at": m.get("modified_at", ""),
        })
    return [m for m in out if m["name"]]