beaver_project/app-instance/backend/beaver/tasks/models.py

"""Models for internal task tracking and user acceptance."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Literal


ValidationStatus = Literal["accepted", "rejected", "insufficient_evidence", "validator_error"]

VALIDATION_STATUSES = {"accepted", "rejected", "insufficient_evidence", "validator_error"}
TASK_OPEN_STATUSES = {"open", "running", "awaiting_acceptance", "needs_revision"}
LEGACY_STATUS_MAP = {
    "validating": "running",
    "awaiting_feedback": "awaiting_acceptance",
    "needs_review": "awaiting_acceptance",
}


@dataclass(slots=True)
class ValidationResult:
    status: ValidationStatus = "rejected"
    score: float = 0.0
    issues: list[str] = field(default_factory=list)
    missing_requirements: list[str] = field(default_factory=list)
    evidence_gaps: list[str] = field(default_factory=list)
    recommended_revision_prompt: str = ""
    validator: str = "heuristic"

    def __init__(
        self,
        *,
        status: ValidationStatus | None = None,
        passed: bool | None = None,
        score: float = 0.0,
        issues: list[str] | None = None,
        missing_requirements: list[str] | None = None,
        evidence_gaps: list[str] | None = None,
        recommended_revision_prompt: str = "",
        validator: str = "heuristic",
    ) -> None:
        if status is not None and status not in VALIDATION_STATUSES:
            raise ValueError(f"unknown validation status: {status}")
        self.status = status or ("accepted" if passed and score >= 0.75 else "rejected")
        self.score = max(0.0, min(1.0, float(score or 0.0)))
        self.issues = list(issues or [])
        self.missing_requirements = list(missing_requirements or [])
        self.evidence_gaps = list(evidence_gaps or [])
        self.recommended_revision_prompt = recommended_revision_prompt
        self.validator = validator

    @property
    def passed(self) -> bool:
        return self.status == "accepted"

    @property
    def accepted(self) -> bool:
        return self.status == "accepted"

    def to_dict(self) -> dict[str, Any]:
        return {
            "status": self.status,
            "passed": self.passed,
            "score": self.score,
            "issues": list(self.issues),
            "missing_requirements": list(self.missing_requirements),
            "evidence_gaps": list(self.evidence_gaps),
            "recommended_revision_prompt": self.recommended_revision_prompt,
            "validator": self.validator,
            "accepted": self.accepted,
        }

    @classmethod
    def from_dict(cls, payload: dict[str, Any] | None) -> "ValidationResult | None":
        if not isinstance(payload, dict):
            return None
        raw_status = payload.get("status")
        if "status" in payload and raw_status not in VALIDATION_STATUSES:
            raise ValueError(f"unknown validation status: {raw_status}")
        status: ValidationStatus | None = raw_status if "status" in payload else None
        return cls(
            status=status,
            passed=bool(payload.get("passed")) if "status" not in payload else None,
            score=float(payload.get("score", 0.0) or 0.0),
            issues=[str(item) for item in payload.get("issues") or []],
            missing_requirements=[str(item) for item in payload.get("missing_requirements") or []],
            evidence_gaps=[str(item) for item in payload.get("evidence_gaps") or []],
            recommended_revision_prompt=str(payload.get("recommended_revision_prompt") or ""),
            validator=str(payload.get("validator") or "unknown"),
        )


@dataclass(slots=True)
class TaskRecord:
    task_id: str
    session_id: str
    description: str
    goal: str
    constraints: list[str]
    priority: int
    status: str
    creator: str
    created_at: str
    updated_at: str
    parent_task_id: str | None = None
    closed_at: str | None = None
    close_reason: str | None = None
    satisfaction: float | None = None
    run_ids: list[str] = field(default_factory=list)
    skill_names: list[str] = field(default_factory=list)
    feedback: list[dict[str, Any]] = field(default_factory=list)
    validation_result: dict[str, Any] | None = None
    metadata: dict[str, Any] = field(default_factory=dict)

    @property
    def is_open(self) -> bool:
        return self.status in TASK_OPEN_STATUSES

    @property
    def is_execution_active(self) -> bool:
        return self.status == "running"

    @property
    def requires_user_action(self) -> bool:
        return self.status in {"awaiting_acceptance", "needs_revision"}

    def to_dict(self) -> dict[str, Any]:
        return {
            "task_id": self.task_id,
            "session_id": self.session_id,
            "parent_task_id": self.parent_task_id,
            "description": self.description,
            "goal": self.goal,
            "constraints": list(self.constraints),
            "priority": self.priority,
            "status": self.status,
            "creator": self.creator,
            "created_at": self.created_at,
            "updated_at": self.updated_at,
            "closed_at": self.closed_at,
            "close_reason": self.close_reason,
            "satisfaction": self.satisfaction,
            "run_ids": list(self.run_ids),
            "skill_names": list(self.skill_names),
            "acceptance": list(self.feedback),
            "feedback": list(self.feedback),
            "validation_result": self.validation_result,
            "metadata": dict(self.metadata),
        }

    @classmethod
    def from_dict(cls, payload: dict[str, Any]) -> "TaskRecord":
        return cls(
            task_id=str(payload["task_id"]),
            session_id=str(payload["session_id"]),
            parent_task_id=_optional_str(payload.get("parent_task_id")),
            description=str(payload.get("description") or ""),
            goal=str(payload.get("goal") or payload.get("description") or ""),
            constraints=[str(item) for item in payload.get("constraints") or []],
            priority=int(payload.get("priority", 0) or 0),
            status=LEGACY_STATUS_MAP.get(str(payload.get("status") or "open"), str(payload.get("status") or "open")),
            creator=str(payload.get("creator") or "main-agent"),
            created_at=str(payload.get("created_at") or ""),
            updated_at=str(payload.get("updated_at") or ""),
            closed_at=_optional_str(payload.get("closed_at")),
            close_reason=_optional_str(payload.get("close_reason")),
            satisfaction=_optional_float(payload.get("satisfaction")),
            run_ids=[str(item) for item in payload.get("run_ids") or []],
            skill_names=[str(item) for item in payload.get("skill_names") or []],
            feedback=[
                _normalize_acceptance_entry(dict(item))
                for item in (payload.get("acceptance") or payload.get("feedback") or [])
                if isinstance(item, dict)
            ],
            validation_result=dict(payload["validation_result"]) if isinstance(payload.get("validation_result"), dict) else None,
            metadata=dict(payload.get("metadata") or {}),
        )


@dataclass(slots=True)
class TaskEvent:
    event_id: str
    task_id: str
    session_id: str
    event_type: str
    created_at: str
    run_id: str | None = None
    payload: dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        return {
            "event_id": self.event_id,
            "task_id": self.task_id,
            "session_id": self.session_id,
            "run_id": self.run_id,
            "event_type": self.event_type,
            "created_at": self.created_at,
            "payload": dict(self.payload),
        }

    @classmethod
    def from_dict(cls, payload: dict[str, Any]) -> "TaskEvent":
        return cls(
            event_id=str(payload["event_id"]),
            task_id=str(payload["task_id"]),
            session_id=str(payload["session_id"]),
            run_id=_optional_str(payload.get("run_id")),
            event_type=str(payload.get("event_type") or ""),
            created_at=str(payload.get("created_at") or ""),
            payload=dict(payload.get("payload") or {}),
        )


@dataclass(slots=True)
class MainAgentDecision:
    mode: str
    reason: str
    starts_new_task: bool = False
    closes_task: bool = False
    abandons_task: bool = False
    short_title: str | None = None
    action: str = ""

    @property
    def is_task(self) -> bool:
        return self.mode == "task"


def _optional_str(value: Any) -> str | None:
    if value in (None, ""):
        return None
    return str(value)


def _optional_float(value: Any) -> float | None:
    if value in (None, ""):
        return None
    return float(value)


def _normalize_acceptance_entry(entry: dict[str, Any]) -> dict[str, Any]:
    if entry.get("acceptance_type") is None and entry.get("feedback_type") is not None:
        feedback_type = str(entry.get("feedback_type") or "")
        entry["acceptance_type"] = "accept" if feedback_type == "satisfied" else feedback_type
    if entry.get("feedback_type") is None and entry.get("acceptance_type") is not None:
        acceptance_type = str(entry.get("acceptance_type") or "")
        entry["feedback_type"] = "satisfied" if acceptance_type == "accept" else acceptance_type
    return entry