docent-python 0.1.21a0__py3-none-any.whl → 0.1.23a0__py3-none-any.whl

docent/judges/impl.py

@@ -1,67 +1,275 @@
-import json
+import random
+import re
 from abc import ABC, abstractmethod
-from typing import Any
+from contextlib import nullcontext
+from typing import Any, Sequence
 
+import anyio
+import yaml
+from pydantic_core import to_jsonable_python
+from tqdm.auto import tqdm
+
+from docent._llm_util.data_models.exceptions import ValidationFailedException
 from docent._llm_util.data_models.llm_output import LLMOutput
-from docent._llm_util.data_models.simple_svc import BaseLLMService, SimpleLLMService
+from docent._llm_util.llm_svc import BaseLLMService
 from docent._log_util import get_logger
 from docent.data_models.agent_run import AgentRun
-from docent.judges.types import JudgeResult, ResultType, Rubric
-from docent.judges.util.parse_output import parse_and_validate_llm_output
-from docent.judges.util.voting import find_modal_result, get_agreement_keys
+from docent.data_models.chat.message import (
+    AssistantMessage,
+    ChatMessage,
+    ToolMessage,
+    UserMessage,
+)
+from docent.data_models.chat.tool import ToolInfo
+from docent.judges.types import JudgeResult, JudgeVariant, ResultType, Rubric
+from docent.judges.util.parse_output import parse_and_validate_output_str
+from docent.judges.util.voting import (
+    JudgeOutputDistribution,
+    compute_output_distributions,
+    find_modal_result,
+    get_agreement_keys,
+)
+from docent.trace import agent_run_context, agent_run_metadata
 
 logger = get_logger(__name__)
 
 
 class BaseJudge(ABC):
-    def __init__(self, cfg: Rubric, llm_svc: BaseLLMService):
+    def __init__(
+        self, cfg: Rubric, llm_svc: BaseLLMService, docent_collection_id: str | None = None
+    ):
         self.cfg = cfg
         self.llm_svc = llm_svc
+        self.docent_collection_id = docent_collection_id
 
     @abstractmethod
-    async def __call__(self, agent_run: AgentRun, *args: Any, **kwargs: Any) -> JudgeResult | None:
+    async def __call__(self, agent_run: AgentRun) -> JudgeResult | None:
         """Returns None if all rollouts failed to produce a valid output."""
 
+    @abstractmethod
+    async def estimate_output_distrs(
+        self, agent_run: AgentRun, **kwargs: Any
+    ) -> None | tuple[dict[str, JudgeOutputDistribution], dict[str, Any]]:
+        """Estimate the output distribution of each output key."""
 
-class MajorityVotingJudge(BaseJudge):
-    """Rolls out the judge multiple times, then uses majority voting to determine the final result."""
-
-    def __init__(
-        self,
-        cfg: Rubric,
-        n_rollouts_per_input: int,
-        llm_svc: BaseLLMService = SimpleLLMService(),
-    ):
-        super().__init__(cfg, llm_svc)
-        self.n_rollouts_per_input = n_rollouts_per_input
-
-    async def __call__(
-        self,
-        agent_run: AgentRun,
-        max_concurrency: int = 10,
-    ) -> JudgeResult | None:
+    def _get_validation_callback(self, agent_run: AgentRun):
         async def _validation_callback(batch_index: int, llm_output: LLMOutput):
