Spaces:
Paused
Paused
| """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}""" | |
| 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, | |
| ) | |