src/agents/recipe_planner.py

"""Recipe planner agent: propose dishes + generate step-by-step recipe.

Uses openbmb/MiniCPM4.1-8B (text-only) as the primary planner.
Falls back to the shared vision model (MiniCPM-V-4.6) when the planner
model is unavailable (e.g. insufficient RAM on the Space).
"""
from __future__ import annotations

import json
import logging
import re

import spaces
import torch

from src import config
from src.pipeline import DishOption, Recipe, RecipeStep

log = logging.getLogger(__name__)

_PROPOSE_PROMPT = (config.PROMPTS_DIR / "planner_propose.txt").read_text(encoding="utf-8")
_RECIPE_PROMPT = (config.PROMPTS_DIR / "planner_recipe.txt").read_text(encoding="utf-8")


# ---------------------------------------------------------------------------
# JSON extraction helpers
# ---------------------------------------------------------------------------

def _extract_json(text: str) -> dict:
    """Robustly extract the first JSON object from raw model output."""
    text = text.strip()
    try:
        return json.loads(text)
    except Exception:
        pass
    # Markdown code-block
    m = re.search(r"```(?:json)?\s*(\{.*?\})\s*```", text, re.DOTALL)
    if m:
        try:
            return json.loads(m.group(1))
        except Exception:
            pass
    # First {...} block with minor auto-fixes
    m = re.search(r"\{.*\}", text, re.DOTALL)
    if m:
        candidate = m.group(0)
        candidate = candidate.replace("'", '"')
        candidate = re.sub(r",\s*([}\]])", r"\1", candidate)
        try:
            return json.loads(candidate)
        except Exception:
            pass
    log.warning("Could not extract JSON from output (first 300 chars): %.300s", text)
    return {}


# ---------------------------------------------------------------------------
# Inference dispatcher
# ---------------------------------------------------------------------------

def _infer(prompt: str, max_new_tokens: int = 1024, temperature: float = 0.0) -> str:
    """Run text inference.

    Primary: the dedicated MiniCPM4.1-8B planner Modal endpoint (transformers
    4.x). Falls back to the local vision model (text-only) if the endpoint is
    unavailable or returns nothing.
    """
    try:
        import modal
        cls = modal.Cls.from_name(config.PLANNER_MODAL_APP, config.PLANNER_MODAL_CLS)
        out = cls().infer.remote(prompt, max_new_tokens=max_new_tokens, temperature=temperature)
        if out and out.strip():
            return out
        log.warning("Planner endpoint returned empty — falling back to vision model.")
    except Exception as exc:
        log.warning("Planner endpoint call failed: %s — falling back to vision model.", exc)

    # Fallback: use the vision model in text-only mode
    log.warning("Using vision model as text fallback.")
    from src.agents.mise_en_place import model as vis_model, processor as vis_proc

    messages = [{"role": "user", "content": [{"type": "text", "text": prompt}]}]
    inputs = vis_proc.apply_chat_template(
        messages,
        add_generation_prompt=True,
        tokenize=True,
        return_dict=True,
        return_tensors="pt",
        enable_thinking=False,
    )
    device = vis_model.device
    inputs = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in inputs.items()}
    for k, v in inputs.items():
        if isinstance(v, torch.Tensor) and torch.is_floating_point(v):
            inputs[k] = v.to(dtype=torch.bfloat16)

    with torch.no_grad():
        generated_ids = vis_model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False)

    trimmed = [out[len(inp):] for inp, out in zip(inputs["input_ids"], generated_ids)]
    return vis_proc.batch_decode(trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]


# ---------------------------------------------------------------------------
# Public agent functions
# ---------------------------------------------------------------------------

@spaces.GPU(duration=90)
def propose_dishes(ingredients: list[str]) -> list[DishOption]:
    """Given detected ingredients, return up to 3 dish proposals."""
    try:
        prompt = _PROPOSE_PROMPT.replace("{ingredients}", ", ".join(ingredients))
        raw = _infer(prompt, max_new_tokens=512, temperature=0.7)
        log.info("propose_dishes raw: %.500s", raw)
        data = _extract_json(raw)
        options = data.get("options", [])
        return [
            DishOption(name=str(o.get("name", "Dish")), why=str(o.get("why", "")))
            for o in options[:3]
            if o.get("name")
        ] or [DishOption(name="Simple Stir-fry", why="Quick and adaptable to most ingredients.")]
    except Exception as exc:
        log.warning("propose_dishes failed: %s", exc)
        return [DishOption(name="Simple Stir-fry", why="Quick and adaptable to most ingredients.")]


@spaces.GPU(duration=120)
def plan_recipe(dish_name: str, ingredients: list[str]) -> Recipe:
    """Generate a full step-by-step recipe for the chosen dish."""
    try:
        prompt = (
            _RECIPE_PROMPT
            .replace("{dish_name}", dish_name)
            .replace("{ingredients}", ", ".join(ingredients))
        )
        raw = _infer(prompt, max_new_tokens=1024, temperature=0.0)
        log.info("plan_recipe raw: %.800s", raw)
        data = _extract_json(raw)

        raw_steps = data.get("steps", [])
        steps = []
        for i, s in enumerate(raw_steps, start=1):
            if not s.get("instruction"):
                continue
            tip_val = s.get("tip")
            steps.append(RecipeStep(
                n=int(s.get("n", i)),
                instruction=str(s["instruction"]),
                duration=str(s.get("duration", "5 min")),
                tip=str(tip_val) if tip_val and str(tip_val).lower() not in ("null", "none") else None,
                visual=str(s.get("visual", "")),
            ))

        return Recipe(
            name=str(data.get("name", dish_name)),
            cuisine=str(data.get("cuisine", "International")),
            servings=int(data.get("servings", 2)),
            total_time_minutes=int(data.get("total_time_minutes", 30)),
            final_dish_visual=str(data.get("final_dish_visual", "")),
            steps=steps or [RecipeStep(n=1, instruction="Prepare and cook ingredients to taste.", duration="20 min")],
        )
    except Exception as exc:
        log.warning("plan_recipe failed: %s", exc)
        return Recipe(
            name=dish_name,
            steps=[RecipeStep(n=1, instruction="Prepare and cook ingredients to taste.", duration="20 min")],
        )