comet-ml--opik/sdks/python/tests/e2e/verifiers.py

import base64
import json
import re
from typing import Any, Callable, Dict, Iterable, List, Literal, Optional, Set, Union
from unittest import mock

import pytest
import opik
from opik import Attachment, Prompt, ChatPrompt, synchronization
from opik.api_objects import rest_helpers
from opik.api_objects.attachment import decoder_helpers
from opik.api_objects.dataset import dataset_item
from opik.rest_api import ExperimentPublic, FeedbackScore, FeedbackScorePublic
from opik.rest_api.types import (
    attachment as rest_api_attachment,
    span_public,
    trace_public,
)
from opik.types import ErrorInfoDict, FeedbackScoreDict, TraceSource
from opik import url_helpers
from .. import testlib

InputType = Union[
    Dict[str, Any],
    List[Any],
    str,
]

OutputType = InputType


def _try_get__dict__(instance: Any) -> Optional[Dict[str, Any]]:
    if instance is None:
        return None

    if hasattr(instance, "model_dump"):
        return instance.model_dump()

    return instance.__dict__


def _try_build_set(iterable: Optional[Iterable[Any]]) -> Optional[Set[Any]]:
    if iterable is None:
        return iterable

    return set(iterable)


def _retry_until_assertions_pass(check: Callable[[], None]) -> None:
    """Run ``check`` repeatedly until its asserts pass, or surface the failure.

    ``check`` is the same body used to verify a trace/span — including
    ``assert``s. We poll it because traces/spans are eventually-consistent in
    ClickHouse: the initial create lands fast, but a later ``update`` (e.g. a
    message replayed from the offline queue, or a follow-up SDK call) can take
    longer to ingest. Asserting on a single read-once snapshot was racy under
    xdist load — the existence-only `is not None` poll passed too early.

    Re-using the assertion body as the predicate (instead of a parallel
    "matches expected" helper) keeps the comparison logic in one place — the
    asserts already encode what counts as a match (incl. ``mock.ANY`` via
    ``assert_equal``, set-equality for tags, etc.).

    Errors during polling are tolerated (``allow_errors=True``) to match the
    rest of this file — a 404 on a not-yet-ingested entity, a transient
    network blip, or any other ``Exception`` simply triggers a retry. This
    does NOT mask real bugs: ``synchronization.until`` is bounded by
    ``max_try_seconds``, and on timeout we re-run ``check()`` so the actual
    exception (with its traceback) reaches pytest.

    ``pytest_deepassert.equal`` (used by ``assert_equal``) raises
    ``pytest.fail.Exception`` (a.k.a. ``_pytest.outcomes.Failed``), which
    inherits from ``BaseException``, NOT from ``Exception`` — so
    ``synchronization.until``'s ``except Exception`` would not catch it and
    the retry would degenerate to a single attempt. Convert it to a return
    value here so the polling actually works for value-comparable fields.
    """

    def _attempt() -> bool:
        try:
            check()
        except pytest.fail.Exception:
            return False
        return True

    if synchronization.until(_attempt, allow_errors=True):
        return
    # Timeout — re-run so the original error reaches pytest.
    check()


def verify_trace(
    opik_client: opik.Opik,
    trace_id: str,
    name: str = mock.ANY,  # type: ignore
    metadata: Dict[str, Any] = mock.ANY,  # type: ignore
    input: InputType = mock.ANY,  # type: ignore
    output: Optional[OutputType] = mock.ANY,  # type: ignore
    tags: Union[List[str], Set[str]] = mock.ANY,  # type: ignore
    feedback_scores: List[FeedbackScoreDict] = mock.ANY,  # type: ignore
    project_name: Optional[str] = mock.ANY,  # type: ignore
    error_info: Optional[ErrorInfoDict] = mock.ANY,  # type: ignore
    thread_id: Optional[str] = mock.ANY,  # type: ignore
    guardrails_validations: Optional[List[Dict[str, Any]]] = mock.ANY,  # type: ignore
    source: Optional[TraceSource] = mock.ANY,  # type: ignore
    environment: Optional[str] = mock.ANY,  # type: ignore
    comments: List[str] = mock.ANY,  # type: ignore
):
    def _check() -> None:
        trace = opik_client.get_trace_content(id=trace_id)
        assert trace is not None, f"Failed to get trace with id {trace_id}."

        testlib.assert_equal(name, trace.name)
        testlib.assert_equal(input, trace.input)
        testlib.assert_equal(output, trace.output)
        testlib.assert_equal(metadata, trace.metadata)
        testlib.assert_equal(source, trace.source)

        if environment is not mock.ANY:
            testlib.assert_equal(environment, trace.environment)

        if tags is not mock.ANY:
            testlib.assert_equal(_try_build_set(tags), _try_build_set(trace.tags))

        if error_info is not mock.ANY:
            testlib.assert_equal(error_info, _try_get__dict__(trace.error_info))

        assert thread_id == trace.thread_id, f"{trace.thread_id} != {thread_id}"

        if project_name is not mock.ANY:
            trace_project = opik_client.get_project(trace.project_id)
            assert trace_project.name == project_name

        if feedback_scores is not mock.ANY:
            _assert_feedback_scores(
                item_id=trace_id,
                feedback_scores=trace.feedback_scores,
                expected_feedback_scores=feedback_scores,
            )

        if guardrails_validations is not mock.ANY:
            if trace.guardrails_validations is None:
                assert guardrails_validations is None, (
                    f"Expected guardrails validation to be None, but got {guardrails_validations}"
                )
                return

            actual_guardrails_validations = (
                []
                if trace.guardrails_validations is None
                else trace.guardrails_validations
            )
            assert len(actual_guardrails_validations) == len(guardrails_validations), (
                f"Expected amount of trace guardrails validation ({len(guardrails_validations)}) is not equal to actual amount ({len(actual_guardrails_validations)})"
            )

            actual_guardrails_validations = [
                guardrail.model_dump() for guardrail in trace.guardrails_validations
            ]

            sorted_actual_guardrails_validations = sorted(
                actual_guardrails_validations, key=lambda item: item["span_id"]
            )
            sorted_expected_guardrails_validations = sorted(
                guardrails_validations, key=lambda item: item["span_id"]
            )
            for actual_guardrail, expected_guardrail in zip(
                sorted_actual_guardrails_validations,
                sorted_expected_guardrails_validations,
            ):
                testlib.assert_dicts_equal(actual_guardrail, expected_guardrail)

        if comments is not mock.ANY:
            actual_comments = [c.text for c in (trace.comments or [])]
            assert actual_comments == comments, (
                f"trace comments differ: expected {comments}, got {actual_comments}"
            )

    _retry_until_assertions_pass(_check)