-            parse_and_validate_llm_output(llm_output, self.cfg.output_schema, agent_run)
-
-        prompt = [{"role": "user", "content": self.cfg.materialize_system_prompt(agent_run)}]
+            validated_output = self._validate_first_response_tag_or_entire_output(
+                llm_output.first_text or "", agent_run
+            )
+            if validated_output is None:
+                raise ValidationFailedException(
+                    "Validation failed", failed_output=llm_output.first_text
+                )
+
+        return _validation_callback
+
+    async def one_rollout(
+        self, agent_run: AgentRun
+    ) -> tuple[dict[str, Any] | None, dict[str, Any] | None]:
+        with agent_run_context() if self.docent_collection_id is not None else nullcontext():
+            if self.cfg.rollout_type == "single_turn":
+                output, metadata = await self.one_single_turn_rollout(agent_run)
+            elif self.cfg.rollout_type == "multi_turn":
+                output, metadata = await self.one_multi_turn_rollout(
+                    agent_run, max_turns=10, max_steps_per_turn=5
+                )
+            else:
+                raise ValueError(f"Invalid rollout type: {self.cfg.rollout_type}")
+
+            if self.docent_collection_id is not None:
+                agent_run_metadata(
+                    {
+                        "agent_run_id": agent_run.id,
+                        "judge_output": output,
+                        "judge_rollout_metadata": to_jsonable_python(metadata),
+                    }
+                )
+
+        return output, metadata
+
+    def _validate_first_response_tag_or_entire_output(
+        self, output_str: str, agent_run: AgentRun
+    ) -> dict[str, Any] | None:
+        """Validate the first <response> tag in the output string.
+        For backward compatibility, also try to validate the entire output as JSON, for
+            old system prompts that don't ask for <response> tags.
+
+        Args:
+            output_str: The output string to validate
+            agent_run: The agent run to validate against
+
+        Returns:
+            The validated output if successful, None otherwise
+        """
+        response_matches = re.findall(r"<response>(.*?)</response>", output_str, re.DOTALL)
+
+        # Try to validate any match; take the first
+        for response_text in response_matches:
+            try:
+                validated_output = parse_and_validate_output_str(
+                    response_text, self.cfg.output_schema, agent_run
+                )
+                return validated_output
+            except ValidationFailedException:
+                continue  # Try the next match if validation fails
+
+        # Try to validate the entire output as JSON
+        # But only if the output _didn't_ contain a <response>...</response> tag
+        if not response_matches:
+            try:
+                validated_output = parse_and_validate_output_str(
+                    output_str, self.cfg.output_schema, agent_run
+                )
+                return validated_output
+            except ValidationFailedException:
+                pass
+
+        return None
+
+    ########################
+    # Single turn rollouts #
+    ########################
+
+    async def one_single_turn_rollout(
+        self, agent_run: AgentRun
+    ) -> tuple[dict[str, Any] | None, dict[str, Any] | None]:
+        prompt = [UserMessage(content=self.cfg.materialize_system_prompt(agent_run))]
         outputs = await self.llm_svc.get_completions(
-            inputs=[prompt for _ in range(self.n_rollouts_per_input)],
+            inputs=[prompt],
             model_options=[self.cfg.judge_model],
             max_new_tokens=16384,
             timeout=180.0,
             use_cache=False,
-            validation_callback=_validation_callback,
-            max_concurrency=max_concurrency,
+            validation_callback=self._get_validation_callback(agent_run),
         )
+        output_str = outputs[0].first_text
 
