AutoDiscovery API Reference

The parts of the autodiscovery package that implement the search described in How It Works. Each entry shows its source.

Search loop

autodiscovery.run

run_mcts(root, nodes_by_level, dataset_paths, log_dirname, work_dir, *, model_name, belief_model_name, vision_model, embedding_model, embedding_dimensions=None, max_iterations=100, branching_factor=8, max_rounds=100000, selection_method=None, allow_generate_experiments=False, n_belief_samples=5, k_parents=10, temperature=1.0, belief_temperature=1.0, reasoning_effort='medium', belief_reasoning_effort='low', implicit_bayes_posterior=False, surprisal_width=0.2, user_query=None, belief_mode='categorical', use_binary_reward=True, run_dedupe=True, experiment_first=False, code_timeout=30 * 60, n_warmstart=0, use_online_beliefs=False, evidence_weight=1.0, kl_scale=20.0, reward_mode='belief_and_kl', warmstart_experiments=None, backend='process', bucket_path=None, datalib_dirs=None, batch_size=1, n_threads=1, agent_usage_mode='per_response')

Run AutoDS exploration. In MCTS, root node level=0 is a dummy node with no experiment, level=1 is the first real node with the dataset loading experiment, levels > 1 are the actual MCTS nodes with hypotheses and experiments.

Parameters:

Name	Type	Description	Default
root		Root MCTSNode to continue from.	required
nodes_by_level		Dictionary to store nodes by level.	required
dataset_paths		List of paths to dataset files.	required
log_dirname		Directory to save logs and MCTS nodes.	required
work_dir		Working directory for agents.	required
model_name		LLM model for agents, as litellm's <provider>/<model>.	required
embedding_model		Deduplication embedding model, as litellm's <provider>/<model>.	required
embedding_dimensions		Optional deduplication embedding dimensions.	None
belief_model_name		LLM model name for belief distribution agent.	required
max_iterations		Maximum number of MCTS iterations.	100
branching_factor		Maximum number of children per node.	8
max_rounds		Maximum number of rounds for the group chat.	100000
selection_method		Function to select nodes in MCTS (default is UCB1).	None
allow_generate_experiments		Whether to allow nodes to generate new experiments on demand.	False
n_belief_samples		Number of samples for belief distribution evaluation.	5
k_parents		Number of parent levels to include in logs (None for all).	10
temperature		Temperature setting for all agents (except belief agent).	1.0
belief_temperature		Temperature setting for the belief agent.	1.0
reasoning_effort		Reasoning effort for non-belief agent model calls.	'medium'
belief_reasoning_effort		Reasoning effort for belief-model calls.	'low'
implicit_bayes_posterior		Whether to use the belief samples with evidence as the direct posterior or to use a Bayesian update that explicitly combines it with the prior.	False
surprisal_width		Minimum difference in mean prior and posterior probabilities required to count as a surprisal.	0.2
user_query		Custom user query to condition experiment generation during exploration.	None
belief_mode		Belief elicitation mode (boolean, categorical, categorical_numeric, or probability).	'categorical'
use_binary_reward		Whether to use binary reward for MCTS instead of a continuous reward (belief change).	True
run_dedupe		Whether to deduplicate nodes before saving to JSON and CSV.	True
experiment_first		If True, an experiment will be generated before its hypothesis.	False
code_timeout		Timeout for code execution in seconds (default is 30 minutes).	30 * 60
n_warmstart		Number of warmstart experiments to run after data loading but before MCTS selection.	0
use_online_beliefs		Whether to use online beliefs (i.e., beliefs updated with evidence from previous nodes).	False
evidence_weight		Weight for the experimental evidence for posterior calculation.	1.0
kl_scale		Normalization factor for KL divergence in reward calculation.	20.0
reward_mode		Mode for reward calculation (belief, kl, or belief_and_kl).	'belief_and_kl'
warmstart_experiments		Path to JSON file with warmstart experiments to run after data loading but before MCTS selection.	None
backend		Code execution backend (local, process, or modal).	'process'
bucket_path		GCS bucket path for Modal sandbox (e.g., gs://example-bucket/discoverybench/).	None
datalib_dirs		Read-only data-library directories available to the run's code.	None
vision_model		Model used for image analysis in code execution.	required
batch_size		Number of nodes to select and expand per iteration.	1
n_threads		Number of threads to use for parallel node expansion.	1
agent_usage_mode	str	Tracking mode for agents chat usage. per_response records usage from each AG2 model response. summary_delta records usage from AG2 usage-summary deltas.	'per_response'

Source code in packages/autodiscovery/src/autodiscovery/run.py

def run_mcts(
    root,
    nodes_by_level,
    dataset_paths,
    log_dirname,
    work_dir,
    *,
    model_name,
    belief_model_name,
    vision_model,
    embedding_model,
    embedding_dimensions=None,
    max_iterations=100,
    branching_factor=8,
    max_rounds=100000,
    selection_method=None,
    allow_generate_experiments=False,
    n_belief_samples=5,
    k_parents=10,
    temperature=1.0,
    belief_temperature=1.0,
    reasoning_effort="medium",
    belief_reasoning_effort="low",
    implicit_bayes_posterior=False,
    surprisal_width=0.2,
    user_query=None,
    belief_mode="categorical",
    use_binary_reward=True,
    run_dedupe=True,
    experiment_first=False,
    code_timeout=30 * 60,
    n_warmstart=0,
    use_online_beliefs=False,
    evidence_weight=1.0,
    kl_scale=20.0,
    reward_mode="belief_and_kl",
    warmstart_experiments=None,
    backend="process",
    bucket_path=None,
    datalib_dirs=None,
    batch_size=1,
    n_threads=1,
    agent_usage_mode: str = "per_response",
):
    """Run AutoDS exploration. In MCTS, root node level=0 is a dummy node with no experiment, level=1 is the first real node with the dataset loading experiment, levels > 1 are the actual MCTS nodes with hypotheses and experiments.

    Args:
        root: Root MCTSNode to continue from.
        nodes_by_level: Dictionary to store nodes by level.
        dataset_paths: List of paths to dataset files.
        log_dirname: Directory to save logs and MCTS nodes.
        work_dir: Working directory for agents.
        model_name: LLM model for agents, as litellm's ``<provider>/<model>``.
        embedding_model: Deduplication embedding model, as litellm's
            ``<provider>/<model>``.
        embedding_dimensions: Optional deduplication embedding dimensions.
        belief_model_name: LLM model name for belief distribution agent.
        max_iterations: Maximum number of MCTS iterations.
        branching_factor: Maximum number of children per node.
        max_rounds: Maximum number of rounds for the group chat.
        selection_method: Function to select nodes in MCTS (default is UCB1).
        allow_generate_experiments: Whether to allow nodes to generate new experiments on demand.
        n_belief_samples: Number of samples for belief distribution evaluation.
        k_parents: Number of parent levels to include in logs (None for all).
        temperature: Temperature setting for all agents (except belief agent).
        belief_temperature: Temperature setting for the belief agent.
        reasoning_effort: Reasoning effort for non-belief agent model calls.
        belief_reasoning_effort: Reasoning effort for belief-model calls.
        implicit_bayes_posterior: Whether to use the belief samples with evidence as the direct posterior or to use a Bayesian update that explicitly combines it with the prior.
        surprisal_width: Minimum difference in mean prior and posterior probabilities required to count as a surprisal.
        user_query: Custom user query to condition experiment generation during exploration.
        belief_mode: Belief elicitation mode (boolean, categorical, categorical_numeric, or probability).
        use_binary_reward: Whether to use binary reward for MCTS instead of a continuous reward (belief change).
        run_dedupe: Whether to deduplicate nodes before saving to JSON and CSV.
        experiment_first: If True, an experiment will be generated before its hypothesis.
        code_timeout: Timeout for code execution in seconds (default is 30 minutes).
        n_warmstart: Number of warmstart experiments to run after data loading but before MCTS selection.
        use_online_beliefs: Whether to use online beliefs (i.e., beliefs updated with evidence from previous nodes).
        evidence_weight: Weight for the experimental evidence for posterior calculation.
        kl_scale: Normalization factor for KL divergence in reward calculation.
        reward_mode: Mode for reward calculation (belief, kl, or belief_and_kl).
        warmstart_experiments: Path to JSON file with warmstart experiments to run after data loading but before MCTS selection.
        backend: Code execution backend (local, process, or modal).
        bucket_path: GCS bucket path for Modal sandbox (e.g., gs://example-bucket/discoverybench/).
        datalib_dirs: Read-only data-library directories available to the run's code.
        vision_model: Model used for image analysis in code execution.
        batch_size: Number of nodes to select and expand per iteration.
        n_threads: Number of threads to use for parallel node expansion.
        agent_usage_mode: Tracking mode for agents chat usage. ``per_response`` records
            usage from each AG2 model response. ``summary_delta`` records usage from
            AG2 usage-summary deltas.
    """

    def _get_executor_rich_outputs(code_executor_agent) -> list:
        """Return the most recent rich outputs from the code executor, if available."""
        executor = getattr(code_executor_agent, "code_executor", None)
        if executor is None or not hasattr(executor, "get_last_rich_outputs"):
            return []
        return executor.get_last_rich_outputs() or []

    def _write_rich_outputs(level: int, node_idx: int, rich_outputs: list) -> None:
        """Persist rich outputs for a node to the rich_outputs directory."""
        rich_outputs_dir = os.path.join(log_dirname, "rich_outputs")
        os.makedirs(rich_outputs_dir, exist_ok=True)
        output_path = os.path.join(rich_outputs_dir, f"ro_{level}_{node_idx}.json")
        with open(output_path, "w") as f:
            json.dump(rich_outputs, f, indent=2)

    def _set_executor_usage_context(code_executor_agent, node_id: str | None) -> None:
        """Set usage context on executors that support direct image-analysis tracking."""
        executor = getattr(code_executor_agent, "code_executor", None)
        if executor is not None and hasattr(executor, "set_usage_context"):
            executor.set_usage_context(usage_tracker, node_id=node_id)

    # Setup logger
    logger = TreeLogger(log_dirname)

    # Track time
    start_time = time()

    # Create work directory if it doesn't exist
    os.makedirs(work_dir, exist_ok=True)

    usage_tracker = UsageTracker()
    usage_tracker.save_events(log_dirname)

    try:
        if agent_usage_mode == "per_response":
            # LiteLLMAG2Client records each response as it is produced.
            configure_ag2_usage_tracking(usage_tracker)
        elif agent_usage_mode == "summary_delta":
            configure_ag2_usage_tracking(None)
        else:
            raise ValueError(
                f"Unknown agent_usage_mode '{agent_usage_mode}'. "
                "Expected one of ['per_response', 'summary_delta']."
            )

        # For Modal sandbox, files are in GCS and will be mounted directly.
        # For local/process backends, ensure every dataset is reachable from
        # work_dir so that the sandbox (which chdir's there) can find files.
        # We symlink anything that isn't already under work_dir.
        if backend != "modal":
            for dataset_fpath in dataset_paths:
                # Resolve to the real file (handles the webstack's data/ layout).
                abs_src = resolve_local_dataset_source(dataset_fpath)
                # Skip if already inside work_dir (e.g. easy.py already symlinked)
                if abs_src.startswith(os.path.abspath(work_dir) + os.sep):
                    continue
                dst = os.path.join(work_dir, os.path.basename(dataset_fpath))
                if not os.path.lexists(dst):
                    os.symlink(abs_src, dst)

        base_agent_objs = None
        if n_threads <= 1:
            base_agent_objs = get_agents(
                work_dir,
                model_name=model_name,
                temperature=temperature,
                reasoning_effort=reasoning_effort,
                branching_factor=branching_factor,
                user_query=user_query,
                experiment_first=experiment_first,
                code_timeout=code_timeout,
                backend=backend,
                bucket_path=bucket_path,
                dataset_paths=dataset_paths,
                datalib_dirs=datalib_dirs,
                vision_model=vision_model,
                usage_tracker=usage_tracker,
            )

        if selection_method is None:
            # Default selection method is UCB1
            selection_method = default_mcts_selection(exploration_weight=1.0)

        # Store the list of (level, node_idx) tuples for surprising nodes; if resuming, load them from the previous run
        all_surprisals = []
        for level in nodes_by_level:
            for node in nodes_by_level[level]:
                if node.surprising:
                    all_surprisals.append((node.level, node.node_idx))

        # Load warmstart experiments if provided
        _warmstart_experiments = None
        if warmstart_experiments is not None:
            with open(warmstart_experiments) as f:
                _warmstart_experiments = json.load(f)

        # TEMPORARY LOGGING
        TEMP_LOG = []

        total_to_sample = max_iterations
        n_sampled = 0
        iteration_idx = 0
        node_mutation_locks: dict[str, threading.Lock] = {}
        node_mutation_locks_guard = threading.Lock()

        def _get_node_mutation_lock(node_obj: MCTSNode) -> threading.Lock:
            """Return a per-node lock used to serialize mutable node state updates."""
            node_key = node_obj.id or f"node_obj_{id(node_obj)}"
            with node_mutation_locks_guard:
                lock = node_mutation_locks.get(node_key)
                if lock is None:
                    lock = threading.Lock()
                    node_mutation_locks[node_key] = lock
            return lock

        while n_sampled < total_to_sample:
            iteration_idx += 1
            # MCTS SELECTION, EXPANSION, and EXECUTION
            print(f"\n\n######### ITERATION {iteration_idx} #########\n")

            remaining_budget = total_to_sample - n_sampled
            next_nodes = select_nodes(
                selection_method,
                root,
                nodes_by_level,
                n_warmstart,
                return_n=min(batch_size, remaining_budget),
            )[:remaining_budget]
            if not next_nodes:
                break
            print(
                "SAMPLED "
                f"{len(next_nodes)} NODE(S) FOR EXPANSION: "
                f"{[f'{node.level}_{node.node_idx}' for node in next_nodes]}\n"
            )

            def _expand_node(
                inbatch_idx,
                node,
                agent_objs,
                logger_obj,
                get_node_idx,
                update_mcts_lock=None,
                is_threaded=False,
            ):
                print(
                    f"({inbatch_idx + 1}/{len(next_nodes)}): "  # noqa: B023
                    f"EXPANDING NODE {node.level}_{node.node_idx}\n"
                    "==========================\n"
                )

                experiment_generator = agent_objs["experiment_generator"]
                user_proxy_local = agent_objs["user_proxy"]
                parent_node = node
                parent_lock = _get_node_mutation_lock(parent_node)
                created_node = None
                node_committed = False

                try:
                    # Fetch/generate experiments and attach a new child atomically per parent node.
                    with parent_lock:
                        new_experiment, new_query = parent_node.get_next_experiment(
                            experiment_generator=experiment_generator
                        )
                        if new_query is None:
                            if not parent_node.untried_experiments:
                                parent_node.allow_generate_experiments = False
                            print(
                                f"No new experiment generated for node {parent_node.level}_{parent_node.node_idx}. "
                                "Skipping this iteration."
                            )
                            return None

                        new_level = parent_node.level + 1
                        new_node_idx = get_node_idx(new_level)
                        created_node = MCTSNode(
                            level=new_level,
                            node_idx=new_node_idx,
                            hypothesis=new_experiment["hypothesis"],
                            experiment_plan=new_experiment["experiment_plan"],
                            query=new_query,
                            parent=parent_node,
                            allow_generate_experiments=allow_generate_experiments and new_level > 0,
                            untried_experiments=_warmstart_experiments if new_level == 1 else None,
                        )
                    node = created_node

                    # Update logger state
                    logger_obj.level = node.level
                    logger_obj.node_idx = node.node_idx

                    # Load previous explorations (make sure the root is always included)
                    node_context = []
                    if node.level > 1:
                        if root.children:
                            root_context = root.children[0].get_context(include_code_output=True)
                            if root_context is not None:
                                node_context.append(root_context)
                        path_context = node.get_path_context(k=k_parents - 1, skip_root=True) or []
                        node_context.extend(
                            context for context in path_context if context is not None
                        )

                    node_messages = [
                        {
                            "name": "user_proxy",
                            "role": "user",
                            "content": "PREVIOUS EXPLORATION:\n\n" + n,
                        }
                        for n in node_context
                    ]
                    node_messages += [
                        {
                            "name": "user_proxy",
                            "role": "user",
                            "content": node.query
                            + (
                                "\n\nNote for the programmer: Dataset files are present one level above the current "
                                "working directory."
                                if is_threaded
                                else ""
                            ),
                        }
                    ]

                    # Set up the group chat
                    groupchat, chat_manager = setup_group_chat(agent_objs, max_rounds)
                    _, last_message = chat_manager.resume(messages=node_messages)
                    agent_usage_before = None
                    if agent_usage_mode == "summary_delta":
                        agent_usage_before = snapshot_agents_actual_usage(agent_objs)
                    _set_executor_usage_context(agent_objs["code_executor"], node.id)
                    if agent_usage_mode == "per_response":
                        set_ag2_usage_context(node_id=node.id, component="agents.chat")

                    # Track time per node
                    _node_start_time = time()

                    # Execute current experiment and generate new experiments
                    try:
                        user_proxy_local.initiate_chat(
                            recipient=chat_manager,
                            message=last_message,
                            clear_history=False,
                        )
                    finally:
                        if agent_usage_mode == "per_response":
                            clear_ag2_usage_context()

                    if agent_usage_mode == "summary_delta":
                        assert agent_usage_before is not None
                        agent_usage_after = snapshot_agents_actual_usage(agent_objs)
                        usage_tracker.record_agent_usage_deltas(
                            agent_usage_before,
                            agent_usage_after,
                            node_id=node.id,
                            component="agents.chat",
                        )

                    # Store the raw message logs for the node
                    logger_obj.log_node(
                        node.level,
                        node.node_idx,
                        chat_manager.messages_to_string(groupchat.messages),
                    )

                    # Get messages starting from the current query and update the node
                    node.messages = get_msgs_from_latest_query(groupchat.messages)
                    node.read_experiment_from_messages(
                        store_new_experiments=not (
                            node.level == 1 and _warmstart_experiments is not None
                        )
                    )
                    rich_outputs = _get_executor_rich_outputs(agent_objs["code_executor"])
                    _write_rich_outputs(node.level, node.node_idx, rich_outputs)

                    # Calculate beliefs and rewards
                    if node.success and node.level > 1:
                        compute_and_store_reward(
                            node,
                            belief_model_name,
                            belief_temperature,
                            belief_reasoning_effort,
                            n_belief_samples,
                            implicit_bayes_posterior,
                            surprisal_width,
                            belief_mode,
                            use_binary_reward,
                            all_surprisals,
                            nodes_by_level=nodes_by_level,
                            use_online_beliefs=use_online_beliefs,
                            evidence_weight=evidence_weight,
                            kl_scale=kl_scale,
                            reward_mode=reward_mode,
                            TEMP_LOG=TEMP_LOG,
                            usage_tracker=usage_tracker,
                        )

                        if node.success:  # i.e., reward was computed successfully
                            # Print debug information
                            print_node_info(node)

                            # TEMPORARY LOGGING
                            if TEMP_LOG:
                                temp_log_file = os.path.join(log_dirname, "temp_log.json")
                                with open(temp_log_file, "w") as f:
                                    json.dump(TEMP_LOG, f, indent=2)
                                print(f"Temporary log saved to {temp_log_file}")

                    # End time tracking for the node
                    _node_end_time = time()
                    node.time_elapsed = round(_node_end_time - _node_start_time, 2)

                    def _backprop_and_save():
                        # MCTS BACKPROPAGATION
                        node.update_counts(visits=1, reward=node.self_value)
                        # Save the new node and its parents' updated counts to JSON files
                        save_mcts_node(node, log_dirname, to_root=True)
                        usage_tracker.save_events(log_dirname)
                        _update_report(log_dirname, "RUNNING")

                    if update_mcts_lock is not None:
                        with update_mcts_lock:
                            _backprop_and_save()
                    else:
                        _backprop_and_save()
                    node_committed = True
                    return node
                except Exception:
                    if agent_usage_mode == "per_response":
                        clear_ag2_usage_context()
                    # Roll back partially-created nodes so selection logic does not keep revisiting them.
                    if created_node is not None and not node_committed:
                        with parent_lock:
                            if created_node in parent_node.children:
                                parent_node.children.remove(created_node)
                    raise

            if n_threads > 1 and len(next_nodes) > 1:
                # Parallel expansion of nodes
                index_lock = threading.Lock()
                update_mcts_lock = threading.Lock()
                next_node_idx_by_level = defaultdict(
                    int, {level: len(nodes) for level, nodes in nodes_by_level.items()}
                )
                thread_local = threading.local()

                def _get_node_idx(new_level):
                    with index_lock:  # noqa: B023
                        new_node_idx = next_node_idx_by_level[new_level]  # noqa: B023
                        next_node_idx_by_level[new_level] += 1  # noqa: B023
                    return new_node_idx

                def _get_thread_agents():
                    if not hasattr(thread_local, "agent_objs"):  # noqa: B023
                        thread_id = threading.get_ident()
                        thread_work_dir = os.path.join(work_dir, f"thread_{thread_id}")
                        os.makedirs(thread_work_dir, exist_ok=True)
                        thread_local.agent_objs = get_agents(  # noqa: B023
                            thread_work_dir,
                            model_name=model_name,
                            temperature=temperature,
                            reasoning_effort=reasoning_effort,
                            branching_factor=branching_factor,
                            user_query=user_query,
                            experiment_first=experiment_first,
                            code_timeout=code_timeout,
                            backend=backend,
                            bucket_path=bucket_path,
                            dataset_paths=dataset_paths,
                            datalib_dirs=datalib_dirs,
                            vision_model=vision_model,
                            usage_tracker=usage_tracker,
                        )
                    return thread_local.agent_objs  # noqa: B023

                def _expand_node_parallel(inbatch_idx, node):
                    return _expand_node(
                        inbatch_idx,
                        node,
                        _get_thread_agents(),
                        TreeLogger(log_dirname),
                        _get_node_idx,
                        update_mcts_lock=update_mcts_lock,  # noqa: B023
                        is_threaded=True,
                    )

                def _on_expand_error(node_id: str, exc: Exception) -> None:
                    # Keep exploration running when one node expansion fails.
                    print(
                        f"[run_mcts] Failed expanding node {node_id}: "
                        f"{exc.__class__.__name__}: {exc}"
                    )
                    traceback.print_exception(type(exc), exc, exc.__traceback__)

                expanded_nodes = []
                with ThreadPoolExecutor(max_workers=n_threads) as executor:
                    future_labels = {
                        executor.submit(_expand_node_parallel, inbatch_idx, node): (
                            f"{node.level}_{node.node_idx}"
                        )
                        for inbatch_idx, node in enumerate(next_nodes)
                    }

                    expanded_nodes = gather_completed_futures(
                        future_labels,
                        on_error=_on_expand_error,
                    )
            else:
                # Sequential expansion of nodes
                if base_agent_objs is None:
                    base_agent_objs = get_agents(
                        work_dir,
                        model_name=model_name,
                        temperature=temperature,
                        reasoning_effort=reasoning_effort,
                        branching_factor=branching_factor,
                        user_query=user_query,
                        experiment_first=experiment_first,
                        code_timeout=code_timeout,
                        backend=backend,
                        bucket_path=bucket_path,
                        dataset_paths=dataset_paths,
                        datalib_dirs=datalib_dirs,
                        vision_model=vision_model,
                        usage_tracker=usage_tracker,
                    )

                next_node_idx_by_level = defaultdict(
                    int, {level: len(nodes) for level, nodes in nodes_by_level.items()}
                )

                def _get_node_idx(new_level):
                    new_node_idx = next_node_idx_by_level[new_level]  # noqa: B023
                    next_node_idx_by_level[new_level] += 1  # noqa: B023
                    return new_node_idx

                expanded_nodes = []
                for inbatch_idx, node in enumerate(next_nodes):
                    node_id = f"{node.level}_{node.node_idx}"
                    try:
                        new_node = _expand_node(
                            inbatch_idx,
                            node,
                            base_agent_objs,
                            logger,
                            _get_node_idx,
                        )
                    except Exception as exc:
                        # Align sequential behavior with threaded mode by continuing after per-node failures.
                        print(
                            f"[run_mcts] Failed expanding node {node_id}: "
                            f"{exc.__class__.__name__}: {exc}"
                        )
                        traceback.print_exception(type(exc), exc, exc.__traceback__)
                        continue
                    if new_node is not None:
                        expanded_nodes.append(new_node)

            # Count only successful expansions toward experiment budget.
            n_sampled += len(expanded_nodes)

            # Add expanded nodes to the tree
            expanded_nodes.sort(key=lambda n: (n.level, n.node_idx))
            for node in expanded_nodes:
                nodes_by_level[node.level].append(node)
    except KeyboardInterrupt:
        print("\n\n######### EXPLORATION INTERRUPTED! SAVING THE CURRENT STATE... #########\n\n")
    finally:
        clear_ag2_usage_context()
        configure_ag2_usage_tracking(None)

    # End time tracking
    end_time = time()
    time_elapsed = end_time - start_time

    # Save all MCTS nodes
    save_nodes(
        nodes_by_level,
        log_dirname,
        run_dedupe=run_dedupe,
        model=belief_model_name,
        embedding_model=embedding_model,
        embedding_dimensions=embedding_dimensions,
        time_elapsed=time_elapsed,
        usage_tracker=usage_tracker,
    )
    usage_tracker.save_events(log_dirname)
    usage_tracker.save_summary(log_dirname)

compute_and_store_reward(node, belief_model_name, belief_temperature, belief_reasoning_effort, n_belief_samples, implicit_bayes_posterior, surprisal_width, belief_mode, use_binary_reward, all_surprisals=None, nodes_by_level=None, use_online_beliefs=False, evidence_weight=1.0, kl_scale=20.0, reward_mode='belief', TEMP_LOG=None, usage_tracker=None)

Compute node belief metrics, reward, and surprisal flags.

Parameters:

Name	Type	Description	Default
node		Node whose reward will be computed.	required
belief_model_name		Model name for belief elicitation.	required
belief_temperature		Temperature for belief sampling.	required
belief_reasoning_effort		Reasoning effort for belief-model calls.	required
n_belief_samples		Number of belief samples.	required
implicit_bayes_posterior		Whether posterior uses implicit prior knowledge.	required
surprisal_width		Width threshold for binary surprisal.	required
belief_mode		Belief elicitation mode.	required
use_binary_reward		Whether to compute binary reward.	required
all_surprisals		Running list of (level, node_idx) tuples of surprising nodes.	None
nodes_by_level		Tree nodes keyed by level, used to resolve all_surprisals. Required when use_online_beliefs is set and all_surprisals is non-empty.	None
use_online_beliefs		Whether to condition prior on prior surprising nodes.	False
evidence_weight		Weight for current evidence in posterior update.	1.0
kl_scale		Scale factor for KL reward normalization.	20.0
reward_mode		Reward mode combining belief and KL terms.	'belief'
TEMP_LOG		Optional debug log collector.	None
usage_tracker	UsageTracker | None	Optional usage tracker.	None

Source code in packages/autodiscovery/src/autodiscovery/run.py

def compute_and_store_reward(
    node,
    belief_model_name,
    belief_temperature,
    belief_reasoning_effort,
    n_belief_samples,
    implicit_bayes_posterior,
    surprisal_width,
    belief_mode,
    use_binary_reward,
    all_surprisals=None,
    nodes_by_level=None,
    use_online_beliefs=False,
    evidence_weight=1.0,
    kl_scale=20.0,
    reward_mode="belief",
    TEMP_LOG=None,
    usage_tracker: UsageTracker | None = None,
):
    """Compute node belief metrics, reward, and surprisal flags.

    Args:
        node: Node whose reward will be computed.
        belief_model_name: Model name for belief elicitation.
        belief_temperature: Temperature for belief sampling.
        belief_reasoning_effort: Reasoning effort for belief-model calls.
        n_belief_samples: Number of belief samples.
        implicit_bayes_posterior: Whether posterior uses implicit prior knowledge.
        surprisal_width: Width threshold for binary surprisal.
        belief_mode: Belief elicitation mode.
        use_binary_reward: Whether to compute binary reward.
        all_surprisals: Running list of (level, node_idx) tuples of surprising nodes.
        nodes_by_level: Tree nodes keyed by level, used to resolve ``all_surprisals``.
            Required when ``use_online_beliefs`` is set and ``all_surprisals`` is non-empty.
        use_online_beliefs: Whether to condition prior on prior surprising nodes.
        evidence_weight: Weight for current evidence in posterior update.
        kl_scale: Scale factor for KL reward normalization.
        reward_mode: Reward mode combining belief and KL terms.
        TEMP_LOG: Optional debug log collector.
        usage_tracker: Optional usage tracker.
    """
    s_conditioned_prior = None
    evidence_msg = []

    # If there are past surprisal, computed the s-conditioned prior
    if all_surprisals is not None and len(all_surprisals) > 0 and use_online_beliefs:
        # Build evidence message for prior belief elicitation
        if nodes_by_level is None:
            raise ValueError("nodes_by_level is required to condition on prior surprisals")
        prior_nodes = [nodes_by_level[level][idx] for level, idx in all_surprisals]
        evidence_msg = [
            {
                "role": "user",
                "content": "Previous study:\n\n"
                + get_context_string(
                    hyp_exp_query=f"Hypothesis: {prior.hypothesis}",
                    analysis=prior.analysis,
                    review=prior.review,
                    belief_mean=prior.posterior.mean,
                    include_code_output=False,
                ),
            }
            for prior in prior_nodes
        ]
        try:
            pt_prior, s_conditioned_prior, _, _ = calculate_prior_and_posterior_beliefs(
                node,
                model=belief_model_name,
                temperature=belief_temperature,
                reasoning_effort=belief_reasoning_effort,
                n_samples=n_belief_samples,
                implicit_bayes_posterior=implicit_bayes_posterior,
                surprisal_width=surprisal_width,
                belief_mode=belief_mode,
                evidence_msg=evidence_msg,
                usage_tracker=usage_tracker,
                usage_node_id=node.id,
                usage_context_label="surprisal_conditioned",
            )
        except ValueError as e:
            print(f"Error for node {node.id}: {e}")
            node.success = False
            return

        # TEMPORARY LOGGING
        if TEMP_LOG is not None:
            TEMP_LOG.append(
                {
                    "node_id": node.id,
                    "belief_change": None,
                    "kl_divergence": None,
                    "hypothesis": node.hypothesis,
                    "pt_prior": pt_prior.to_dict(),
                    "surprisal_evidence": [e["content"] for e in evidence_msg],
                    "s_conditioned_prior": s_conditioned_prior.to_dict(),
                }
            )

    # Build the evidence message for the current node
    evidence_msg.append(
        {
            "role": "user",
            "content": "Current experiment:\n\n"
            + get_context_string(
                hyp_exp_query=node.query,
                code_output=node.code_output,
                analysis=node.analysis,
                review=node.review,
                include_code_output=False,
            ),
        }
    )

    # Compute the prior and posterior beliefs for the current node
    try:
        prior, posterior, belief_change, kl_divergence = calculate_prior_and_posterior_beliefs(
            node,
            model=belief_model_name,
            temperature=belief_temperature,
            reasoning_effort=belief_reasoning_effort,
            n_samples=n_belief_samples,
            implicit_bayes_posterior=implicit_bayes_posterior,
            surprisal_width=surprisal_width,
            belief_mode=belief_mode,
            prior=s_conditioned_prior,
            evidence_msg=evidence_msg,
            evidence_weight=evidence_weight,
            usage_tracker=usage_tracker,
            usage_node_id=node.id,
            usage_context_label="main",
        )
    except ValueError as e:
        print(f"Error for node {node.id}: {e}")
        node.success = False
        return

    normalized_surprisal = None
    if belief_mode == "boolean_cat" and prior is not None and posterior is not None:
        prior_params = getattr(prior, "prior_params", (0.5, 0.5))
        theoretical_max = _theoretical_max_boolean_cat(
            n_belief_samples, evidence_weight, prior_params=prior_params
        )
        if theoretical_max != 0:
            prior_mean = prior.get_mean_belief(recompute=True)
            posterior_mean = posterior.get_mean_belief(prior=prior, recompute=True)
            normalized_surprisal = (posterior_mean - prior_mean) / theoretical_max

    # TEMPORARY LOGGING
    if TEMP_LOG is not None and len(TEMP_LOG) > 0:
        # Generate the posterior without surprisals
        _, _posterior, _belief_change, _kl_divergence = calculate_prior_and_posterior_beliefs(
            node,
            model=belief_model_name,
            temperature=belief_temperature,
            reasoning_effort=belief_reasoning_effort,
            n_samples=n_belief_samples,
            implicit_bayes_posterior=implicit_bayes_posterior,
            surprisal_width=surprisal_width,
            belief_mode=belief_mode,
            prior=pt_prior,
            evidence_msg=evidence_msg[-1:],
            evidence_weight=evidence_weight,
            usage_tracker=usage_tracker,
            usage_node_id=node.id,
            usage_context_label="offline_diagnostic",
        )

        TEMP_LOG[-1]["current_evidence"] = evidence_msg[-1]["content"]
        TEMP_LOG[-1]["online_posterior"] = posterior.to_dict()
        TEMP_LOG[-1]["belief_change"] = belief_change
        TEMP_LOG[-1]["kl_divergence"] = kl_divergence
        TEMP_LOG[-1]["offline_posterior"] = _posterior.to_dict()
        TEMP_LOG[-1]["offline_belief_change"] = _belief_change
        TEMP_LOG[-1]["offline_kl_divergence"] = _kl_divergence
        TEMP_LOG[-1]["current_surprisals"] = all_surprisals.copy()

        print("\n\n======================= SURPRISAL-CONDITION BELIEFS =======================\n")
        print(
            json.dumps(
                {
                    k: v
                    for k, v in TEMP_LOG[-1].items()
                    if k
                    in ["pt_prior", "s_conditioned_prior", "online_posterior", "offline_posterior"]
                },
                indent=2,
            )
        )

    node.prior = prior
    node.posterior = posterior
    node.belief_change = belief_change
    node.normalized_surprisal = normalized_surprisal
    node.kl_divergence = kl_divergence
    # Compute reward and surprisal
    node.self_value, node.surprising = get_self_value(
        belief_change=node.belief_change,
        kl_divergence=node.kl_divergence,
        binary=use_binary_reward,
        width=surprisal_width,
        kl_scale=kl_scale,
        mode=reward_mode,
    )
    if node.surprising:
        # Store the surprisal
        all_surprisals.append((node.level, node.node_idx))

_theoretical_max_boolean_cat(n_samples, evidence_weight, prior_params=(0.5, 0.5))

Compute the maximum possible mean shift for boolean_cat beliefs.

Parameters:

Name	Type	Description	Default
n_samples	int	Maximum number of prior-stage samples (N).	required
evidence_weight	float	Weight applied to stage-2 (evidence) samples.	required
prior_params	tuple[float, float]	Beta prior parameters (alpha, beta).	(0.5, 0.5)

Returns:

Type	Description
float	Theoretical maximum |mu2 - mu1| over all nodes given N and evidence weight.

Source code in packages/autodiscovery/src/autodiscovery/run.py

def _theoretical_max_boolean_cat(
    n_samples: int, evidence_weight: float, prior_params: tuple[float, float] = (0.5, 0.5)
) -> float:
    """Compute the maximum possible mean shift for boolean_cat beliefs.

    Args:
        n_samples: Maximum number of prior-stage samples (N).
        evidence_weight: Weight applied to stage-2 (evidence) samples.
        prior_params: Beta prior parameters (alpha, beta).

    Returns:
        Theoretical maximum |mu2 - mu1| over all nodes given N and evidence weight.
    """
    n = float(n_samples)
    w = float(evidence_weight)
    if n <= 0 or w <= 0:
        return 0.0

    alpha, beta = prior_params
    s = alpha + beta
    d = min(alpha, beta)
    t = w * n

    # Unconstrained (relaxed) optimum in u = S + n space.
    u_star = d + math.sqrt(d * (d + t))
    # Clamp to feasible interval u in [S, S + N] because n in [0, N].
    u_opt = min(max(u_star, s), s + n)
    return (t * (u_opt - d)) / (u_opt * (u_opt + t))

Tree and node selection

autodiscovery.mcts

MCTSNode

Source code in packages/autodiscovery/src/autodiscovery/mcts.py

class MCTSNode:
    _creation_counter = 0

    def __init__(
        self,
        level=None,
        node_idx=None,
        hypothesis=None,
        query=None,
        parent_idx=None,
        parent=None,
        untried_experiments=None,
        allow_generate_experiments=False,
        experiment_plan=None,
        code=None,
        code_output=None,
        analysis=None,
        review=None,
        creation_idx=None,
        created_at=None,
    ):
        # Tree attributes
        self.creation_idx = creation_idx if creation_idx is not None else MCTSNode._creation_counter
        MCTSNode._creation_counter += 1  # Used to replay MCTS from log files
        self.created_at = created_at if created_at is not None else datetime.now(UTC).isoformat()
        self.time_elapsed = None
        self.level = level
        self.node_idx = node_idx
        self.id = (
            f"node_{self.level}_{self.node_idx}"
            if self.level is not None and self.node_idx is not None
            else None
        )
        self.children = []
        self.parent = parent  # MCTSNode
        if self.parent is not None:
            self.parent_idx = self.parent.node_idx
            self.parent_id = self.parent.id
            self.parent.children.append(self)
        else:
            self.parent_idx = parent_idx
            self.parent_id = (
                f"node_{self.level - 1}_{self.parent_idx}" if self.parent_idx is not None else None
            )
        self.success = None

        # Agent attributes
        self.query = query
        self.messages = []
        self.untried_experiments = (
            copy.deepcopy(untried_experiments) if untried_experiments is not None else []
        )
        self.tried_experiments = []  # Track all tried experiments
        self.allow_generate_experiments = allow_generate_experiments

        # Experiment attributes
        self.hypothesis = hypothesis
        self.experiment_plan = experiment_plan
        self.code = code
        self.code_output = code_output
        self.analysis = analysis
        self.review = review

        # MCTS attributes
        self.visits = 0  # Visits to this node or its children
        self.value = 0.0  # Number of surprising hypotheses
        self.self_value = 0.0  # Value of this node only (not aggregated from children)

        # Belief attributes
        self.surprising: bool | None = None
        self.prior = None
        self.posterior = None
        self.belief_change: float | None = None  # Change in belief from prior to posterior
        self.normalized_surprisal: float | None = None
        self.kl_divergence: float | None = None

    def init_from_dict(self, data):
        """Initialize node attributes from a dictionary."""
        # Tree attributes
        self.creation_idx = data.get("creation_idx", MCTSNode._creation_counter)
        self.created_at = data.get("created_at", self.created_at)
        self.time_elapsed = data.get("time_elapsed", self.time_elapsed)
        self.id = data.get("id", None)
        if self.id is not None:
            self.level, self.node_idx = get_node_level_idx(self.id)
        else:
            self.level = data["level"]
            self.node_idx = data["node_idx"]
            self.id = f"node_{self.level}_{self.node_idx}"
        self.parent_id = data.get("parent_id", self.parent_id)
        if self.parent_id is not None:
            _, self.parent_idx = get_node_level_idx(self.parent_id)
        else:
            self.parent_idx = data.get("parent_idx", self.parent_idx)
            if self.parent_idx is not None:
                self.parent_id = f"node_{self.level - 1}_{self.parent_idx}"
        self.success = data.get("success", self.success)

        # Agent attributes
        self.query = data.get("query", "N/A")
        self.messages = data.get("messages", self.messages)
        self.untried_experiments = data.get("untried_experiments", self.untried_experiments)
        # self.tried_experiments = data.get('tried_experiments', self.tried_experiments)
        self.allow_generate_experiments = self.allow_generate_experiments and self.level > 0

        # Experiment attributes
        self.hypothesis = data.get("hypothesis", self.hypothesis)
        self.experiment_plan = data.get("experiment_plan", self.experiment_plan)
        self.code = data.get("code", self.code)
        self.code_output = data.get("code_output", self.code_output)
        self.analysis = data.get("analysis", self.analysis)
        self.review = data.get("review", self.review)

        # MCTS attributes
        self.visits = data.get("visits", self.visits)
        self.value = data.get("value", self.value)
        self.self_value = data.get("self_value", self.self_value)

        # Belief attributes
        self.surprising = data.get("surprising", self.surprising)
        from autodiscovery.beliefs import (
            BELIEF_MODE_TO_CLS,
        )  # Import here to avoid circular import issues

        if "prior" in data and data["prior"]:
            belief_cls = BELIEF_MODE_TO_CLS[data["prior"]["_type"]]
            self.prior = belief_cls.DistributionFormat(**data["prior"])
        if "posterior" in data and data["posterior"]:
            belief_cls = BELIEF_MODE_TO_CLS[data["posterior"]["_type"]]
            self.posterior = belief_cls.DistributionFormat(**data["posterior"])
        self.belief_change = data.get("belief_change", self.belief_change)
        self.normalized_surprisal = data.get("normalized_surprisal", self.normalized_surprisal)
        self.kl_divergence = data.get("kl_divergence", self.kl_divergence)

    def get_next_experiment(self, experiment_generator=None, n_retry=3):
        """Returns the next untried experiment. If none left and generating experiments is allowed, generates more using
        the experiment generator agent.
        """
        new_experiment, new_query = None, None

        if n_retry > 0:
            if self.untried_experiments:
                idx = random.randrange(len(self.untried_experiments))
                new_experiment = self.untried_experiments.pop(idx)
                self.tried_experiments.append(new_experiment)
            elif self.allow_generate_experiments and experiment_generator is not None:
                # Generate new experiments on-demand, providing all previous experiments as context
                _messages = self.messages + [
                    {
                        "role": "user",
                        "content": f"Generate new experiments given these previously attempted experiments: {json.dumps(self.tried_experiments)}",
                    }
                ]
                _reply = experiment_generator.generate_reply(messages=_messages)
                try:
                    experiments = try_loading_dict(_reply).get("experiments", [])
                except (json.JSONDecodeError, TypeError):
                    experiments = []
                self.untried_experiments = experiments.copy()
                if self.untried_experiments:
                    idx = random.randrange(len(self.untried_experiments))
                    new_experiment = self.untried_experiments.pop(idx)
                    self.tried_experiments.append(new_experiment)

            if new_experiment is not None:
                with contextlib.suppress(Exception):
                    new_query = get_query_from_experiment(new_experiment)
            if new_query is None:
                return self.get_next_experiment(
                    experiment_generator=experiment_generator, n_retry=n_retry - 1
                )

        return new_experiment, new_query

    def has_untried_experiments(self):
        return bool(self.untried_experiments) or self.allow_generate_experiments

    def to_dict(self):
        return {
            "id": self.id,
            "success": self.success,
            "parent_id": self.parent_id,
            "creation_idx": self.creation_idx,
            "created_at": self.created_at,
            "time_elapsed": self.time_elapsed,
            "visits": self.visits,
            "value": self.value,
            "self_value": self.self_value,
            "surprising": self.surprising,
            "belief_change": self.belief_change,
            "normalized_surprisal": self.normalized_surprisal,
            "kl_divergence": self.kl_divergence,
            "prior": self.prior.to_dict() if self.prior else None,
            "posterior": self.posterior.to_dict() if self.posterior else None,
            "hypothesis": self.hypothesis,
            "experiment_plan": self.experiment_plan,
            "code": self.code,
            "code_output": self.code_output,
            "analysis": self.analysis,
            "review": self.review,
            "untried_experiments": self.untried_experiments,
            "tried_experiments": self.tried_experiments,
            "query": self.query,
            "messages": self.messages,
        }

    def read_experiment_from_messages(self, store_new_experiments=False):
        """Extracts experiment details from messages and updates the node's attributes."""
        latest_experiment = None
        was_revised = False
        latest_programmer = None
        latest_code_executor = None
        latest_analyst = None
        latest_reviewer = None
        latest_reviewer_feedback = "N/A"
        latest_reviewer_success = False
        latest_experiment_generator = None

        for msg in reversed(self.messages):
            if not latest_experiment and msg.get("name") in ["user_proxy", "experiment_reviser"]:
                latest_experiment = msg.get("content")
                if msg.get("name") == "experiment_reviser":
                    was_revised = True
            elif not latest_programmer and msg.get("name") == "experiment_programmer":
                latest_programmer = try_loading_dict(msg.get("content")).get("code", "N/A")
            elif not latest_code_executor and msg.get("name") == "code_executor":
                latest_code_executor = msg.get("content")
            elif not latest_analyst and msg.get("name") in [
                "experiment_analyst",
                "experiment_code_analyst",
            ]:
                latest_analyst = try_loading_dict(msg.get("content")).get("analysis", "N/A")
            elif not latest_reviewer and msg.get("name") == "experiment_reviewer":
                latest_reviewer = try_loading_dict(msg.get("content"))
                latest_reviewer_feedback = latest_reviewer.get("feedback", "N/A")
                if latest_reviewer_feedback == "":
                    latest_reviewer_feedback = "N/A"
                latest_reviewer_success = latest_reviewer.get("success", False)
            elif not latest_experiment_generator and msg.get("name") == "experiment_generator":
                latest_experiment_generator = try_loading_dict(msg.get("content")).get(
                    "experiments", []
                )

            if (
                latest_experiment
                and latest_programmer
                and latest_code_executor
                and latest_analyst
                and latest_reviewer
            ):
                break

        if was_revised:
            latest_experiment_obj = try_loading_dict(latest_experiment)
            # Change what the query should now be based on the revised experiment
            self.query = get_query_from_experiment(latest_experiment_obj)
        else:
            latest_experiment_obj = get_experiment_from_query(
                latest_experiment
            )  # assuming it is a query string

        self.hypothesis = latest_experiment_obj.get("hypothesis", "N/A")
        self.experiment_plan = latest_experiment_obj.get("experiment_plan", "N/A")
        self.code = latest_programmer
        self.code_output = latest_code_executor
        self.analysis = latest_analyst
        self.review = latest_reviewer_feedback
        self.success = latest_reviewer_success

        # Store new experiments into untried_experiments
        if store_new_experiments and latest_experiment_generator:
            self.untried_experiments += latest_experiment_generator

    def get_context(self, include_code_output=False) -> None | str:
        """Returns the node's hypothesis, experiment, output, analysis, and review."""
        if len(self.messages) == 0:
            return None
        context_str = get_context_string(
            self.query,
            self.code_output,
            self.analysis,
            self.review,
            include_code_output=include_code_output,
        )
        return context_str

    def get_path_context(self, k: int | None = None, skip_root=True) -> list:
        """Returns messages from the node to the root

        Args:
            k: Optional maximum number of parent levels to include. If None, includes all parents.
            skip_root: If True, skips the root node in the context.
        """
        if self.parent is None:
            return []

        current_context = self.parent.get_context()
        context: list = []
        if current_context is not None and not (skip_root and self.parent.level <= 1):
            context = [current_context]

        k_remaining = None if k is None else k - 1
        if self.parent is not None and (k_remaining is None or k_remaining > 0):
            parent_context = self.parent.get_path_context(k=k_remaining, skip_root=skip_root)
            if parent_context:
                context = parent_context + context

        return context

    def update_counts(self, visits: int = 1, reward: float = 0):
        """Update the visit count and value of this node and its parents."""
        self.visits += visits
        self.value += reward
        if self.parent is not None:
            self.parent.update_counts(visits=visits, reward=reward)

get_context(include_code_output=False)

Returns the node's hypothesis, experiment, output, analysis, and review.

Source code in packages/autodiscovery/src/autodiscovery/mcts.py

def get_context(self, include_code_output=False) -> None | str:
    """Returns the node's hypothesis, experiment, output, analysis, and review."""
    if len(self.messages) == 0:
        return None
    context_str = get_context_string(
        self.query,
        self.code_output,
        self.analysis,
        self.review,
        include_code_output=include_code_output,
    )
    return context_str

get_next_experiment(experiment_generator=None, n_retry=3)

Returns the next untried experiment. If none left and generating experiments is allowed, generates more using the experiment generator agent.

Source code in packages/autodiscovery/src/autodiscovery/mcts.py

def get_next_experiment(self, experiment_generator=None, n_retry=3):
    """Returns the next untried experiment. If none left and generating experiments is allowed, generates more using
    the experiment generator agent.
    """
    new_experiment, new_query = None, None

    if n_retry > 0:
        if self.untried_experiments:
            idx = random.randrange(len(self.untried_experiments))
            new_experiment = self.untried_experiments.pop(idx)
            self.tried_experiments.append(new_experiment)
        elif self.allow_generate_experiments and experiment_generator is not None:
            # Generate new experiments on-demand, providing all previous experiments as context
            _messages = self.messages + [
                {
                    "role": "user",
                    "content": f"Generate new experiments given these previously attempted experiments: {json.dumps(self.tried_experiments)}",
                }
            ]
            _reply = experiment_generator.generate_reply(messages=_messages)
            try:
                experiments = try_loading_dict(_reply).get("experiments", [])
            except (json.JSONDecodeError, TypeError):
                experiments = []
            self.untried_experiments = experiments.copy()
            if self.untried_experiments:
                idx = random.randrange(len(self.untried_experiments))
                new_experiment = self.untried_experiments.pop(idx)
                self.tried_experiments.append(new_experiment)

        if new_experiment is not None:
            with contextlib.suppress(Exception):
                new_query = get_query_from_experiment(new_experiment)
        if new_query is None:
            return self.get_next_experiment(
                experiment_generator=experiment_generator, n_retry=n_retry - 1
            )

    return new_experiment, new_query

get_path_context(k=None, skip_root=True)

Returns messages from the node to the root

Parameters:

Name	Type	Description	Default
k	int | None	Optional maximum number of parent levels to include. If None, includes all parents.	None
skip_root		If True, skips the root node in the context.	True

Source code in packages/autodiscovery/src/autodiscovery/mcts.py

def get_path_context(self, k: int | None = None, skip_root=True) -> list:
    """Returns messages from the node to the root

    Args:
        k: Optional maximum number of parent levels to include. If None, includes all parents.
        skip_root: If True, skips the root node in the context.
    """
    if self.parent is None:
        return []

    current_context = self.parent.get_context()
    context: list = []
    if current_context is not None and not (skip_root and self.parent.level <= 1):
        context = [current_context]

    k_remaining = None if k is None else k - 1
    if self.parent is not None and (k_remaining is None or k_remaining > 0):
        parent_context = self.parent.get_path_context(k=k_remaining, skip_root=skip_root)
        if parent_context:
            context = parent_context + context

    return context

init_from_dict(data)

Initialize node attributes from a dictionary.

Source code in packages/autodiscovery/src/autodiscovery/mcts.py

def init_from_dict(self, data):
    """Initialize node attributes from a dictionary."""
    # Tree attributes
    self.creation_idx = data.get("creation_idx", MCTSNode._creation_counter)
    self.created_at = data.get("created_at", self.created_at)
    self.time_elapsed = data.get("time_elapsed", self.time_elapsed)
    self.id = data.get("id", None)
    if self.id is not None:
        self.level, self.node_idx = get_node_level_idx(self.id)
    else:
        self.level = data["level"]
        self.node_idx = data["node_idx"]
        self.id = f"node_{self.level}_{self.node_idx}"
    self.parent_id = data.get("parent_id", self.parent_id)
    if self.parent_id is not None:
        _, self.parent_idx = get_node_level_idx(self.parent_id)
    else:
        self.parent_idx = data.get("parent_idx", self.parent_idx)
        if self.parent_idx is not None:
            self.parent_id = f"node_{self.level - 1}_{self.parent_idx}"
    self.success = data.get("success", self.success)

    # Agent attributes
    self.query = data.get("query", "N/A")
    self.messages = data.get("messages", self.messages)
    self.untried_experiments = data.get("untried_experiments", self.untried_experiments)
    # self.tried_experiments = data.get('tried_experiments', self.tried_experiments)
    self.allow_generate_experiments = self.allow_generate_experiments and self.level > 0

    # Experiment attributes
    self.hypothesis = data.get("hypothesis", self.hypothesis)
    self.experiment_plan = data.get("experiment_plan", self.experiment_plan)
    self.code = data.get("code", self.code)
    self.code_output = data.get("code_output", self.code_output)
    self.analysis = data.get("analysis", self.analysis)
    self.review = data.get("review", self.review)

    # MCTS attributes
    self.visits = data.get("visits", self.visits)
    self.value = data.get("value", self.value)
    self.self_value = data.get("self_value", self.self_value)

    # Belief attributes
    self.surprising = data.get("surprising", self.surprising)
    from autodiscovery.beliefs import (
        BELIEF_MODE_TO_CLS,
    )  # Import here to avoid circular import issues

    if "prior" in data and data["prior"]:
        belief_cls = BELIEF_MODE_TO_CLS[data["prior"]["_type"]]
        self.prior = belief_cls.DistributionFormat(**data["prior"])
    if "posterior" in data and data["posterior"]:
        belief_cls = BELIEF_MODE_TO_CLS[data["posterior"]["_type"]]
        self.posterior = belief_cls.DistributionFormat(**data["posterior"])
    self.belief_change = data.get("belief_change", self.belief_change)
    self.normalized_surprisal = data.get("normalized_surprisal", self.normalized_surprisal)
    self.kl_divergence = data.get("kl_divergence", self.kl_divergence)

read_experiment_from_messages(store_new_experiments=False)

Extracts experiment details from messages and updates the node's attributes.

Source code in packages/autodiscovery/src/autodiscovery/mcts.py

def read_experiment_from_messages(self, store_new_experiments=False):
    """Extracts experiment details from messages and updates the node's attributes."""
    latest_experiment = None
    was_revised = False
    latest_programmer = None
    latest_code_executor = None
    latest_analyst = None
    latest_reviewer = None
    latest_reviewer_feedback = "N/A"
    latest_reviewer_success = False
    latest_experiment_generator = None

    for msg in reversed(self.messages):
        if not latest_experiment and msg.get("name") in ["user_proxy", "experiment_reviser"]:
            latest_experiment = msg.get("content")
            if msg.get("name") == "experiment_reviser":
                was_revised = True
        elif not latest_programmer and msg.get("name") == "experiment_programmer":
            latest_programmer = try_loading_dict(msg.get("content")).get("code", "N/A")
        elif not latest_code_executor and msg.get("name") == "code_executor":
            latest_code_executor = msg.get("content")
        elif not latest_analyst and msg.get("name") in [
            "experiment_analyst",
            "experiment_code_analyst",
        ]:
            latest_analyst = try_loading_dict(msg.get("content")).get("analysis", "N/A")
        elif not latest_reviewer and msg.get("name") == "experiment_reviewer":
            latest_reviewer = try_loading_dict(msg.get("content"))
            latest_reviewer_feedback = latest_reviewer.get("feedback", "N/A")
            if latest_reviewer_feedback == "":
                latest_reviewer_feedback = "N/A"
            latest_reviewer_success = latest_reviewer.get("success", False)
        elif not latest_experiment_generator and msg.get("name") == "experiment_generator":
            latest_experiment_generator = try_loading_dict(msg.get("content")).get(
                "experiments", []
            )

        if (
            latest_experiment
            and latest_programmer
            and latest_code_executor
            and latest_analyst
            and latest_reviewer
        ):
            break

    if was_revised:
        latest_experiment_obj = try_loading_dict(latest_experiment)
        # Change what the query should now be based on the revised experiment
        self.query = get_query_from_experiment(latest_experiment_obj)
    else:
        latest_experiment_obj = get_experiment_from_query(
            latest_experiment
        )  # assuming it is a query string

    self.hypothesis = latest_experiment_obj.get("hypothesis", "N/A")
    self.experiment_plan = latest_experiment_obj.get("experiment_plan", "N/A")
    self.code = latest_programmer
    self.code_output = latest_code_executor
    self.analysis = latest_analyst
    self.review = latest_reviewer_feedback
    self.success = latest_reviewer_success

    # Store new experiments into untried_experiments
    if store_new_experiments and latest_experiment_generator:
        self.untried_experiments += latest_experiment_generator

update_counts(visits=1, reward=0)

Update the visit count and value of this node and its parents.

Source code in packages/autodiscovery/src/autodiscovery/mcts.py

def update_counts(self, visits: int = 1, reward: float = 0):
    """Update the visit count and value of this node and its parents."""
    self.visits += visits
    self.value += reward
    if self.parent is not None:
        self.parent.update_counts(visits=visits, reward=reward)

ucb1(node, exploration_weight=1.0)

Upper Confidence Bound 1 calculation for node selection

Source code in packages/autodiscovery/src/autodiscovery/mcts.py

def ucb1(node, exploration_weight=1.0):
    """Upper Confidence Bound 1 calculation for node selection"""
    if node.visits == 0:
        return float("inf")
    exploitation_term = node.value / node.visits
    exploration_term = 0
    if node.parent:
        exploration_term = np.sqrt(2 * np.log(node.parent.visits) / node.visits)
    return exploitation_term + exploration_weight * exploration_term

progressive_widening(k, alpha, exploration_weight=1.0)

Create a progressive widening selection function.

Parameters:

Name	Type	Description	Default
k		Progressive widening constant.	required
alpha		Progressive widening exponent.	required
exploration_weight		Exploration weight for UCB1 selection method.	1.0

Returns:

Type	Description
	A callable function that accepts a node and returns the selected node.

Source code in packages/autodiscovery/src/autodiscovery/mcts.py

def progressive_widening(k, alpha, exploration_weight=1.0):
    """Create a progressive widening selection function.

    Args:
        k: Progressive widening constant.
        alpha: Progressive widening exponent.
        exploration_weight: Exploration weight for UCB1 selection method.

    Returns:
        A callable function that accepts a `node` and returns the selected node.
    """

    def select(node, nodes_by_level, return_n=1):
        return_nodes = []
        queue = [node]

        while len(return_nodes) < return_n and queue:
            current_node = queue.pop(0)
            num_visits = current_node.visits
            num_children = len(current_node.children)
            if (
                num_children == 0 or num_children < k * (num_visits**alpha)
            ) and current_node.has_untried_experiments():
                return_nodes.append(current_node)
                if len(return_nodes) >= return_n:
                    break

            if current_node.children:
                sorted_children = sorted(
                    current_node.children, key=lambda n: ucb1(n, exploration_weight), reverse=True
                )
                queue += sorted_children

        return return_nodes

    return select

progressive_widening_all(k, alpha, exploration_weight=1.0)

Create a progressive widening selection function with an alternative implementation that selects the node which has the highest UCB1 value from the entire set of nodes.

Parameters:

Name	Type	Description	Default
k		Progressive widening constant.	required
alpha		Progressive widening exponent.	required
exploration_weight		Exploration weight for UCB1 selection method.	1.0

Returns:

Type	Description
	A callable function that accepts a node and returns the selected node.

Source code in packages/autodiscovery/src/autodiscovery/mcts.py

def progressive_widening_all(k, alpha, exploration_weight=1.0):
    """Create a progressive widening selection function with an alternative implementation that selects the node
    which has the highest UCB1 value from the entire set of nodes.

    Args:
        k: Progressive widening constant.
        alpha: Progressive widening exponent.
        exploration_weight: Exploration weight for UCB1 selection method.

    Returns:
        A callable function that accepts a `node` and returns the selected node.
    """

    def select(node, nodes_by_level, return_n=1):
        return_nodes = []
        all_nodes = [n for level, nodes in nodes_by_level.items() if level > 0 for n in nodes]
        # Sort all nodes by UCB1 value
        all_nodes_sorted = sorted(
            all_nodes, key=lambda n: ucb1(n, exploration_weight), reverse=True
        )

        for _node in all_nodes_sorted:
            num_visits = _node.visits
            num_children = len(_node.children)
            # If we can add a new child based on the progressive widening condition
            if (
                num_children == 0 or num_children < k * (num_visits**alpha)
            ) and _node.has_untried_experiments():
                return_nodes.append(_node)
                if len(return_nodes) >= return_n:
                    break

        return return_nodes

    return select

default_mcts_selection(exploration_weight)

Create a selection function that returns nodes with available experiments by UCB1 score.

Parameters:

Name	Type	Description	Default
exploration_weight		Exploration weight for UCB1 selection.	required

Returns:

Type	Description
	Callable selection function.

Source code in packages/autodiscovery/src/autodiscovery/mcts.py

def default_mcts_selection(exploration_weight):
    """Create a selection function that returns nodes with available experiments by UCB1 score.

    Args:
        exploration_weight: Exploration weight for UCB1 selection.

    Returns:
        Callable selection function.
    """

    def select(node, nodes_by_level, return_n=1):
        return_nodes = []
        queue = [node]

        while len(return_nodes) < return_n and queue:
            current_node = queue.pop(0)
            if current_node.has_untried_experiments():
                return_nodes.append(current_node)
            if current_node.children:
                sorted_children = sorted(
                    current_node.children, key=lambda n: ucb1(n, exploration_weight), reverse=True
                )
                queue += sorted_children

        return return_nodes

    return select

ucb1_recursive(exploration_weight=1.0)

Create a UCB1 traversal selection function.

Parameters:

Name	Type	Description	Default
exploration_weight		Exploration weight for UCB1 selection method.	1.0

Returns:

Type	Description
	A callable function that accepts a node and returns the selected node.

Source code in packages/autodiscovery/src/autodiscovery/mcts.py

def ucb1_recursive(exploration_weight=1.0):
    """Create a UCB1 traversal selection function.

    Args:
        exploration_weight: Exploration weight for UCB1 selection method.

    Returns:
        A callable function that accepts a `node` and returns the selected node.
    """

    def select(node, nodes_by_level, return_n=1):
        # Sort self and children by UCB1 value
        all_nodes = [node] + node.children
        sorted_nodes = sorted(all_nodes, key=lambda n: ucb1(n, exploration_weight), reverse=True)

        return_nodes = []
        for best_node in sorted_nodes:
            if len(return_nodes) >= return_n:
                break
            if best_node.has_untried_experiments():
                if best_node is node:
                    return_nodes.append(best_node)
                else:
                    return_nodes.extend(
                        select(best_node, nodes_by_level, return_n=return_n - len(return_nodes))
                    )

        return return_nodes

    return select

beam_search(branching_factor, beam_width, log_dirname=None)

Create a beam search selection function.

Parameters:

Name	Type	Description	Default
branching_factor		Maximum number of children per node	required
beam_width		Number of nodes to keep in beam	required

Returns:

Type	Description
	A callable function that accepts a root node and returns selected nodes.

Source code in packages/autodiscovery/src/autodiscovery/mcts.py

def beam_search(branching_factor, beam_width, log_dirname=None):
    """Create a beam search selection function.

    Args:
        branching_factor: Maximum number of children per node
        beam_width: Number of nodes to keep in beam

    Returns:
        A callable function that accepts a root node and returns selected nodes.
    """

    def select(root, nodes_by_level, return_n=1):
        return_nodes = []
        beam = [root]

        while beam and len(return_nodes) < return_n:
            for node in beam:
                if node.has_untried_experiments() and len(node.children) < branching_factor:
                    return_nodes.append(node)
                    if len(return_nodes) >= return_n:
                        break
            if len(return_nodes) >= return_n:
                break

            candidates = []
            for node in beam:
                children = node.children
                if not children:
                    continue
                if len(children) > branching_factor:
                    candidates.extend(
                        sorted(children, key=lambda n: ucb1(n), reverse=True)[:branching_factor]
                    )
                else:
                    candidates.extend(children)

            if not candidates:
                break

            if len(candidates) > beam_width:
                beam = sorted(candidates, key=lambda n: ucb1(n), reverse=True)[:beam_width]
            else:
                beam = candidates

        if log_dirname:
            beam_state = [{"level": node.level, "node_idx": node.node_idx} for node in beam]
            level = beam[0].level if beam else 0
            with open(os.path.join(log_dirname, f"beam_level_{level}.json"), "w") as f:
                json.dump(beam_state, f, indent=2)

        return return_nodes

    return select

autodiscovery.mcts_utils

select_nodes(selection_method, root, nodes_by_level, n_warmstart=0, return_n=1)

Select the next nodes to expand in MCTS using the provided selection method.

Parameters:

Name	Type	Description	Default
selection_method		Function to select nodes in MCTS.	required
root		Root MCTSNode to select from.	required
nodes_by_level		Dictionary of nodes by level.	required
n_warmstart		Number of warmstart experiments to run after data loading but before MCTS selection.	0
return_n		Number of nodes to return.	1

Returns:

Type	Description
	List of selected MCTSNode instances for expansion. Empty when nothing is left to expand.

Source code in packages/autodiscovery/src/autodiscovery/mcts_utils.py

def select_nodes(selection_method, root, nodes_by_level, n_warmstart=0, return_n=1):
    """Select the next nodes to expand in MCTS using the provided selection method.

    Args:
        selection_method: Function to select nodes in MCTS.
        root: Root MCTSNode to select from.
        nodes_by_level: Dictionary of nodes by level.
        n_warmstart: Number of warmstart experiments to run after data loading but before MCTS selection.
        return_n: Number of nodes to return.

    Returns:
        List of selected MCTSNode instances for expansion. Empty when nothing is left to expand.
    """
    # If the data loader node hasn't been executed, return the root. This is not run in batch mode.
    # Once the root has nothing left to try (its data-loader expansion failed), return nothing so
    # the caller stops instead of selecting the root forever.
    if len(nodes_by_level[1]) == 0:
        return [root] if root.has_untried_experiments() else []

    # If there are warmstart experiments left to run, select the data loader node.
    n_children_at_data_loader = len(nodes_by_level[2])
    if (n_warmstart - n_children_at_data_loader) > 0:
        return [nodes_by_level[1][0]] * min(return_n, n_warmstart - n_children_at_data_loader)

    # Otherwise, use the selection policy to select the next nodes.
    try:
        selected = selection_method(root, nodes_by_level, return_n=return_n)
    except TypeError:
        selected = selection_method(root, nodes_by_level)
    return_nodes = selected if isinstance(selected, list) else [selected]
    if len(return_nodes) == 0:
        return []
    # Repeat return_nodes sequentially if needed to fill the batch
    return_nodes = (
        return_nodes * (return_n // len(return_nodes))
        + return_nodes[: return_n % len(return_nodes)]
    )
    return return_nodes

get_self_value(belief_change, kl_divergence, binary=True, width=0.2, kl_scale=20.0, mode='belief_and_kl')

Get self value for a node based on its belief.

Parameters:

Name	Type	Description	Default
belief_change	float	Change in belief from prior to posterior.	required
kl_divergence	float	KL divergence between prior and posterior beliefs.	required
binary	bool	Whether the surprisal reward is binary or continuous.	True
width	float	Surprisal width for belief change.	0.2
kl_scale	float	Normalization factor for KL divergence.	20.0
mode	str	Mode to use for self value calculation. Choices: "belief", "kl", "both".	'belief_and_kl'

Returns:

Name	Type	Description
float		Self value based on the belief type.
bool		Whether it is a surprisal.

Source code in packages/autodiscovery/src/autodiscovery/mcts_utils.py

def get_self_value(
    belief_change,
    kl_divergence,
    binary=True,
    width=0.2,
    kl_scale=20.0,
    mode: Literal["belief", "kl", "belief_and_kl"] = "belief_and_kl",
):
    """Get self value for a node based on its belief.

    Args:
        belief_change (float): Change in belief from prior to posterior.
        kl_divergence (float): KL divergence between prior and posterior beliefs.
        binary (bool): Whether the surprisal reward is binary or continuous.
        width (float): Surprisal width for belief change.
        kl_scale (float): Normalization factor for KL divergence.
        mode (str): Mode to use for self value calculation. Choices: "belief", "kl", "both".

    Returns:
        float: Self value based on the belief type.
        bool: Whether it is a surprisal.
    """
    if mode == "belief":
        if binary:
            return float(belief_change >= width), bool(belief_change >= width)
        else:
            # Continuous reward normalized by the surprisal width
            return belief_change / width, bool((belief_change / width) >= 1.0)
    elif mode == "kl":
        # KL divergence reward normalized by the KL scale
        if binary:
            return float(kl_divergence >= kl_scale), bool(kl_divergence >= kl_scale)
        else:
            return kl_divergence / kl_scale, bool((kl_divergence / kl_scale) >= 1.0)
    elif mode == "belief_and_kl":
        # Satisfy both modes
        belief_value, is_surprising_belief = get_self_value(
            belief_change, kl_divergence, binary, width, kl_scale, mode="belief"
        )
        kl_value, is_surprising_kl = get_self_value(
            belief_change, kl_divergence, binary, width, kl_scale, mode="kl"
        )
        # Combine both values
        combined_value = max(belief_value, kl_value)
        is_surprising = bool(is_surprising_belief or is_surprising_kl)
        return combined_value, is_surprising

    raise ValueError(f"Invalid mode: {mode}. Choose from 'belief', 'kl', or 'belief_and_kl'.")

get_context_string(hyp_exp_query, code_output=None, analysis=None, review=None, belief_mean=None, include_code_output=False)

Source code in packages/autodiscovery/src/autodiscovery/mcts_utils.py

def get_context_string(
    hyp_exp_query,
    code_output=None,
    analysis=None,
    review=None,
    belief_mean=None,
    include_code_output=False,
):
    # Format the experiment to include as context in, e.g., an LLM call.
    context_str = hyp_exp_query
    if include_code_output and code_output is not None:
        context_str += f"\n\nCode Output:\n{code_output}"
    if analysis is not None:
        context_str += f"\n\nAnalysis:\n{analysis}"
    if review is not None:
        context_str += f"\n\nReview:\n{review}"
    if belief_mean is not None:
        context_str += f"\n\nBelief about this hypothesis (range 0-1: definitely false -> uncertain -> definitely true): {belief_mean:.4f}"

    return context_str

load_mcts_from_json(json_obj_or_file_or_dir, args=None, replay_mcts=True)

Load and reconstruct MCTS nodes from a JSON object, log file, or directory.

Parameters:

Name	Type	Description	Default
json_obj_or_file_or_dir		Loaded JSON object or a path to the mcts_nodes.json file or to a directory with mcts_node_*.json files.	required

Returns:

Name	Type	Description
root		Root MCTSNode
nodes_by_level		Dictionary mapping levels to lists of MCTSNodes

Source code in packages/autodiscovery/src/autodiscovery/mcts_utils.py

def load_mcts_from_json(json_obj_or_file_or_dir, args=None, replay_mcts=True):
    """Load and reconstruct MCTS nodes from a JSON object, log file, or directory.

    Args:
        json_obj_or_file_or_dir: Loaded JSON object or a path to the mcts_nodes.json file or to a directory with mcts_node_*.json files.

    Returns:
        root: Root MCTSNode
        nodes_by_level: Dictionary mapping levels to lists of MCTSNodes
    """
    from autodiscovery.mcts import MCTSNode  # Import here to avoid circular import issues

    node_data = get_nodes(json_obj_or_file_or_dir)

    # Initialize tree data structures
    nodes_by_level = defaultdict(list)
    node_map = {}  # Map (level, idx) to node objects for linking

    # Iterate over the nodes in level order and build the tree
    node_data.sort(key=lambda x: (int(x["id"].split("_")[1]), int(x["id"].split("_")[2])))
    for data in node_data:
        # Create an empty node and initialize from dict (parent links added in second pass)
        node = MCTSNode(
            allow_generate_experiments=args.allow_generate_experiments if args else True
        )
        node.init_from_dict(data)
        # Add to data structures
        nodes_by_level[node.level].append(node)
        node_map[(node.level, node.node_idx)] = node
        # Link to parent
        if node.parent_id is not None:
            parent_level = node.level - 1
            parent_idx = node.parent_idx
            try:
                node.parent = node_map[(parent_level, parent_idx)]
                node.parent.children.append(node)
            except KeyError:
                assert (parent_level, parent_idx) == (0, 0), (
                    f"Parent node ({parent_level}, {parent_idx}) not found in node_map."
                )

    # Create root node if it does not exist
    if (0, 0) not in node_map:
        node = MCTSNode(level=0, node_idx=0, creation_idx=0)
        nodes_by_level[0].append(node)  # Figure out creation_idx use
        node_map[(0, 0)] = node
        # Link root to the tree
        node.children = [node_map[(1, 0)]]
        node_map[(1, 0)].parent = node

    assert len(node_map) == MCTSNode._creation_counter
    root = node_map[(0, 0)]

    # Fix tried/untried experiments
    for node in node_map.values():
        _tried_experiments, _untried_experiments = [], []
        cur_untried_experiments = set(
            list(map(get_query_from_experiment, node.untried_experiments))
        )
        for child in node.children:
            # Keep only children in tried experiments
            _tried_experiments.append(get_experiment_from_query(child.query))
            # Remove child from untried experiments if exists
            if child.query in cur_untried_experiments:
                cur_untried_experiments.remove(child.query)
        _untried_experiments = list(map(get_experiment_from_query, list(cur_untried_experiments)))
        node.tried_experiments = _tried_experiments
        node.untried_experiments = _untried_experiments

    if replay_mcts:
        # Replay MCTS to assign correct visits and values in order of creation_idx
        _nodes = sorted(node_map.values(), key=lambda x: x.creation_idx)
        # Reset visits and value
        for _node in _nodes:
            _node.visits = 0
            _node.value = 0
        # Backpropagate visits and values
        for _node in _nodes:
            _node.update_counts(visits=1, reward=_node.self_value)

    return root, nodes_by_level

Beliefs

autodiscovery.beliefs

BeliefTrueFalseCat

Source code in packages/autodiscovery/src/autodiscovery/beliefs.py

class BeliefTrueFalseCat:
    score_per_category = {
        "definitely_false": 0.0,
        "maybe_false": 0.25,
        "uncertain": 0.5,
        "maybe_true": 0.75,
        "definitely_true": 1.0,
    }

    class DistributionFormat:
        """A distribution of beliefs about the hypothesis using categorical buckets (Categorical).

        Attributes:
            n: Number of samples used to compute the distribution
            definitely_true: Number of "definitely true" responses
            maybe_true: Number of "maybe true" responses
            uncertain: Number of "uncertain" responses
            maybe_false: Number of "maybe false" responses
            definitely_false: Number of "definitely false" responses
            mean: Mean belief probability (optional, computed if not provided)
            prior_params: Parameters for the prior Beta distribution (alpha, beta)
        """

        def __init__(
            self,
            n: float = Field(..., description="Number of samples used to compute the distribution"),
            definitely_true: float = Field(
                ..., description='Number of "definitely true" responses'
            ),
            maybe_true: float = Field(..., description='Number of "maybe true" responses'),
            uncertain: float = Field(..., description='Number of "uncertain" responses'),
            maybe_false: float = Field(..., description='Number of "maybe false" responses'),
            definitely_false: float = Field(
                ..., description='Number of "definitely false" responses'
            ),
            mean: float | None = None,
            prior_params: tuple[float, float] = (0.5, 0.5),
            **kwargs,
        ):
            self.n = n
            self.definitely_true = definitely_true
            self.maybe_true = maybe_true
            self.uncertain = uncertain
            self.maybe_false = maybe_false
            self.definitely_false = definitely_false
            self.mean = mean
            self._empirical_mean = 0.5
            self.prior_params = prior_params  # Parameters for the prior Beta distribution

        def __repr__(self):
            return (
                f"BeliefTrueFalseCat.DistributionFormat(n={self.n}, definitely_true={self.definitely_true}, "
                f"maybe_true={self.maybe_true}, uncertain={self.uncertain}, "
                f"maybe_false={self.maybe_false}, definitely_false={self.definitely_false})"
            )

        def to_dict(self):
            return {
                "_type": "boolean_cat",
                "prior_params": self.prior_params,
                "n": self.n,
                "definitely_true": self.definitely_true,
                "maybe_true": self.maybe_true,
                "uncertain": self.uncertain,
                "maybe_false": self.maybe_false,
                "definitely_false": self.definitely_false,
                "_empirical_mean": self._empirical_mean,
                "mean": self.mean,
            }

        def get_mean_belief(self, prior=None, recompute=False) -> float:
            """Get the mean of the prior/posterior belief distribution.

            Args:
                prior (BeliefTrueFalseCat.DistributionFormat): Prior distribution format object.
                recompute (bool): Whether to recompute the mean even if it is already set.

            Returns:
                float: The mean belief probability.
            """
            if self.mean is None or recompute:
                # Compute the mean belief using the Beta distribution
                alpha1, alpha2 = BeliefTrueFalseCat.get_beta_params_from_cat_samples(
                    self.definitely_true,
                    self.maybe_true,
                    self.uncertain,
                    self.maybe_false,
                    self.definitely_false,
                )
                if self.n > 0:
                    self._empirical_mean = alpha1 / self.n
                self.mean = (self.prior_params[0] + alpha1) / (self.n + sum(self.prior_params))

                if prior is not None:
                    # Bayesian update: Beta(n_true + a, n_false + b) where a and b are prior parameters
                    prior_alpha1, prior_alpha2 = (
                        BeliefTrueFalseCat.get_beta_params_from_cat_samples(
                            prior.definitely_true,
                            prior.maybe_true,
                            prior.uncertain,
                            prior.maybe_false,
                            prior.definitely_false,
                        )
                    )
                    post_alpha = prior_alpha1 + prior.prior_params[0]
                    # post_beta = prior_alpha2 + prior.prior_params[1]
                    self.mean = (alpha1 + post_alpha) / (self.n + prior.n + sum(prior.prior_params))
            return self.mean

        def update(
            self,
            definitely_true: int = 0,
            maybe_true: int = 0,
            uncertain: int = 0,
            maybe_false: int = 0,
            definitely_false: int = 0,
            distr=None,
            normalize: bool = False,
        ):
            """Update the distribution with new counts."""
            if distr is not None:
                self.definitely_true += distr.definitely_true
                self.maybe_true += distr.maybe_true
                self.uncertain += distr.uncertain
                self.maybe_false += distr.maybe_false
                self.definitely_false += distr.definitely_false
            else:
                self.definitely_true += definitely_true
                self.maybe_true += maybe_true
                self.uncertain += uncertain
                self.maybe_false += maybe_false
                self.definitely_false += definitely_false
            n = (
                distr.n
                if distr is not None
                else (definitely_true + maybe_true + uncertain + maybe_false + definitely_false)
            )
            if normalize:
                total = self.n + n
                self.definitely_true /= total / self.n
                self.maybe_true /= total / self.n
                self.uncertain /= total / self.n
                self.maybe_false /= total / self.n
                self.definitely_false /= total / self.n
            else:
                self.n += n
            # Reset mean
            _ = self.get_mean_belief(recompute=True)

        def get_params(self) -> tuple[float, float]:
            """Get the parameters of the Beta distribution.

            Returns:
                Tuple[float, float]: Parameters (alpha, beta) of the Beta distribution.
            """
            alpha1, alpha2 = BeliefTrueFalseCat.get_beta_params_from_cat_samples(
                self.definitely_true,
                self.maybe_true,
                self.uncertain,
                self.maybe_false,
                self.definitely_false,
            )
            return self.prior_params[0] + alpha1, self.prior_params[1] + alpha2

    class ResponseFormat(BaseModel):
        """Belief about the support for the hypothesis.

        Attributes:
            belief (str): Belief about the support for the hypothesis. Choices are:
                "definitely true": Hypothesis is definitely true.
                "maybe true": Hypothesis may be true.
                "uncertain": Hypothesis is equally likely to be true or false (e.g., because of relevant but contradictory evidence).
                "maybe false": Hypothesis may be false.
                "definitely false": Hypothesis is definitely false.
                "cannot comment": Not enough information to comment on the hypothesis (e.g., due to lack of domain knowledge or lack of relevant evidence).
        """

        belief: str = Field(
            ...,
            description="Belief about the hypothesis",
            choices=[
                "definitely true",
                "maybe true",
                "uncertain",
                "maybe false",
                "definitely false",
                "cannot comment",
            ],
        )

    @staticmethod
    def parse_response(
        response: list[dict], prior_params: tuple[float, float] = (0.5, 0.5), weight: float = 1.0
    ) -> BeliefTrueFalseCat.DistributionFormat:
        """Parse the response from the LLM into a DistributionFormat.

        Args:
            response (dict): The response from the LLM containing belief counts.
            prior_params (Tuple[float, float]): Parameters for the prior Beta distribution (alpha, beta).
            weight (float): Weight to apply to the counts (default is 1.0).

        Returns:
            BeliefTrueFalseCat.DistributionFormat: Parsed distribution format.
        """
        cannot_comment = sum(1 for _res in response if _res["belief"] == "cannot comment")
        definitely_true = weight * sum(
            1 for _res in response if _res["belief"] == "definitely true"
        )
        maybe_true = weight * sum(1 for _res in response if _res["belief"] == "maybe true")
        uncertain = weight * sum(1 for _res in response if _res["belief"] == "uncertain")
        maybe_false = weight * sum(1 for _res in response if _res["belief"] == "maybe false")
        definitely_false = weight * sum(
            1 for _res in response if _res["belief"] == "definitely false"
        )
        n = weight * (len(response) - cannot_comment)  # Exclude responses with "cannot comment"

        return BeliefTrueFalseCat.DistributionFormat(
            n=n,
            definitely_true=definitely_true,
            maybe_true=maybe_true,
            uncertain=uncertain,
            maybe_false=maybe_false,
            definitely_false=definitely_false,
            prior_params=prior_params,
        )

    @staticmethod
    def kl_divergence(
        distr1: BeliefTrueFalseCat.DistributionFormat,
        distr2: BeliefTrueFalseCat.DistributionFormat,
    ) -> float:
        """Compute the KL divergence between two distributions.

        Args:
            distr1 (BeliefTrueFalseCat.DistributionFormat): First distribution.
            distr2 (BeliefTrueFalseCat.DistributionFormat): Second distribution.

        Returns:
            float: KL divergence between the two distributions.
        """
        alpha1, beta1 = distr1.get_params()
        alpha2, beta2 = distr2.get_params()
        term1 = betaln(alpha2, beta2) - betaln(alpha1, beta1)
        term2 = (alpha1 - alpha2) * psi(alpha1)
        term3 = (beta1 - beta2) * psi(beta1)
        term4 = (alpha2 - alpha1 + beta2 - beta1) * psi(alpha1 + beta1)
        return term1 + term2 + term3 + term4

    @staticmethod
    def get_beta_params_from_cat_samples(
        definitely_true: float,
        maybe_true: float,
        uncertain: float,
        maybe_false: float,
        definitely_false: float,
    ) -> tuple[float, float]:
        """Convert categorical counts into parameters for a Beta distribution.

        Args:
            definitely_true: Count of "definitely true" responses.
            maybe_true: Count of "maybe true" responses.
            uncertain: Count of "uncertain" responses.
            maybe_false: Count of "maybe false" responses.
            definitely_false: Count of "definitely false" responses.

        Returns:
            Tuple[float, float]: Parameters (alpha, beta) for the Beta distribution.
        """
        total = definitely_true + maybe_true + uncertain + maybe_false + definitely_false
        alpha = (
            definitely_true * BeliefTrueFalseCat.score_per_category["definitely_true"]
            + maybe_true * BeliefTrueFalseCat.score_per_category["maybe_true"]
            + uncertain * BeliefTrueFalseCat.score_per_category["uncertain"]
            + maybe_false * BeliefTrueFalseCat.score_per_category["maybe_false"]
            + definitely_false * BeliefTrueFalseCat.score_per_category["definitely_false"]
        )
        beta = total - alpha
        return alpha, beta

DistributionFormat

A distribution of beliefs about the hypothesis using categorical buckets (Categorical).

Attributes:

Name	Type	Description
n		Number of samples used to compute the distribution
definitely_true		Number of "definitely true" responses
maybe_true		Number of "maybe true" responses
uncertain		Number of "uncertain" responses
maybe_false		Number of "maybe false" responses
definitely_false		Number of "definitely false" responses
mean		Mean belief probability (optional, computed if not provided)
prior_params		Parameters for the prior Beta distribution (alpha, beta)

Source code in packages/autodiscovery/src/autodiscovery/beliefs.py

class DistributionFormat:
    """A distribution of beliefs about the hypothesis using categorical buckets (Categorical).

    Attributes:
        n: Number of samples used to compute the distribution
        definitely_true: Number of "definitely true" responses
        maybe_true: Number of "maybe true" responses
        uncertain: Number of "uncertain" responses
        maybe_false: Number of "maybe false" responses
        definitely_false: Number of "definitely false" responses
        mean: Mean belief probability (optional, computed if not provided)
        prior_params: Parameters for the prior Beta distribution (alpha, beta)
    """

    def __init__(
        self,
        n: float = Field(..., description="Number of samples used to compute the distribution"),
        definitely_true: float = Field(
            ..., description='Number of "definitely true" responses'
        ),
        maybe_true: float = Field(..., description='Number of "maybe true" responses'),
        uncertain: float = Field(..., description='Number of "uncertain" responses'),
        maybe_false: float = Field(..., description='Number of "maybe false" responses'),
        definitely_false: float = Field(
            ..., description='Number of "definitely false" responses'
        ),
        mean: float | None = None,
        prior_params: tuple[float, float] = (0.5, 0.5),
        **kwargs,
    ):
        self.n = n
        self.definitely_true = definitely_true
        self.maybe_true = maybe_true
        self.uncertain = uncertain
        self.maybe_false = maybe_false
        self.definitely_false = definitely_false
        self.mean = mean
        self._empirical_mean = 0.5
        self.prior_params = prior_params  # Parameters for the prior Beta distribution

    def __repr__(self):
        return (
            f"BeliefTrueFalseCat.DistributionFormat(n={self.n}, definitely_true={self.definitely_true}, "
            f"maybe_true={self.maybe_true}, uncertain={self.uncertain}, "
            f"maybe_false={self.maybe_false}, definitely_false={self.definitely_false})"
        )

    def to_dict(self):
        return {
            "_type": "boolean_cat",
            "prior_params": self.prior_params,
            "n": self.n,
            "definitely_true": self.definitely_true,
            "maybe_true": self.maybe_true,
            "uncertain": self.uncertain,
            "maybe_false": self.maybe_false,
            "definitely_false": self.definitely_false,
            "_empirical_mean": self._empirical_mean,
            "mean": self.mean,
        }

    def get_mean_belief(self, prior=None, recompute=False) -> float:
        """Get the mean of the prior/posterior belief distribution.

        Args:
            prior (BeliefTrueFalseCat.DistributionFormat): Prior distribution format object.
            recompute (bool): Whether to recompute the mean even if it is already set.

        Returns:
            float: The mean belief probability.
        """
        if self.mean is None or recompute:
            # Compute the mean belief using the Beta distribution
            alpha1, alpha2 = BeliefTrueFalseCat.get_beta_params_from_cat_samples(
                self.definitely_true,
                self.maybe_true,
                self.uncertain,
                self.maybe_false,
                self.definitely_false,
            )
            if self.n > 0:
                self._empirical_mean = alpha1 / self.n
            self.mean = (self.prior_params[0] + alpha1) / (self.n + sum(self.prior_params))

            if prior is not None:
                # Bayesian update: Beta(n_true + a, n_false + b) where a and b are prior parameters
                prior_alpha1, prior_alpha2 = (
                    BeliefTrueFalseCat.get_beta_params_from_cat_samples(
                        prior.definitely_true,
                        prior.maybe_true,
                        prior.uncertain,
                        prior.maybe_false,
                        prior.definitely_false,
                    )
                )
                post_alpha = prior_alpha1 + prior.prior_params[0]
                # post_beta = prior_alpha2 + prior.prior_params[1]
                self.mean = (alpha1 + post_alpha) / (self.n + prior.n + sum(prior.prior_params))
        return self.mean

    def update(
        self,
        definitely_true: int = 0,
        maybe_true: int = 0,
        uncertain: int = 0,
        maybe_false: int = 0,
        definitely_false: int = 0,
        distr=None,
        normalize: bool = False,
    ):
        """Update the distribution with new counts."""
        if distr is not None:
            self.definitely_true += distr.definitely_true
            self.maybe_true += distr.maybe_true
            self.uncertain += distr.uncertain
            self.maybe_false += distr.maybe_false
            self.definitely_false += distr.definitely_false
        else:
            self.definitely_true += definitely_true
            self.maybe_true += maybe_true
            self.uncertain += uncertain
            self.maybe_false += maybe_false
            self.definitely_false += definitely_false
        n = (
            distr.n
            if distr is not None
            else (definitely_true + maybe_true + uncertain + maybe_false + definitely_false)
        )
        if normalize:
            total = self.n + n
            self.definitely_true /= total / self.n
            self.maybe_true /= total / self.n
            self.uncertain /= total / self.n
            self.maybe_false /= total / self.n
            self.definitely_false /= total / self.n
        else:
            self.n += n
        # Reset mean
        _ = self.get_mean_belief(recompute=True)

    def get_params(self) -> tuple[float, float]:
        """Get the parameters of the Beta distribution.

        Returns:
            Tuple[float, float]: Parameters (alpha, beta) of the Beta distribution.
        """
        alpha1, alpha2 = BeliefTrueFalseCat.get_beta_params_from_cat_samples(
            self.definitely_true,
            self.maybe_true,
            self.uncertain,
            self.maybe_false,
            self.definitely_false,
        )
        return self.prior_params[0] + alpha1, self.prior_params[1] + alpha2

get_mean_belief(prior=None, recompute=False)

Get the mean of the prior/posterior belief distribution.

Parameters:

Name	Type	Description	Default
prior	DistributionFormat	Prior distribution format object.	None
recompute	bool	Whether to recompute the mean even if it is already set.	False

Returns:

Name	Type	Description
float	float	The mean belief probability.

Source code in packages/autodiscovery/src/autodiscovery/beliefs.py

def get_mean_belief(self, prior=None, recompute=False) -> float:
    """Get the mean of the prior/posterior belief distribution.

    Args:
        prior (BeliefTrueFalseCat.DistributionFormat): Prior distribution format object.
        recompute (bool): Whether to recompute the mean even if it is already set.

    Returns:
        float: The mean belief probability.
    """
    if self.mean is None or recompute:
        # Compute the mean belief using the Beta distribution
        alpha1, alpha2 = BeliefTrueFalseCat.get_beta_params_from_cat_samples(
            self.definitely_true,
            self.maybe_true,
            self.uncertain,
            self.maybe_false,
            self.definitely_false,
        )
        if self.n > 0:
            self._empirical_mean = alpha1 / self.n
        self.mean = (self.prior_params[0] + alpha1) / (self.n + sum(self.prior_params))

        if prior is not None:
            # Bayesian update: Beta(n_true + a, n_false + b) where a and b are prior parameters
            prior_alpha1, prior_alpha2 = (
                BeliefTrueFalseCat.get_beta_params_from_cat_samples(
                    prior.definitely_true,
                    prior.maybe_true,
                    prior.uncertain,
                    prior.maybe_false,
                    prior.definitely_false,
                )
            )
            post_alpha = prior_alpha1 + prior.prior_params[0]
            # post_beta = prior_alpha2 + prior.prior_params[1]
            self.mean = (alpha1 + post_alpha) / (self.n + prior.n + sum(prior.prior_params))
    return self.mean

get_params()

Get the parameters of the Beta distribution.

Returns:

Type	Description
tuple[float, float]	Tuple[float, float]: Parameters (alpha, beta) of the Beta distribution.

Source code in packages/autodiscovery/src/autodiscovery/beliefs.py

def get_params(self) -> tuple[float, float]:
    """Get the parameters of the Beta distribution.

    Returns:
        Tuple[float, float]: Parameters (alpha, beta) of the Beta distribution.
    """
    alpha1, alpha2 = BeliefTrueFalseCat.get_beta_params_from_cat_samples(
        self.definitely_true,
        self.maybe_true,
        self.uncertain,
        self.maybe_false,
        self.definitely_false,
    )
    return self.prior_params[0] + alpha1, self.prior_params[1] + alpha2

update(definitely_true=0, maybe_true=0, uncertain=0, maybe_false=0, definitely_false=0, distr=None, normalize=False)

Update the distribution with new counts.

Source code in packages/autodiscovery/src/autodiscovery/beliefs.py

def update(
    self,
    definitely_true: int = 0,
    maybe_true: int = 0,
    uncertain: int = 0,
    maybe_false: int = 0,
    definitely_false: int = 0,
    distr=None,
    normalize: bool = False,
):
    """Update the distribution with new counts."""
    if distr is not None:
        self.definitely_true += distr.definitely_true
        self.maybe_true += distr.maybe_true
        self.uncertain += distr.uncertain
        self.maybe_false += distr.maybe_false
        self.definitely_false += distr.definitely_false
    else:
        self.definitely_true += definitely_true
        self.maybe_true += maybe_true
        self.uncertain += uncertain
        self.maybe_false += maybe_false
        self.definitely_false += definitely_false
    n = (
        distr.n
        if distr is not None
        else (definitely_true + maybe_true + uncertain + maybe_false + definitely_false)
    )
    if normalize:
        total = self.n + n
        self.definitely_true /= total / self.n
        self.maybe_true /= total / self.n
        self.uncertain /= total / self.n
        self.maybe_false /= total / self.n
        self.definitely_false /= total / self.n
    else:
        self.n += n
    # Reset mean
    _ = self.get_mean_belief(recompute=True)

ResponseFormat

Bases: BaseModel

Belief about the support for the hypothesis.

Attributes:

Name	Type	Description
belief	str	Belief about the support for the hypothesis. Choices are: "definitely true": Hypothesis is definitely true. "maybe true": Hypothesis may be true. "uncertain": Hypothesis is equally likely to be true or false (e.g., because of relevant but contradictory evidence). "maybe false": Hypothesis may be false. "definitely false": Hypothesis is definitely false. "cannot comment": Not enough information to comment on the hypothesis (e.g., due to lack of domain knowledge or lack of relevant evidence).

Source code in packages/autodiscovery/src/autodiscovery/beliefs.py

class ResponseFormat(BaseModel):
    """Belief about the support for the hypothesis.

    Attributes:
        belief (str): Belief about the support for the hypothesis. Choices are:
            "definitely true": Hypothesis is definitely true.
            "maybe true": Hypothesis may be true.
            "uncertain": Hypothesis is equally likely to be true or false (e.g., because of relevant but contradictory evidence).
            "maybe false": Hypothesis may be false.
            "definitely false": Hypothesis is definitely false.
            "cannot comment": Not enough information to comment on the hypothesis (e.g., due to lack of domain knowledge or lack of relevant evidence).
    """

    belief: str = Field(
        ...,
        description="Belief about the hypothesis",
        choices=[
            "definitely true",
            "maybe true",
            "uncertain",
            "maybe false",
            "definitely false",
            "cannot comment",
        ],
    )

get_beta_params_from_cat_samples(definitely_true, maybe_true, uncertain, maybe_false, definitely_false) staticmethod

Convert categorical counts into parameters for a Beta distribution.

Parameters:

Name	Type	Description	Default
definitely_true	float	Count of "definitely true" responses.	required
maybe_true	float	Count of "maybe true" responses.	required
uncertain	float	Count of "uncertain" responses.	required
maybe_false	float	Count of "maybe false" responses.	required
definitely_false	float	Count of "definitely false" responses.	required

Returns:

Type	Description
tuple[float, float]	Tuple[float, float]: Parameters (alpha, beta) for the Beta distribution.

Source code in packages/autodiscovery/src/autodiscovery/beliefs.py

@staticmethod
def get_beta_params_from_cat_samples(
    definitely_true: float,
    maybe_true: float,
    uncertain: float,
    maybe_false: float,
    definitely_false: float,
) -> tuple[float, float]:
    """Convert categorical counts into parameters for a Beta distribution.

    Args:
        definitely_true: Count of "definitely true" responses.
        maybe_true: Count of "maybe true" responses.
        uncertain: Count of "uncertain" responses.
        maybe_false: Count of "maybe false" responses.
        definitely_false: Count of "definitely false" responses.

    Returns:
        Tuple[float, float]: Parameters (alpha, beta) for the Beta distribution.
    """
    total = definitely_true + maybe_true + uncertain + maybe_false + definitely_false
    alpha = (
        definitely_true * BeliefTrueFalseCat.score_per_category["definitely_true"]
        + maybe_true * BeliefTrueFalseCat.score_per_category["maybe_true"]
        + uncertain * BeliefTrueFalseCat.score_per_category["uncertain"]
        + maybe_false * BeliefTrueFalseCat.score_per_category["maybe_false"]
        + definitely_false * BeliefTrueFalseCat.score_per_category["definitely_false"]
    )
    beta = total - alpha
    return alpha, beta

kl_divergence(distr1, distr2) staticmethod

Compute the KL divergence between two distributions.

Parameters:

Name	Type	Description	Default
distr1	DistributionFormat	First distribution.	required
distr2	DistributionFormat	Second distribution.	required

Returns:

Name	Type	Description
float	float	KL divergence between the two distributions.

Source code in packages/autodiscovery/src/autodiscovery/beliefs.py

@staticmethod
def kl_divergence(
    distr1: BeliefTrueFalseCat.DistributionFormat,
    distr2: BeliefTrueFalseCat.DistributionFormat,
) -> float:
    """Compute the KL divergence between two distributions.

    Args:
        distr1 (BeliefTrueFalseCat.DistributionFormat): First distribution.
        distr2 (BeliefTrueFalseCat.DistributionFormat): Second distribution.

    Returns:
        float: KL divergence between the two distributions.
    """
    alpha1, beta1 = distr1.get_params()
    alpha2, beta2 = distr2.get_params()
    term1 = betaln(alpha2, beta2) - betaln(alpha1, beta1)
    term2 = (alpha1 - alpha2) * psi(alpha1)
    term3 = (beta1 - beta2) * psi(beta1)
    term4 = (alpha2 - alpha1 + beta2 - beta1) * psi(alpha1 + beta1)
    return term1 + term2 + term3 + term4

parse_response(response, prior_params=(0.5, 0.5), weight=1.0) staticmethod

Parse the response from the LLM into a DistributionFormat.

Parameters:

Name	Type	Description	Default
response	dict	The response from the LLM containing belief counts.	required
prior_params	Tuple[float, float]	Parameters for the prior Beta distribution (alpha, beta).	(0.5, 0.5)
weight	float	Weight to apply to the counts (default is 1.0).	1.0

Returns:

Type	Description
DistributionFormat	BeliefTrueFalseCat.DistributionFormat: Parsed distribution format.

Source code in packages/autodiscovery/src/autodiscovery/beliefs.py

@staticmethod
def parse_response(
    response: list[dict], prior_params: tuple[float, float] = (0.5, 0.5), weight: float = 1.0
) -> BeliefTrueFalseCat.DistributionFormat:
    """Parse the response from the LLM into a DistributionFormat.

    Args:
        response (dict): The response from the LLM containing belief counts.
        prior_params (Tuple[float, float]): Parameters for the prior Beta distribution (alpha, beta).
        weight (float): Weight to apply to the counts (default is 1.0).

    Returns:
        BeliefTrueFalseCat.DistributionFormat: Parsed distribution format.
    """
    cannot_comment = sum(1 for _res in response if _res["belief"] == "cannot comment")
    definitely_true = weight * sum(
        1 for _res in response if _res["belief"] == "definitely true"
    )
    maybe_true = weight * sum(1 for _res in response if _res["belief"] == "maybe true")
    uncertain = weight * sum(1 for _res in response if _res["belief"] == "uncertain")
    maybe_false = weight * sum(1 for _res in response if _res["belief"] == "maybe false")
    definitely_false = weight * sum(
        1 for _res in response if _res["belief"] == "definitely false"
    )
    n = weight * (len(response) - cannot_comment)  # Exclude responses with "cannot comment"

    return BeliefTrueFalseCat.DistributionFormat(
        n=n,
        definitely_true=definitely_true,
        maybe_true=maybe_true,
        uncertain=uncertain,
        maybe_false=maybe_false,
        definitely_false=definitely_false,
        prior_params=prior_params,
    )

get_belief(hypothesis, *, model, evidence=None, belief_mode='boolean', n_samples=5, temperature=None, reasoning_effort=None, use_llm_prior=False, explicit_prior=None, n_retries=3, weight=1.0, usage_tracker=None, usage_component='belief', usage_node_id=None, usage_metadata=None)

Get belief distribution for a hypothesis with optional evidence.

Parameters:

Name	Type	Description	Default
hypothesis	str	The hypothesis to evaluate	required
evidence	list[dict[str, str]] | None	Optional evidence messages to condition the belief	None
model	str	The LLM model to use, as litellm's <provider>/<model>	required
belief_mode	str	The belief mode to use for parsing responses (e.g., BeliefTrueFalse, BeliefCategorical)	'boolean'
n_samples	int	Number of samples to draw from the LLM	5
temperature	float | None	Temperature for sampling	None
reasoning_effort	str | None	Reasoning effort for supported models.	None
use_llm_prior	bool	Whether to use implicit Bayesian posterior	False
explicit_prior		Optional prior distribution to use for Bayesian updates	None
n_retries		Number of retries for querying the LLM in case of errors	3
weight	float	Weight to apply to the empirical distribution	1.0
usage_tracker	UsageTracker | None	Optional usage tracker.	None
usage_component	str	Component label used for usage tracking.	'belief'
usage_node_id	str | None	Optional node id used for usage tracking.	None
usage_metadata	dict | None	Optional metadata attached to usage events.	None

Source code in packages/autodiscovery/src/autodiscovery/beliefs.py

def get_belief(
    hypothesis: str,
    *,
    model: str,
    evidence: list[dict[str, str]] | None = None,
    belief_mode: str = "boolean",
    n_samples: int = 5,
    temperature: float | None = None,
    reasoning_effort: str | None = None,
    use_llm_prior: bool = False,
    explicit_prior=None,
    n_retries=3,
    weight: float = 1.0,
    usage_tracker: UsageTracker | None = None,
    usage_component: str = "belief",
    usage_node_id: str | None = None,
    usage_metadata: dict | None = None,
):
    """Get belief distribution for a hypothesis with optional evidence.

    Args:
        hypothesis: The hypothesis to evaluate
        evidence: Optional evidence messages to condition the belief
        model: The LLM model to use, as litellm's ``<provider>/<model>``
        belief_mode: The belief mode to use for parsing responses (e.g., BeliefTrueFalse, BeliefCategorical)
        n_samples: Number of samples to draw from the LLM
        temperature: Temperature for sampling
        reasoning_effort: Reasoning effort for supported models.
        use_llm_prior: Whether to use implicit Bayesian posterior
        explicit_prior: Optional prior distribution to use for Bayesian updates
        n_retries: Number of retries for querying the LLM in case of errors
        weight: Weight to apply to the empirical distribution
        usage_tracker: Optional usage tracker.
        usage_component: Component label used for usage tracking.
        usage_node_id: Optional node id used for usage tracking.
        usage_metadata: Optional metadata attached to usage events.
    """
    belief_cls = BELIEF_MODE_TO_CLS.get(belief_mode)
    if belief_cls is None:
        raise ValueError(
            f"Unknown belief_mode '{belief_mode}'; expected one of {list(BELIEF_MODE_TO_CLS.keys())}"
        )

    # Construct the system prompt based on whether we are eliciting prior, implicit posterior, or explicit posterior beliefs
    _system_msgs = [
        "You are a research scientist skilled at analyzing scientific hypotheses. Your task is to provide your belief about the given hypothesis."
    ]
    if evidence is not None:
        # posterior belief
        _system_msgs.append(
            "Use the provided evidence collected from running experiments to help make your decision. Carefully consider each piece of evidence and decide whether and how any of them affects your belief about the current hypothesis. Note that evidence from previous studies may have an indirect bearing on the hypothesis, so think about how they might relate to the hypothesis even if they do not directly test it."
        )
    # else:  # prior belief
    if use_llm_prior:
        # implicit posterior
        _system_msgs.append(
            "Use your prior knowledge of the research domain to help in your assessment of the hypothesis."
        )
    else:
        # explicit posterior
        assert evidence is not None
        _system_msgs.append(
            "Disregard any prior beliefs you have about the hypothesis and focus only on the provided evidence."
        )
    system_prompt = {"role": "system", "content": " ".join(_system_msgs)}

    hypothesis_msg = {
        "role": "user",
        "content": f"Hypothesis: {hypothesis}\n\nCarefully reason before making your assessment.",
    }

    all_msgs = [system_prompt]
    if evidence is not None:
        all_msgs += evidence
    all_msgs.append(hypothesis_msg)

    distribution, mean_belief = None, None
    for attempt in range(n_retries):
        try:
            response = query_llm(
                all_msgs,
                model=model,
                n_samples=n_samples,
                temperature=temperature,
                reasoning_effort=reasoning_effort,
                response_format=belief_cls.ResponseFormat,
                usage_tracker=usage_tracker,
                usage_component=usage_component,
                usage_agent_name="belief_agent",
                usage_node_id=usage_node_id,
                usage_metadata=usage_metadata,
            )
            prior_params_or_none = {}
            if explicit_prior is not None:
                prior_params_or_none["prior_params"] = explicit_prior.get_params()
            distribution = belief_cls.parse_response(
                response, weight=weight, **prior_params_or_none
            )
            # Compute and store the mean belief
            mean_belief = distribution.get_mean_belief()
            break
        except Exception as e:
            if attempt == n_retries - 1:
                print(
                    f"Querying LLM: ERROR: {e}\nMax retries reached. Returning empty distribution."
                )
                return None, None
            else:
                print(f"Querying LLM: ERROR: {e}\nRetrying ({attempt + 1}/{n_retries})...")

    return distribution, mean_belief

calculate_prior_and_posterior_beliefs(node, *, model, n_samples=4, temperature=None, reasoning_effort=None, implicit_bayes_posterior=False, surprisal_width=0.2, belief_mode='boolean', evidence_msg=None, prior=None, evidence_weight=1.0, usage_tracker=None, usage_node_id=None, usage_context_label='default')

Calculate prior and posterior belief distributions for a hypothesis.

Parameters:

Name	Type	Description	Default
node		MCTSNode instance containing node information and messages or a dictionary with node data	required
n_samples		Number of samples to draw from the LLM	4
model		The LLM model to use for querying, as litellm's <provider>/<model>	required
temperature		Temperature for sampling	None
reasoning_effort		Reasoning effort for supported models.	None
implicit_bayes_posterior		Whether to use implicit Bayesian posterior	False
surprisal_width		Minimum difference in mean prior and posterior probabilities required to count as a surprisal	0.2
belief_mode		The belief mode to use for parsing responses (e.g., "boolean", "categorical")	'boolean'
evidence_msg		Optional evidence messages to condition the posterior belief	None
prior		Optional pre-computed prior distribution to use for posterior calculation	None
evidence_weight		Weight to apply to the evidence when calculating the posterior belief	1.0
usage_tracker	UsageTracker | None	Optional usage tracker.	None
usage_node_id	str | None	Optional node id used for usage tracking.	None
usage_context_label	str	Context label used for usage tracking.	'default'

Source code in packages/autodiscovery/src/autodiscovery/beliefs.py

def calculate_prior_and_posterior_beliefs(
    node,
    *,
    model,
    n_samples=4,
    temperature=None,
    reasoning_effort=None,
    implicit_bayes_posterior=False,
    surprisal_width=0.2,
    belief_mode="boolean",
    evidence_msg=None,
    prior=None,
    evidence_weight=1.0,
    usage_tracker: UsageTracker | None = None,
    usage_node_id: str | None = None,
    usage_context_label: str = "default",
):
    """Calculate prior and posterior belief distributions for a hypothesis.

    Args:
        node: MCTSNode instance containing node information and messages or a dictionary with node data
        n_samples: Number of samples to draw from the LLM
        model: The LLM model to use for querying, as litellm's ``<provider>/<model>``
        temperature: Temperature for sampling
        reasoning_effort: Reasoning effort for supported models.
        implicit_bayes_posterior: Whether to use implicit Bayesian posterior
        surprisal_width: Minimum difference in mean prior and posterior probabilities required to count as a surprisal
        belief_mode: The belief mode to use for parsing responses (e.g., "boolean", "categorical")
        evidence_msg: Optional evidence messages to condition the posterior belief
        prior: Optional pre-computed prior distribution to use for posterior calculation
        evidence_weight: Weight to apply to the evidence when calculating the posterior belief
        usage_tracker: Optional usage tracker.
        usage_node_id: Optional node id used for usage tracking.
        usage_context_label: Context label used for usage tracking.
    """
    # MODEL_CTXT_LIMITS = {
    #     "o4-mini": 200_000,
    #     "gpt-4o": 128_000,
    # }
    belief_cls = BELIEF_MODE_TO_CLS.get(belief_mode)
    if belief_cls is None:
        raise ValueError(
            f"Unknown belief_mode '{belief_mode}'; expected one of {list(BELIEF_MODE_TO_CLS.keys())}"
        )

    if type(node) is MCTSNode:
        hypothesis = node.hypothesis
        query = node.query  # Contains the hypothesis and experiment plan
        code_output = node.code_output
        analysis = node.analysis
        review = node.review
    else:
        hypothesis = node["hypothesis"]
        query = node.get("query", "N/A")
        code_output = node.get("code_output", "N/A")
        analysis = node.get("analysis", "N/A")
        review = node.get("review", "N/A")

    if hypothesis is None:
        return None, None, None, None

    if evidence_msg is None:
        evidence_msg = [
            {
                "role": "user",
                "content": get_context_string(
                    query, code_output, analysis, review, include_code_output=True
                ),
            }
        ]

    if prior is None:
        prior, _ = get_belief(
            hypothesis=hypothesis,
            evidence=None,
            model=model,
            belief_mode=belief_mode,
            n_samples=n_samples,
            temperature=temperature,
            reasoning_effort=reasoning_effort,
            use_llm_prior=True,
            usage_tracker=usage_tracker,
            usage_component=f"belief.{usage_context_label}.prior",
            usage_node_id=usage_node_id,
        )

    posterior, _ = get_belief(
        hypothesis=hypothesis,
        evidence=evidence_msg,
        model=model,
        belief_mode=belief_mode,
        n_samples=n_samples,
        temperature=temperature,
        reasoning_effort=reasoning_effort,
        use_llm_prior=implicit_bayes_posterior,
        explicit_prior=prior,
        weight=evidence_weight,
        usage_tracker=usage_tracker,
        usage_component=f"belief.{usage_context_label}.posterior",
        usage_node_id=usage_node_id,
    )

    if prior is None or posterior is None:
        raise ValueError("Belief distribution could not be computed.")

    belief_change = abs(posterior.mean - prior.mean)
    kl_divergence = belief_cls.kl_divergence(posterior, prior)

    return prior, posterior, belief_change, kl_divergence

Agents and proposals

autodiscovery.agents

get_agents(work_dir, *, model_name, vision_model, temperature=None, reasoning_effort=None, branching_factor=3, user_query=None, experiment_first=False, code_timeout=30 * 60, backend='process', bucket_path=None, dataset_paths=None, datalib_dirs=None, usage_tracker=None)

Build and return the conversational agents used by AutoDiscovery.

Parameters:

Name	Type	Description	Default
work_dir		Working directory for code execution.	required
model_name		Model for AG2 conversational agents, as litellm's <provider>/<model>.	required
temperature		Sampling temperature for non-reasoning models.	None
reasoning_effort		Reasoning effort for compatible models.	None
branching_factor		Number of experiment candidates to request.	3
user_query		Optional user query injected into generator prompts.	None
experiment_first		Whether generator returns experiment-first outputs.	False
code_timeout		Timeout in seconds for code execution.	30 * 60
backend		Code execution backend (local, process, or modal).	'process'
bucket_path		Optional GCS bucket path for Modal datasets.	None
dataset_paths		Optional dataset paths (reserved for future use).	None
datalib_dirs		Read-only data-library directories the code may read, as absolute paths. The modal backend mounts each at the same path.	None
vision_model		Vision model for plot analysis, as litellm's <provider>/<model>.	required
usage_tracker	UsageTracker | None	Optional usage tracker for direct image-analysis calls.	None

Returns:

Type	Description
dict[str, ConversableAgent]	Dictionary mapping agent name to agent instance.

Source code in packages/autodiscovery/src/autodiscovery/agents.py

def get_agents(
    work_dir,
    *,
    model_name,
    vision_model,
    temperature=None,
    reasoning_effort=None,
    branching_factor=3,
    user_query=None,
    experiment_first=False,
    code_timeout=30 * 60,
    backend="process",
    bucket_path=None,
    dataset_paths=None,
    datalib_dirs=None,
    usage_tracker: UsageTracker | None = None,
) -> dict[str, ConversableAgent]:
    """Build and return the conversational agents used by AutoDiscovery.

    Args:
        work_dir: Working directory for code execution.
        model_name: Model for AG2 conversational agents, as litellm's
            ``<provider>/<model>``.
        temperature: Sampling temperature for non-reasoning models.
        reasoning_effort: Reasoning effort for compatible models.
        branching_factor: Number of experiment candidates to request.
        user_query: Optional user query injected into generator prompts.
        experiment_first: Whether generator returns experiment-first outputs.
        code_timeout: Timeout in seconds for code execution.
        backend: Code execution backend (local, process, or modal).
        bucket_path: Optional GCS bucket path for Modal datasets.
        dataset_paths: Optional dataset paths (reserved for future use).
        datalib_dirs: Read-only data-library directories the code may read, as
            absolute paths. The modal backend mounts each at the same path.
        vision_model: Vision model for plot analysis, as litellm's
            ``<provider>/<model>``.
        usage_tracker: Optional usage tracker for direct image-analysis calls.

    Returns:
        Dictionary mapping agent name to agent instance.
    """
    llm_config = get_llm_config(
        model_name=model_name,
        temperature=temperature,
        reasoning_effort=reasoning_effort,
    )

    # Create token limit transform.
    # `model` must be set explicitly: MessageTokenLimiter defaults to "gpt-3.5-turbo-0613"
    # and silently caps max_tokens_per_message to that model's 4096-token limit.
    # The value only drives tokenizer choice and the cap lookup, so any large-context OpenAI model works.
    token_limit_capability = transform_messages.TransformMessages(
        transforms=[
            transforms.MessageTokenLimiter(
                max_tokens_per_message=10_000, min_tokens=12_000, model="gpt-4o"
            )
        ]
    )

    # Experiment Generator
    _user_query_or_empty = f"{user_query}\n\n" if user_query is not None else ""
    datalib_dirs = list(datalib_dirs or [])
    # Empty without a data library, so those prompts are unchanged.
    _datalib_note = datalib.prompt_note(datalib_dirs)

    experiment_generator = ConversableAgent(
        name="experiment_generator",
        llm_config={
            **llm_config,
            "response_format": ExperimentList if not experiment_first else ExperimentHypothesisList,
        },
        system_message=(
            "You are a research scientist who is interested in doing open-ended, data-driven research using the provided dataset(s). "
            f"{_user_query_or_empty}"
            f"Be creative and think of new and interesting verifiable {'experiments' if experiment_first else 'hypotheses'} and corresponding {'hypotheses' if experiment_first else 'experiments'}. "
            "The hypothesis should be a falsifiable statement that can be sufficiently tested by an experiment using the provided data. "
            "Explain in natural language what this experiment plan is so that a programmer can implement it (do not provide the code yourself). "
            "Remember, you are interested in open-ended research, so your proposals may be exploratory in nature and may have only an indirect connection to the previous explorations provided. "
            f"{_datalib_note}"
            "Here are some instructions that you must follow:\n"
            "1. Strictly use only the dataset(s) provided and do not simulate dummy/synthetic data or columns that cannot be derived from the existing columns.\n"
            "2. Each hypothesis (and experiment plan) should be creative, independent, and self-contained.\n"
            "3. Use the prior experiments/hypotheses as inspiration to think of interesting and creative new experiments/hypotheses. However, do not repeat the same experiments/hypotheses.\n\n"
            "Here is a possible approach to coming up with a new hypothesis and experiment plan:\n"
            "1. Find an interesting context: this could be a specific subset of the data. E.g., if the dataset has multiple categorical variables, you could split the data based on specific values of such variables, which would then allow you to validate a hypothesis in the specific contexts defined by the values of those variables.\n"
            "2. Find interesting variables: these could be the columns in the dataset that you find interesting or relevant to the context. You are allowed and encouraged to create composite variables derived from the existing variables.\n"
            "3. Find interesting relationships: these are interactions between the variables that you find interesting or relevant to the context. You are encouraged to propose experiments involving complex predictive or causal models.\n"
            "4. You must require that your proposed hypotheses are verifiable using robust statistical tests. Remember, your programmer can install python packages via pip which can allow it to write code for complex statistical analyses.\n"
            "5. Multiple datasets: If you are provided with more than one dataset, then try to also propose hypotheses that utilize contexts, variables, and relationships across datasets, e.g., this may involve using join or similar operations.\n\n"
            "Generally, in typical data-driven research, you will need to explore and visualize the data for possible high-level insights, clean, transform, or derive new variables from the dataset to be suited for the investigation, deep-dive into specific parts of the data for fine-grained analysis, perform data modeling, and run statistical tests. "
            f"Now, generate exactly {branching_factor} new hypotheses with their experiment plans."
        ),
        human_input_mode="NEVER",
    )

    if backend == "process":
        # The process backend places an `install-package` script on PATH
        # that installs into a per-cell temp directory via uv.
        install_snippet = """\nimport subprocess

def install(package):
    subprocess.check_call(["install-package", package])\n\n\n"""
    else:
        install_snippet = """\nimport subprocess
import sys

def install(package):
    subprocess.check_call([sys.executable, "-m", "pip", "install", "--quiet", package])\n\n\n"""

    # Experiment Programmer
    experiment_programmer = ConversableAgent(
        name="experiment_programmer",
        llm_config={**llm_config, "response_format": ExperimentCode},
        system_message=(
            "You are a scientific experiment programmer proficient in writing python code given an experiment plan. "
            "Your code will be included in a python file that is executed and any relevant results should be printed to standard out or presented using plt.show appropriately. "
            "Make sure you provide python code in the proper format to execute. "
            "Ensure your code is clean and concise, and include debug statements only when they are absolutely necessary. "
            "Use only the dataset given and do not assume any other files are available. "
            f"{_datalib_note}"
            "The state is not preserved between code blocks, so do not assume any variables or imports from previous code blocks. "
            "Import any libraries you need to use. Always attempt to import a library before installing it (it may already be installed). "
            "If you need to install a library, use the following code example:"
            f"{install_snippet}"
            "When installing python packages, use the --quiet option to minimize unnecessary output."
            "Prefer using installed libraries over installing new libraries whenever possible. "
            "If possible, instead of downgrading library versions, try to adapt your code to work with a more updated version that is already installed. "
            "Never attempt to create a new environment. Always use the current environment. "
            "If the code requires generating plots, use plt.show (not plt.savefig).  "
            "Avoid printing the whole data structure to the console directly if it is large; instead, print concise results that are directly relevant to the experiment. "
            "You are allowed 6 total attempts to run the code, including debugging attempts.\n\n"
            "Debugging instructions:\n"
            "1. Only debug if you are either unsure about the executability or validity of the code (i.e., whether it satisfies the proposed experiment).\n"
            '2. If the code you are writing is intended for debugging, the first line of your code must be "# [debug]" only.\n'
            '3. DO NOT use "[debug]" anywhere else in your code.\n'
            "4. DO NOT combine any debug code and the actual experiment implementation code; keep them separate.\n"
            "5. For each experiment, you are allowed to debug at most 3 times.\n"
            "6. As much as possible, minimize the number of debugging steps you use."
        ),
        human_input_mode="NEVER",
    )

    # Experiment Analyst
    experiment_analyst = ConversableAgent(
        name="experiment_code_analyst",
        llm_config={**llm_config, "response_format": ExperimentAnalyst},
        system_message=(
            "You are a research scientist responsible for evaluating the code execution output for a scientific experiment written by a programmer. "
            "If no code was executed, there was an error, or the code fails silently, return the success status as **false**. "
            'If the code includes a line "# [debug]" i.e "[debug]" as a comment, strictly treat this as a debugging experiment. '
            "In such cases, strictly return the success status as **false**, provide information that it was a debug code execution, "
            "give feedback and request the experiment to be retried with the new information. "
            "Otherwise, analyze the results and provide a short summary of the code output."
        ),
        human_input_mode="NEVER",
    )

    # Experiment Reviewer
    experiment_reviewer = ConversableAgent(
        name="experiment_reviewer",
        llm_config={**llm_config, "response_format": ExperimentReviewer},
        system_message=(
            "You are a research scientist responsible for holistically reviewing the entire experiment pipeline, i.e., the generated code, the output, and the analysis w.r.t. the original experiment plan. "
            "Assess whether the experiment was faithfully implemented, i.e., whether the implementation follows the experiment plan without significant deviation and whether the hypothesis was in fact tested sufficiently. "
            "If you find issues or inconsistencies in any part of the experiment pipeline, return the success status as **false** and provide feedback about what is wrong. "
            "Otherwise, return the success status as **true** and provide a summary of the hypothesis, experiment results, and findings."
        ),
        human_input_mode="NEVER",
    )

    # Experiment Reviser
    experiment_reviser = ConversableAgent(
        name="experiment_reviser",
        llm_config={**llm_config, "response_format": Experiment},
        system_message=(
            "You are a research scientist revisiting the most recent experiment, which could not be conducted correctly due to issues in the code or the formulation of the experiment plan,"
            "as indicated by the reviewer. Your goal is to revise this failed experiment plan by addressing the issues and limitations pointed out by the reviewer. "
            "The revised experiment plan should still aim to validate the most recent hypothesis. "
            "Do not provide the code yourself but explain in natural language what the experiment should do for a programmer. "
            "Strictly use only the dataset provided and do not create synthetic data or columns that cannot be derived from the given columns. "
            f"{_datalib_note}"
            "The experiment should be creative, independent, and self-contained. "
            "Generally, in typical data-driven research, you will need to explore and visualize the data for possible high-level insights, clean, transform, or derive new variables from the dataset to be suited for the investigation, deep-dive into specific parts of the data for fine-grained analysis, perform data modeling, and run statistical tests."
        ),
        human_input_mode="NEVER",
    )

    ## Code Executor Setup
    modal_working_dir = None  # Track working directory for Modal sandbox

    if backend == "modal":
        # Use Modal sandbox for code execution
        if not bucket_path:
            raise ValueError("bucket_path is required when backend is 'modal'")

        import modal
        from asta_sandbox import CloudShare
        from asta_sandbox.backends.modal_ephemeral import ModalEphemeralExecutor
        from autodiscovery_modal.ipython_session import build_sandbox_image

        # Parse bucket path
        bucket_name, key_prefix = parse_bucket_path(bucket_path)

        # Same path the run's data has in the job container, so the sandbox shows the
        # store's layout unchanged.
        modal_mount_path = store_mount_path(key_prefix)
        modal_working_dir = modal_mount_path

        # Get Modal configuration from environment
        app_name = os.environ.get("MODAL_APP_NAME", "asta-autodiscovery")
        secret_name = os.environ.get("MODAL_BUCKET_SECRET", "example-bucket-secret")
        bucket_endpoint_url = os.environ.get("GCS_ENDPOINT_URL", "https://storage.googleapis.com")

        sandbox_image = build_sandbox_image(
            extra_packages=[
                "numpy",
                "pandas",
                "matplotlib",
                "matplotlib-inline",
                "seaborn",
                "scikit-learn",
                "scipy",
                "statsmodels",
            ]
        )

        cloud_share = CloudShare(
            dest=modal_mount_path,
            bucket=bucket_name,
            key_prefix=key_prefix,
            read_only=True,
            bucket_endpoint_url=bucket_endpoint_url,
            modal_secret=modal.Secret.from_name(secret_name),
        )

        # sandbox_timeout_s covers startup + execution + teardown, so it must exceed
        # code_timeout (the process-level limit). The buffer gives time for Modal sandbox
        # startup and for the finally-block terminate() to run after the process times out.
        _SANDBOX_OVERHEAD_S = 60
        modal_executor = ModalEphemeralExecutor(
            app_name=app_name,
            image=sandbox_image,
            environment={"DATASET_ROOT": modal_mount_path},
            sandbox_timeout_s=code_timeout + _SANDBOX_OVERHEAD_S,
        )
        # Each data-library directory, read-only, at the same path it has in the
        # job container, so the paths the agents were given work here too. Like
        # the run's data, it is in the bucket under the key its path implies.
        datalib_shares = [
            CloudShare(
                dest=path,
                bucket=bucket_name,
                key_prefix=datalib.store_key_prefix(path, STORE_MOUNT_ROOT),
                read_only=True,
                bucket_endpoint_url=bucket_endpoint_url,
                modal_secret=modal.Secret.from_name(secret_name),
            )
            for path in datalib_dirs
        ]
        _run_async(modal_executor.add_shares(cloud_share, *datalib_shares))

        executor = SandboxCodeExecutor(
            modal_executor,
            timeout=code_timeout,
            vision_model=vision_model,
            usage_tracker=usage_tracker,
        )
        print(
            f"Using Modal sandbox with bucket gs://{bucket_name}/{key_prefix} mounted at {modal_mount_path}"
        )
        print(f"Working directory will be: {modal_working_dir}")
    elif backend == "process":
        # Use isolated subprocess for code execution
        from code_execution import ProcessIPythonBackend

        process_backend = ProcessIPythonBackend(cwd=work_dir)
        executor = SandboxCodeExecutor(
            _IPythonBackendAdapter(process_backend),
            timeout=code_timeout,
            vision_model=vision_model,
            usage_tracker=usage_tracker,
        )
        print(f"Using process backend with work_dir: {work_dir}")
    elif backend == "local":
        # Local: the job's own Python environment, with no per-cell venv.
        # It still runs each cell in a child process (use_subprocess), which is
        # what LocalCommandLineCodeExecutor used to give us and what makes the
        # timeout enforceable; the difference from `process` is only that the
        # cell sees the job's installed packages rather than a curated venv.
        from code_execution import LocalIPythonBackend

        executor = SandboxCodeExecutor(
            _IPythonBackendAdapter(LocalIPythonBackend(cwd=work_dir), use_subprocess=True),
            timeout=code_timeout,
            vision_model=vision_model,
            usage_tracker=usage_tracker,
        )
        print(f"Using local backend with work_dir: {work_dir}")
    else:
        raise ValueError(f"unknown code execution backend: {backend!r}")

    # Create an agent with code executor configuration.
    code_executor = ConversableAgent(
        "code_executor",
        llm_config=False,
        code_execution_config={"executor": executor},
        human_input_mode="NEVER",
    )

    # Every backend gets the same transform: the agent's code, prefixed only by a
    # chdir into the directory the data is mounted at. Figure interpretation is not
    # injected here -- the executor returns figures as rich outputs and analyses them
    # in this process.
    sandbox_working_dir = code_transform_working_dir(backend, work_dir, modal_working_dir)
    transform_messages_capability = transform_messages.TransformMessages(
        transforms=[SimpleCodeBlockTransform(working_dir=sandbox_working_dir)]
    )
    transform_messages_capability.add_to_agent(code_executor)

    user_proxy = UserProxyAgent(
        name="user_proxy",
        description="Responsible for providing the initial query",
        code_execution_config=False,
        human_input_mode="NEVER",
    )

    agents = [
        experiment_generator,
        experiment_programmer,
        experiment_analyst,
        experiment_reviewer,
        experiment_reviser,
        code_executor,
        user_proxy,
    ]

    for agent in agents:
        if agent.llm_config is not False:
            agent.register_model_client(LiteLLMAG2Client)

    # Apply token limit to all agents
    for agent in agents:
        token_limit_capability.add_to_agent(agent)

    agents_dict = {agent.name: agent for agent in agents}
    return agents_dict

autodiscovery.transitions

SpeakerSelector

Source code in packages/autodiscovery/src/autodiscovery/transitions.py

class SpeakerSelector:
    def __init__(self):
        self.code_failure_count = 0
        self.experiment_revision_count = 0

    def select_next_speaker(self, last_speaker: Agent, groupchat: GroupChat) -> Agent | None:
        """Define a customized speaker selection function for the data exploration workflow.

        Args:
            last_speaker: The previous speaker in the conversation
            groupchat: The GroupChat instance containing conversation history

        Returns:
            The next agent to speak or None to end the conversation
        """
        messages = groupchat.messages

        if last_speaker.name == "user_proxy":
            return groupchat.agent_by_name("experiment_programmer")

        elif last_speaker.name == "experiment_programmer":
            return groupchat.agent_by_name("code_executor")

        elif last_speaker.name == "code_executor":
            return groupchat.agent_by_name("experiment_code_analyst")

        elif last_speaker.name == "experiment_code_analyst":
            # Check if experiment failed based on structured response
            content = messages[-1].get("content", "")
            try:
                response = json.loads(content)
            except json.JSONDecodeError:
                # If JSON parsing fails, treat it as an error
                response = {"success": False, "analysis": "Error parsing response"}
            if not response.get("success", False) and self.code_failure_count < 6:
                self.code_failure_count += 1
                return groupchat.agent_by_name("experiment_programmer")
            else:
                self.code_failure_count = 0
                return groupchat.agent_by_name("experiment_reviewer")

        elif last_speaker.name == "experiment_reviewer":
            content = messages[-1].get("content", "")
            try:
                response = json.loads(content)
            except json.JSONDecodeError:
                response = {"success": False, "feedback": "Error parsing reviewer response"}
            if not response.get("success", True) and self.experiment_revision_count < 1:
                self.experiment_revision_count += 1
                return groupchat.agent_by_name("experiment_reviser")
            else:
                self.experiment_revision_count = 0
                return groupchat.agent_by_name("experiment_generator")

        if last_speaker.name == "experiment_reviser":
            return groupchat.agent_by_name("experiment_programmer")

        elif last_speaker.name == "experiment_generator":
            return None

        return None

select_next_speaker(last_speaker, groupchat)

Define a customized speaker selection function for the data exploration workflow.

Parameters:

Name	Type	Description	Default
last_speaker	Agent	The previous speaker in the conversation	required
groupchat	GroupChat	The GroupChat instance containing conversation history	required

Returns:

Type	Description
Agent | None	The next agent to speak or None to end the conversation

Source code in packages/autodiscovery/src/autodiscovery/transitions.py

def select_next_speaker(self, last_speaker: Agent, groupchat: GroupChat) -> Agent | None:
    """Define a customized speaker selection function for the data exploration workflow.

    Args:
        last_speaker: The previous speaker in the conversation
        groupchat: The GroupChat instance containing conversation history

    Returns:
        The next agent to speak or None to end the conversation
    """
    messages = groupchat.messages

    if last_speaker.name == "user_proxy":
        return groupchat.agent_by_name("experiment_programmer")

    elif last_speaker.name == "experiment_programmer":
        return groupchat.agent_by_name("code_executor")

    elif last_speaker.name == "code_executor":
        return groupchat.agent_by_name("experiment_code_analyst")

    elif last_speaker.name == "experiment_code_analyst":
        # Check if experiment failed based on structured response
        content = messages[-1].get("content", "")
        try:
            response = json.loads(content)
        except json.JSONDecodeError:
            # If JSON parsing fails, treat it as an error
            response = {"success": False, "analysis": "Error parsing response"}
        if not response.get("success", False) and self.code_failure_count < 6:
            self.code_failure_count += 1
            return groupchat.agent_by_name("experiment_programmer")
        else:
            self.code_failure_count = 0
            return groupchat.agent_by_name("experiment_reviewer")

    elif last_speaker.name == "experiment_reviewer":
        content = messages[-1].get("content", "")
        try:
            response = json.loads(content)
        except json.JSONDecodeError:
            response = {"success": False, "feedback": "Error parsing reviewer response"}
        if not response.get("success", True) and self.experiment_revision_count < 1:
            self.experiment_revision_count += 1
            return groupchat.agent_by_name("experiment_reviser")
        else:
            self.experiment_revision_count = 0
            return groupchat.agent_by_name("experiment_generator")

    if last_speaker.name == "experiment_reviser":
        return groupchat.agent_by_name("experiment_programmer")

    elif last_speaker.name == "experiment_generator":
        return None

    return None

autodiscovery.structured_outputs

ExperimentList

Bases: BaseModel

A collection of experiments.

Attributes:

Name	Type	Description
experiments	list[Experiment]	List of Experiment objects

Source code in packages/autodiscovery/src/autodiscovery/structured_outputs.py

class ExperimentList(BaseModel):
    """A collection of experiments.

    Attributes:
        experiments (list[Experiment]): List of Experiment objects
    """

    experiments: list[Experiment]

Experiment

Bases: BaseModel

Represents an experiment with a hypothesis and corresponding experiment plan.

Attributes:

Name	Type	Description
hypothesis	str	A natural-language hypothesis representing an assertion about the world
experiment_plan	ExperimentPlan	The structured experiment plan to verify the hypothesis

Source code in packages/autodiscovery/src/autodiscovery/structured_outputs.py

class Experiment(BaseModel):
    """Represents an experiment with a hypothesis and corresponding experiment plan.

    Attributes:
        hypothesis (str): A natural-language hypothesis representing an assertion about the world
        experiment_plan (ExperimentPlan): The structured experiment plan to verify the hypothesis
    """

    hypothesis: str
    experiment_plan: ExperimentPlan

autodiscovery.deduplication

dedupe(nodes_or_json_path, *, model, embedding_model, n_samples=10, merge_threshold=0.7, seed=42, rep_mode='biggest', embedding_dimensions=None, n_nodes=None, verbose=False, log_comparisons_fname=None, usage_tracker=None)

Deduplicate hypotheses using embeddings + LLM similarity checks.

Parameters:

Name	Type	Description	Default
nodes_or_json_path		Nodes list or path to nodes JSON.	required
n_samples		Number of LLM samples per comparison.	10
merge_threshold		Merge threshold for LLM votes.	0.7
seed		Random seed.	42
rep_mode		Representative selection mode for merged clusters.	'biggest'
model		LLM model, as litellm's <provider>/<model>.	required
embedding_model		Embedding model, as litellm's <provider>/<model>.	required
embedding_dimensions		Optional embedding dimensions override.	None
n_nodes		Optional cap on nodes processed.	None
verbose		Whether to print verbose details.	False
log_comparisons_fname		Optional JSON path to log LLM comparisons.	None
usage_tracker	UsageTracker | None	Optional usage tracker.	None

Returns:

Type	Description
	Tuple of (deduplicated nodes list, duplicates mapping).

Source code in packages/autodiscovery/src/autodiscovery/deduplication.py

def dedupe(
    nodes_or_json_path,
    *,
    model,
    embedding_model,
    n_samples=10,
    merge_threshold=0.7,
    seed=42,
    rep_mode="biggest",
    embedding_dimensions=None,
    n_nodes=None,
    verbose=False,
    log_comparisons_fname=None,
    usage_tracker: UsageTracker | None = None,
):
    """Deduplicate hypotheses using embeddings + LLM similarity checks.

    Args:
        nodes_or_json_path: Nodes list or path to nodes JSON.
        n_samples: Number of LLM samples per comparison.
        merge_threshold: Merge threshold for LLM votes.
        seed: Random seed.
        rep_mode: Representative selection mode for merged clusters.
        model: LLM model, as litellm's ``<provider>/<model>``.
        embedding_model: Embedding model, as litellm's ``<provider>/<model>``.
        embedding_dimensions: Optional embedding dimensions override.
        n_nodes: Optional cap on nodes processed.
        verbose: Whether to print verbose details.
        log_comparisons_fname: Optional JSON path to log LLM comparisons.
        usage_tracker: Optional usage tracker.

    Returns:
        Tuple of (deduplicated nodes list, duplicates mapping).
    """
    random.seed(seed)
    np.random.seed(seed)

    from autodiscovery.mcts_utils import get_nodes  # Importing here to avoid circular import issues

    nodes_list = get_nodes(nodes_or_json_path)[:n_nodes]
    # Remove nodes without hypotheses
    nodes_list = [node for node in nodes_list if node.get("hypothesis", None) is not None]

    dedup_hyp, hyp_to_index, orig_to_dedup = [], {}, []

    # Deduplicate hypotheses by exact match
    for node in nodes_list:
        hyp = node["hypothesis"]  # Hypothesis string
        if hyp not in hyp_to_index:
            hyp_to_index[hyp] = len(dedup_hyp)
            dedup_hyp.append(hyp)
        orig_to_dedup.append(hyp_to_index[hyp])
    n_dedup = len(dedup_hyp)

    # Generate embeddings for deduplicated hypotheses
    embeds = np.array(
        get_embedding(
            dedup_hyp,
            model=embedding_model,
            dimensions=embedding_dimensions,
            n_attempts=3,
            usage_tracker=usage_tracker,
        )
    )

    # Initialize assignment structures
    clusters = {i: [i] for i in range(n_dedup)}
    cluster_assignment = {i: i for i in range(n_dedup)}
    hac_to_current = {i: i for i in range(n_dedup)}
    cluster_rep = {i: i for i in range(n_dedup)}

    # Map dedup indices back to original node indices (for duplicates output).
    dedup_to_orig = {i: [] for i in range(n_dedup)}
    for orig_idx, dedup_idx in enumerate(orig_to_dedup):
        dedup_to_orig[dedup_idx].append(orig_idx)

    def build_output_from_clusters(in_clusters, in_cluster_rep):
        """Build deduplicated nodes and duplicates mapping from clusters."""
        deduped_nodes, duplicates = [], {}
        for cluster_id, cluster in in_clusters.items():
            rep_dedup_idx = in_cluster_rep[cluster_id]
            rep_orig_idx = dedup_to_orig[rep_dedup_idx][0]
            node_copy = copy.deepcopy(nodes_list[rep_orig_idx])
            dup_orig_indices = []
            for dedup_idx in cluster:
                for orig_idx in dedup_to_orig[dedup_idx]:
                    if orig_idx != rep_orig_idx:
                        dup_orig_indices.append(orig_idx)
            node_copy["duplicate_nodes"] = [nodes_list[i]["id"] for i in dup_orig_indices]
            duplicates[node_copy["id"]] = node_copy["duplicate_nodes"]
            deduped_nodes.append(node_copy)
        return deduped_nodes, duplicates

    if n_dedup < 2:
        return build_output_from_clusters(clusters, cluster_rep)

    # Perform HAC over LM embeddings and get the linkage matrix
    linkage_matrix = linkage(embeds, method="ward")

    # Iterate through the linkage matrix to additionally merge clusters based on LLM decisions
    pbar = tqdm(linkage_matrix, desc="Deduplicating")
    pbar.set_postfix({"n_clusters": len(clusters)})
    llm_comparisons = []
    for r, row in enumerate(pbar):
        hac_node_id = n_dedup + r
        left_hac, right_hac = int(row[0]), int(row[1])
        left_current = hac_to_current.get(left_hac)
        right_current = hac_to_current.get(right_hac)
        if left_current is None or right_current is None or left_current == right_current:
            hac_to_current[hac_node_id] = (
                left_current if left_current is not None else right_current
            )
            continue
        rep_left, rep_right = cluster_rep[left_current], cluster_rep[right_current]
        # struct_left, struct_right = dedup_struct_hyp[rep_left], dedup_struct_hyp[rep_right]
        struct_left, struct_right = dedup_hyp[rep_left], dedup_hyp[rep_right]
        # Get the LLM merge decision
        llm_decision = get_llm_merge_decision(
            struct_left,
            struct_right,
            n_samples=n_samples,
            threshold=merge_threshold,
            model=model,
            usage_tracker=usage_tracker,
        )
        if verbose:
            print(f"""\n\n
Cluster left (size: {len(clusters[left_current])}):
{struct_left}

Cluster right (size: {len(clusters[right_current])}):
{struct_right}

LLM Decision: {"Merge" if llm_decision else "Do not merge"}\n\n""")
        if log_comparisons_fname is not None:
            llm_comparisons.append(
                {
                    "left_size": len(clusters[left_current]),
                    "right_size": len(clusters[right_current]),
                    "left_hypothesis": struct_left,
                    "right_hypothesis": struct_right,
                    "llm_decision": llm_decision,
                }
            )

        if llm_decision:
            if rep_mode == "random":
                new_rep = random.choice([rep_left, rep_right])
            elif rep_mode == "biggest":
                new_rep = (
                    rep_left
                    if len(clusters[left_current]) >= len(clusters[right_current])
                    else rep_right
                )
            else:
                raise NotImplementedError
            merged_cluster_id = min(left_current, right_current)
            other_cluster_id = max(left_current, right_current)
            clusters[merged_cluster_id] += clusters[other_cluster_id]
            for idx in clusters[merged_cluster_id]:
                cluster_assignment[idx] = merged_cluster_id
            cluster_rep[merged_cluster_id] = new_rep
            del clusters[other_cluster_id]
            del cluster_rep[other_cluster_id]
            hac_to_current[hac_node_id] = merged_cluster_id
        else:
            hac_to_current[hac_node_id] = None

        # Update pbar with number of clusters
        pbar.set_postfix({"n_clusters": len(clusters)})

    if log_comparisons_fname is not None:
        with open(log_comparisons_fname, "w") as f:
            json.dump(llm_comparisons, f, indent=2)
        print(f"LLM comparisons logged to {log_comparisons_fname}")

    # Return final deduplicated nodes and the list of clusters
    return build_output_from_clusters(clusters, cluster_rep)