`src/llm.py`

51 lines
"""
Thin Ollama client. No cloud APIs.

Uses /api/chat with format="json" so the model returns parsable JSON.
Falls back to /api/generate for classify (string label) calls.
"""
from __future__ import annotations

import json
import os

import httpx

OLLAMA_HOST = os.environ.get("OLLAMA_HOST", "http://localhost:11434")


def chat_json(model: str, system: str, user: str, *, num_ctx: int = 32768) -> dict:
    """Chat call constrained to JSON output. Returns parsed dict."""
    with httpx.Client(timeout=600) as c:
        r = c.post(
            f"{OLLAMA_HOST}/api/chat",
            json={
                "model": model,
                "format": "json",
                "stream": False,
                "options": {"temperature": 0, "num_ctx": num_ctx},
                "messages": [
                    {"role": "system", "content": system},
                    {"role": "user", "content": user},
                ],
            },
        )
        r.raise_for_status()
    body = r.json()["message"]["content"]
    return json.loads(body)


def generate_text(model: str, prompt: str, *, num_ctx: int = 8192) -> str:
    with httpx.Client(timeout=600) as c:
        r = c.post(
            f"{OLLAMA_HOST}/api/generate",
            json={
                "model": model,
                "prompt": prompt,
                "stream": False,
                "options": {"temperature": 0, "num_ctx": num_ctx},
            },
        )
        r.raise_for_status()
    return r.json()["response"].strip()