-        # Process each rollout independently
-        indep_results: list[dict[str, Any]] = []
-        for output in outputs:
-            if validated_output := parse_and_validate_llm_output(
-                output, self.cfg.output_schema, agent_run
-            ):
-                indep_results.append(validated_output)
+        # Extract all <response>...</response> tags from the current message
+        validated_output = self._validate_first_response_tag_or_entire_output(
+            output_str or "", agent_run
+        )
+        if validated_output is not None:
+            return validated_output, {"full_output": output_str}
+        else:
+            return None, None
+
+    #######################
+    # Multi-turn rollouts #
+    #######################
+
+    async def one_multi_turn_rollout(
+        self, agent_run: AgentRun, max_turns: int, max_steps_per_turn: int
+    ) -> tuple[dict[str, Any] | None, dict[str, Any] | None]:
+        msgs = [UserMessage(content=self.cfg.materialize_system_prompt(agent_run))]
+        for _ in range(max_turns):
+            msgs = await self.agent_one_turn(msgs, max_steps_per_turn=max_steps_per_turn)
+
+            last_msg_content = msgs[-1].text if msgs else None
+            # Extract all <response>...</response> tags from the current message
+            # Return if we find a valid response; otherwise, continue
+            validated_output = self._validate_first_response_tag_or_entire_output(
+                last_msg_content or "", agent_run
+            )
+            if validated_output is not None:
+                # When returning, strip out the system message, which duplicates the agent run
+                # content many times.
+                return validated_output, {"rollout_messages": msgs[1:]}
+
+        # No <response>...</response> tags with valid JSON,so return None
+        return None, None
+
+    async def agent_one_turn(self, init_msgs: Sequence[ChatMessage], max_steps_per_turn: int):
+        """Given a list of messages, run one turn of the agent.
+        The agent may invoke tools, so we loop until there are no more to handle.
+        """
+
+        msgs = list(init_msgs)  # Shallow copy is fine
+        for _ in range(max_steps_per_turn):
+            last_msg = msgs[-1]
+            if last_msg.role == "system" or last_msg.role == "user" or last_msg.role == "tool":
+                outputs = await self.llm_svc.get_completions(
+                    inputs=[msgs],
+                    model_options=[self.cfg.judge_model],
+                    tools=[
+                        ToolInfo(
+                            name="step_finished",
+                            description="Call this tool to indicate that you have finished one step in the decision procedure",
+                        )
+                    ],
+                    max_new_tokens=16384,
+                    timeout=180.0,
+                    use_cache=False,
+                )
+                output = outputs[0].first
+                if output is None:
+                    # FIXME(mengk): handle empty completion
+                    raise ValueError("Empty completion in agent one turn")
+                new_assistant_msg = AssistantMessage(
+                    content=output.text or "", tool_calls=output.tool_calls
+                )
+                msgs.append(new_assistant_msg)
+            elif last_msg.role == "assistant":
+                if last_msg.tool_calls is not None:
+                    msgs.extend(
+                        [
+                            ToolMessage(
+                                content="Step completed",
+                                tool_call_id=tool_call.id,
+                            )
+                            for tool_call in last_msg.tool_calls
+                        ]
+                    )
+                else:
+                    break  # Terminate if there are no more tool calls to handle
+            else:
+                raise ValueError(f"Unknown message role: {last_msg.role}")
+        return msgs
+
+
+class SingleRolloutJudge(BaseJudge):
+    """Rolls out the judge once."""
+
+    def __init__(self, cfg: Rubric, llm_svc: BaseLLMService):
+        super().__init__(cfg, llm_svc)
 
+    async def __call__(self, agent_run: AgentRun) -> JudgeResult | None:
+        output, metadata = await self.one_rollout(agent_run)
+        if output is None:
+            return None
+        else:
+            return JudgeResult(
+                agent_run_id=agent_run.id,
+                rubric_id=self.cfg.id,
+                rubric_version=self.cfg.version,
+                output=output,
+                result_metadata={"rollout_metadata": metadata},
+                result_type=ResultType.DIRECT_RESULT,
+            )
+
+
+class MajorityVotingJudge(BaseJudge):
+    """Rolls out the judge multiple times, then uses majority voting to determine the final result."""
+
+    def __init__(
+        self, cfg: Rubric, llm_svc: BaseLLMService, docent_collection_id: str | None = None
+    ):
+        super().__init__(cfg, llm_svc, docent_collection_id)
+
+    async def __call__(self, agent_run: AgentRun) -> JudgeResult | None:
+        indep_results: list[dict[str, Any]] = []
+        indep_rollout_metadata: list[dict[str, Any] | None] = []
+
+        async def _execute():
+            result, metadata = await self.one_rollout(agent_run)
+            if result is not None:
+                indep_results.append(result)
+                indep_rollout_metadata.append(metadata)
+
+        # Run rollouts concurrently
+        async with anyio.create_task_group() as tg:
+            for _ in range(self.cfg.n_rollouts_per_input):
+                tg.start_soon(_execute)
         if not indep_results:
             return None
 
