arena/validator_graph.py

"""LangGraph Validator workflow for prompt-first Runs."""

from __future__ import annotations

import os
from dataclasses import dataclass
from typing import Protocol

from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import END, START, StateGraph
from pydantic import BaseModel, Field
from typing_extensions import TypedDict

from .run_models import CodebaseArchive
from .validation_models import Criteria, ValidationCheck, ValidationCheckResult, ValidationReport
from .validator_plan import ValidatorContext, ValidatorPlan


class ValidatorExecutor(Protocol):
    def run(self, archive_pointer: str, command: str) -> tuple[bool, str]:
        ...


class RubricReview(BaseModel):
    score: float = Field(ge=0.0, le=1.0)
    feedback: str = Field(default="")
    risks: list[str] = Field(default_factory=list)


class RubricReviewer(Protocol):
    def review(
        self,
        *,
        prompt: str,
        criteria: list[Criteria],
        check_results: list[ValidationCheckResult],
    ) -> RubricReview:
        ...


class DeterministicValidatorExecutor:
    def run(self, archive_pointer: str, command: str) -> tuple[bool, str]:
        if "fail" in command.lower():
            return False, "deterministic validator marked command as failed"
        return True, "deterministic validator marked command as passed"


class DeterministicRubricReviewer:
    def review(
        self,
        *,
        prompt: str,
        criteria: list[Criteria],
        check_results: list[ValidationCheckResult],
    ) -> RubricReview:
        if not check_results:
            return RubricReview(score=0.5, feedback="No checks to evaluate.")
        passed = sum(1 for r in check_results if r.passed)
        return RubricReview(
            score=passed / len(check_results),
            feedback=f"Deterministic review: {passed}/{len(check_results)} checks passed.",
        )


class OpenRouterRubricReviewer:
    def __init__(self, model: str | None = None) -> None:
        from langchain_openrouter import ChatOpenRouter

        model_name = (model or os.getenv("RUBRIC_MODEL") or os.getenv("DEEPAGENT_MODEL") or "openrouter/owl-alpha")
        for prefix in ("openrouter:", "openrouter/"):
            if model_name.startswith(prefix):
                model_name = model_name[len(prefix):]
        self._model = ChatOpenRouter(
            model=model_name,
            temperature=0.0,
            max_tokens=768,
        )

    def review(
        self,
        *,
        prompt: str,
        criteria: list[Criteria],
        check_results: list[ValidationCheckResult],
    ) -> RubricReview:
        import json as _json
        from langchain_core.messages import HumanMessage, SystemMessage

        checks_text = "\n".join(
            f"- {r.check_id}: {'PASSED' if r.passed else 'FAILED'}; output: {r.output[:300]}"
            for r in check_results
        ) or "- No checks were run."
        criteria_text = "\n".join(f"- {c.text}" for c in criteria) or "- None provided"

        instruct = (
            "You are a code review rubric evaluator. Evaluate whether the codebase "
            "meets the user's prompt and criteria.\n\n"
            "Return a JSON object with exactly these keys:\n"
            '  "rubric_score": float (0.0 to 1.0)\n'
            '  "rubric_feedback": string (concise, actionable)\n'
            '  "rubric_risks": list of strings\n\n'
            "Do not fabricate files or behaviors. Only output the JSON object, no other text."
        )
        query = (
            f"User prompt: {prompt}\n\n"
            f"Criteria:\n{criteria_text}\n\n"
            f"Validation check results:\n{checks_text}"
        )
        response = self._model.invoke([
            SystemMessage(content=instruct),
            HumanMessage(content=query),
        ])
        text = ""
        if hasattr(response, "content"):
            text = str(response.content)
        elif isinstance(response, str):
            text = response
        elif isinstance(response, dict):
            text = response.get("content", _json.dumps(response))
        else:
            text = str(response)

        for candidate in _json_match(text):
            try:
                data = _json.loads(candidate)
                return RubricReview(
                    score=float(data.get("rubric_score", data.get("score", 0))),
                    feedback=str(data.get("rubric_feedback", data.get("feedback", ""))),
                    risks=[str(r) for r in data.get("rubric_risks", data.get("risks", []))],
                )
            except (ValueError, TypeError, _json.JSONDecodeError):
                continue

        return RubricReview(
            score=0.0,
            feedback=text[:512] if text else "No rubric feedback produced",
            risks=["Rubric LLM did not return valid JSON"],
        )


def _json_match(text: str):
    import re
    yield text
    m = re.search(r'\{[\s\S]*\}', text)
    if m:
        yield m.group()


class ValidatorState(TypedDict, total=False):
    run_id: str
    context: ValidatorContext
    codebase_archive: CodebaseArchive
    plan: ValidatorPlan
    results: list[ValidationCheckResult]
    stagehand_results: list[ValidationCheckResult]
    rubric_score: float
    rubric_feedback: str
    rubric_risks: list[str]
    report: ValidationReport


class StagehandExecutor(Protocol):
    def run(self, url: str, instruction: str) -> tuple[bool, str]:
        ...


class NoopStagehandExecutor:
    def run(self, url: str, instruction: str) -> tuple[bool, str]:
        return False, "Stagehand is not configured. Set BROWSERBASE_API_KEY and MODEL_API_KEY."


def _create_stagehand_executor() -> StagehandExecutor:
    try:
        from .stagehand_validator import run_stagehand_check

        class _Executor:
            def run(self, url: str, instruction: str) -> tuple[bool, str]:
                return run_stagehand_check(url=url, instruction=instruction)

        return _Executor()
    except ImportError:
        return NoopStagehandExecutor()


