tests/observability/test_contracts.py

from __future__ import annotations

import typing

import pytest
from pydantic import ValidationError

from app.observability.contracts import (
    METRIC_ATTRIBUTE_KEYS,
    ExperimentIdentity,
    OperationObservation,
    RetrievalOutcome,
    TelemetryMode,
    metric_attributes,
)


# --- RetrievalOutcome -------------------------------------------------


def test_retrieval_outcome_distinguishes_empty_from_error():
    empty = RetrievalOutcome.success([])
    failed = RetrievalOutcome.failure("collection_unavailable")

    assert empty.status == "success_empty"
    assert empty.empty_result is True
    assert empty.retrieval_error is False

    assert failed.status == "error_fallback"
    assert failed.empty_result is False
    assert failed.retrieval_error is True


def test_retrieval_outcome_success_with_candidates_is_not_empty():
    outcome = RetrievalOutcome.success([{"chunk_id": "c1"}, {"chunk_id": "c2"}])

    assert outcome.status == "success"
    assert outcome.empty_result is False
    assert outcome.retrieval_error is False
    assert outcome.candidates == [{"chunk_id": "c1"}, {"chunk_id": "c2"}]


def test_retrieval_outcome_failure_carries_error_type_and_no_candidates():
    outcome = RetrievalOutcome.failure("timeout")

    assert outcome.error_type == "timeout"
    assert outcome.candidates == []


# --- metric_attributes / METRIC_ATTRIBUTE_KEYS -------------------------


def test_metric_attributes_remove_high_cardinality_values():
    safe = metric_attributes(
        {
            "pipeline.version": "rag-naive-v1",
            "query.category": "methodology",
            "query.id": "q-001",
            "project.id": "private-project",
        }
    )
    assert safe == {
        "pipeline.version": "rag-naive-v1",
        "query.category": "methodology",
    }


def test_metric_attributes_drops_non_scalar_values_even_for_safe_keys():
    safe = metric_attributes(
        {
            "pipeline.version": "rag-naive-v1",
            "operation": ["retrieval", "context_assembly"],
            "status": {"nested": "dict"},
        }
    )
    assert safe == {"pipeline.version": "rag-naive-v1"}


def test_metric_attributes_handles_empty_input():
    assert metric_attributes({}) == {}


def test_metric_attributes_never_mutates_input():
    original = {"pipeline.version": "v1", "query.id": "q-1"}
    snapshot = dict(original)
    metric_attributes(original)
    assert original == snapshot


# --- TelemetryMode -------------------------------------------------


def test_telemetry_mode_allows_only_the_three_canonical_values():
    assert typing.get_args(TelemetryMode) == ("disabled", "local", "full")


# --- ExperimentIdentity -------------------------------------------------


def test_experiment_identity_defaults_to_all_unset():
    identity = ExperimentIdentity()

    assert identity.experiment_id is None
    assert identity.run_id is None
    assert identity.pipeline_version is None
    assert identity.query_set_version is None
    assert identity.corpus_version is None
    assert identity.query_id is None
    assert identity.query_category is None
    assert identity.repetition is None
    assert identity.cache_state is None
    assert identity.git_commit is None
    assert identity.embedding_model is None
    assert identity.cerebras_model is None
    assert identity.chroma_version is None
    assert identity.cognee_version is None


def test_experiment_identity_accepts_all_canonical_fields():
    identity = ExperimentIdentity(
        experiment_id="exp-1",
        run_id="run-1",
        pipeline_version="rag-naive-v1",
        query_set_version="qs-1",
        corpus_version="corpus-1",
        query_id="q-001",
        query_category="methodology",
        repetition=3,
        cache_state="warm",
        git_commit="abc123",
        embedding_model="all-MiniLM-L6-v2",
        cerebras_model="gemma-4-31b",
        chroma_version="0.5.0",
        cognee_version="1.2.2",
    )

    assert identity.pipeline_version == "rag-naive-v1"
    assert identity.repetition == 3


def test_experiment_identity_rejects_unknown_fields():
    with pytest.raises(ValidationError):
        ExperimentIdentity(unknown_field="nope")


# --- OperationObservation -------------------------------------------------


def test_operation_observation_requires_operation_and_status():
    with pytest.raises(ValidationError):
        OperationObservation()


def test_operation_observation_minimal_construction_has_safe_defaults():
    obs = OperationObservation(operation="retrieval", status="success")

    assert obs.operation == "retrieval"
    assert obs.status == "success"
    assert obs.subsystem is None
    assert obs.consumer is None
    assert obs.duration_ms is None
    assert obs.stage_durations_ms == {}
    assert obs.counts == {}
    assert obs.trace_id is None
    assert obs.experiment is None
    assert obs.evaluation_run is False
    assert obs.attributes == {}


def test_operation_observation_carries_stage_durations_counts_and_trace_id():
    obs = OperationObservation(
        operation="retrieval",
        subsystem="chroma",
        consumer="chat",
        status="success",
        duration_ms=123.4,
        stage_durations_ms={"vector_search_ms": 88.0, "collection_lookup_ms": 12.0},
        counts={"raw_candidate_count": 20, "selected_candidate_count": 5},
        trace_id="trace-abc",
    )

    assert obs.stage_durations_ms["vector_search_ms"] == 88.0
    assert obs.counts["selected_candidate_count"] == 5
    assert obs.trace_id == "trace-abc"


def test_operation_observation_can_embed_experiment_identity_and_retrieval_outcome():
    identity = ExperimentIdentity(pipeline_version="rag-naive-v1", query_category="methodology")
    outcome = RetrievalOutcome.success([{"chunk_id": "c1"}])

    obs = OperationObservation(
        operation="retrieval",
        status=outcome.status,
        experiment=identity,
        retrieval=outcome,
        evaluation_run=True,
    )

    assert obs.experiment.pipeline_version == "rag-naive-v1"
    assert obs.retrieval.status == "success"
    assert obs.evaluation_run is True


def test_operation_observation_missing_data_uses_sentinel_not_zero():
    # Per the canonical contract: missing observations must never be silently
    # coerced to zero. A "not instrumented" duration must stay None, not 0.
    obs = OperationObservation(operation="memory", status="not_instrumented")

    assert obs.status == "not_instrumented"
    assert obs.duration_ms is None