@@ -74,6 +282,11 @@ class MajorityVotingJudge(BaseJudge):
         )
         final_output = indep_results[final_max_idx]
 
+        # Compute the distribution of the output across the agreement keys
+        final_output_distributions = compute_output_distributions(
+            indep_results, self.cfg.output_schema, agreement_keys
+        )
+
         return JudgeResult(
             agent_run_id=agent_run.id,
             rubric_id=self.cfg.id,
@@ -85,53 +298,122 @@ class MajorityVotingJudge(BaseJudge):
                 "final_results": indep_results,
                 "final_agt_key_modes_and_counts": final_agt_key_modes_and_counts,
                 "final_max_idx": final_max_idx,
+                "final_output_distributions": final_output_distributions,
+                "final_rollout_metadata": indep_rollout_metadata,
             },
             result_type=ResultType.DIRECT_RESULT,
         )
 
+    async def estimate_output_distrs(
+        self, agent_run: AgentRun, *, n_initial_rollouts_to_sample: int, **kwargs: Any
+    ) -> None | tuple[dict[str, JudgeOutputDistribution], dict[str, Any]]:
+        if self.cfg.n_rollouts_per_input > n_initial_rollouts_to_sample:
+            raise ValueError(
+                "n_initial_rollouts_to_sample must be greater than or equal to cfg.n_rollouts_per_input"
+            )
+
+        indep_results: list[dict[str, Any]] = []
+        indep_rollout_metadata: list[dict[str, Any] | None] = []
+        pbar = tqdm(total=n_initial_rollouts_to_sample, desc="Independent rollouts", leave=False)
+
+        async def _execute():
+            result, metadata = await self.one_rollout(agent_run)
+            if result is not None:
+                indep_results.append(result)
+                indep_rollout_metadata.append(metadata)
+            pbar.update(1)
+
+        # Run rollouts concurrently
+        async with anyio.create_task_group() as tg:
+            for _ in range(n_initial_rollouts_to_sample):
+                tg.start_soon(_execute)
+
+        pbar.close()
+
+        if not indep_results:
+            return None
+
+        # Compute the probability vector for each agreement key
+        distributions = compute_output_distributions(
+            indep_results, self.cfg.output_schema, get_agreement_keys(self.cfg.output_schema)
+        )
+
+        return distributions, {
+            "first_step_rollouts": indep_results,
+            "first_step_rollout_metadata": indep_rollout_metadata,
+        }
+
 
 class MultiReflectionJudge(BaseJudge):
     """Rolls out the judge multiple times, then uses reflection to determine the final result."""
 
     def __init__(
-        self,
-        cfg: Rubric,
-        n_rollouts_per_input: int,
-        llm_svc: BaseLLMService = SimpleLLMService(),
+        self, cfg: Rubric, llm_svc: BaseLLMService, docent_collection_id: str | None = None
     ):
