src/planner/llm_base.py

from __future__ import annotations

import json
import re
from dataclasses import dataclass
from typing import Any, Dict, Optional, Tuple

from transformers import pipeline

_JSON_RE = re.compile(r"\{.*\}", re.DOTALL)


@dataclass(frozen=True)
class LLMConfig:
    model_name: str = "google/flan-t5-base"
    max_new_tokens: int = 512
    temperature: float = 0.2
    top_p: float = 0.95
    max_retries: int = 3


class LocalLLM:
    """
    Local HF LLM wrapper. Uses text2text-generation (T5-style instruction models).
    Designed for stable structured JSON output.
    """

    def __init__(self, config: Optional[LLMConfig] = None):
        self.config = config or LLMConfig()
        self.pipe = pipeline(
            "text2text-generation",
            model=self.config.model_name,
        )

    def generate(self, prompt: str) -> str:
        out = self.pipe(
            prompt,
            max_new_tokens=self.config.max_new_tokens,
            do_sample=True,
            temperature=self.config.temperature,
            top_p=self.config.top_p,
            num_return_sequences=1,
        )[0]["generated_text"]
        return out.strip()


def extract_json(text: str) -> Optional[str]:
    """
    Extract the first JSON object from model output.
    Handles cases where the model wraps JSON in explanations.
    """
    text = text.strip()

    if "```" in text:
        blocks = re.findall(
            r"```(?:json)?\s*(\{.*?\})\s*```",
            text,
            flags=re.DOTALL | re.IGNORECASE,
        )
        if blocks:
            return blocks[0].strip()

    match = _JSON_RE.search(text)
    if not match:
        return None
    return match.group(0).strip()


def safe_json_loads(s: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
    try:
        return json.loads(s), None
    except Exception as exc:
        return None, str(exc)


def generate_validated_json(
    llm: LocalLLM,
    prompt: str,
    validator_fn,
    max_retries: int = 3,
) -> Dict[str, Any]:
    """
    Generate JSON with retries.
    validator_fn must raise on invalid JSON or invalid schema.
    """
    last_error = None
    last_raw = None

    for attempt in range(1, max_retries + 1):
        raw = llm.generate(prompt)
        last_raw = raw

        js = extract_json(raw)
        if js is None:
            last_error = f"No JSON found in output (attempt {attempt}). Raw: {raw[:200]}..."
            continue

        data, err = safe_json_loads(js)
        if data is None:
            last_error = f"JSON parse error (attempt {attempt}): {err}"
            continue

        try:
            validator_fn(data)
            return data
        except Exception as exc:
            last_error = f"Schema validation error (attempt {attempt}): {exc}"
            continue

    raise RuntimeError(
        "Failed to generate valid JSON after retries.\n"
        f"Last error: {last_error}\n"
        f"Last raw output: {last_raw}"
    )