def verify_span(
    opik_client: opik.Opik,
    span_id: str,
    trace_id: str,
    parent_span_id: Optional[str],
    name: str = mock.ANY,  # type: ignore
    metadata: Dict[str, Any] = mock.ANY,  # type: ignore
    input: InputType = mock.ANY,  # type: ignore
    output: Optional[OutputType] = mock.ANY,  # type: ignore
    tags: Union[List[str], Set[str]] = mock.ANY,  # type: ignore
    type: str = mock.ANY,  # type: ignore
    feedback_scores: List[FeedbackScoreDict] = mock.ANY,  # type: ignore
    project_name: Optional[str] = mock.ANY,
    model: Optional[str] = mock.ANY,  # type: ignore
    provider: Optional[str] = mock.ANY,  # type: ignore
    error_info: Optional[ErrorInfoDict] = mock.ANY,  # type: ignore
    total_cost: Optional[float] = mock.ANY,  # type: ignore
    source: Optional[TraceSource] = mock.ANY,  # type: ignore
    environment: Optional[str] = mock.ANY,  # type: ignore
    comments: List[str] = mock.ANY,  # type: ignore
):
    def _check() -> None:
        span = opik_client.get_span_content(id=span_id)
        assert span is not None, f"Failed to get span with id {span_id}."

        assert span.trace_id == trace_id, f"{span.trace_id} != {trace_id}"

        if parent_span_id is None:
            assert span.parent_span_id is None, (
                f"{span.parent_span_id} != {parent_span_id}"
            )
        else:
            assert span.parent_span_id == parent_span_id, (
                f"{span.parent_span_id} != {parent_span_id}"
            )

        testlib.assert_equal(name, span.name)
        testlib.assert_equal(type, span.type)

        testlib.assert_equal(input, span.input)
        testlib.assert_equal(output, span.output)
        testlib.assert_equal(metadata, span.metadata)
        testlib.assert_equal(source, span.source)

        if environment is not mock.ANY:
            testlib.assert_equal(environment, span.environment)

        if tags is not mock.ANY:
            testlib.assert_equal(_try_build_set(tags), _try_build_set(span.tags))

        if error_info is not mock.ANY:
            testlib.assert_equal(error_info, _try_get__dict__(span.error_info))

        assert span.model == model, f"{span.model} != {model}"
        assert span.provider == provider, f"{span.provider} != {provider}"
        assert span.total_estimated_cost == total_cost, (
            f"{span.total_estimated_cost} != {total_cost}"
        )

        if project_name is not mock.ANY:
            span_project = opik_client.get_project(span.project_id)
            assert span_project.name == project_name

        if feedback_scores is not mock.ANY:
            _assert_feedback_scores(
                item_id=span_id,
                feedback_scores=span.feedback_scores,
                expected_feedback_scores=feedback_scores,
            )

        if comments is not mock.ANY:
            actual_comments = [c.text for c in (span.comments or [])]
            assert actual_comments == comments, (
                f"span comments differ: expected {comments}, got {actual_comments}"
            )

    _retry_until_assertions_pass(_check)


def verify_dataset(
    opik_client: opik.Opik,
    name: str,
    description: str = mock.ANY,
    dataset_items: List[dataset_item.DatasetItem] = mock.ANY,
    project_name: Optional[str] = None,
):
    assert synchronization.until(
        lambda: opik_client.get_dataset(name=name) is not None,
        allow_errors=True,
    ), f"Failed to get dataset with name {name}."

    actual_dataset = opik_client.get_dataset(name=name, project_name=project_name)
    assert actual_dataset.description == description
    assert actual_dataset.name == name
    if project_name is not None:
        assert actual_dataset.project_name == project_name

    actual_dataset_items = list(
        actual_dataset.__internal_api__stream_items_as_dataclasses__()
    )
    assert len(actual_dataset_items) == len(dataset_items), (
        f"Amount of actual dataset items ({len(actual_dataset_items)}) is not the same as of expected ones ({len(dataset_items)})"
    )

    actual_dataset_items_dicts = [item.model_dump() for item in actual_dataset_items]
    expected_dataset_items_dicts = [item.model_dump() for item in dataset_items]

    sorted_actual_items = sorted(
        actual_dataset_items_dicts, key=lambda item: json.dumps(item, sort_keys=True)
    )
    sorted_expected_items = sorted(
        expected_dataset_items_dicts, key=lambda item: json.dumps(item, sort_keys=True)
    )

    for actual_item, expected_item in zip(sorted_actual_items, sorted_expected_items):
        testlib.assert_dicts_equal(actual_item, expected_item, ignore_keys=["id"])


