From 1d431e3a05b94c66e2984ed42ae349e3dca60957 Mon Sep 17 00:00:00 2001 From: Leonid Pershin Date: Sun, 23 Aug 2026 09:28:06 +0300 Subject: [PATCH] Align Assistent debug patch extraction with string generate flags and unfenced JSON. Co-authored-by: Cursor --- src/gpu_rent/debug_assistent.py | 64 ++++++++++++++++++++----- src/gpu_rent/debug_assistent_session.py | 10 +--- tests/test_debug_api.py | 11 +++++ 3 files changed, 66 insertions(+), 19 deletions(-) diff --git a/src/gpu_rent/debug_assistent.py b/src/gpu_rent/debug_assistent.py index dfc79a1..7975595 100644 --- a/src/gpu_rent/debug_assistent.py +++ b/src/gpu_rent/debug_assistent.py @@ -52,6 +52,7 @@ _PATCH_KEYS = frozenset( "vision_from", "vision_slots", "variants", + "checkpoint", } ) _FENCE_RE = re.compile(r"```(?:json)?\s*([\s\S]*?)```", re.IGNORECASE) @@ -343,12 +344,39 @@ def _compact_api(name: str, data: Any) -> Any: return {"keys": list(data.keys())[:20]} +def _generate_flag_on(obj: dict[str, Any] | None) -> bool: + if not isinstance(obj, dict): + return False + g = obj.get("generate") + if g is True or g == 1: + return True + if isinstance(g, str) and g.strip().lower() in {"true", "1", "yes", "on"}: + return True + acts = obj.get("actions") + return isinstance(acts, list) and "generate" in [str(a) for a in acts] + + +def _normalize_extracted_patch(obj: dict[str, Any]) -> dict[str, Any]: + patch = dict(obj) + if _generate_flag_on(patch): + patch["generate"] = True + if isinstance(patch.get("ask"), str): + one = patch["ask"].strip() + if one: + patch["ask"] = [one] + else: + patch.pop("ask", None) + return patch + + def extract_assistent_patch(text: str | None) -> dict[str, Any]: """Pull last fenced JSON patch from Assistent reply (mirrors SA.extractPatch).""" if not text: return {"prose": text or "", "patch": None} - last_patch: dict[str, Any] | None = None - prose = text + last_any: dict[str, Any] | None = None + last_any_span: tuple[int, int] | None = None + last_term: dict[str, Any] | None = None + last_term_span: tuple[int, int] | None = None for match in _FENCE_RE.finditer(text): try: obj = json.loads(match.group(1).strip()) @@ -358,15 +386,29 @@ def extract_assistent_patch(text: str | None) -> dict[str, Any]: continue if not any(k in obj and obj[k] is not None for k in _PATCH_KEYS): continue - patch = dict(obj) - acts = [str(a) for a in patch["actions"]] if isinstance(patch.get("actions"), list) else [] - if patch.get("generate") is True or "generate" in acts: - patch["generate"] = True - if isinstance(patch.get("ask"), str): - patch["ask"] = [patch["ask"]] - last_patch = patch - prose = (text[: match.start()] + text[match.end() :]).strip() - return {"prose": prose, "patch": last_patch} + patch = _normalize_extracted_patch(obj) + last_any = patch + last_any_span = (match.start(), match.end()) + if _generate_flag_on(patch) or patch.get("look_at") is not None or patch.get("ask"): + last_term = patch + last_term_span = last_any_span + elif isinstance(patch.get("prompt"), str) and len(patch["prompt"].strip()) >= 48: + last_term = patch + last_term_span = last_any_span + chosen = last_term or last_any + span = last_term_span or last_any_span + if chosen and span: + prose = (text[: span[0]] + text[span[1] :]).strip() + return {"prose": prose, "patch": chosen} + brace = text.rfind("{") + if brace >= 0: + try: + obj = json.loads(text[brace:].strip()) + except (json.JSONDecodeError, TypeError, ValueError): + obj = None + if isinstance(obj, dict) and any(k in obj and obj[k] is not None for k in _PATCH_KEYS): + return {"prose": text[:brace].strip(), "patch": _normalize_extracted_patch(obj)} + return {"prose": text, "patch": None} def _summarize_patch(patch: dict[str, Any] | None) -> dict[str, Any] | None: diff --git a/src/gpu_rent/debug_assistent_session.py b/src/gpu_rent/debug_assistent_session.py index ec91086..0c289a3 100644 --- a/src/gpu_rent/debug_assistent_session.py +++ b/src/gpu_rent/debug_assistent_session.py @@ -36,6 +36,7 @@ from gpu_rent.debug_assistent import ( _swarm_base, collect_assistent_deep, extract_assistent_patch, + _generate_flag_on, ) SESSION_TTL_SEC = 45 * 60 @@ -245,14 +246,7 @@ def analyze_exact_merge( ) -> dict[str, Any]: """Simulate client mergeExactParamsForGenerate for diagnostics.""" patch = patch if isinstance(patch, dict) else None - acts = patch.get("actions") if patch else None - wants_gen = bool( - patch - and ( - patch.get("generate") is True - or (isinstance(acts, list) and "generate" in [str(a) for a in acts]) - ) - ) + wants_gen = _generate_flag_on(patch) defaults = resolve_exact_profile_defaults(exact, profile_name=krea_profile) if recommended and isinstance(recommended, dict): for k in EXACT_GENERATE_PARAM_KEYS: diff --git a/tests/test_debug_api.py b/tests/test_debug_api.py index 4503bb5..b614e2b 100644 --- a/tests/test_debug_api.py +++ b/tests/test_debug_api.py @@ -190,6 +190,17 @@ def test_extract_assistent_patch(): assert got2["patch"]["steps"] == 8 assert got2["patch"]["cfg"] == 2 + str_flag = extract_assistent_patch( + 'ok\n```json\n{"prompt":"a slender redhead in leather","generate":"true"}\n```' + ) + assert str_flag["patch"]["generate"] is True + + raw = extract_assistent_patch( + 'done\n{"prompt":"a fiery redhead, street light, 85mm","generate":true}' + ) + assert raw["patch"]["generate"] is True + assert raw["prose"] == "done" + def test_chat_eval_mocked_http(monkeypatch, tmp_path): from gpu_rent import debug_assistent