@dataclass
class ValidatorGraph:
    executor: ValidatorExecutor | None = None
    stagehand_executor: StagehandExecutor | None = None
    rubric_reviewer: RubricReviewer | None = None
    rubric_enabled: bool = True

    def __post_init__(self) -> None:
        self._executor = self.executor or DeterministicValidatorExecutor()
        self._stagehand_executor = self.stagehand_executor or _create_stagehand_executor()
        self._rubric_reviewer = self.rubric_reviewer or DeterministicRubricReviewer()
        self._checkpointer = MemorySaver()
        self._graph = self._build_graph()

    def validate(
        self,
        *,
        run_id: str,
        codebase_archive: CodebaseArchive,
        plan: ValidatorPlan,
        prompt: str = "",
        criteria: list[Criteria] | None = None,
    ) -> ValidationReport:
        context = _resolve_context(plan, prompt=prompt, criteria=criteria)
        effective_plan = plan if (plan.context.prompt or plan.context.criteria) else _with_context(plan, context)
        try:
            state = self._graph.invoke(
                {
                    "run_id": run_id,
                    "context": effective_plan.context,
                    "codebase_archive": codebase_archive,
                    "plan": effective_plan,
                },
                config={"configurable": {"thread_id": f"validator-{run_id}"}},
            )
            return state["report"]
        finally:
            close = getattr(self._executor, "close", None)
            if callable(close):
                close()

    def _build_graph(self):
        graph = StateGraph(ValidatorState)
        graph.add_node("run_validation_checks", self._run_validation_checks)
        graph.add_node("run_stagehand_checks", self._run_stagehand_checks)
        if self.rubric_enabled:
            graph.add_node("rubric_review", self._rubric_review)
        graph.add_node("finalize_report", self._finalize_report)

        graph.add_edge(START, "run_validation_checks")
        graph.add_edge("run_validation_checks", "run_stagehand_checks")
        if self.rubric_enabled:
            graph.add_edge("run_stagehand_checks", "rubric_review")
            graph.add_edge("rubric_review", "finalize_report")
        else:
            graph.add_edge("run_stagehand_checks", "finalize_report")
        graph.add_edge("finalize_report", END)
        return graph.compile(checkpointer=self._checkpointer, name="validator-graph")

    def _run_validation_checks(self, state: ValidatorState) -> ValidatorState:
        archive_pointer = state["codebase_archive"].pointer
        results = []
        for step in state["plan"].command_steps:
            check = step.check
            results.append(
                _validation_result(check, self._executor.run(archive_pointer, check.command))
            )
        return {"results": results}

    def _run_stagehand_checks(self, state: ValidatorState) -> ValidatorState:
        stagehand_steps = state["plan"].stagehand_steps
        if not stagehand_steps:
            return {"stagehand_results": []}
        results = []
        for step in stagehand_steps:
            check = step.check
            passed, output = self._stagehand_executor.run(
                url=check.command,
                instruction=check.name,
            )
            results.append(ValidationCheckResult(
                check_id=check.id,
                passed=passed,
                output=output,
            ))
        return {"stagehand_results": results}

    def _rubric_review(self, state: ValidatorState) -> ValidatorState:
        if state["plan"].rubric_step is None:
            return {}
        result = self._rubric_reviewer.review(
            prompt=state["context"].prompt,
            criteria=list(state["context"].criteria),
            check_results=state["results"],
        )
        return {
            "rubric_score": result.score,
            "rubric_feedback": result.feedback,
            "rubric_risks": result.risks,
        }

    def _finalize_report(self, state: ValidatorState) -> ValidatorState:
        checks = state["plan"].checks
        results = list(state["results"]) + list(state.get("stagehand_results", []))
        passed = all(result.passed for result in results) if results else True
        risks = _risks(checks, results)

        rubric_feedback = state.get("rubric_feedback", "")
        if rubric_feedback:
            risks = list(risks) + state.get("rubric_risks", [])

        if not checks:
            summary = "No validation checks configured."
        else:
            passed_count = sum(1 for result in results if result.passed)
            summary = f"Validation checks passed: {passed_count}/{len(results)}."

        rubric_score = state.get("rubric_score")
        if rubric_score is not None:
            summary = f"{summary} Rubric score: {rubric_score:.2f}."

        return {
            "report": ValidationReport(
                run_id=state["run_id"],
                passed=passed,
                results=results,
                summary=summary,
                risks=risks,
                rubric_score=rubric_score,
                rubric_feedback=rubric_feedback,
            )
        }


def _resolve_context(
    plan: ValidatorPlan,
    *,
    prompt: str,
    criteria: list[Criteria] | None,
) -> ValidatorContext:
    if plan.context.prompt or plan.context.criteria:
        return plan.context
    if prompt or criteria:
        return ValidatorContext(prompt=prompt, criteria=tuple(criteria or ()))
    return plan.context


def _with_context(plan: ValidatorPlan, context: ValidatorContext) -> ValidatorPlan:
    return ValidatorPlan(
        context=context,
        command_steps=plan.command_steps,
        stagehand_steps=plan.stagehand_steps,
        rubric_step=plan.rubric_step,
    )


def _validation_result(
    check: ValidationCheck,
    result: tuple[bool, str],
) -> ValidationCheckResult:
    passed, output = result
    return ValidationCheckResult(
        check_id=check.id,
        passed=passed,
        output=output,
    )


def _risks(
    checks: list[ValidationCheck],
    results: list[ValidationCheckResult],
) -> list[str]:
    names_by_id = {check.id: check.name for check in checks}
    failed_names = [names_by_id.get(result.check_id, result.check_id) for result in results if not result.passed]
    if not failed_names:
        return []
    return [f"Failed checks: {', '.join(failed_names)}"]