-        super().__init__(cfg, llm_svc)
-        self.n_rollouts_per_input = n_rollouts_per_input
-
-    async def __call__(
-        self,
-        agent_run: AgentRun,
-        max_concurrency: int = 10,
-    ) -> JudgeResult | None:
-        rubric = self.cfg
-
-        async def _validation_callback(batch_index: int, llm_output: LLMOutput):
-            parse_and_validate_llm_output(llm_output, rubric.output_schema, agent_run)
-
-        # Run several independent rollouts
-        prompt = [{"role": "user", "content": self.cfg.materialize_system_prompt(agent_run)}]
+        super().__init__(cfg, llm_svc, docent_collection_id)
+
+    async def one_rollout_second_stage(
+        self, agent_run: AgentRun, first_stage_results: list[dict[str, Any]]
+    ) -> tuple[dict[str, Any] | None, dict[str, Any] | None]:
+        """Reflect on the results of the first stage of rollouts.
+        TODO(mengk): this is only done in a single-turn way. We should generalize this to multi-turn.
+        """
+
+        # Construct *single* reflection prompt
+        first_stage_results_text = "\n\n".join(
+            [
+                f"Rollout {j+1}:\n{yaml.dump(r, width=float('inf'))}"
+                for j, r in enumerate(first_stage_results)
+            ]
+        )
+        reflection_instruction = (
+            f"We have sampled a judge {len(first_stage_results)} times to get {len(first_stage_results)} independent answers to the same rubric evaluation:\n"
+            f"{first_stage_results_text}\n\n"
+            f"Please reflect on these answers. Consider all the information and evidence presented. "
+            f"Return a final answer in the same JSON format as before."
+        )
+        reflection_prompt = [
+            # Original system prompt
+            {"role": "system", "content": self.cfg.materialize_system_prompt(agent_run)},
+            # Additional reflection instruction as a user message (kind of awkward)
+            {"role": "user", "content": reflection_instruction},
+        ]
+
+        # Ask the judge to reflect on the others' results
         outputs = await self.llm_svc.get_completions(
-            inputs=[prompt for _ in range(self.n_rollouts_per_input)],
-            model_options=[rubric.judge_model],
+            inputs=[reflection_prompt],
+            model_options=[self.cfg.judge_model],
             max_new_tokens=16384,
             timeout=180.0,
             use_cache=False,
-            validation_callback=_validation_callback,
-            max_concurrency=max_concurrency,
+            validation_callback=self._get_validation_callback(agent_run),
         )
+        output_str = outputs[0].first_text
 
-        # Process each rollout
-        indep_results: list[dict[str, Any]] = []
-        for output in outputs:
-            if output.first_text is None:
-                continue
-            if v_output := parse_and_validate_llm_output(output, rubric.output_schema, agent_run):
-                indep_results.append(v_output)
+        validated_output = self._validate_first_response_tag_or_entire_output(
+            output_str or "", agent_run
+        )
+        if validated_output is not None:
+            return validated_output, None
+        else:
+            return None, None
 
+    async def __call__(self, agent_run: AgentRun) -> JudgeResult | None:
+        rubric = self.cfg
+
+        indep_results: list[dict[str, Any]] = []
+        indep_rollout_metadata: list[dict[str, Any] | None] = []
+
+        async def _execute():
+            result, metadata = await self.one_rollout(agent_run)
+            if result is not None:
+                indep_results.append(result)
+                indep_rollout_metadata.append(metadata)
+
+        # Stage 1: run rollouts concurrently
+        async with anyio.create_task_group() as tg:
+            for _ in range(self.cfg.n_rollouts_per_input):
+                tg.start_soon(_execute)
         if not indep_results:
             return None
 
@@ -141,61 +423,28 @@ class MultiReflectionJudge(BaseJudge):
             indep_results, agreement_keys
         )
 
-        def _get_reflection_prompt(cur_index: int):
-            # Current result
-            result = indep_results[cur_index]
-            # Get other results (excluding the current one)
-            other_results = [r for j, r in enumerate(indep_results) if j != cur_index]
-
-            # Create the reflection message
-            other_results_text = "\n\n".join(
-                [f"Answer {j+1}:\n{json.dumps(r, indent=2)}" for j, r in enumerate(other_results)]
-            )
-
-            reflection_instruction = (
-                f"Here are {len(other_results)} other independent answers to the same rubric evaluation:\n\n"
-                f"{other_results_text}\n\n"
-                f"Please reflect on these other answers and your own answer. "
-                f"Consider if any of them have identified important aspects you missed, or if there are disagreements that should be resolved. "
-                f"Then provide your final answer in the same JSON format as before."
-            )
-
-            # Construct the multi-message prompt
-            # 1. Original user message
-            # 2. Assistant message with the rollout's result
-            # 3. New user message asking for reflection
-            return [
-                *prompt,  # Original user message(s)
-                {"role": "assistant", "content": json.dumps(result, indent=2)},
-                {"role": "user", "content": reflection_instruction},
-            ]
-
-        final_results = indep_results.copy()  # Shallow copy
+        # Stage 2: reflect on the results
+        # Shallow copies are fine
+        final_results = indep_results.copy()
+        final_rollout_metadata = indep_rollout_metadata.copy()
         if len(indep_results) > 1:
-            # Ask the judge to reflect on the others' results
-            reflection_outputs = await self.llm_svc.get_completions(
-                inputs=[_get_reflection_prompt(i) for i in range(len(indep_results))],
-                model_options=[rubric.judge_model],
-                max_new_tokens=16384,
-                timeout=180.0,
-                use_cache=False,
-                validation_callback=_validation_callback,
-                max_concurrency=max_concurrency,
-            )
+            candidate_final_results: list[dict[str, Any]] = []
+            candidate_final_rollout_metadata: list[dict[str, Any] | None] = []
+
+            async def _execute_second_stage():
+                result, metadata = await self.one_rollout_second_stage(agent_run, indep_results)
+                if result is not None:
+                    candidate_final_results.append(result)
+                    candidate_final_rollout_metadata.append(metadata)
 
-            # Process reflection outputs in the same way as the initial rollouts
-            reflected_results: list[dict[str, Any]] = []
-            for output in reflection_outputs:
-                if output.first_text is None:
-                    continue
-                if v_output := parse_and_validate_llm_output(
-                    output, rubric.output_schema, agent_run
-                ):
-                    reflected_results.append(v_output)
+            async with anyio.create_task_group() as tg:
+                for _ in range(self.cfg.n_rollouts_per_input):
+                    tg.start_soon(_execute_second_stage)
 
             # Use reflected results if we got any, otherwise fall back to original results
-            if reflected_results:
-                final_results = reflected_results
+            if candidate_final_results:
+                final_results = candidate_final_results
+                final_rollout_metadata = candidate_final_rollout_metadata
             else:
                 logger.warning("No reflected results found, falling back to original results")
 
@@ -213,10 +462,126 @@ class MultiReflectionJudge(BaseJudge):
                 "final_results": final_results,
                 "final_agt_key_modes_and_counts": final_agt_key_modes_and_counts,
                 "final_max_idx": final_max_idx,
+                "final_rollout_metadata": final_rollout_metadata,
                 # Also include initial measurements
                 "indep_results": indep_results,
                 "indep_max_idx": indep_max_idx,
                 "indep_agt_key_modes_and_counts": indep_agt_key_modes_and_counts,
+                "indep_rollout_metadata": indep_rollout_metadata,
             },
             result_type=ResultType.DIRECT_RESULT,
         )
