fix(sampling): make sample final and add sample_index to majority-vote hook payloads

BaseMBRDSampling._sample fans out number_of_samples concurrent calls to
BaseSamplingStrategy._sample, each receiving the same sampling_id and
starting their own _subsample_iteration loop from subsample_index=0.
This caused all branches to emit the same iteration numbers under one
sampling_id, making (sampling_id, iteration) non-unique for consumers
of SAMPLING_ITERATION and SAMPLING_REPAIR hooks.

Fix by introducing sample_index: int | None = None on both payload
classes and threading it through the emit helpers and _subsample_iteration
so each fan-out branch carries a distinct 0-based ordinal. Non-fan-out
strategies leave sample_index=None; no existing call sites change.

- Renamed _sample_impl to _sample
- Marked sample method as @Final enforces that subclasses override _sample
  rather than sample()
- Add sample_index field to SamplingIterationPayload and
  SamplingRepairPayload (mellea/plugins/hooks/sampling.py)
- Add sample_index kwarg to _emit_sampling_iteration and
  _emit_sampling_repair, forwarded to the payload
  (mellea/core/sampling.py)
- Add sample_index param to BaseSamplingStrategy._sample and
  _subsample_iteration; forward to both _emit_* calls
  (mellea/stdlib/sampling/base.py)
- Pass sample_index=i in the BaseMBRDSampling fan-out loop
  (mellea/stdlib/sampling/majority_voting.py)
- Emit mellea.sampling.sample_index span-event attribute in
  SamplingTracingPlugin.on_iteration and on_repair when not None
  (mellea/telemetry/tracing_plugins.py)

Tests added:
- test_majority_vote_iteration_sample_index_is_unique: e2e regression
  proving (sampling_id, sample_index, iteration) is unique across all
  branches with number_of_samples=3, loop_budget=2
- test_majority_vote_repair_sample_index_matches_branch: repair events
  carry the same sample_index as the failed iteration that triggered them
- test_sample_index_defaults_to_none / test_sample_index_construction
  on both payload classes
- test_sampling_iteration_includes/omits_sample_index_when_set/none
- test_sampling_repair_includes/omits_sample_index_when_set/none

Assisted-by: IBM Bob
Signed-off-by: Vishal V <VishalV@ibm.com>

Author: Vishal V

docs/docs/community/building-extensions.md

@@ -13,11 +13,11 @@ Components, Sampling Strategies, and Backends without modifying the core library
 
 Choose the pathway that fits the scope of your work:
 
