workspace/tests/test_retrieval.py

from __future__ import annotations

from pathlib import Path

from dovla_cil.data.datasets import CILDataset
from dovla_cil.generation.pipeline import generate_cil_dataset
from dovla_cil.retrieval.embeddings import cosine_similarity, embed_observation_language
from dovla_cil.retrieval.eval import RetrievalEvalQuery, evaluate_retrieval_baselines
from dovla_cil.retrieval.index import CILRetrievalIndex
from dovla_cil.retrieval.prompting import build_retrieval_prompt
from dovla_cil.retrieval.retriever import CriticGatedRetriever, RetrievalConditionedPolicyWrapper
from dovla_cil.tasks.library import built_in_toy_tasks
from dovla_cil.transfercritic.schema import TransferContext


def _make_dataset(tmp_path: Path) -> CILDataset:
    generate_cil_dataset(
        backend="toy",
        tasks=built_in_toy_tasks()[:2],
        out_dir=tmp_path,
        num_states_per_task=1,
        k=4,
        seed=21,
        shard_size=8,
        inline_observations=True,
    )
    return CILDataset(tmp_path)


def test_embedding_index_over_tiny_cil_dataset(tmp_path: Path) -> None:
    dataset = _make_dataset(tmp_path)
    index = CILRetrievalIndex.from_dataset(tmp_path, dim=32)
    record = dataset[0]
    query = embed_observation_language(record.observation_inline, record.instruction, dim=32)

    hits = index.query(query, top_k=3)

    assert len(index) == len(dataset)
    assert len(query) == 32
    assert len(hits) == 3
    assert hits[0].similarity >= hits[-1].similarity
    assert cosine_similarity(query, query) > 0.999


def test_retriever_modes_and_same_state_filter(tmp_path: Path) -> None:
    dataset = _make_dataset(tmp_path)
    index = CILRetrievalIndex.from_dataset(tmp_path, dim=32)
    first = dataset[0]
    retriever = CriticGatedRetriever(index, dim=32)

    same_state = retriever.retrieve(
        first.observation_inline,
        first.instruction,
        k=3,
        mode="nearest_neighbor",
        same_state_group_id=first.group_id,
    )
    success_only = retriever.retrieve(first.observation_inline, first.instruction, k=3, mode="success_only")
    contrastive = retriever.retrieve(
        first.observation_inline,
        first.instruction,
        k=3,
        mode="success_failure_contrastive",
    )

    assert same_state.examples
    assert all(example.item.group_id == first.group_id for example in same_state.examples)
    assert all(example.item.success for example in success_only.examples)
    assert contrastive.examples
    assert any(example.role == "positive_successful" for example in contrastive.examples)


def test_critic_gated_retrieval_uses_optional_critic(tmp_path: Path) -> None:
    dataset = _make_dataset(tmp_path)
    index = CILRetrievalIndex.from_dataset(tmp_path, dim=32)
    first = dataset[0]

    class MockCritic:
        def score_atom(self, atom, selected_atoms, context):
            del selected_atoms, context
            return 10.0 if atom.reward_summary.get("success", 0.0) else 0.0

    retriever = CriticGatedRetriever(
        index,
        critic=MockCritic(),
        transfer_context=TransferContext(benchmark_name="CausalStress"),
        dim=32,
    )
    result = retriever.retrieve(first.observation_inline, first.instruction, k=3, mode="critic_gated")

    assert result.examples
    assert result.examples[0].gate_score >= result.examples[-1].gate_score


def test_retrieval_conditioned_policy_wrapper(tmp_path: Path) -> None:
    dataset = _make_dataset(tmp_path)
    index = CILRetrievalIndex.from_dataset(tmp_path, dim=32)
    first = dataset[0]
    retriever = CriticGatedRetriever(index, dim=32)

    class DummyPolicy:
        def forward_policy(self, observation, instruction, retrieved_examples=None):
            del observation, instruction
            return {"retrieved": len(retrieved_examples or [])}

    wrapper = RetrievalConditionedPolicyWrapper(DummyPolicy(), retriever, k=2)
    output = wrapper.policy(first.observation_inline, first.instruction)

    assert output["retrieved"] == 2
    assert len(wrapper.last_retrieved_examples) == 2
    prompt = build_retrieval_prompt(first.instruction, wrapper.last_retrieved_examples)
    assert "Retrieved exemplars" in prompt


def test_retrieval_eval_baselines(tmp_path: Path) -> None:
    dataset = _make_dataset(tmp_path)
    index = CILRetrievalIndex.from_dataset(tmp_path, dim=32)
    retriever = CriticGatedRetriever(index, dim=32)
    queries = [
        RetrievalEvalQuery(
            observation=record.observation_inline or {},
            instruction=record.instruction,
            group_id=record.group_id,
        )
        for record in dataset.records[:2]
    ]

    report = evaluate_retrieval_baselines(retriever, queries, k=3)

    assert set(report) == {
        "no_retrieval",
        "nearest_neighbor",
        "success_only",
        "success_failure_contrastive",
        "critic_gated",
    }
    assert report["nearest_neighbor"]["retrieval_coverage"] == 1.0