ss-tools/backend/src/services/dataset_review/clarification_engine.py

# #region ClarificationEngine [C:4] [TYPE Module] [SEMANTICS pydantic, dataset, clarification, finding, resolution]
# @BRIEF Manage one-question-at-a-time clarification state, deterministic answer persistence, and readiness/finding updates.
# @LAYER Domain
# @RELATION DEPENDS_ON -> [DatasetReviewSessionRepository]
# @RELATION DEPENDS_ON -> [ClarificationSession]
# @RELATION DEPENDS_ON -> [ClarificationQuestion]
# @RELATION DEPENDS_ON -> [ClarificationAnswer]
# @RELATION DEPENDS_ON -> [ValidationFinding]
# @RELATION DISPATCHES -> [ClarificationHelpers]
# @PRE  Target session contains a persisted clarification aggregate in the current ownership scope.
# @POST  Active clarification payload exposes one highest-priority unresolved question, and each recorded answer is persisted before pointer/readiness mutation.
# @SIDE_EFFECT  Persists clarification answers, question/session states, and related readiness/finding changes.
# @DATA_CONTRACT  Input[DatasetReviewSession|ClarificationAnswerCommand] -> Output[ClarificationStateResult]
# @INVARIANT  Only one active clarification question may exist at a time; skipped and expert-review items remain unresolved and visible.
# @RATIONALE  Original 635-line file exceeded INV_7 (400-line module limit). Extracted pure helpers into _helpers sub-module.
# @REJECTED  Keeping all clarification logic in one file because it exceeded the fractal limit.

from __future__ import annotations

from dataclasses import dataclass, field
from datetime import datetime

from src.core.logger import belief_scope, logger
from src.models.auth import User
from src.models.dataset_review import (
    AnswerKind,
    ClarificationAnswer,
    ClarificationQuestion,
    ClarificationSession,
    ClarificationStatus,
    DatasetReviewSession,
    QuestionState,
    ReadinessState,
    RecommendedAction,
    SessionPhase,
    ValidationFinding,
)
from src.services.dataset_review.clarification_pkg._helpers import (
    build_impact_summary,
    count_remaining_questions,
    count_resolved_questions,
    derive_readiness_state,
    derive_recommended_action,
    normalize_answer_value,
    select_next_open_question,
    upsert_clarification_finding,
)
from src.services.dataset_review.repositories.session_repository import (
    DatasetReviewSessionRepository,
)


# #region ClarificationQuestionPayload [C:2] [TYPE Class]
# @BRIEF Typed active-question payload returned to the API layer.
@dataclass
class ClarificationQuestionPayload:
    question_id: str
    clarification_session_id: str
    topic_ref: str
    question_text: str
    why_it_matters: str
    current_guess: str | None
    priority: int
    state: QuestionState
    options: list[dict[str, object]] = field(default_factory=list)


# #endregion ClarificationQuestionPayload


# #region ClarificationStateResult [C:2] [TYPE Class]
# @BRIEF Clarification state result carrying the current session, active payload, and changed findings.
@dataclass
class ClarificationStateResult:
    clarification_session: ClarificationSession
    current_question: ClarificationQuestionPayload | None
    session: DatasetReviewSession
    changed_findings: list[ValidationFinding] = field(default_factory=list)


# #endregion ClarificationStateResult


# #region ClarificationAnswerCommand [C:2] [TYPE Class]
# @BRIEF Typed answer command for clarification state mutation.
@dataclass
class ClarificationAnswerCommand:
    session: DatasetReviewSession
    question_id: str
    answer_kind: AnswerKind
    answer_value: str | None
    user: User


# #endregion ClarificationAnswerCommand