def verify_dashboard(
    opik_client: opik.Opik,
    dashboard_id: str,
    name: str = mock.ANY,
    type: str = mock.ANY,
    version: int = mock.ANY,
    section_count: int = mock.ANY,
    expected_widget_configs: Optional[Dict[str, Dict[str, Any]]] = None,
    expected_widget_positions: Optional[Dict[str, Dict[str, int]]] = None,
):
    assert synchronization.until(
        lambda: opik_client.get_dashboard(dashboard_id),
        allow_errors=True,
    ), f"Failed to get dashboard with id {dashboard_id}."

    fetched = opik_client.get_dashboard(dashboard_id)
    config = fetched.config

    testlib.assert_equal(name, fetched.name)
    testlib.assert_equal(type, fetched.type)
    if version is not mock.ANY:
        assert config["version"] == version
    if section_count is not mock.ANY:
        assert len(config["sections"]) == section_count

    configs_by_type: Dict[str, List[Dict[str, Any]]] = {}
    for section in config["sections"]:
        for widget in section.get("widgets", []):
            configs_by_type.setdefault(widget["type"], []).append(widget["config"])
    if expected_widget_configs is not None:
        for widget_type, expected_config in expected_widget_configs.items():
            assert widget_type in configs_by_type, (
                f"Expected widget of type {widget_type!r} not found in dashboard. "
                f"Present types: {sorted(configs_by_type)}"
            )
            actual_config = configs_by_type[widget_type][0]
            actual_subset = {key: actual_config.get(key) for key in expected_config}
            testlib.assert_equal(expected_config, actual_subset)

    # Every widget has a matching layout item and vice versa.
    for section in config["sections"]:
        section_widget_ids = {widget["id"] for widget in section.get("widgets", [])}
        layout_ids = {item["i"] for item in section.get("layout", [])}
        assert section_widget_ids == layout_ids, (
            f"Widget/layout id mismatch in section {section.get('id')!r}"
        )

    if expected_widget_positions is not None:
        layout_by_id = {
            item.id: item for section in fetched.sections for item in section.layout
        }
        for widget_id, expected_pos in expected_widget_positions.items():
            actual_item = layout_by_id[widget_id]
            actual_subset = {key: getattr(actual_item, key) for key in expected_pos}
            testlib.assert_equal(expected_pos, actual_subset)


def verify_experiment(
    opik_client: opik.Opik,
    id: str,
    experiment_name: str,
    experiment_metadata: Optional[Dict[str, Any]],
    feedback_scores_amount: int,
    traces_amount: int,
    prompts: Optional[List[Prompt]] = None,
    experiment_scores: Optional[Dict[str, float]] = None,
    experiment_tags: Optional[List[str]] = None,
    dataset_version_id: Optional[str] = mock.ANY,  # type: ignore
    project_name: Optional[str] = mock.ANY,
):
    rest_client = (
        opik_client._rest_client
    )  # temporary solution until backend prepares proper endpoints

    rest_client.datasets.find_dataset_items_with_experiment_items

    assert synchronization.until(
        lambda: rest_client.experiments.get_experiment_by_id(id) is not None,
        allow_errors=True,
    ), f"Failed to get experiment with id {id}."

    experiment_content = rest_client.experiments.get_experiment_by_id(id)

    _verify_experiment_metadata(experiment_content, experiment_metadata)

    assert experiment_content.name == experiment_name, (
        f"{experiment_content.name} != {experiment_name}"
    )

    testlib.assert_equal(project_name, experiment_content.project_name)

    actual_scores_count = (
        0
        if experiment_content.feedback_scores is None
        else len(experiment_content.feedback_scores)
    )
    assert actual_scores_count == feedback_scores_amount, (
        f"{actual_scores_count} != {feedback_scores_amount}"
    )

    actual_trace_count = (
        0 if experiment_content.trace_count is None else experiment_content.trace_count
    )
    assert actual_trace_count == traces_amount, (
        f"{actual_trace_count} != {traces_amount}"
    )

    _verify_experiment_prompts(experiment_content, prompts)

    _verify_experiment_scores(experiment_content, experiment_scores)

    testlib.assert_equal(expected=experiment_tags, actual=experiment_content.tags)

    if dataset_version_id is not mock.ANY:
        assert experiment_content.dataset_version_id == dataset_version_id, (
            f"Expected dataset_version_id {dataset_version_id}, "
            f"got {experiment_content.dataset_version_id}"
        )


