nextcraft/apps/ai-service/ai_service/grading/engine.py

"""GradingEngine — rubric scoring over real process traces (REQ-3-004).

The Wave-2 composition of the grading stack:

    trace completeness gate (G-4, FIRST — nothing is sent to any LLM
    when the gate trips) → TraceStore.get_trace → compute_digest (D-028)
    → prompts.grading.render_trace_digest → D-020 4-layer structured
    defense (agents/structured.py — REUSED, composed, never duplicated)
    → RubricScore validation → GradeStore persistence → GradeRecord.

Gate verdicts are FIRST-CLASS RESULTS, not exceptions (G-4 is binding):
    UNGRADABLE_TRACE_INCOMPLETE — seq gaps in the store OR the trace is
        flagged by TraceIntegrityMap (INCOMPLETE_FLOODED). The gap list /
        flag reason is surfaced in `scores` for API rendering, and the
        record is PERSISTED like any grade so a learner sees why no
        credential can be issued for this trace — the gate outcome is
        durable and auditable, not a transient error string.
    UNGRADABLE_EMPTY_TRACE      — no events stored for the pair.
On either verdict `scores.criteria` is empty and the LLM is never called.

DI (D-027/D-032 house pattern): the engine receives trace_store,
grade_store, integrity and provider through the constructor and knows
NOTHING of FastAPI — api/ composes it (Task 3-3-01). `model` is injected
alongside the provider so tests script the mock against the production
wiring without touching Settings.

Boundary (D-027): grading/ imports llm/ (provider protocol + Message
type), telemetry/ (store + integrity map), prompts/ (rubric text) and
ONLY agents.structured — the sanctioned shared D-020 defense. We import
the MODULE directly (`from ..agents.structured import structured_completion`)
rather than the `agents` package, mirroring how agents/base.py consumes
it (same direct-module import): that keeps the dependency surface to
exactly the two names the engine needs (structured_completion,
StructuredOutputError) and avoids executing agents/__init__ re-exports
(BaseAgent, registry, session store) that grading has no business
loading — a side-effect-hygiene choice that keeps this import line
grep-auditable as "the one agents dependency". grading/ never imports api/.

RubricScore placement (documented decision): the validated output model
lives HERE, not in prompts/. The pydantic model is the engine's return
CONTRACT (the shape GradeStore.scores must hold), while prompts/grading.py
is pure prompt text + its mirror schema HINT string — the same split as
agents/assessor.py (model + hint) but with the model owned by the engine
module that validates it. Prompt files hold text, per the prompts/ house
style; engine files hold typed contracts.
"""

import logging
from datetime import UTC, datetime
from typing import Final

from pydantic import BaseModel, ConfigDict, Field, field_validator

from ..agents.structured import StructuredOutputError, structured_completion
from ..llm.base import LLMProvider
from ..llm.types import Message
from ..prompts.grading import (
    RUBRIC_CRITERIA,
    RUBRIC_SCORE_SCHEMA_HINT,
    SYSTEM_PROMPT,
    render_trace_digest,
)
from ..telemetry.ingest import TraceIntegrityMap
from ..telemetry.store import TraceStore
from .features import compute_digest
from .store import GradeRecord, GradeStore

logger = logging.getLogger(__name__)

#: First-class gate verdicts (G-4). GradeRecord.verdict values; non-empty
#: by store contract. Rubric verdicts (mastered/developing/not_yet) ride in
#: `scores.verdict` — `record.verdict` stays the machine-readable outcome.
VERDICT_UNGRADABLE_INCOMPLETE: Final = "UNGRADABLE_TRACE_INCOMPLETE"
VERDICT_UNGRADABLE_EMPTY: Final = "UNGRADABLE_EMPTY_TRACE"
#: record.verdict for a successfully LLM-graded trace (the rubric verdict
#: travels inside scores); keeps verdict non-empty for every persisted row.
VERDICT_GRADED: Final = "GRADED"

#: Gate-detail keys surfaced in GradeRecord.scores (tests assert on these).
_INTEGRITY_FLAG_KEY: Final = "integrity_flag"
_MISSING_SEQS_KEY: Final = "missing_seqs"


class RubricScore(BaseModel):
    """Validated LLM output: per-criterion 0-4 scores + strengths + gaps + verdict.

    The D-020 schema for the grader: `structured_completion` parses the
    model reply into THIS shape (layer 3), retrying once with the
    validation error fed back (layer 4). Exact criteria set + 0-4 ranges
    are enforced here, so the scores dict persisted to GradeStore is
    always rubric-shaped no matter what the model produced.
    """

    model_config = ConfigDict(extra="forbid")

    criteria: dict[str, int]
    strengths: list[str] = Field(min_length=1, max_length=2)
    gaps: list[str] = Field(min_length=1, max_length=2)
    verdict: str  # "mastered" | "developing" | "not_yet"

    @field_validator("criteria")
    @classmethod
    def _criteria_rubric_shaped(cls, value: dict[str, int]) -> dict[str, int]:
        """Exact criteria keys (no extras, no omissions) and 0-4 scores."""
        expected = set(RUBRIC_CRITERIA)
        got = set(value)
        if got != expected:
            raise ValueError(
                f"criteria keys must be exactly {sorted(expected)}, got {sorted(got)}"
            )
        for key, score in value.items():
            if not 0 <= score <= 4:
                raise ValueError(f"criterion {key!r} must be within 0-4, got {score}")
        return value

    @field_validator("verdict")
    @classmethod
    def _verdict_known(cls, value: str) -> str:
        allowed = {"mastered", "developing", "not_yet"}
        if value not in allowed:
            raise ValueError(f"verdict must be one of {sorted(allowed)}, got {value!r}")
        return value