+
+    async def estimate_output_distrs(
+        self,
+        agent_run: AgentRun,
+        *,
+        n_initial_rollouts_to_sample: int,
+        n_combinations_to_sample: int,
+        n_reflection_rollouts_to_sample: int,
+        **kwargs: Any,
+    ) -> None | tuple[dict[str, JudgeOutputDistribution], dict[str, Any]]:
+        if self.cfg.n_rollouts_per_input > n_initial_rollouts_to_sample:
+            raise ValueError(
+                "n_initial_rollouts_to_sample must be greater than or equal to cfg.n_rollouts_per_input"
+            )
+        if self.cfg.n_rollouts_per_input > n_reflection_rollouts_to_sample:
+            raise ValueError(
+                "n_reflection_rollouts_to_sample must be greater than or equal to cfg.n_rollouts_per_input"
+            )
+
+        first_step_rollouts: list[dict[str, Any]] = []
+        first_step_rollout_metadata: list[dict[str, Any] | None] = []
+        first_step_combinations: list[list[dict[str, Any]]] = []
+        second_step_rollouts: list[list[dict[str, Any]]] = []
+        second_step_rollout_metadata: list[list[dict[str, Any] | None]] = []
+
+        ##########
+        # Step 1 #
+        ##########
+
+        pbar_first = tqdm(
+            total=n_initial_rollouts_to_sample, desc="Stage 1: Initial rollouts", leave=False
+        )
+
+        async def _execute_first_stage():
+            result, metadata = await self.one_rollout(agent_run)
+            if result is not None:
+                first_step_rollouts.append(result)
+                first_step_rollout_metadata.append(metadata)
+            pbar_first.update(1)
+
+        # Collect rollouts of the first stage
+        async with anyio.create_task_group() as tg_first:
+            for _ in range(n_initial_rollouts_to_sample):
+                tg_first.start_soon(_execute_first_stage)
+
+        pbar_first.close()
+
+        if len(first_step_rollouts) < self.cfg.n_rollouts_per_input:
+            raise ValueError("Not enough first step rollouts to sample combinations")
+
+        # Sample random k-sized combinations of the first step rollouts
+        for _ in range(n_combinations_to_sample):
+            combination = random.sample(first_step_rollouts, self.cfg.n_rollouts_per_input)
+            first_step_combinations.append(combination)
+            second_step_rollouts.append([])
+            second_step_rollout_metadata.append([])
+
+        ##########
+        # Step 2 #
+        ##########
+
+        pbar_second = tqdm(
+            total=n_combinations_to_sample, desc="Stage 2: Combinations", leave=False
+        )
+
+        async with anyio.create_task_group() as tg_second:
+
+            async def _execute_second_stage(i: int, combination: list[dict[str, Any]]):
+                pbar_third = tqdm(
+                    total=n_reflection_rollouts_to_sample,
+                    desc=f"Stage 2: Combination {i+1}/{n_combinations_to_sample}",
+                    leave=False,
+                )
+
+                async def _execute_second_stage_inner():
+                    result, metadata = await self.one_rollout_second_stage(agent_run, combination)
+                    if result is not None:
+                        second_step_rollouts[i].append(result)
+                        second_step_rollout_metadata[i].append(metadata)
+                    pbar_third.update(1)
+
+                async with anyio.create_task_group() as tg:
+                    for _ in range(n_reflection_rollouts_to_sample):
+                        tg.start_soon(_execute_second_stage_inner)
+
+                pbar_third.close()
+                pbar_second.update(1)
+
+            for i, combination in enumerate(first_step_combinations):
+                tg_second.start_soon(_execute_second_stage, i, combination)
+
+        pbar_second.close()
+
+        output_distributions = compute_output_distributions(
+            [sublist for el in second_step_rollouts for sublist in el],
+            self.cfg.output_schema,
+            get_agreement_keys(self.cfg.output_schema),
+        )
+
+        return output_distributions, {
+            "first_step_rollouts": first_step_rollouts,
+            "first_step_rollout_metadata": first_step_rollout_metadata,
+            "first_step_combinations": first_step_combinations,
+            "second_step_rollouts": second_step_rollouts,
+            "second_step_rollout_metadata": second_step_rollout_metadata,
+        }
+
+
+def build_judge(rubric: Rubric, llm_svc: BaseLLMService, docent_collection_id: str | None = None):
+    if rubric.judge_variant == JudgeVariant.MAJORITY:
+        return MajorityVotingJudge(rubric, llm_svc, docent_collection_id)
+    elif rubric.judge_variant == JudgeVariant.MULTI_REFLECT:
+        return MultiReflectionJudge(rubric, llm_svc, docent_collection_id)
+    raise ValueError(f"Invalid variant: {rubric.judge_variant}")