def verify_experiment_items_completed(
    opik_client: opik.Opik,
    experiment_id: str,
    *,
    expected_completed_dataset_item_ids: Set[str],
    timeout: float = 15,
) -> None:
    """
    Assert that the set of dataset item ids with at least one fully-
    completed experiment item matches ``expected_completed_dataset_item_ids``
    exactly.

    Goes through the same predicate ``opik.evaluate_resume`` uses
    (``trace.output`` is set only when the engine's happy-path-only line
    runs), so this check correctly excludes both task-side failures
    (task raised before producing output) and scoring-side failures
    (output stripped by the engine when scoring did not finish).
    """
    from opik.evaluation.resume import context as resume_context

    # The experiment record is eventually consistent: looking it up right
    # after ``evaluate()`` returns can hit a 404, or a transient API/network
    # error. Fetch the handle inside the poll (matching ``verify_experiment``
    # / ``verify_test_suite_result``), so a fresh handle is also taken on
    # each iteration rather than a stale one being cached across the poll.
    def _completed_ids_or_none() -> Optional[Set[str]]:
        experiment = opik_client.get_experiment_by_id(experiment_id)
        return {
            item.dataset_item_id
            for item in experiment.get_items()
            if resume_context.is_trial_fully_completed(item)
        }

    last_observed: Set[str] = set()

    def _converged() -> bool:
        nonlocal last_observed
        last_observed = _completed_ids_or_none() or set()
        return last_observed == expected_completed_dataset_item_ids

    assert synchronization.until(
        _converged, max_try_seconds=timeout, allow_errors=True
    ), (
        f"Expected completed dataset item ids "
        f"{sorted(expected_completed_dataset_item_ids)}; got "
        f"{sorted(last_observed)} after {timeout}s"
    )


def verify_attachments(
    opik_client: opik.Opik,
    entity_type: Literal["trace", "span"],
    entity_id: str,
    attachments: Dict[str, Attachment],
    data_sizes: Dict[str, int],
    timeout: int = 10,
) -> None:
    attachment_list = _wait_for_attachments_list(
        opik_client=opik_client,
        entity_type=entity_type,
        entity_id=entity_id,
        expected_size=len(attachments),
        timeout=timeout,
    )
    url_override = opik_client._config.url_override

    for attachment in attachment_list:
        expected_attachment = attachments.get(attachment.file_name, None)
        assert expected_attachment is not None, (
            f"Attachment {attachment.file_name} not found in expected attachments"
        )

        assert attachment.file_size == data_sizes[expected_attachment.file_name], (
            f"Wrong size for attachment {attachment.file_name}: {attachment.file_size} != {data_sizes[expected_attachment.file_name]}"
        )

        assert attachment.mime_type == expected_attachment.content_type, (
            f"Wrong content type for attachment {attachment.file_name}: {attachment.mime_type} != {expected_attachment.content_type}"
        )

        if not url_helpers.is_aws_presigned_url(attachment.link):
            assert attachment.link.startswith(url_override), (
                f"Wrong link for attachment {attachment.file_name}: {attachment.link} does not start with {url_override}"
            )


def verify_auto_extracted_attachments(
    opik_client: opik.Opik,
    entity_type: Literal["trace", "span"],
    entity_id: str,
    expected_sizes: List[int],
    timeout: int = 10,
) -> None:
    attachment_list = _wait_for_attachments_list(
        opik_client=opik_client,
        entity_type=entity_type,
        entity_id=entity_id,
        expected_size=len(expected_sizes),
        timeout=timeout,
    )

    file_name_pattern = re.compile(decoder_helpers.ATTACHMENT_FILE_NAME_REGEX)

    for attachment in attachment_list:
        assert attachment.file_size in expected_sizes, (
            f"Wrong size for attachment {attachment.file_name}: {attachment.file_size} not in {expected_sizes}"
        )
        assert file_name_pattern.match(attachment.file_name), (
            f"Wrong file name for attachment {attachment.file_name} - it does not match {file_name_pattern.pattern}"
        )


def _wait_for_attachments_list(
    opik_client: opik.Opik,
    entity_type: Literal["trace", "span"],
    entity_id: str,
    expected_size: int,
    timeout: int,
) -> List[Attachment]:
    assert synchronization.until(
        lambda: (
            _get_trace_or_span(
                opik_client, entity_type=entity_type, entity_id=entity_id
            )
            is not None
        ),
        allow_errors=True,
        max_try_seconds=timeout,
    ), f"Failed to get {entity_type} with id {entity_id}."

    trace_or_span = _get_trace_or_span(
        opik_client, entity_type=entity_type, entity_id=entity_id
    )
    url_override = opik_client._config.url_override
    url_override_path = base64.b64encode(url_override.encode("utf-8")).decode("utf-8")

    assert synchronization.until(
        lambda: (
            len(
                _get_attachments(
                    opik_client=opik_client,
                    project_id=trace_or_span.project_id,
                    entity_type=entity_type,
                    entity_id=entity_id,
                    url_override_path=url_override_path,
                )
            )
            == expected_size
        ),
        allow_errors=True,
        max_try_seconds=timeout,
    ), f"Failed to get all expected attachments for {entity_type} with id {entity_id}."

    return _get_attachments(
        opik_client=opik_client,
        project_id=trace_or_span.project_id,
        entity_type=entity_type,
        entity_id=entity_id,
        url_override_path=url_override_path,
    )


