Initial commit

opc/llm/retry.py

@@ -0,0 +1,290 @@
+"""Unified retry helper for LLM calls that expect structured (JSON) output.
+
+The orchestration layers of OPC frequently ask an LLM to produce a JSON object
+that is then parsed, type-validated, and fed into downstream business logic
+(work-item assignments, gate harness decisions, approval reviews, memory
+extraction, permission classifications, etc.).
+
+Historically, each call site re-implemented its own markdown-fence stripping,
+`json.loads` wrapping, and silent fallback. When the model occasionally
+produced malformed JSON or an unrecognised enum value, the whole orchestration
+step would either:
+
+* raise an unhandled exception (breaking the run), or
+* silently fall back to a degraded default (masking the real failure and
+  producing the wrong answer).
+
+This helper replaces that ad-hoc handling with a unified 3-attempt retry loop
+that feeds the model's previous error back into the next prompt so it can
+self-correct. It mirrors the retry pattern already used by
+the company runtime coordination builder and the recruiter:
+append a `retry_feedback` list to the payload on each retry, then re-invoke
+`LLMProvider.simple_chat`.
+
+Call sites that need more than a simple JSON parse can also pass a
+`validator` callable; if it returns a non-empty string, that string is treated
+as the error message for the next retry.
+"""
+
+from __future__ import annotations
+
+import json
+from typing import Any, Awaitable, Callable
+
+from loguru import logger
+
+from opc.llm.provider import LLMProvider
+
+
+Validator = Callable[[Any], "str | None"]
+
+
+class LLMRetryError(RuntimeError):
+    """Raised when the retry loop exhausts all attempts without success."""
+
+    def __init__(self, label: str, attempts: int, last_error: str, last_raw: str | None) -> None:
+        super().__init__(
+            f"[{label}] LLM JSON retry exhausted after {attempts} attempts. "
+            f"Last error: {last_error}"
+        )
+        self.label = label
+        self.attempts = attempts
+        self.last_error = last_error
+        self.last_raw = last_raw
+
+
+def _strip_fences(raw: str) -> str:
+    """Remove surrounding ```...``` fences, tolerant of missing newlines."""
+    text = str(raw or "").strip()
+    if text.startswith("```"):
+        parts = text.split("\n", 1)
+        text = parts[1] if len(parts) == 2 else text[3:]
+        if text.endswith("```"):
+            text = text[:-3]
+    return text.strip()
+
+
+def _compose_prompt(
+    payload: Any,
+    retry_feedback: list[str],
+) -> str:
+    """Build the prompt body for an attempt.
+
+    For dict payloads we copy and inject `retry_feedback` so the model can see
+    exactly why its previous attempt was rejected. For string payloads we
+    append a labelled section so we don't silently mutate the caller's text.
+    """
+    if isinstance(payload, dict):
+        attempt_payload = dict(payload)
+        if retry_feedback:
+            attempt_payload["retry_feedback"] = list(retry_feedback)
+        return json.dumps(attempt_payload, ensure_ascii=False)
+    text = str(payload or "")
+    if retry_feedback:
+        feedback_block = "\n".join(f"- {item}" for item in retry_feedback)
+        return (
+            f"{text}\n\n"
+            "# retry_feedback (your previous attempt failed; correct and retry)\n"
+            f"{feedback_block}"
+        )
+    return text
+
+
+async def call_llm_json_with_retry(
+    llm: LLMProvider,
+    *,
+    system: str,
+    payload: Any,
+    task_type: str = "quick_tasks",
+    validator: Validator | None = None,
+    max_attempts: int = 3,
+    label: str = "llm_json",
+    require_object: bool = True,
+) -> Any:
+    """Call `llm.simple_chat` expecting a JSON object, retrying on failure.
+
+    Args:
+        llm: LLMProvider instance.
+        system: System prompt describing the JSON contract.
+        payload: dict (preferred) or string. When a dict, the helper injects
+            a `retry_feedback` key on retries so the model can see its
+            previous error. When a string, the feedback is appended as a
+            trailing section.
+        task_type: routing hint forwarded to `simple_chat`.
+        validator: optional callable `validator(parsed) -> Optional[str]`.
+            Return `None` when the parsed object passes validation; return an
+            error description string to trigger a retry with that feedback.
+        max_attempts: total attempts including the first one (default 3).
+        label: short tag used in log messages and the raised exception.
+        require_object: when True (default) the parsed top-level value must
+            be a dict; otherwise any JSON value (list, scalar) is accepted.
+
+    Returns:
+        The parsed, validated JSON value.
+
+    Raises:
+        LLMRetryError: when all attempts fail. Callers typically catch this
+            and fall back to their heuristic default so the orchestration
+            stays alive.
+    """
+    retry_feedback: list[str] = []
+    last_error: str = "no attempt made"
+    last_raw: str | None = None
+
+    for attempt in range(1, max_attempts + 1):
+        prompt = _compose_prompt(payload, retry_feedback)
+        try:
+            raw = await llm.simple_chat(
+                prompt=prompt,
+                system=system,
+                task_type=task_type,
+            )
+        except Exception as exc:
+            last_error = f"LLM transport error: {exc}"
+            last_raw = None
+            logger.warning(
+                f"[{label}] attempt {attempt}/{max_attempts} transport error: {exc}"
+            )
+            retry_feedback.append(last_error)
+            continue
+
+        last_raw = raw
+        text = _strip_fences(raw)
+        if not text:
+            last_error = "Response was empty after stripping markdown fences."
+            logger.warning(
+                f"[{label}] attempt {attempt}/{max_attempts} empty response."
+            )
+            retry_feedback.append(last_error)
+            continue
+
+        try:
+            data = json.loads(text)
+        except json.JSONDecodeError as exc:
+            snippet = text[:200].replace("\n", "\\n")
+            last_error = (
+                f"Response was not valid JSON: {exc.msg} at char {exc.pos}. "
+                f"Snippet: {snippet}"
+            )
+            logger.warning(
+                f"[{label}] attempt {attempt}/{max_attempts} JSON parse failed: {exc}"
+            )
+            retry_feedback.append(last_error)
+            continue
+
+        if require_object and not isinstance(data, dict):
+            last_error = (
+                f"Top-level response must be a JSON object; got {type(data).__name__}."
+            )
+            logger.warning(f"[{label}] attempt {attempt}/{max_attempts}: {last_error}")
+            retry_feedback.append(last_error)
+            continue
+
+        if validator is not None:
+            try:
+                validation_error = validator(data)
+            except Exception as exc:
+                validation_error = f"Validator raised: {exc}"
+            if validation_error:
+                last_error = str(validation_error)
+                logger.warning(
+                    f"[{label}] attempt {attempt}/{max_attempts} validation failed: {last_error}"
+                )
+                retry_feedback.append(last_error)
+                continue
+
+        return data
+
+    raise LLMRetryError(
+        label=label,
+        attempts=max_attempts,
+        last_error=last_error,
+        last_raw=last_raw,
+    )
+
+
+async def call_llm_json_with_retry_custom(
+    llm: LLMProvider,
+    *,
+    system: str,
+    payload: Any,
+    task_type: str = "quick_tasks",
+    builder: Callable[[Any], Awaitable[Any]] | Callable[[Any], Any],
+    max_attempts: int = 3,
+    label: str = "llm_json_custom",
+) -> Any:
+    """Variant that lets the caller build/validate the final object in one pass.
+
+    The `builder` callable receives the parsed JSON value and must either
+    return the caller's normalized object (any type) or raise an exception
+    with a human-readable message. A raised exception is converted into
+    retry feedback, so the model can see exactly what failed during its
+    last attempt (e.g. dataclass construction, enum coercion, schema
+    normalization). This matches how the work-item assignment planner wraps its
+    `_sanitize_work_item_assignment_packet` call in an exception handler.
+    """
+    retry_feedback: list[str] = []
+    last_error: str = "no attempt made"
+    last_raw: str | None = None
+
+    for attempt in range(1, max_attempts + 1):
+        prompt = _compose_prompt(payload, retry_feedback)
+        try:
+            raw = await llm.simple_chat(
+                prompt=prompt,
+                system=system,
+                task_type=task_type,
+            )
+        except Exception as exc:
+            last_error = f"LLM transport error: {exc}"
+            last_raw = None
+            logger.warning(
+                f"[{label}] attempt {attempt}/{max_attempts} transport error: {exc}"
+            )
+            retry_feedback.append(last_error)
+            continue
+
+        last_raw = raw
+        text = _strip_fences(raw)
+        if not text:
+            last_error = "Response was empty after stripping markdown fences."
+            logger.warning(
+                f"[{label}] attempt {attempt}/{max_attempts} empty response."
+            )
+            retry_feedback.append(last_error)
+            continue
+
+        try:
+            data = json.loads(text)
+        except json.JSONDecodeError as exc:
+            snippet = text[:200].replace("\n", "\\n")
+            last_error = (
+                f"Response was not valid JSON: {exc.msg} at char {exc.pos}. "
+                f"Snippet: {snippet}"
+            )
+            logger.warning(
+                f"[{label}] attempt {attempt}/{max_attempts} JSON parse failed: {exc}"
+            )
+            retry_feedback.append(last_error)
+            continue
+
+        try:
+            result = builder(data)
+            if hasattr(result, "__await__"):
+                result = await result  # type: ignore[assignment]
+        except Exception as exc:
+            last_error = f"Response could not be normalized: {exc}"
+            logger.warning(
+                f"[{label}] attempt {attempt}/{max_attempts} build failed: {exc}"
+            )
+            retry_feedback.append(last_error)
+            continue
+
+        return result
+
+    raise LLMRetryError(
+        label=label,
+        attempts=max_attempts,
+        last_error=last_error,
+        last_raw=last_raw,
+    )