`scripts/_ollama.py`

90 lines
#!/usr/bin/env python3
"""Shared local-Ollama helper for the linkedin-voice-agent skill.

Zero external dependencies (stdlib urllib). Runs on Mac2/Mac1 Ollama at $0.
Mirrors the call+fallback+think-strip pattern used by the kartiseira skill so
behaviour is consistent across Steve's persona-voice tools.
"""
import json
import os
import re
import sys
import urllib.request

DEFAULT_MODEL = os.environ.get("LVA_MODEL", "qwen3:14b")
BASE = os.environ.get("OLLAMA_URL", "http://localhost:11434")
# Preference order if the requested model isn't pulled locally.
FALLBACKS = ("qwen3", "hermes3", "gemma3", "qwen2.5", "llama3.1", "llama3")


def _available_models():
    try:
        with urllib.request.urlopen(f"{BASE}/api/tags", timeout=8) as r:
            data = json.load(r)
        return {m["name"].split(":")[0]: m["name"] for m in data.get("models", [])}
    except Exception:
        return {}


def _resolve_model(want):
    avail = _available_models()
    if not avail:
        return want  # let the call fail loudly if ollama is truly down
    if want in avail.values():
        return want
    for fam in FALLBACKS:
        if fam in avail:
            return avail[fam]
    return next(iter(avail.values()))


def ask(prompt, model=None, temperature=0.8, num_predict=1400):
    """Return (text, model_used). Strips <think> blocks. Raises on hard failure."""
    model = _resolve_model(model or DEFAULT_MODEL)
    body = json.dumps({
        "model": model,
        "prompt": prompt,
        "stream": False,
        "options": {"temperature": temperature, "num_predict": num_predict},
    }).encode()
    req = urllib.request.Request(f"{BASE}/api/generate", data=body,
                                 headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=240) as r:
        out = json.load(r).get("response", "")
    out = re.sub(r"<think>.*?</think>", "", out, flags=re.S | re.I).strip()
    return out, model


def extract_json(text):
    """Pull the first JSON object/array out of a possibly-chatty model reply."""
    text = text.strip()
    m = re.search(r"```(?:json)?\s*(.+?)```", text, flags=re.S)
    if m:
        text = m.group(1).strip()
    # Match whichever bracket appears FIRST — a top-level object may contain
    # arrays (and vice-versa), so trying "[" unconditionally would grab an
    # inner array out of an object.
    pairs = [(o, c) for o, c in (("[", "]"), ("{", "}")) if text.find(o) != -1]
    pairs.sort(key=lambda p: text.find(p[0]))
    for opener, closer in pairs:
        i = text.find(opener)
        if i == -1:
            continue
        depth = 0
        for j in range(i, len(text)):
            if text[j] == opener:
                depth += 1
            elif text[j] == closer:
                depth -= 1
                if depth == 0:
                    try:
                        return json.loads(text[i:j + 1])
                    except Exception:
                        break
    return None


if __name__ == "__main__":
    txt, used = ask(sys.argv[1] if len(sys.argv) > 1 else "Say hi in one line.")
    print(f"[model: {used}]\n{txt}")