class GradingEngine:
    """Scores a (learner_id, task_id) trace into a persisted GradeRecord."""

    def __init__(
        self,
        trace_store: TraceStore,
        grade_store: GradeStore,
        integrity: TraceIntegrityMap,
        provider: LLMProvider,
        *,
        model: str = "gemma4:31b",
    ) -> None:
        self._trace_store = trace_store
        self._grade_store = grade_store
        self._integrity = integrity
        self._provider = provider
        self._model = model

    async def grade(self, learner_id: str, task_id: str) -> GradeRecord:
        """Grade one trace; persist latest-state (GradeStore upserts); return it.

        Gate FIRST (G-4): the LLM is only ever reached from the fully
        guarded path — no gate state can be masked by an LLM error.
        """
        record = await self._grade(learner_id, task_id)
        self._grade_store.save(record)
        return record

    # ------------------------------------------------------------------ core

    async def _grade(self, learner_id: str, task_id: str) -> GradeRecord:
        # -- G-4 gate FIRST: integrity flag OR seq gaps. Ordering matters:
        # gaps() returns [] for an EMPTY trace, so the empty check below is
        # reachable only when no rows exist at all; a gapped or flooded
        # trace can never fall through to the LLM path.
        if self._integrity.is_incomplete(learner_id, task_id):
            reason = self._integrity.reason(learner_id, task_id) or "unknown"
            gaps = self._trace_store.gaps(learner_id, task_id)
            logger.info(
                "grade gate (G-4): %s/%s integrity-flagged (%s) — ungradable",
                learner_id,
                task_id,
                reason,
            )
            return self._ungradable(
                learner_id,
                task_id,
                detail={_INTEGRITY_FLAG_KEY: reason, _MISSING_SEQS_KEY: gaps},
                verdict=VERDICT_UNGRADABLE_INCOMPLETE,
            )
        gaps = self._trace_store.gaps(learner_id, task_id)
        if gaps:
            logger.info(
                "grade gate (G-4): %s/%s seq gaps %s — ungradable",
                learner_id,
                task_id,
                gaps,
            )
            return self._ungradable(
                learner_id,
                task_id,
                detail={_INTEGRITY_FLAG_KEY: None, _MISSING_SEQS_KEY: gaps},
                verdict=VERDICT_UNGRADABLE_INCOMPLETE,
            )

        trace = self._trace_store.get_trace(learner_id, task_id)
        if not trace:
            logger.info(
                "grade gate: %s/%s empty trace — ungradable", learner_id, task_id
            )
            return self._ungradable(
                learner_id,
                task_id,
                detail={_INTEGRITY_FLAG_KEY: None, _MISSING_SEQS_KEY: []},
                verdict=VERDICT_UNGRADABLE_EMPTY,
            )

        # -- Guarded path: digest (D-028) → prompt → D-020 4-layer defense.
        digest = compute_digest(trace)
        messages = [
            Message(role="system", content=SYSTEM_PROMPT),
            Message(role="user", content=render_trace_digest(digest)),
        ]
        try:
            rubric = await structured_completion(
                self._provider,
                messages,
                model=self._model,
                schema=RubricScore,
                schema_hint=RUBRIC_SCORE_SCHEMA_HINT,
            )
        except StructuredOutputError as exc:
            # The trace was gradable but the model failed to produce valid
            # JSON within the D-020 budget (two attempts). Raise — the API
            # layer maps this to a 502 (assessor precedent). Persisting a
            # fabricated or partial grade here would violate the no-silent-
            # fallback rule: no credential-worthy record without a validated
            # RubricScore.
            raise StructuredOutputError(f"grading LLM failed for {task_id}: {exc}") from exc

        logger.debug(
            "graded %s/%s: %s (model=%s)",
            learner_id,
            task_id,
            rubric.verdict,
            self._model,
        )
        return GradeRecord(
            learner_id=learner_id,
            task_id=task_id,
            variant_seed=None,  # null until P4 (D-029)
            digest=digest.model_dump(),
            scores=rubric.model_dump(),
            verdict=VERDICT_GRADED,
            model=self._model,
            created_at=datetime.now(tz=UTC),
        )

    # ------------------------------------------------------------ gate record

    @staticmethod
    def _ungradable(
        learner_id: str,
        task_id: str,
        *,
        detail: dict,
        verdict: str,
    ) -> GradeRecord:
        """Build a gate record: no digest (nothing was graded), gate detail
        surfaced in `scores` (the store allows an empty scores dict, but
        G-4 requires the gap list / flag reason surfaced — the detail IS the
        verdict's payload), model="none" (no LLM was involved; provenance
        stays honest).
        """
        return GradeRecord(
            learner_id=learner_id,
            task_id=task_id,
            variant_seed=None,
            digest={},
            scores=detail,
            verdict=verdict,
            model="none",
            created_at=datetime.now(tz=UTC),
        )