Retry a rejected report grouping instead of falling back

composer/spec/source/report/grouping.py

@@ -1,24 +1,30 @@
 """LLM-driven grouping of inferred properties into high-level audit claims.
 
-A single structured LLM call takes the `FormalizedProperty` list and partitions it into high-level
+A structured LLM call takes the `FormalizedProperty` list and partitions it into high-level
 `PropertyGroup`s (the "P-NN" headings) — each property in exactly one group, while the rules those
 properties are formalized by may surface under several groups. Each group's status is rolled up from
 its members' rules' verdicts. Groups are identified by the slug the LLM assigns — a per-run snapshot.
 
+A response that misses the schema costs the report every heading it has, so one thing stands
+between a malformed answer and the fallback: the call is retried once with the rejection appended.
+
 A single ``general`` fallback group (every property in one group) is used by `build` when the LLM
 call raises, validation rejects the grouping, or the grouping covers no properties.
 """
+import logging
 from typing import Iterable
 
 from langchain_core.language_models.chat_models import BaseChatModel
 from langchain_core.messages import HumanMessage, SystemMessage
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, ValidationError
 
 from composer.templates.loader import load_jinja_template
 from composer.spec.source.report.schema import (
     FormalizedProperty, GroupStatus, Outcome, PropertyGroup, PropertyKey, RuleRef,
 )
 
+_log = logging.getLogger(__name__)
+
 FALLBACK_SLUG = "general"
 FALLBACK_TITLE = "General"
 
@@ -78,19 +84,41 @@ async def call_grouping_llm(
     llm: BaseChatModel,
     contract_name: str,
     properties: list[FormalizedProperty],
+    max_attempts: int = 2,
 ) -> GroupingResult:
     """One structured LLM call: the property list in, a `GroupingResult` out, via langchain's
-    `with_structured_output`. The model + token budget come from the passed `llm`."""
+    `with_structured_output`. The model + token budget come from the passed `llm`.
+
+    A response that does not match the schema is retried once with the rejection appended, the
+    way `spec/prioritize.py` corrects a ranking. Without it a single malformed response costs the
+    report every heading it has, and the caller's fallback is meant for a grouping that could not
+    be obtained at all, not for one the model would have got right on a second look. The
+    correction is appended to the original request rather than sent as a follow-up turn, so the
+    retry stays one self-contained user message."""
     system = load_jinja_template("autoprove_report_grouping_system.j2")
     user = load_jinja_template(
         "autoprove_report_grouping_prompt.j2",
         contract_name=contract_name,
         properties=properties,
     )
     bound = llm.with_structured_output(GroupingResult)
-    result = await bound.ainvoke([SystemMessage(system), HumanMessage(user)])
-    assert isinstance(result, GroupingResult)
-    return result
+
+    prompt = user
+    for attempt in range(max_attempts):
+        try:
+            result = await bound.ainvoke([SystemMessage(system), HumanMessage(prompt)])
+        except ValidationError as e:
+            if attempt == max_attempts - 1:
+                raise
+            _log.warning("report grouping rejected (attempt %d): %s", attempt + 1, e)
+            prompt = (
+                f"{user}\n\nA previous attempt was rejected:\n\n{e}\n\n"
+                "Return the corrected grouping as structured output matching the schema."
+            )
+            continue
+        assert isinstance(result, GroupingResult)
+        return result
+    raise AssertionError("unreachable: the loop returns or raises")
 
 
 def build_groups(

tests/test_autoprove_report.py

@@ -14,6 +14,7 @@
 import pathlib
 
 import pytest
+from pydantic import ValidationError as PydanticValidationError
 from prover_output_utility.models import NodeStatus
 from prover_output_utility import ProverOutputAPI
 from langchain_core.language_models import BaseChatModel
@@ -30,6 +31,7 @@
 from composer.spec.source.report import build
 from composer.spec.source.report.collect import ReportComponentInput, collect
 from composer.spec.source.report.coverage import ValidationError, validate
+from composer.spec.source.report import grouping
 from composer.spec.source.report.grouping import (
     FALLBACK_SLUG, GroupingResult, PropertyGroupDraft, aggregate_status,
     build_fallback_grouping, build_groups,
@@ -592,6 +594,64 @@ async def test_build_groups_properties(tmp_path):
     assert report.coverage.property_coverage_complete is True
 
 
+class _FlakyStructuredModel(_StructuredStubModel):
+    """Raises the given exceptions on successive calls, then returns `output`. Records how many
+    times the structured binding was invoked."""
+    failures: list[Exception]
+    calls: list[str] = []
+
+    def with_structured_output(self, schema, **kwargs) -> Runnable:  # type: ignore[override]
+        out, failures, calls = self.output, self.failures, self.calls
+
+        def _invoke(messages):
+            calls.append(messages[-1].content)
+            if failures:
+                raise failures.pop(0)
+            return out
+
+        return RunnableLambda(_invoke)
+
+
+def _schema_error() -> PydanticValidationError:
+    try:
+        GroupingResult.model_validate({"groups": 17})
+    except PydanticValidationError as e:
+        return e
+    raise AssertionError("expected a validation error")
+
+
+@pytest.mark.asyncio
+async def test_grouping_retries_once_with_the_rejection_appended():
+    good = GroupingResult(groups=[PropertyGroupDraft(
+        slug="g", title="G", description="d", members=[("C", "p1")])])
+    llm = _FlakyStructuredModel(output=good, failures=[_schema_error()], calls=[])
+
+    result = await grouping.call_grouping_llm(
+        llm=llm, contract_name="C", properties=[_fp("C", "p1", [])],
+    )
+
+    assert [g.slug for g in result.groups] == ["g"]
+    assert len(llm.calls) == 2
+    # The retry carries the rejection, and still stands alone as one request.
+    assert "A previous attempt was rejected" in llm.calls[1]
+    assert "A previous attempt was rejected" not in llm.calls[0]
+
+
+@pytest.mark.asyncio
+async def test_grouping_gives_up_after_the_retry_so_the_caller_can_fall_back():
+    llm = _FlakyStructuredModel(
+        output=GroupingResult(groups=[]),
+        failures=[_schema_error(), _schema_error()],
+        calls=[],
+    )
+
+    with pytest.raises(PydanticValidationError):
+        await grouping.call_grouping_llm(
+            llm=llm, contract_name="C", properties=[_fp("C", "p1", [])],
+        )
+    assert len(llm.calls) == 2
+
+
 @pytest.mark.asyncio
 async def test_build_empty_grouping_falls_back(tmp_path):
     gen = _gen({"p1": ["r1"], "p2": ["r2"]})