def _get_attachments(
    opik_client: opik.Opik,
    entity_type: Literal["trace", "span"],
    entity_id: str,
    project_id: str,
    url_override_path: str,
) -> List[rest_api_attachment.Attachment]:
    return opik_client._rest_client.attachments.attachment_list(
        project_id=project_id,
        entity_type=entity_type,
        entity_id=entity_id,
        path=url_override_path,
    ).content


def _get_trace_or_span(
    opik_client: opik.Opik,
    entity_type: Literal["trace", "span"],
    entity_id: str,
) -> Union[trace_public.TracePublic, span_public.SpanPublic]:
    if entity_type == "trace":
        return opik_client.get_trace_content(id=entity_id)
    elif entity_type == "span":
        return opik_client.get_span_content(id=entity_id)
    else:
        raise ValueError(f"Invalid entity type: {entity_type}")


def _verify_experiment_metadata(
    experiment_content: ExperimentPublic,
    metadata: Optional[Dict[str, Any]],
):
    experiment_metadata = experiment_content.metadata
    if experiment_content.metadata is not None:
        experiment_metadata = {**experiment_content.metadata}
        # SDK-managed keys that ``evaluate()`` writes into the
        # experiment config; tests assert on user-supplied keys only.
        experiment_metadata.pop("prompt", None)
        experiment_metadata.pop("prompts", None)
        experiment_metadata.pop("_opik_resume", None)
        # Tests that supplied no ``experiment_config`` expect
        # ``metadata=None``; after stripping SDK keys we may be left
        # with an empty dict for the same case.
        if not experiment_metadata:
            experiment_metadata = None

    assert experiment_metadata == metadata, f"{experiment_metadata} != {metadata}"


def verify_experiment_traces_have_opik_prompts(
    opik_client: opik.Opik,
    trace_ids: List[str],
    prompts: List[Prompt],
) -> None:
    """Verify each trace produced by an experiment carries `opik_prompts`
    in its metadata. Backend needs this for trace-level prompt linkage."""
    assert trace_ids, "Expected at least one trace_id to verify prompt injection"

    expected_prompt_keys = sorted(
        (p.__internal_api__prompt_id__, p.commit) for p in prompts
    )

    for trace_id in trace_ids:
        assert synchronization.until(
            lambda: opik_client.get_trace_content(id=trace_id) is not None,
            allow_errors=True,
        ), f"Failed to get trace {trace_id}"

        trace_content = opik_client.get_trace_content(id=trace_id)
        assert trace_content.metadata is not None, (
            f"Trace {trace_id} has no metadata; expected opik_prompts"
        )
        actual_prompts = trace_content.metadata.get("opik_prompts")
        assert actual_prompts, (
            f"Trace {trace_id} metadata is missing 'opik_prompts': "
            f"{trace_content.metadata}"
        )
        actual_keys = sorted(
            (p.get("id"), (p.get("version") or {}).get("commit"))
            for p in actual_prompts
        )
        assert actual_keys == expected_prompt_keys, (
            f"Trace {trace_id} opik_prompts mismatch: "
            f"actual={actual_keys}, expected={expected_prompt_keys}"
        )


def _verify_experiment_prompts(
    experiment_content: ExperimentPublic,
    prompts: Optional[List[Prompt]],
):
    if prompts is None:
        return

    # asserting Prompt vs Experiment.prompt_version
    experiment_content_prompt_versions = sorted(
        experiment_content.prompt_versions, key=lambda x: x.id
    )
    prompts = sorted(prompts, key=lambda x: x.__internal_api__version_id__)

    for i, prompt in enumerate(prompts):
        assert (
            experiment_content_prompt_versions[i].id
            == prompt.__internal_api__version_id__
        ), (
            f"{experiment_content_prompt_versions[i].id} != {prompt.__internal_api__version_id__}"
        )
        assert (
            experiment_content_prompt_versions[i].prompt_id
            == prompt.__internal_api__prompt_id__
        ), (
            f"{experiment_content_prompt_versions[i].prompt_id} != {prompt.__internal_api__prompt_id__}"
        )

        assert experiment_content_prompt_versions[i].commit == prompt.commit, (
            f"{experiment_content_prompt_versions[i].commit} != {prompt.commit}"
        )

    # check that experiment config/metadata contains Prompt's template
    experiment_prompts = experiment_content.metadata["prompts"]

    for prompt in prompts:
        assert experiment_prompts[prompt.name] == prompt.prompt, (
            f"{experiment_prompts[prompt.name]} != {prompt.prompt}"
        )


def _verify_experiment_scores(
    experiment_content: ExperimentPublic,
    experiment_scores: Optional[Dict[str, float]],
):
    """Verify experiment-level scores match expected values."""
    if experiment_scores is None:
        return

    actual_experiment_scores = experiment_content.experiment_scores

    assert actual_experiment_scores is not None, (
        f"Expected experiment_scores to be set, but got None. "
        f"Experiment ID: {experiment_content.id}, "
        f"Expected scores: {experiment_scores}"
    )

    # Create a dict of actual scores for easy comparison
    actual_scores_dict = {score.name: score.value for score in actual_experiment_scores}

    assert len(actual_scores_dict) == len(experiment_scores), (
        f"Expected {len(experiment_scores)} experiment scores, "
        f"but got {len(actual_scores_dict)}. "
        f"Expected: {experiment_scores}, "
        f"Actual: {actual_scores_dict}"
    )

    for expected_name, expected_value in experiment_scores.items():
        assert expected_name in actual_scores_dict, (
            f"Expected experiment score '{expected_name}' not found. "
            f"Available scores: {list(actual_scores_dict.keys())}"
        )

        assert actual_scores_dict[expected_name] == expected_value, (
            f"Expected experiment score '{expected_name}' to have value {expected_value}, "
            f"but got {actual_scores_dict[expected_name]}"
        )


