tests/test_vlm_annotation.py

from __future__ import annotations

from dovla_cil.data.schema import ActionChunk, FailureInfo, RewardInfo, StructuredEffect
from dovla_cil.tasks.library import ToyTaskLibrary
from dovla_cil.vlm.annotation import VLMFailureAnnotator
from dovla_cil.vlm.client import VLMClient


def test_mock_vlm_annotation_works(monkeypatch, tmp_path) -> None:
    monkeypatch.setenv("OPENCLAUDE_MOCK", "1")
    task = ToyTaskLibrary().get_by_id("toy_pick_object_among_distractors")
    local_failure = _failure("wrong_target")
    annotator = VLMFailureAnnotator(
        client=VLMClient(api_key="test-secret", model="mock-model"),
        cache_path=tmp_path / "annotations.json",
    )

    annotated = annotator.annotate_failure(
        task=task,
        instruction=task.instruction,
        action=_action(candidate_type="wrong_target", target="blue_mug"),
        effect=_effect(),
        reward=_reward(success=False),
        local_failure=local_failure,
    )

    assert annotated.type == "wrong_target"
    assert "Mock VLM explanation" in (annotated.language_explanation or "")
    metadata = annotated.metadata["semantic_annotation"]["vlm_annotation"]
    assert metadata["source"] == "vlm"
    assert metadata["suggested_failure_type"] == "wrong_target"
    assert "test-secret" not in str(annotated.to_dict())


def test_annotation_cache_hit_avoids_client_call(tmp_path) -> None:
    task = ToyTaskLibrary().get_by_id("toy_lift_can")
    client = CountingAnnotationClient()
    annotator = VLMFailureAnnotator(client=client, cache_path=tmp_path / "cache.json")
    kwargs = {
        "task": task,
        "instruction": task.instruction,
        "action": _action(candidate_type="no_motion", target="can"),
        "effect": _effect(),
        "reward": _reward(success=False),
        "local_failure": _failure("no_motion"),
    }

    first = annotator.annotate_failure(**kwargs)
    second = annotator.annotate_failure(**kwargs)

    assert client.calls == 1
    assert first.language_explanation == second.language_explanation
    assert second.metadata["semantic_annotation"]["vlm_annotation"]["cache_hit"] is True


def test_invalid_vlm_json_falls_back_to_local_explanation(tmp_path) -> None:
    task = ToyTaskLibrary().get_by_id("toy_lift_can")
    local_failure = _failure("no_motion")
    annotator = VLMFailureAnnotator(client=InvalidAnnotationClient(), cache_path=tmp_path / "cache.json")

    annotated = annotator.annotate_failure(
        task=task,
        instruction=task.instruction,
        action=_action(candidate_type="noop", target="can"),
        effect=_effect(),
        reward=_reward(success=False),
        local_failure=local_failure,
    )

    assert annotated.type == local_failure.type
    assert annotated.language_explanation == local_failure.language_explanation
    metadata = annotated.metadata["semantic_annotation"]["vlm_annotation"]
    assert metadata["source"] == "local_fallback"
    assert metadata["fallback_error"]


def test_vlm_annotation_cannot_change_reward_or_local_failure_type(tmp_path) -> None:
    task = ToyTaskLibrary().get_by_id("toy_pick_object_among_distractors")
    reward = _reward(success=False)
    local_failure = _failure("wrong_target")
    annotator = VLMFailureAnnotator(client=OverridingAnnotationClient(), cache_path=tmp_path / "cache.json")

    annotated = annotator.annotate_failure(
        task=task,
        instruction=task.instruction,
        action=_action(candidate_type="wrong_target", target="blue_mug"),
        effect=_effect(),
        reward=reward,
        local_failure=local_failure,
    )

    assert reward.progress == 0.25
    assert reward.success is False
    assert reward.terminal_success is False
    assert annotated.type == "wrong_target"
    assert annotated.metadata["semantic_annotation"]["suggested_failure_type"] == "success"


class CountingAnnotationClient:
    def __init__(self) -> None:
        self.calls = 0

    def chat_json(self, system: str, user: str, schema_hint=None) -> dict:
        del system, user, schema_hint
        self.calls += 1
        return {
            "failure_type": "no_motion",
            "explanation": "The action produced no task-relevant motion.",
            "avoidance_hint": "Choose an action that contacts the can.",
            "confidence": 0.9,
        }


class InvalidAnnotationClient:
    def chat_json(self, system: str, user: str, schema_hint=None) -> dict:
        del system, user, schema_hint
        return {"failure_type": "no_motion", "confidence": "high"}


class OverridingAnnotationClient:
    def chat_json(self, system: str, user: str, schema_hint=None) -> dict:
        del system, user, schema_hint
        return {
            "failure_type": "success",
            "explanation": "The VLM incorrectly claims this succeeded.",
            "avoidance_hint": "No hint.",
            "confidence": 1.0,
        }


def _action(*, candidate_type: str, target: str) -> ActionChunk:
    return ActionChunk(
        representation="semantic",
        values=[{"command": "grasp", "object": target}],
        skill_type="grasp",
        metadata={
            "candidate_type": candidate_type,
            "intended_target": target,
            "intended_relation": "grasped",
            "difficulty": 0.3,
        },
    )


def _effect() -> StructuredEffect:
    before = {
        "objects": {
            "red_mug": {"position": [0.0, 0.0, 0.03], "grasped": False, "lifted": False},
            "blue_mug": {"position": [0.4, 0.0, 0.03], "grasped": False, "lifted": False},
        },
        "robot": {"eef_position": [0.0, 0.0, 0.2], "gripper": "open", "held_object": None},
    }
    after = {
        "objects": {
            "red_mug": {"position": [0.0, 0.0, 0.03], "grasped": False, "lifted": False},
            "blue_mug": {"position": [0.6, 0.0, 0.03], "grasped": False, "lifted": False},
        },
        "robot": {"eef_position": [0.4, 0.0, 0.2], "gripper": "closed", "held_object": None},
    }
    return StructuredEffect(
        object_pose_delta={"blue_mug": [0.2, 0.0, 0.0]},
        contact_events=[{"type": "touch", "object": "blue_mug"}],
        relation_before={"grasped(red_mug)": False},
        relation_after={"grasped(red_mug)": False},
        grasp_success=False,
        moved_objects=["blue_mug"],
        symbolic_before=before,
        symbolic_after=after,
    )


def _reward(*, success: bool) -> RewardInfo:
    return RewardInfo(
        progress=1.0 if success else 0.25,
        success=success,
        terminal_success=success,
        dense_components={"partial": 0.25},
    )


def _failure(failure_type: str) -> FailureInfo:
    return FailureInfo(
        type=failure_type,
        symbolic_reason=f"local {failure_type}",
        language_explanation=f"Local explanation for {failure_type}.",
    )