diff --git a/composer/spec/source/report/grouping.py b/composer/spec/source/report/grouping.py index cb952edd..8b92df00 100644 --- a/composer/spec/source/report/grouping.py +++ b/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,9 +84,17 @@ 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", @@ -88,9 +102,23 @@ async def call_grouping_llm( 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( diff --git a/tests/test_autoprove_report.py b/tests/test_autoprove_report.py index 542a9d49..28aa15ea 100644 --- a/tests/test_autoprove_report.py +++ b/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"]})