def verify_optimization(
    opik_client: opik.Opik,
    optimization_id: str,
    name: str = mock.ANY,  # type: ignore
    dataset_name: Optional[str] = mock.ANY,  # type: ignore
    status: Optional[str] = mock.ANY,  # type: ignore
    objective_name: Optional[str] = mock.ANY,  # type: ignore
    project_name: Optional[str] = None,
) -> None:
    assert synchronization.until(
        lambda: opik_client.get_optimization_by_id(optimization_id) is not None,
        allow_errors=True,
    ), f"Failed to get optimization with id {optimization_id}."

    optimization = opik_client.get_optimization_by_id(optimization_id)

    optimization_content = optimization.fetch_content()

    assert optimization_content.name == name, f"{optimization_content.name} != {name}"

    assert optimization_content.dataset_name == dataset_name, (
        f"{optimization_content.dataset_name} != {dataset_name}"
    )

    assert optimization_content.status == status, (
        f"{optimization_content.status} != {status}"
    )

    assert optimization_content.objective_name == objective_name, (
        f"{optimization_content.objective_name} != {objective_name}"
    )

    if project_name is not None:
        project_id = rest_helpers.resolve_project_id_by_name(
            rest_client=opik_client.rest_client, project_name=project_name
        )
        assert optimization_content.project_id == project_id, (
            f"{optimization_content.project_id} != {project_id}"
        )


def verify_thread(
    opik_client: opik.Opik,
    thread_id: str,
    project_name: Optional[str] = None,
    feedback_scores: List[FeedbackScoreDict] = mock.ANY,  # type: ignore
) -> None:
    assert synchronization.until(
        lambda: (
            len(
                opik_client.search_threads(
                    project_name=project_name, filter_string=f'id = "{thread_id}"'
                )
            )
            == 1
        ),
    ), f"Failed to get thread with id '{thread_id}'."
    threads = opik_client.search_threads(
        project_name=project_name,
        filter_string=f'id = "{thread_id}"',
    )
    assert len(threads) == 1

    thread = threads[0]
    assert thread.id == thread_id

    def _get_feedback_scores() -> Optional[List[Union[FeedbackScore]]]:
        return opik_client.search_threads(
            project_name=project_name,
            filter_string=f'id = "{thread_id}"',
        )[0].feedback_scores

    if feedback_scores is not mock.ANY:
        # wait for feedback scores to propagate
        assert synchronization.until(lambda: _get_feedback_scores() is not None), (
            f"Failed to get feedback scores for thread with id '{thread_id}'."
        )

        actual_feedback_scores = _get_feedback_scores()
        _assert_feedback_scores(
            item_id=thread_id,
            feedback_scores=actual_feedback_scores,
            expected_feedback_scores=feedback_scores,
        )


def _assert_feedback_scores(
    item_id: str,
    feedback_scores: Optional[List[Union[FeedbackScore, FeedbackScorePublic]]],
    expected_feedback_scores: Optional[List[FeedbackScoreDict]],
) -> None:
    if feedback_scores is None:
        assert expected_feedback_scores is None, (
            f"Expected feedback scores to be None, but got {expected_feedback_scores}"
        )
        return

    actual_feedback_scores = [] if feedback_scores is None else feedback_scores
    assert len(actual_feedback_scores) == len(expected_feedback_scores), (
        f"Expected amount of feedback scores ({len(expected_feedback_scores)}) is not equal to actual amount ({len(actual_feedback_scores)})"
    )

    actual_feedback_scores: List[FeedbackScoreDict] = [
        FeedbackScoreDict(
            category_name=score.category_name,
            id=item_id,
            name=score.name,
            reason=score.reason.strip() if score.reason is not None else None,
            value=score.value,
        )
        for score in feedback_scores
    ]

    sorted_actual_feedback_scores = sorted(
        actual_feedback_scores, key=lambda item: json.dumps(item, sort_keys=True)
    )
    sorted_expected_feedback_scores = sorted(
        expected_feedback_scores, key=lambda item: json.dumps(item, sort_keys=True)
    )
    for actual_score, expected_score in zip(
        sorted_actual_feedback_scores, sorted_expected_feedback_scores
    ):
        testlib.assert_dicts_equal(actual_score, expected_score, ignore_keys=["value"])
        assert expected_score["value"] == pytest.approx(
            actual_score["value"], abs=0.0001
        )