-| Pathway | When to use |
-| ------- | ----------- |
-| **Core repository** | General-purpose additions that benefit all users — open an issue first to discuss placement |
-| **Your own repo** (`mellea-` prefix) | Application-specific or domain-specific libraries |
-| **[mellea-contribs](https://github.com/generative-computing/mellea-contribs)** | Experimental or specialized components not yet ready for the standard library |
+| Pathway                                                                        | When to use                                                                                 |
+| ------------------------------------------------------------------------------ | ------------------------------------------------------------------------------------------- |
+| **Core repository**                                                            | General-purpose additions that benefit all users — open an issue first to discuss placement |
+| **Your own repo** (`mellea-` prefix)                                           | Application-specific or domain-specific libraries                                           |
+| **[mellea-contribs](https://github.com/generative-computing/mellea-contribs)** | Experimental or specialized components not yet ready for the standard library               |
 
 > **Note:** For general-purpose Components, Requirements, or Sampling Strategies,
 > open an issue before submitting a PR. This avoids duplication and ensures
@@ -143,12 +143,14 @@ For a full walkthrough of the Component protocol and templating system, see
 
 A [`SamplingStrategy`](../reference/glossary#sampling-strategy) controls how Mellea
 generates and validates outputs — for example, rejection sampling, best-of-n, or
-beam search. Subclass `SamplingStrategy` and implement `sample`:
+beam search. Subclass `SamplingStrategy` and implement `_sample`:
+
+> **Note:** `sample()` is `@final` — it owns the sampling lifecycle (hooks, budget
+> management). Override `_sample()` instead; `sample()` calls it automatically.
 
 ```python
-import asyncio
 from mellea.core.backend import Backend
-from mellea.core.base import Component, Context, ModelOutputThunk, S
+from mellea.core.base import CBlock, Component, Context, ModelOutputThunk, S
 from mellea.core.requirement import Requirement
 from mellea.core.sampling import SamplingResult, SamplingStrategy
 
@@ -164,36 +166,42 @@ class BestOfNStrategy(SamplingStrategy):
         """
         self.n = n
 
-    async def sample(
+    async def _sample(
         self,
-        action: Component[S],
+        action: Component[S] | CBlock | ModelOutputThunk,
         context: Context,
         backend: Backend,
-        requirements: list[Requirement] | None,
+        requirements: list[Requirement],
         *,
+        effective_loop_budget: int,
         validation_ctx: Context | None = None,
         format: type | None = None,
         model_options: dict | None = None,
         tool_calls: bool = False,
+        sampling_id: str,
+        **kwargs,
     ) -> SamplingResult[S]:
         """Generate N candidates and return the best one.
 
         Args:
-            action: The component to generate a response for.
+            action: The component, block, or thunk to generate a response for.
             context: The current session context.
             backend: The backend used for generation.
             requirements: Requirements to validate each candidate against.
+            effective_loop_budget: Loop budget after hook modification.
             validation_ctx: Optional context override for validation.
             format: Structured output format, if any.
             model_options: Model options to pass to the backend.
             tool_calls: Whether to enable tool calls during generation.
+            sampling_id: UUID correlating iteration/repair/end hooks to this loop.
+            **kwargs: Additional keyword arguments (e.g., `show_progress`).
 
         Returns:
             SamplingResult containing the selected candidate and validation details.
         """
         generations: list[ModelOutputThunk[S]] = []
         contexts: list[Context] = []
-        actions: list[Component[S]] = []
+        actions: list[Component[S] | CBlock | ModelOutputThunk] = []
         validations: list[list[tuple[Requirement, object]]] = []
 
         for _ in range(self.n):

mellea/core/sampling.py

@@ -14,7 +14,7 @@
 import abc
 import uuid
 from collections.abc import Sequence
-from typing import Generic
+from typing import Generic, final
 
 from .backend import Backend, BaseModelSubclass
 from .base import (
@@ -125,8 +125,17 @@ class SamplingStrategy(abc.ABC):
 
     This class provides a template for creating concrete sampling strategies that can be used to generate model outputs based on given instructions.
     It allows setting custom validation and generation functions through properties.
+
+    Attributes:
+        loop_budget: Maximum number of generate/validate cycles. Defaults to `1`.
+        requirements: Global requirements evaluated on every sample. When set,
+            overrides per-call requirements. Defaults to `None`.
     """
 
+    loop_budget: int = 1
+    requirements: list[Requirement] | None = None
+
+    @final
     async def sample(
         self,
         action: Component[S] | CBlock | ModelOutputThunk,
@@ -143,7 +152,7 @@ async def sample(
         """Concrete wrapper: owns the sampling lifecycle and fires loop start/end hooks.
 
         Mints a `sampling_id`, dispatches `sampling_loop_start` (which may modify
-        `loop_budget`), delegates to `_sample_impl`, and dispatches
+        `loop_budget`), delegates to `_sample`, and dispatches
         `sampling_loop_end` on every exit path — success, budget exhaustion, and
         raised exceptions.
 
@@ -156,7 +165,7 @@ async def sample(
             format: output format for structured outputs.
             model_options: model options to pass to the backend during generation / validation.
             tool_calls: True if tool calls should be used during this sampling strategy.
-            **kwargs: Additional keyword arguments forwarded to `_sample_impl`.
+            **kwargs: Additional keyword arguments forwarded to `_sample`.
 
         Returns:
             SamplingResult[S]: A result object indicating the success or failure of the sampling process.
@@ -174,8 +183,7 @@ async def sample(
 
         try:
             reqs = self._merge_requirements(requirements)
-            loop_budget = self._get_loop_budget()
-            effective_loop_budget = loop_budget
+            effective_loop_budget = self.loop_budget
 
             # --- sampling_loop_start hook ---
             if has_plugins(HookType.SAMPLING_LOOP_START):
@@ -187,7 +195,7 @@ async def sample(
                     action=action,
                     context=context,
                     requirements=reqs,
-                    loop_budget=loop_budget,
+                    loop_budget=self.loop_budget,
                 )
                 _, start_payload = await invoke_hook(
                     HookType.SAMPLING_LOOP_START, start_payload, backend=backend
@@ -202,7 +210,7 @@ async def sample(
                     f"{effective_loop_budget}; must be >= 1."
                 )
 
-            s_result = await self._sample_impl(
+            s_result = await self._sample(
                 action=action,
                 context=context,
                 backend=backend,
@@ -255,7 +263,7 @@ async def sample(
                 )
 
     @abc.abstractmethod
-    async def _sample_impl(
+    async def _sample(
         self,
         action: Component[S] | CBlock | ModelOutputThunk,
         context: Context,
@@ -301,6 +309,8 @@ async def _emit_sampling_iteration(
         result: ModelOutputThunk,
         validation_results: list[tuple[Requirement, ValidationResult]],
         backend: Backend,
+        *,
+        sample_index: int | None = None,
     ) -> None:
         """Emit the sampling-iteration hook payload if any plugin is registered."""
         from ..plugins.manager import has_plugins, invoke_hook
@@ -316,6 +326,7 @@ async def _emit_sampling_iteration(
             sampling_id=sampling_id,
             strategy_name=type(self).__name__,
             iteration=iteration,
+            sample_index=sample_index,
             action=action,
             result=result,
             validation_results=validation_results,
@@ -335,6 +346,8 @@ async def _emit_sampling_repair(
         repair_action: SampleActionType,
         repair_context: Context,
         backend: Backend,
+        *,
+        sample_index: int | None = None,
     ) -> None:
         """Emit the sampling-repair hook payload if any plugin is registered."""
         from ..plugins.manager import has_plugins, invoke_hook
@@ -354,6 +367,7 @@ async def _emit_sampling_repair(
             repair_action=repair_action,
             repair_context=repair_context,
             repair_iteration=repair_iteration,
+            sample_index=sample_index,
         )
         await invoke_hook(HookType.SAMPLING_REPAIR, repair_payload, backend=backend)
 
@@ -370,18 +384,9 @@ def _merge_requirements(
         Returns:
             Deduplicated list of requirements to use for this sampling call.
         """
-        strategy_reqs: list[Requirement] | None = getattr(self, "requirements", None)
         reqs: list[Requirement] = []
-        if strategy_reqs is not None:
-            reqs += strategy_reqs
+        if self.requirements is not None:
+            reqs += self.requirements
         elif call_requirements is not None:
             reqs += call_requirements
         return list(set(reqs))
-
-    def _get_loop_budget(self) -> int:
-        """Return the strategy's configured loop budget (default 1).
-
-        Returns:
-            The loop budget from `self.loop_budget` if set, otherwise `1`.
-        """
-        return getattr(self, "loop_budget", 1)

mellea/plugins/hooks/sampling.py

@@ -41,6 +41,9 @@ class SamplingIterationPayload(MelleaBasePayload):
         strategy_name: Class name of the sampling strategy (e.g. `"RejectionSamplingStrategy"`).
         iteration: 1-based iteration number within the sampling loop. There is no guarantee that
             iteration will be monotonically increasing when concurrency is enabled.
+        sample_index: Optional 0-based index identifying the outer concurrent sample branch —
+            set by fan-out strategies such as majority voting; `None` for strategies without an
+            outer fan-out.
         action: The `Component` used for this attempt.
 
         result: The `ModelOutputThunk` produced by this attempt.
@@ -54,6 +57,7 @@ class SamplingIterationPayload(MelleaBasePayload):
     sampling_id: str = ""
     strategy_name: str = ""
     iteration: int = 0
+    sample_index: int | None = None
     action: Any = None
     result: Any = None
     validation_results: list[tuple[Any, Any]] = []
@@ -75,6 +79,9 @@ class SamplingRepairPayload(MelleaBasePayload):
         repair_action: The repaired `Component` to use for the next attempt.
         repair_context: The `Context` to use for the next attempt.
         repair_iteration: 1-based iteration at which the repair was triggered.
+        sample_index: Optional 0-based index identifying the outer concurrent sample branch —
+            set by fan-out strategies such as majority voting; `None` for strategies without an
+            outer fan-out.
     """
 
     sampling_id: str = ""
@@ -85,6 +92,7 @@ class SamplingRepairPayload(MelleaBasePayload):
     repair_action: Any = None
     repair_context: Any = None
     repair_iteration: int = 0
+    sample_index: int | None = None
 
 
 class SamplingLoopEndPayload(MelleaBasePayload):

mellea/stdlib/sampling/base.py

@@ -199,7 +199,7 @@ def select_from_failure(
         """
         ...
 
-    async def _sample_impl(
+    async def _sample(
         self,
         action: Component[S] | CBlock | ModelOutputThunk,
         context: Context,
@@ -213,6 +213,7 @@ async def _sample_impl(
         tool_calls: bool = False,
         sampling_id: str,
         show_progress: bool = True,
+        sample_index: int | None = None,
         **kwargs,
     ) -> SamplingResult[S]:
         """Execute the rejection-sampling loop.
@@ -229,6 +230,7 @@ async def _sample_impl(
             tool_calls: True if tool calls should be used during this sampling strategy.
             sampling_id: UUID correlating iteration/repair/end hooks for this loop.
             show_progress: if true, a tqdm progress bar is used. Otherwise, messages will still be sent to flog.
+            sample_index: Optional 0-based index of this branch within a fan-out strategy (e.g. majority voting). `None` for strategies without an outer fan-out.
             **kwargs: Additional keyword arguments forwarded by `SamplingStrategy.sample()`.
 
         Returns:
@@ -279,6 +281,7 @@ async def _sample_impl(
                     model_options=model_options,
                     tool_calls=tool_calls,
                     sampling_id=sampling_id,
+                    sample_index=sample_index,
                 )
                 for idx in range(self.concurrency_budget)
             ]
@@ -406,6 +409,7 @@ async def _subsample_iteration(
         model_options: dict | None = None,
         tool_calls: bool = False,
         sampling_id: str,
+        sample_index: int | None = None,
     ) -> AsyncGenerator[_SamplingResultSlice[S], Any]:
         """Run one concurrent subsample: up to `iterations` generate/validate/repair attempts.
 
@@ -476,6 +480,7 @@ async def _subsample_iteration(
                     result=result,
                     validation_results=constraint_scores,
                     backend=backend,
+                    sample_index=sample_index,
                 )
 
                 if not all_validations_passed:
@@ -530,6 +535,7 @@ async def _subsample_iteration(
                     repair_action=next_action,
                     repair_context=next_context,
                     backend=backend,
+                    sample_index=sample_index,
                 )
 
 

mellea/stdlib/sampling/budget_forcing.py

@@ -102,7 +102,7 @@ def _get_repair_type() -> str:
         """Return the repair-type label for telemetry hooks."""
         return "budget_forcing"
 
-    async def _sample_impl(
+    async def _sample(
         self,
         action: Component[S] | CBlock | ModelOutputThunk,
         context: Context,

mellea/stdlib/sampling/majority_voting.py

@@ -134,7 +134,7 @@ def maybe_apply_weighted(self, scr: np.ndarray) -> np.ndarray:
 
         return scr
 
-    async def _sample_impl(
+    async def _sample(
         self,
         action: Component[S] | CBlock | ModelOutputThunk,
         context: Context,
@@ -152,7 +152,7 @@ async def _sample_impl(
     ) -> SamplingResult[S]:
         """Samples using majority voting.
 
-        Fan-out calls `_sample_impl` on the parent class directly so that each
+        Fan-out calls `_sample` on the parent class directly so that each
         inner sample does not emit its own enclosing `sampling` span — there is
         exactly one `sampling` span for the top-level call, emitted by the
         `SamplingStrategy.sample` wrapper.
@@ -174,12 +174,12 @@ async def _sample_impl(
         Returns:
             SamplingResult[S]: A result object indicating the success or failure of the sampling process.
         """
-        # execute sampling concurrently — call _sample_impl directly on the
+        # execute sampling concurrently — call _sample directly on the
         # parent so each inner sample does not fire its own lifecycle hooks.
         tasks: list[asyncio.Task[SamplingResult]] = []
         for i in range(self.number_of_samples):
             task = asyncio.create_task(
-                super()._sample_impl(
+                super()._sample(
                     action,
                     context,
                     backend,
@@ -190,6 +190,7 @@ async def _sample_impl(
                     tool_calls=tool_calls,
                     sampling_id=sampling_id,
                     show_progress=show_progress,
+                    sample_index=i,
                 )
             )
             tasks.append(task)

mellea/stdlib/sampling/sofai.py

@@ -579,7 +579,7 @@ async def _generate_and_validate(
     # Main Sample Method
     # =========================================================================
 
-    async def _sample_impl(
+    async def _sample(
         self,
         action: Component[S] | CBlock | ModelOutputThunk,
         context: Context,