# #region ClarificationEngine [C:4] [TYPE Class]
# @BRIEF Provide deterministic one-question-at-a-time clarification selection and answer persistence.
# @RELATION DEPENDS_ON -> [DatasetReviewSessionRepository]
# @RELATION CALLS -> [ClarificationHelpers]
# @PRE  Repository is bound to the current request transaction scope.
# @POST  Returned clarification state is persistence-backed and aligned with session readiness/recommended action.
# @SIDE_EFFECT  Mutates clarification answers, session flags, and related clarification findings.
class ClarificationEngine:
    # region ClarificationEngine_init [TYPE Function]
    # @PURPOSE: Bind repository dependency for clarification persistence operations.
    def __init__(self, repository: DatasetReviewSessionRepository) -> None:
        self.repository = repository

    # endregion ClarificationEngine_init

    # region build_question_payload [TYPE Function]
    # @PURPOSE: Return the one active highest-priority clarification question payload.
    # @PRE  Session contains unresolved clarification state or a resumable clarification session.
    # @POST  Returns exactly one active/open question payload or None when no unresolved question remains.
    # @SIDE_EFFECT  Normalizes the active-question pointer and clarification status in persistence.
    def build_question_payload(
        self, session: DatasetReviewSession,
    ) -> ClarificationQuestionPayload | None:
        with belief_scope("ClarificationEngine.build_question_payload"):
            clarification_session = self._get_latest_clarification_session(session)
            if clarification_session is None:
                logger.reason("No clarification session found", extra={"session_id": session.session_id})
                return None

            active_questions = [
                q for q in clarification_session.questions if q.state == QuestionState.OPEN
            ]
            active_questions.sort(key=lambda item: (-int(item.priority), item.created_at, item.question_id))

            if not active_questions:
                clarification_session.current_question_id = None
                clarification_session.status = ClarificationStatus.COMPLETED
                session.readiness_state = derive_readiness_state(session, clarification_session)
                session.recommended_action = derive_recommended_action(session, clarification_session)
                if session.current_phase == SessionPhase.CLARIFICATION:
                    session.current_phase = SessionPhase.REVIEW
                self.repository.db.commit()
                logger.reflect("No unresolved clarification question remains", extra={"session_id": session.session_id})
                return None

            selected_question = active_questions[0]
            clarification_session.current_question_id = selected_question.question_id
            clarification_session.status = ClarificationStatus.ACTIVE
            session.readiness_state = ReadinessState.CLARIFICATION_ACTIVE
            session.recommended_action = RecommendedAction.ANSWER_NEXT_QUESTION
            session.current_phase = SessionPhase.CLARIFICATION

            logger.reason("Selected active clarification question", extra={"session_id": session.session_id, "question_id": selected_question.question_id, "priority": selected_question.priority})
            self.repository.db.commit()

            payload = ClarificationQuestionPayload(
                question_id=selected_question.question_id,
                clarification_session_id=selected_question.clarification_session_id,
                topic_ref=selected_question.topic_ref,
                question_text=selected_question.question_text,
                why_it_matters=selected_question.why_it_matters,
                current_guess=selected_question.current_guess,
                priority=selected_question.priority,
                state=selected_question.state,
                options=[
                    {"option_id": o.option_id, "question_id": o.question_id, "label": o.label, "value": o.value, "is_recommended": o.is_recommended, "display_order": o.display_order}
                    for o in sorted(selected_question.options, key=lambda item: (item.display_order, item.label, item.option_id))
                ],
            )
            logger.reflect("Clarification payload built", extra={"session_id": session.session_id, "question_id": payload.question_id, "option_count": len(payload.options)})
            return payload

    # endregion build_question_payload

    # region record_answer [TYPE Function]
    # @PURPOSE: Persist one clarification answer before any pointer/readiness mutation.
    # @PRE  Target question belongs to the session's active clarification session and is still open.
    # @POST  Answer row is persisted before current-question pointer advances.
    # @SIDE_EFFECT  Inserts answer row, mutates question/session states, updates clarification findings, and commits.
    def record_answer(self, command: ClarificationAnswerCommand) -> ClarificationStateResult:
        with belief_scope("ClarificationEngine.record_answer"):
            session = command.session
            clarification_session = self._get_latest_clarification_session(session)
            if clarification_session is None:
                logger.explore("Cannot record clarification answer because no clarification session exists", extra={"session_id": session.session_id})
                raise ValueError("Clarification session not found")

            question = self._find_question(clarification_session, command.question_id)
            if question is None:
                logger.explore("Cannot record clarification answer for foreign or missing question", extra={"session_id": session.session_id, "question_id": command.question_id})
                raise ValueError("Clarification question not found")

            if question.answer is not None:
                logger.explore("Rejected duplicate clarification answer submission", extra={"session_id": session.session_id, "question_id": command.question_id})
                raise ValueError("Clarification question already answered")

            if clarification_session.current_question_id and clarification_session.current_question_id != question.question_id:
                logger.explore("Rejected answer for non-active clarification question", extra={"session_id": session.session_id, "question_id": question.question_id, "current_question_id": clarification_session.current_question_id})
                raise ValueError("Only the active clarification question can be answered")

            normalized_answer_value = normalize_answer_value(command.answer_kind, command.answer_value, question)

            logger.reason("Persisting clarification answer before state advancement", extra={"session_id": session.session_id, "question_id": question.question_id, "answer_kind": command.answer_kind.value})
            persisted_answer = ClarificationAnswer(
                question_id=question.question_id,
                answer_kind=command.answer_kind,
                answer_value=normalized_answer_value,
                answered_by_user_id=command.user.id,
                impact_summary=build_impact_summary(question, command.answer_kind, normalized_answer_value),
            )
            self.repository.db.add(persisted_answer)
            self.repository.db.flush()

            changed_finding = upsert_clarification_finding(
                session=session, question=question, answer_kind=command.answer_kind,
                answer_value=normalized_answer_value, db_session=self.repository.db,
            )

            if command.answer_kind == AnswerKind.SELECTED or command.answer_kind == AnswerKind.CUSTOM:
                question.state = QuestionState.ANSWERED
            elif command.answer_kind == AnswerKind.SKIPPED:
                question.state = QuestionState.SKIPPED
            elif command.answer_kind == AnswerKind.EXPERT_REVIEW:
                question.state = QuestionState.EXPERT_REVIEW

            question.updated_at = datetime.utcnow()
            self.repository.db.flush()

            clarification_session.resolved_count = count_resolved_questions(clarification_session)
            clarification_session.remaining_count = count_remaining_questions(clarification_session)
            clarification_session.summary_delta = self.summarize_progress(clarification_session)
            clarification_session.updated_at = datetime.utcnow()

            next_question = select_next_open_question(clarification_session)
            clarification_session.current_question_id = next_question.question_id if next_question else None
            clarification_session.status = ClarificationStatus.ACTIVE if next_question else ClarificationStatus.COMPLETED
            if clarification_session.status == ClarificationStatus.COMPLETED:
                clarification_session.completed_at = datetime.utcnow()

            session.readiness_state = derive_readiness_state(session, clarification_session)
            session.recommended_action = derive_recommended_action(session, clarification_session)
            session.current_phase = SessionPhase.CLARIFICATION if clarification_session.current_question_id else SessionPhase.REVIEW
            self.repository.bump_session_version(session)

            self.repository.db.commit()
            self.repository.db.refresh(session)

            logger.reflect("Clarification answer recorded and session advanced", extra={"session_id": session.session_id, "question_id": question.question_id, "next_question_id": clarification_session.current_question_id, "readiness_state": session.readiness_state.value, "remaining_count": clarification_session.remaining_count})

            return ClarificationStateResult(
                clarification_session=clarification_session,
                current_question=self.build_question_payload(session),
                session=session,
                changed_findings=[changed_finding] if changed_finding else [],
            )

    # endregion record_answer

    # region summarize_progress [TYPE Function]
    # @PURPOSE: Produce a compact progress summary for pause/resume and completion UX.
    def summarize_progress(self, clarification_session: ClarificationSession) -> str:
        resolved = count_resolved_questions(clarification_session)
        remaining = count_remaining_questions(clarification_session)
        return f"{resolved} resolved, {remaining} unresolved"

    # endregion summarize_progress

    # region _get_latest_clarification_session [TYPE Function]
    # @PURPOSE: Select the latest clarification session for the current dataset review aggregate.
    def _get_latest_clarification_session(self, session: DatasetReviewSession) -> ClarificationSession | None:
        if not session.clarification_sessions:
            return None
        ordered = sorted(session.clarification_sessions, key=lambda item: (item.started_at, item.clarification_session_id), reverse=True)
        return ordered[0]

    # endregion _get_latest_clarification_session

    # region _find_question [TYPE Function]
    # @PURPOSE: Resolve a clarification question from the active clarification aggregate.
    def _find_question(self, clarification_session: ClarificationSession, question_id: str) -> ClarificationQuestion | None:
        for q in clarification_session.questions:
            if q.question_id == question_id:
                return q
        return None

    # endregion _find_question


# #endregion ClarificationEngine

# #endregion ClarificationEngine