def verify_prompt_version(
    prompt: Prompt,
    *,
    name: Any = mock.ANY,  # type: ignore
    template: Any = mock.ANY,  # type: ignore
    type: Any = mock.ANY,  # type: ignore
    metadata: Any = mock.ANY,  # type: ignore
    version_id: Any = mock.ANY,  # type: ignore
    prompt_id: Any = mock.ANY,  # type: ignore
    commit: Any = mock.ANY,  # type: ignore
    project_name: Any = mock.ANY,
    environments: Optional[Iterable[str]] = None,
) -> None:
    testlib.assert_equal(name, prompt.name)
    testlib.assert_equal(project_name, prompt.project_name)
    testlib.assert_equal(template, prompt.prompt)
    testlib.assert_equal(type, prompt.type)
    testlib.assert_equal(metadata, prompt.metadata)
    assert version_id == prompt.__internal_api__version_id__, (
        f"{prompt.__internal_api__version_id__} != {version_id}"
    )
    assert prompt_id == prompt.__internal_api__prompt_id__, (
        f"{prompt.__internal_api__prompt_id__} != {prompt_id}"
    )
    assert commit == prompt.commit, f"{prompt.commit} != {commit}"
    if environments is not None:
        verify_prompt_environments(prompt, contains=environments)


def verify_chat_prompt_version(
    chat_prompt: ChatPrompt,
    *,
    name: Any = mock.ANY,  # type: ignore
    messages: Any = mock.ANY,  # type: ignore
    type: Any = mock.ANY,  # type: ignore
    metadata: Any = mock.ANY,  # type: ignore
    version_id: Any = mock.ANY,  # type: ignore
    prompt_id: Any = mock.ANY,  # type: ignore
    commit: Any = mock.ANY,  # type: ignore
    project_name: Any = mock.ANY,  # type: ignore
    environments: Optional[Iterable[str]] = None,
) -> None:
    """
    Verifies that a ChatPrompt has the expected properties.

    This verifier checks all the same fields as verify_prompt_version but adapted for ChatPrompt:
    - messages instead of template
    - template_structure field is always verified to be "chat"
    """
    testlib.assert_equal(name, chat_prompt.name)
    testlib.assert_equal(messages, chat_prompt.template)
    testlib.assert_equal(type, chat_prompt.type)
    testlib.assert_equal(metadata, chat_prompt.metadata)
    testlib.assert_equal(project_name, chat_prompt.project_name)
    assert version_id == chat_prompt.__internal_api__version_id__, (
        f"{chat_prompt.__internal_api__version_id__} != {version_id}"
    )
    assert prompt_id == chat_prompt.__internal_api__prompt_id__, (
        f"{chat_prompt.__internal_api__prompt_id__} != {prompt_id}"
    )
    assert commit == chat_prompt.commit, f"{chat_prompt.commit} != {commit}"
    if environments is not None:
        verify_prompt_environments(chat_prompt, contains=environments)


def verify_prompt_environments(
    prompt: Union[Prompt, ChatPrompt],
    *,
    contains: Optional[Iterable[str]] = None,
    excludes: Optional[Iterable[str]] = None,
    exactly: Optional[Iterable[str]] = None,
) -> None:
    """Verify the environment ownership of a prompt version.

    Environment order is not guaranteed by the backend, so checks are
    set-based:
    - ``contains``: every listed environment must be present
    - ``excludes``: none of the listed environments may be present
    - ``exactly``: the environment set must match exactly (use an empty
      iterable to assert ownership was cleared)
    """
    actual: Set[str] = set(prompt.environments or [])

    if exactly is not None:
        expected = set(exactly)
        assert actual == expected, f"environments {actual} != expected {expected}"

    if contains is not None:
        missing = set(contains) - actual
        assert not missing, f"environments {actual} missing {missing}"

    if excludes is not None:
        present = set(excludes) & actual
        assert not present, f"environments {actual} should not contain {present}"


def verify_opik_prompt_entry(
    entry: Dict[str, Any],
    *,
    name: str,
    template_structure: str,
) -> None:
    """Verify a single entry inside metadata['opik_prompts']."""
    assert entry["name"] == name, f"Expected name {name!r}, got {entry.get('name')!r}"
    assert entry.get("template_structure") == template_structure, (
        f"Expected template_structure {template_structure!r}, got {entry.get('template_structure')!r}"
    )
    assert "id" in entry, f"Missing 'id' in opik_prompts entry for {name}"
    version = entry.get("version", {})
    assert "template" in version, (
        f"Missing 'version.template' in opik_prompts entry for {name}"
    )
    assert "commit" in version, (
        f"Missing 'version.commit' in opik_prompts entry for {name}"
    )
    # Sequential version identifier (e.g. "v1"). Mask versions don't carry one,
    # but the injected entries in this happy-path test come from non-mask
    # `create_prompt` / `create_chat_prompt` calls, which always have it.
    assert version.get("version_number", "").startswith("v"), (
        f"Expected 'version.version_number' starting with 'v' in opik_prompts entry for {name}, "
        f"got {version.get('version_number')!r}"
    )


def verify_dataset_filtered_items(
    opik_client: opik.Opik,
    dataset_name: str,
    filter_string: str,
    expected_count: int,
    expected_inputs: Set[str],
    project_name: Optional[str] = None,
) -> None:
    """
    Verifies that filtering dataset items with filter_string returns the expected results.

    Args:
        opik_client: The Opik client instance
        dataset_name: The name of the dataset to retrieve
        filter_string: The filter string to apply
        expected_count: Expected number of items matching the filter
        expected_inputs: Set of expected question strings from input field
        project_name: The name of the project to retrieve the dataset from, if any
    """
    dataset = opik_client.get_dataset(name=dataset_name, project_name=project_name)

    filtered_items = dataset.get_items(filter_string=filter_string)
    assert len(filtered_items) == expected_count, (
        f"Expected {expected_count} items, got {len(filtered_items)}"
    )

    if expected_count > 0:
        inputs = {item["input"]["question"] for item in filtered_items}
        assert inputs == expected_inputs, (
            f"Input mismatch: {inputs} != {expected_inputs}"
        )


