JARVIS/backend/app/agents/verifier.py

from __future__ import annotations

from typing import Any, cast

from pydantic import BaseModel, Field

from app.agents.schemas.task import AgentTask, TaskResult, TaskResultStatus, VerificationStatus
from app.agents.state import AgentState


class VerificationVerdict(BaseModel):
    status: VerificationStatus
    summary: str | None = None
    evidence: list[dict[str, Any]] = Field(default_factory=list)


def normalize_task_result(
    task_result: TaskResult | dict[str, Any],
    *,
    default_task_id: str | None = None,
) -> TaskResult:
    payload = task_result.model_dump(mode="json") if isinstance(task_result, TaskResult) else dict(task_result or {})
    normalized_status = payload.get("status")
    if normalized_status not in {"completed", "failed", "blocked", "passed", "skipped"}:
        normalized_status = "failed"
    return TaskResult(
        task_id=str(payload.get("task_id") or default_task_id or "unknown-task"),
        status=cast(TaskResultStatus, normalized_status),
        summary=payload.get("summary"),
        evidence=list(payload.get("evidence") or []),
        owner_agent_id=payload.get("owner_agent_id"),
        parent_task_id=payload.get("parent_task_id"),
        child_task_ids=list(payload.get("child_task_ids") or []),
        thread_id=payload.get("thread_id"),
        message_id=payload.get("message_id"),
        message_index=payload.get("message_index") if isinstance(payload.get("message_index"), int) else None,
        interrupt_records=list(payload.get("interrupt_records") or []),
        recovery_records=list(payload.get("recovery_records") or []),
        budget_snapshot=payload.get("budget_snapshot") if isinstance(payload.get("budget_snapshot"), dict) else None,
        next_action=payload.get("next_action"),
        output_data=payload.get("output_data") if isinstance(payload.get("output_data"), dict) else None,
    )


def verify_task_result(
    *,
    task: AgentTask | dict[str, Any] | None = None,
    result: TaskResult | dict[str, Any] | None = None,
    summary: str | None = None,
    evidence: list[dict[str, Any]] | None = None,
    status: VerificationStatus | None = None,
) -> VerificationVerdict:
    normalized_result = result.model_dump() if isinstance(result, TaskResult) else dict(result or {})
    normalized_task = task.model_dump() if isinstance(task, AgentTask) else dict(task or {})
    normalized_summary = summary or normalized_result.get("summary") or normalized_task.get("result_summary")
    normalized_evidence = list(evidence or normalized_result.get("evidence") or normalized_task.get("evidence") or [])

    if status is not None:
        return VerificationVerdict(status=status, summary=normalized_summary, evidence=normalized_evidence)

    normalized_status = normalized_result.get("status")
    if normalized_status in {"passed", "failed", "skipped"}:
        inferred_status = normalized_status
    elif normalized_status == "completed":
        inferred_status = "passed"
    elif normalized_status == "blocked":
        inferred_status = "skipped"
    elif normalized_result.get("success") is True:
        inferred_status = "passed"
    elif normalized_result.get("success") is False:
        inferred_status = "failed"
    elif normalized_summary or normalized_evidence:
        inferred_status = "skipped"
    else:
        inferred_status = "failed"
        normalized_summary = "No verification input available."

    return VerificationVerdict(
        status=inferred_status,
        summary=normalized_summary,
        evidence=normalized_evidence,
    )


def apply_verification_verdict(state: AgentState, verdict: VerificationVerdict) -> AgentState:
    next_state = dict(state)
    next_state["verification_status"] = verdict.status
    next_state["verification_summary"] = verdict.summary
    next_state["verification_evidence"] = list(verdict.evidence)
    return AgentState(**next_state)


__all__ = ["VerificationVerdict", "apply_verification_verdict", "normalize_task_result", "verify_task_result"]