beaver_project/app-instance/backend/beaver/skills/learning/pipeline.py

"""Manual skill learning pipeline orchestration."""

from __future__ import annotations

from typing import Any, Callable

from beaver.engine.providers import ProviderBundle
from beaver.memory.skills import SkillDraftEvalReport, SkillDraftSafetyReport, SkillLearningCandidate, SkillLearningStore
from beaver.skills.drafts import DraftService
from beaver.skills.learning.eval import SkillDraftEvaluator
from beaver.skills.learning.replay import ReplayRunner
from beaver.skills.learning.service import SkillLearningService
from beaver.skills.learning.safety import SkillDraftSafetyChecker
from beaver.skills.publisher import SkillPublisher
from beaver.skills.reviews import ReviewService
from beaver.skills.specs import SkillDraft, SkillReviewRecord, SkillReviewState, SkillSpec, SkillVersion

_REJECTABLE_DRAFT_STATUSES = {
    SkillReviewState.DRAFT.value,
    SkillReviewState.IN_REVIEW.value,
    SkillReviewState.APPROVED.value,
}


class SkillLearningPipelineService:
    """Coordinates candidate -> draft -> review -> publish lifecycle."""

    def __init__(
        self,
        *,
        learning_store: SkillLearningStore,
        learning_service: SkillLearningService,
        draft_service: DraftService,
        review_service: ReviewService,
        publisher: SkillPublisher,
        safety_checker: SkillDraftSafetyChecker | None = None,
        evaluator: SkillDraftEvaluator | None = None,
    ) -> None:
        self.learning_store = learning_store
        self.learning_service = learning_service
        self.draft_service = draft_service
        self.review_service = review_service
        self.publisher = publisher
        self.safety_checker = safety_checker or SkillDraftSafetyChecker()
        self.evaluator = evaluator

    def list_candidates(self, status: str | None = None) -> list[SkillLearningCandidate]:
        return self.learning_store.list_learning_candidates(status=status)

    def get_candidate(self, candidate_id: str) -> SkillLearningCandidate:
        for candidate in self.learning_store.list_learning_candidates():
            if candidate.candidate_id == candidate_id:
                return candidate
        raise ValueError(f"Unknown learning candidate: {candidate_id}")

    async def synthesize_draft(
        self,
        candidate_id: str,
        *,
        provider_bundle: ProviderBundle,
    ) -> SkillDraft:
        draft = await self.learning_service.synthesize_draft(candidate_id, provider_bundle)
        self.mark_draft_synthesized(candidate_id, draft)
        return draft

    async def regenerate_draft(
        self,
        candidate_id: str,
        *,
        provider_bundle: ProviderBundle,
    ) -> SkillDraft:
        self.learning_store.transition_learning_candidate(
            candidate_id,
            "synthesizing",
            event_type="draft_synthesis_started",
            last_error=None,
        )
        return await self.synthesize_draft(candidate_id, provider_bundle=provider_bundle)

    def mark_candidate_queued(self, candidate_id: str) -> SkillLearningCandidate:
        return self._require_updated(
            self.learning_store.transition_learning_candidate(
                candidate_id,
                "queued",
                event_type="candidate_queued",
                last_error=None,
            ),
            candidate_id,
        )

    def mark_candidate_synthesizing(self, candidate_id: str) -> SkillLearningCandidate:
        return self._require_updated(
            self.learning_store.transition_learning_candidate(
                candidate_id,
                "synthesizing",
                event_type="draft_synthesis_started",
                last_error=None,
            ),
            candidate_id,
        )

    def mark_draft_synthesized(self, candidate_id: str, draft: SkillDraft) -> SkillLearningCandidate:
        candidate = self.get_candidate(candidate_id)
        evidence = dict(candidate.evidence)
        evidence["draft_id"] = draft.draft_id
        evidence["draft_skill_name"] = draft.skill_name
        return self._require_updated(
            self.learning_store.transition_learning_candidate(
                candidate_id,
                "draft_ready",
                event_type="draft_synthesis_completed",
                evidence=evidence,
                draft_id=draft.draft_id,
                draft_skill_name=draft.skill_name,
                risk_level=candidate.risk_level,
                last_error=None,
                payload={"draft_id": draft.draft_id, "skill_name": draft.skill_name},
            ),
            candidate_id,
        )

    def mark_candidate_failed(
        self,
        candidate_id: str,
        error: str,
        *,
        retry_count: int,
        terminal: bool,
    ) -> SkillLearningCandidate:
        return self._require_updated(
            self.learning_store.transition_learning_candidate(
                candidate_id,
                "failed" if terminal else "open",
                event_type="failed",
                retry_count=retry_count,
                last_error=error,
                payload={"error": error, "terminal": terminal, "retry_count": retry_count},
            ),
            candidate_id,
        )

    def mark_candidate_superseded(self, candidate_id: str, reason: str) -> SkillLearningCandidate:
        return self._require_updated(
            self.learning_store.transition_learning_candidate(
                candidate_id,
                "superseded",
                event_type="superseded",
                last_error=reason,
                payload={"reason": reason},
            ),
            candidate_id,
        )

    def list_drafts(self, skill_name: str | None = None) -> list[SkillDraft]:
        return self.draft_service.list_drafts(skill_name)

    def get_draft(self, skill_name: str, draft_id: str) -> SkillDraft:
        draft = self.draft_service.get_draft(skill_name, draft_id)
        if draft is None:
            raise ValueError(f"Draft not found: {skill_name}/{draft_id}")
        return draft

    def submit_review(
        self,
        skill_name: str,
        draft_id: str,
        *,
        requested_by: str = "system",
        notes: str = "",
    ) -> SkillReviewRecord:
        draft = self.get_draft(skill_name, draft_id)
        if draft.status != SkillReviewState.DRAFT.value:
            raise ValueError("Draft must be in draft status before review submission")
        safety = self.get_safety_report(skill_name, draft_id)
        if safety is not None and (not safety.passed or safety.risk_level == "critical"):
            raise ValueError("Draft cannot enter review because safety check failed")
        review = self.review_service.submit_for_review(
            skill_name,
            draft_id,
            reviewer_request=notes,
            requested_by=requested_by,
        )
        self._mark_candidate_by_draft(
            skill_name,
            draft_id,
            "review_pending",
            "review_submitted",
            last_error=None,
        )
        return review

    def approve(
        self,
        skill_name: str,
        draft_id: str,
        *,
        reviewer: str = "system",
        notes: str = "",
    ) -> SkillReviewRecord:
        draft = self.get_draft(skill_name, draft_id)
        if draft.status != SkillReviewState.IN_REVIEW.value:
            raise ValueError("Draft must be in review before approval")
        safety = self.get_safety_report(skill_name, draft_id)
        if safety is not None and (not safety.passed or safety.risk_level == "critical"):
            raise ValueError("Draft cannot be approved because safety check failed")
        review = self.review_service.approve(skill_name, draft_id, reviewer=reviewer, notes=notes)
        self._mark_candidate_by_draft(skill_name, draft_id, "approved", "approved")
        return review

    def reject(
        self,
        skill_name: str,
        draft_id: str,
        *,
        reviewer: str = "system",
        notes: str = "",
    ) -> SkillReviewRecord:
        draft = self.get_draft(skill_name, draft_id)
        if draft.status not in _REJECTABLE_DRAFT_STATUSES:
            raise ValueError("Draft is not rejectable from its current status")
        review = self.review_service.reject(skill_name, draft_id, reviewer=reviewer, notes=notes)
        self._mark_candidate_by_draft(skill_name, draft_id, "rejected", "rejected")
        return review

    def publish(
        self,
        skill_name: str,
        draft_id: str,
        *,
        publisher: str = "system",
        notes: str = "",
        confirm_high_risk: bool = False,
    ) -> SkillVersion | SkillSpec:
        draft = self.get_draft(skill_name, draft_id)
        self._validate_publish_gates(draft, confirm_high_risk=confirm_high_risk)
        if draft.proposal_kind == "retire_skill":
            result = self.publisher.apply_retire_proposal(skill_name, draft_id, actor=publisher, notes=notes)
        else:
            result = self.publisher.publish(skill_name, draft_id, publisher=publisher, notes=notes)
        self._mark_candidate_by_draft(skill_name, draft_id, "published", "published")
        return result

    def rollback(
        self,
        skill_name: str,
        target_version: str,
        *,
        actor: str = "system",
        reason: str = "",
    ) -> SkillSpec:
        return self.publisher.rollback(skill_name, target_version, actor=actor, reason=reason or "manual rollback")

    def disable(
        self,
        skill_name: str,
        *,
        actor: str = "system",
        reason: str = "",
    ) -> SkillSpec:
        return self.publisher.disable(skill_name, actor=actor, reason=reason or "manual disable")

    def reviews_for_draft(self, skill_name: str, draft_id: str) -> list[SkillReviewRecord]:
        return self.review_service.store.list_reviews(skill_name, draft_id=draft_id)

    def check_safety(self, skill_name: str, draft_id: str) -> SkillDraftSafetyReport:
        draft = self.get_draft(skill_name, draft_id)
        report = self.safety_checker.check(draft)
        self.learning_store.write_safety_report(report)
        status = (
            "safety_failed"
            if not report.passed or report.risk_level == "critical"
            else self._candidate_status_for_draft(draft)
        )
        current = self._candidate_by_draft(skill_name, draft_id)
        if current is not None and current.status == "eval_failed" and status != "safety_failed":
            status = "eval_failed"
        self._mark_candidate_by_draft(
            skill_name,
            draft_id,
            status,
            "safety_checked",
            safety_report_id=report.report_id,
            risk_level=report.risk_level,
            last_error="; ".join(report.blocked_reasons) if status == "safety_failed" else None,
        )
        return report

    def get_safety_report(self, skill_name: str, draft_id: str) -> SkillDraftSafetyReport | None:
        return self.learning_store.get_safety_report(skill_name, draft_id)

    def get_eval_report(self, skill_name: str, draft_id: str) -> SkillDraftEvalReport | None:
        return self.learning_store.get_eval_report(skill_name, draft_id)

    async def evaluate_draft(
        self,
        candidate_id: str,
        skill_name: str,
        draft_id: str,
        *,
        provider_bundle: ProviderBundle | None,
        replay_runner: ReplayRunner | None = None,
        progress_callback: Callable[[dict[str, Any]], None] | None = None,
    ) -> SkillDraftEvalReport:
        draft = self.get_draft(skill_name, draft_id)
        candidate = self.get_candidate(candidate_id)
        evaluator = self.evaluator or SkillDraftEvaluator(self.learning_service.run_store)
        report = await evaluator.evaluate(
            candidate=candidate,
            draft=draft,
            provider_bundle=provider_bundle,
            replay_runner=replay_runner,
            progress_callback=progress_callback,
        )
        self.learning_store.write_eval_report(report)
        if report.status == "skipped_provider_unavailable":
            status = self._candidate_status_for_draft(draft)
            error = "eval skipped: provider unavailable"
        elif report.passed:
            status = self._candidate_status_for_draft(draft)
            error = None
        else:
            status = "eval_failed"
            error = "eval failed"
        current = self._candidate_by_draft(skill_name, draft_id)
        if current is not None and current.status == "safety_failed" and status == "draft_ready":
            status = "safety_failed"
            error = current.last_error
        self.learning_store.transition_learning_candidate(
            candidate_id,
            status,
            event_type="eval_completed",
            eval_report_id=report.report_id,
            eval_progress={
                "phase": "completed",
                "completed_arms": len(report.cases) * 2 if report.mode == "replay" else 0,
                "total_arms": len(report.cases) * 2 if report.mode == "replay" else 0,
                "completed_cases": len(report.cases),
                "total_cases": len(report.cases),
            },
            last_error=error,
            payload=report.to_dict(),
        )
        return report

    def mark_eval_progress(self, candidate_id: str, progress: dict[str, Any]) -> SkillLearningCandidate:
        return self._require_updated(
            self.learning_store.update_learning_candidate(
                candidate_id,
                eval_progress=dict(progress),
            ),
            candidate_id,
        )

    def mark_eval_failed(self, candidate_id: str, error: str) -> SkillLearningCandidate:
        candidate = self.get_candidate(candidate_id)
        progress = dict(candidate.eval_progress)
        progress["phase"] = "failed"
        return self._require_updated(
            self.learning_store.transition_learning_candidate(
                candidate_id,
                "eval_failed",
                eval_progress=progress,
                event_type="eval_failed",
                last_error=error,
                payload={"error": error},
            ),
            candidate_id,
        )

    def _validate_publish_gates(self, draft: SkillDraft, *, confirm_high_risk: bool) -> None:
        reviews = self.reviews_for_draft(draft.skill_name, draft.draft_id)
        if not any(review.status in {SkillReviewState.IN_REVIEW.value, SkillReviewState.APPROVED.value} for review in reviews):
            raise ValueError("Draft must be submitted for review before publish")
        safety = self.get_safety_report(draft.skill_name, draft.draft_id)
        if safety is None:
            raise ValueError("Draft requires a passing safety report before publish")
        if not safety.passed:
            raise ValueError("Draft safety report did not pass")
        if safety.risk_level == "critical":
            raise ValueError("Critical risk drafts cannot be published")
        if safety.risk_level == "high" and not confirm_high_risk:
            raise ValueError("High risk draft publish requires confirm_high_risk=true")
        eval_report = self.get_eval_report(draft.skill_name, draft.draft_id)
        if eval_report is not None and eval_report.status != "skipped_provider_unavailable" and not eval_report.passed:
            raise ValueError("Draft eval report did not pass")
        if eval_report is not None and eval_report.mode == "replay":
            if eval_report.confidence == "low":
                raise ValueError("Draft replay eval has low confidence and requires revision before publish")
            if eval_report.blocked_coverage >= 1.0:
                raise ValueError("Draft replay eval blocked all important tool calls")
            preservation = eval_report.preservation_report or {}
            if preservation.get("passed") is False:
                raise ValueError("Draft preservation check did not pass")

    def _mark_candidate_by_draft(
        self,
        skill_name: str,
        draft_id: str,
        status: str,
        event_type: str,
        **updates: object,
    ) -> SkillLearningCandidate | None:
        candidate = self._candidate_by_draft(skill_name, draft_id)
        if candidate is None:
            return None
        if candidate.status in {"safety_failed", "eval_failed"} and status in {"review_pending", "approved"}:
            return candidate
        return self.learning_store.transition_learning_candidate(
            candidate.candidate_id,
            status,
            event_type=event_type,
            **updates,
        )

    def _candidate_by_draft(self, skill_name: str, draft_id: str) -> SkillLearningCandidate | None:
        for candidate in self.learning_store.list_learning_candidates():
            if candidate.draft_skill_name == skill_name and candidate.draft_id == draft_id:
                return candidate
        return None

    @staticmethod
    def _candidate_status_for_draft(draft: SkillDraft) -> str:
        if draft.status == SkillReviewState.APPROVED.value:
            return "approved"
        if draft.status == SkillReviewState.IN_REVIEW.value:
            return "review_pending"
        return "draft_ready"

    @staticmethod
    def _require_updated(candidate: SkillLearningCandidate | None, candidate_id: str) -> SkillLearningCandidate:
        if candidate is None:
            raise ValueError(f"Unknown learning candidate: {candidate_id}")
        return candidate


def model_to_dict(value: Any) -> dict[str, Any]:
    if hasattr(value, "to_dict"):
        return value.to_dict()
    if isinstance(value, dict):
        return dict(value)
    raise TypeError(f"Cannot convert {type(value).__name__} to dict")