llm.py

#!/usr/bin/env python
"""LLM backends for the Web Research Agent (#13): local Ollama (default, free/private) or Claude.

On-demand only — these are plain function calls, nothing runs in the background. Errors log the
exception TYPE (never repr) so a key can't leak into logs.
"""
import json
import os
import urllib.request

BACKEND      = os.environ.get("RESEARCH_LLM_BACKEND", "ollama").lower()
OLLAMA_URL   = os.environ.get("RESEARCH_OLLAMA_URL", "http://127.0.0.1:11434")   # set to your Spark/fleet host
OLLAMA_MODEL = os.environ.get("RESEARCH_OLLAMA_MODEL", "qwen2.5:7b")
ANTHROPIC_KEY   = (os.environ.get("ANTHROPIC_API_KEY") or "").strip() or None
ANTHROPIC_MODEL = os.environ.get("ANTHROPIC_MODEL", "claude-haiku-4-5")


def _ollama(system, user, temperature, max_tokens, model=None):
    body = json.dumps({
        "model": model or OLLAMA_MODEL,
        "messages": [{"role": "system", "content": system}, {"role": "user", "content": user}],
        "stream": False, "options": {"temperature": temperature, "num_predict": max_tokens},
    }).encode()
    req = urllib.request.Request(OLLAMA_URL + "/api/chat", data=body,
                                 headers={"Content-Type": "application/json"})
    return json.loads(urllib.request.urlopen(req, timeout=180).read())["message"]["content"].strip()


def _claude(system, user, temperature, max_tokens, model=None):
    body = json.dumps({
        "model": model or ANTHROPIC_MODEL, "max_tokens": max_tokens, "temperature": temperature,
        "system": system, "messages": [{"role": "user", "content": user}],
    }).encode()
    req = urllib.request.Request("https://api.anthropic.com/v1/messages", data=body, headers={
        "x-api-key": ANTHROPIC_KEY, "anthropic-version": "2023-06-01", "content-type": "application/json"})
    data = json.loads(urllib.request.urlopen(req, timeout=180).read())
    return "".join(b.get("text", "") for b in data.get("content", [])).strip()


def chat(system: str, user: str, temperature: float = 0.3, max_tokens: int = 900, model: str = None) -> str:
    """One completion. Routes to Claude if RESEARCH_LLM_BACKEND=claude (and a key is set), else Ollama."""
    fn = _claude if (BACKEND == "claude" and ANTHROPIC_KEY) else _ollama
    return fn(system, user, temperature, max_tokens, model)


def chat_json(system: str, user: str, temperature: float = 0.0, max_tokens: int = 500, model: str = None):
    """Completion parsed as JSON. Tolerant: strips code fences and pulls the first {...} or [...] block."""
    raw = chat(system + "\nReply with ONLY valid JSON, no prose, no code fences.",
               user, temperature, max_tokens, model)
    return _parse_json(raw)


def _parse_json(raw: str):
    s = raw.strip()
    if s.startswith("```"):
        s = s.split("```", 2)[1] if s.count("```") >= 2 else s.strip("`")
        if s.lower().startswith("json"):
            s = s[4:]
    s = s.strip()
    try:
        return json.loads(s)
    except Exception:
        # pull first array/object
        for op, cl in (("[", "]"), ("{", "}")):
            i, j = s.find(op), s.rfind(cl)
            if 0 <= i < j:
                try:
                    return json.loads(s[i:j + 1])
                except Exception:
                    pass
    return None


def backend_label() -> str:
    return "Claude" if (BACKEND == "claude" and ANTHROPIC_KEY) else f"Ollama ({OLLAMA_MODEL} @ {OLLAMA_URL})"


if __name__ == "__main__":
    print("backend:", backend_label())
    print(chat("You are terse.", "Say 'research llm ok' and nothing else.", max_tokens=20))
    print("json:", chat_json("You split questions.", "Give 2 search queries for: impact of SLMs on enterprise cost. As a JSON array of strings."))