def verify_traces_annotation_queue(
    opik_client: opik.Opik,
    queue_id: str,
    name: str = mock.ANY,  # type: ignore
    scope: str = mock.ANY,  # type: ignore
    description: Optional[str] = mock.ANY,  # type: ignore
    instructions: Optional[str] = mock.ANY,  # type: ignore
) -> None:
    assert synchronization.until(
        lambda: opik_client.get_traces_annotation_queue(queue_id) is not None,
        allow_errors=True,
    ), f"Failed to get annotation queue with id {queue_id}."

    queue = opik_client.get_traces_annotation_queue(queue_id)

    assert queue.id is not None, "Queue id should not be None"
    testlib.assert_equal(name, queue.name)
    testlib.assert_equal(scope, queue.scope)
    testlib.assert_equal(description, queue.description)
    testlib.assert_equal(instructions, queue.instructions)


def verify_threads_annotation_queue(
    opik_client: opik.Opik,
    queue_id: str,
    name: str = mock.ANY,  # type: ignore
    scope: str = mock.ANY,  # type: ignore
    description: Optional[str] = mock.ANY,  # type: ignore
    instructions: Optional[str] = mock.ANY,  # type: ignore
) -> None:
    assert synchronization.until(
        lambda: opik_client.get_threads_annotation_queue(queue_id) is not None,
        allow_errors=True,
    ), f"Failed to get annotation queue with id {queue_id}."

    queue = opik_client.get_threads_annotation_queue(queue_id)

    assert queue.id is not None, "Queue id should not be None"
    testlib.assert_equal(name, queue.name)
    testlib.assert_equal(scope, queue.scope)
    testlib.assert_equal(description, queue.description)
    testlib.assert_equal(instructions, queue.instructions)


def verify_test_suite_result(
    opik_client: opik.Opik,
    suite_result: Any,
    items_total: int = mock.ANY,  # type: ignore
    items_passed: int = mock.ANY,  # type: ignore
    experiment_items_count: int = mock.ANY,  # type: ignore
    total_feedback_scores: int = mock.ANY,  # type: ignore
    expected_score_names: Optional[Set[str]] = None,
    project_name: Optional[str] = None,
):
    """
    Verify a TestSuiteResult — both in-memory properties and persisted
    experiment data from the backend.

    Args:
        opik_client: The Opik client instance.
        suite_result: The TestSuiteResult returned by suite.run().
        items_total: Expected total number of dataset items in the suite.
        items_passed: Expected number of dataset items that passed.
        experiment_items_count: Expected number of experiment items (traces)
            persisted in the backend. For multi-run suites this is
            dataset_items * runs_per_item.
        total_feedback_scores: Expected total number of feedback scores
            across all experiment items.
        expected_score_names: If provided, the union of all score names
            across all experiment items must equal this set.
        project_name: The project name associated with the test suite.
    """
    if items_total is not mock.ANY:
        assert suite_result.items_total == items_total, (
            f"Expected items_total={items_total}, got {suite_result.items_total}"
        )

    if items_passed is not mock.ANY:
        assert suite_result.items_passed == items_passed, (
            f"Expected items_passed={items_passed}, got {suite_result.items_passed}"
        )

        expected_all_passed = items_passed == (
            items_total if items_total is not mock.ANY else suite_result.items_total
        )
        assert suite_result.all_items_passed is expected_all_passed, (
            f"Expected all_items_passed={expected_all_passed}, "
            f"got {suite_result.all_items_passed}"
        )

    if experiment_items_count is mock.ANY:
        return

    retrieved_experiment = None
    experiment_items: List[Any] = []

    def _experiment_ready() -> bool:
        nonlocal retrieved_experiment, experiment_items
        try:
            retrieved_experiment = opik_client.get_experiment_by_name(
                suite_result.experiment_name, project_name=project_name
            )
        except Exception:
            return False
        experiment_items = retrieved_experiment.get_items()
        return len(experiment_items) == experiment_items_count

    assert synchronization.until(_experiment_ready, max_try_seconds=10), (
        f"Experiment '{suite_result.experiment_name}' did not become "
        f"visible with {experiment_items_count} items in time "
        f"(got {len(experiment_items)})"
    )

    testlib.assert_equal(retrieved_experiment.name, suite_result.experiment_name)

    all_scores = []
    all_score_names: Set[str] = set()
    for exp_item in experiment_items:
        if exp_item.feedback_scores:
            all_scores.extend(exp_item.feedback_scores)
            all_score_names.update(s["name"] for s in exp_item.feedback_scores)
        if exp_item.assertion_results:
            all_scores.extend(exp_item.assertion_results)
            all_score_names.update(ar["value"] for ar in exp_item.assertion_results)

    if total_feedback_scores is not mock.ANY:
        assert len(all_scores) == total_feedback_scores, (
            f"Expected {total_feedback_scores} total feedback scores, "
            f"got {len(all_scores)}"
        )

    if expected_score_names is not None:
        assert expected_score_names == all_score_names, (
            f"Expected score names {expected_score_names}, got {all_score_names}"
        )