Sampling When Every Request Is Different

# ── What uniform sampling costs you here ─────────────────────────
#
#  10,000 turns/day. The distribution is heavily skewed:
#     steps:   p50 = 3    p95 = 7    p99 = 14    max = 41
#     cost:    p50 = $0.012   p99 = $0.31   max = $2.90
#     failures: 0.4% of turns produce a complaint
#
#  Uniform 1% sampling keeps 100 turns:
#     turns with >= 14 steps      expected  1.0   often ZERO
#     turns costing > $0.31       expected  1.0   often ZERO
#     turns that produced a complaint  expected 0.4  usually ZERO
#
#  So the sample tells you a great deal about the median turn, which
#  is the one you had no questions about.

# ── Stratified: keep the classes, sample the middle ─────────────
from opentelemetry.sdk.trace.sampling import (
    Sampler, SamplingResult, Decision, TraceIdRatioBased,
)

BASELINE = 0.02


class AgentSampler(Sampler):
    """Head sampling. Note the hard limit: at span START we know the
    request, not the outcome. Anything outcome-based is tail
    sampling in the Collector -> module 12."""

    def __init__(self):
        self._ratio = TraceIdRatioBased(BASELINE)

    def should_sample(self, parent_context, trace_id, name,
                      kind=None, attributes=None, links=None,
                      trace_state=None):
        a = attributes or {}

        def keep(rule: str, weight: float):
            return SamplingResult(
                Decision.RECORD_AND_SAMPLE,
                {
                    **a,
                    "sampling.rule": rule,
                    # THE important attribute. It lets analysis
                    # reconstruct a population estimate from a
                    # deliberately biased sample.
                    "sampling.weight": weight,
                },
                trace_state,
            )

        # 1. Anything a human flagged. Rare, always wanted.
        if a.get("feedback.present") is True:
            return keep("user_feedback", 1.0)

        # 2. A conversation someone is already investigating.
        if a.get("debug.session") is True:
            return keep("debug_session", 1.0)

        # 3. New prompt or agent version: keep everything until the
        #    rollout is established, then let it fall to baseline.
        if a.get("prompt.version") in ROLLING_OUT:
            return keep("version_rollout", 1.0)

        # 4. A tenant with a contractual SLO.
        if a.get("tenant.tier") == "enterprise":
            return keep("enterprise_tenant", 1.0)

        # 5. First turn of a conversation: cold cache, no history,
        #    systematically different from turn 20.
        if a.get("gen_ai.conversation.turn_index") == 0:
            return keep("first_turn", 1.0)

        # 6. Otherwise the baseline, with the weight recorded so the
        #    sample can be scaled back up.
        r = self._ratio.should_sample(parent_context, trace_id, name,
                                      kind, a, links, trace_state)
        if r.decision != Decision.DROP:
            return keep("baseline", 1.0 / BASELINE)
        return SamplingResult(Decision.DROP)

    def get_description(self):
        return "AgentSampler"


# ── Why sampling.weight matters ─────────────────────────────────
#  Without it, "mean cost per turn" computed over the sample is
#  badly wrong: you over-represent enterprise tenants and first
#  turns by 50x. With it:
#
#     estimated_total = sum(cost_i * weight_i)
#
#  Uniformly sampled traces have weight 1/rate; always-kept classes
#  have weight 1. Record it and your biased sample stays usable for
#  quantitative work. Omit it and every aggregate you compute from
#  traces is silently skewed — which is also the reason cost and
#  token METRICS must be emitted unsampled (module 9).

# ── What must be decided at the TAIL instead ────────────────────
#  keep every turn that ended in an error         -> outcome
#  keep every turn over 10 steps                  -> outcome
#  keep every turn costing more than $0.50        -> outcome
#  keep every turn whose eval score was low       -> much later
#
#  None of these are knowable at span start. Module 12 configures
#  them as tail policies, including the delayed-eval case.