api/enhance.py

"""Prompt enhancement: turn a tester's short request into a precise brief before generation.

It runs through OpenCode with a free OpenCode Zen model (ENHANCE_MODEL, default opencode/deepseek-v4-flash-free),
never through OpenRouter, so it costs nothing and is recorded in the ledger as exempt. It is best effort: when
OpenCode is missing, slow or returns something unusable, generation simply goes ahead with the original text.

The enhancer may only make the request clearer and more complete. It keeps every number, name and constraint the
tester gave, and it must not add goals of its own; the routing tier and the cost estimate are still decided from the
tester's own words.
"""

from __future__ import annotations

import logging
import re
import time
from dataclasses import asdict, dataclass

from .opencode_runner import OpenCodeError, OpenCodeWorkspace, find_opencode
from .zen import zen_config, zen_text

log = logging.getLogger("rosdiff.enhance")

MAX_CHARS = 3000

SCENE_BRIEF = """You improve requests for a MuJoCo scene generator used by robotics students.
Rewrite the request below as a precise scene brief:
- keep every number, unit, object, colour and constraint exactly as given;
- where a size, position or mass is missing, choose a sensible value and state it (metres, kilograms);
- say where each object is relative to the robot (the robot stands at the origin, facing +x);
- mention what should be tested, if the request implies it (e.g. walking up the steps, being pushed);
- do not add unrelated objects or goals, and do not write MJCF or code.
Robot: {robot}.
Reply with only the rewritten brief inside <prompt></prompt>, at most 150 words.

Request:
{text}"""

CODE_BRIEF = """You improve requests for a ROS 2 package generator used by robotics students.
Rewrite the request below as a precise package brief:
- keep every topic, message type, parameter, rate and threshold exactly as given;
- where one is missing, choose the standard ROS 2 choice and state it (topic names, message types, QoS, rates);
- list the node(s), what each subscribes to and publishes, parameters with defaults, and the launch file;
- do not add features the request did not ask for, and do not write code.
Target: ROS 2 {distro}, {language}.
Reply with only the rewritten brief inside <prompt></prompt>, at most 180 words.

Request:
{text}"""


@dataclass
class Enhanced:
    original: str
    text: str  # what generation uses: the enhanced brief, or the original when enhancement was skipped
    used: bool
    model: str | None = None
    note: str | None = None  # why it was skipped
    usage: dict | None = None
    ms: int = 0

    def to_dict(self) -> dict:
        return asdict(self)


def _extract(events: list[dict]) -> str:
    raw = zen_text(events)
    m = re.search(r"<prompt>(.*?)</prompt>", raw, re.S)
    return (m.group(1) if m else raw).strip()


def _plausible(original: str, brief: str) -> str | None:
    """Reason to reject the brief, or None. Every number the tester wrote must survive the rewrite."""
    if len(brief) < 10:
        return "empty answer"
    if len(brief) > MAX_CHARS:
        return "answer too long"
    # Only check standalone numbers (not parts of identifiers like G1, ROS 2)
    nums = re.findall(r"(?<![A-Za-z])\d+(?:\.\d+)?(?![A-Za-z])", original)
    lost = [n for n in nums if not re.search(r"(?<!\d)" + re.escape(n) + r"(?!\d)", brief)]
    if lost:
        return f"the rewrite dropped the value(s) {', '.join(lost[:3])}"
    return None


def enhance(settings, kind: str, text: str, **ctx) -> Enhanced:
    """kind: 'scene' (ctx: robot) or 'code' (ctx: distro, language)."""
    if not settings.prompt_enhance:
        return Enhanced(text, text, False, note="switched off (PROMPT_ENHANCE=false)")
    if not find_opencode(settings):
        return Enhanced(text, text, False, note="OpenCode is not installed on the server")
    if kind == "scene":
        message = SCENE_BRIEF.format(robot=ctx.get("robot") or "none", text=text)
    else:
        lang = "C++" if ctx.get("language") == "cpp" else "Python"
        message = CODE_BRIEF.format(distro=ctx.get("distro", "humble"), language=lang, text=text)
    provider, _, model = settings.enhance_model.partition("/")
    start = time.monotonic()
    try:
        with OpenCodeWorkspace(settings, config=enhance_config(settings), timeout=settings.enhance_timeout) as ws:
            run = ws.run(message, context="", model=model, provider=provider)
    except OpenCodeError as e:
        log.warning("prompt enhancement skipped: %s", e)
        return Enhanced(text, text, False, settings.enhance_model, str(e), ms=int((time.monotonic() - start) * 1000))
    ms = int((time.monotonic() - start) * 1000)
    if run.error:
        log.warning("prompt enhancement failed: %s", run.error)
        return Enhanced(text, text, False, settings.enhance_model, run.error, run.usage(), ms)
    brief = _extract(run.events)
    if why := _plausible(text, brief):
        return Enhanced(text, text, False, settings.enhance_model, why, run.usage(), ms)
    return Enhanced(text, brief, True, settings.enhance_model, None, run.usage(), ms)


def enhance_config(settings) -> dict:
    """Only the free Zen model, with no tools: the enhancer reads and answers, nothing else."""
    return zen_config(settings, settings.enhance_model)


def record(logs, tester: str, generation_id: str | None, category: str, e: Enhanced) -> None:
    """Free calls are still recorded (as exempt), so the ledger shows everything RosDiff asked a model."""
    if not e.model or e.usage is None:
        return
    logs.record_usage(
        tester=tester,
        generation_id=generation_id,
        category=category,
        model=e.model,
        tier="enhance",
        provider="opencode-zen",
        input_tokens=int(e.usage.get("input_tokens") or 0),
        output_tokens=int(e.usage.get("output_tokens") or 0),
        cost_usd=0.0,
        cost_inr=0.0,
        billable=0,
        exempt_reason="opencode_zen_free",
        estimated=0,
        estimate_inr=None,
    )