Skip to content

AutoDiscovery Changes Since NeurIPS

Jump to bottom
dirkraft edited this page Jul 1, 2026 · 1 revision

Changes to the core AutoDiscovery unit since the NeurIPS repo

This page lightly describes what's changed to the core AutoDiscovery unit since the NeurIPS repo (specifically its src/), which now lives down in packages/autodiscovery/src/autodiscovery.

Snapshot as of 2026-07-01. This is a point-in-time comparison and is not kept up to date — more may have changed since this was written. Diff the two paths directly for the current state.

TL;DR

Same discovery algorithm and Bayesian-surprise core you already know from the NeurIPS repo — now it speaks Gemini/Vertex, retries robustly, tracks token cost everywhere, expands nodes in parallel across sandboxed execution backends, and ships an easy CLI + HTML report. The only result-affecting logic changes are in dedup (stricter), beliefs (retry fix, fewer samples), and a couple of MCTS defaults/metrics.

The NeurIPS repo is the research reference implementation. The version in asta-autodiscovery is the same MCTS-based Bayesian-surprise discovery engine, but hardened and generalized into a production package inside the larger Asta app suite. The core science (Bayesian surprise math, MCTS structure) is essentially unchanged — the churn is almost entirely productionization: multi-provider LLM support, parallelism, cost tracking, execution sandboxes, and reporting.

Orientation: the mechanical/packaging shift (ignore when reading diffs)

Every file carries cosmetic churn you can tune out: Black reformatting, from src.* → from autodiscovery.*, Optional[X]/List[X] → X | None/list[X], docstring reflow. The package is now a uv workspace member, invoked as python -m autodiscovery.run instead of the old autodiscovery console script, and version comes from package metadata (the hardcoded __version__ was dropped).

The five substantive themes

1. Multi-provider LLM support (OpenAI-only → Gemini via Vertex)

This is the biggest theme. The default models flipped from OpenAI (o4-mini/gpt-4o) to Gemini (gemini-3.1-pro-preview / gemini-3-flash-preview).

  • agents.py now branches on is_gemini_model() and routes Gemini through Vertex AI's OpenAI-compatible endpoint, using a Vertex access token as the API key.
  • New files vertex_client.py / vertex_config.py: build the Vertex base URL from project/location env vars and provide an OpenAI-client wrapper that refreshes Google ADC credentials before each call.
  • utils.py query path now fans Gemini calls out across a thread pool and applies reasoning-effort to Gemini too.

2. Centralized retry + token/cost tracking (both entirely new)

  • New llm_retry.py: tenacity-based backoff wrapping all LLM calls, with broad retryable-error classification across Google, OpenAI, requests, urllib3, honoring Retry-After. Provider-level max_retries set to 0 so this is the single retry path. Also handles Vertex token refresh mid-run.
  • New llm_usage.py: a UsageTracker threaded through everything (rewards, beliefs, dedup, agent chats, image analysis) that attributes token usage/cost by model/agent/node/component, with reasoning_tokens, and persists llm_usage_events.jsonl + llm_usage_summary.json. New --agent_usage_mode arg (per_response vs summary_delta) controls how AG2 usage is captured.

3. Parallelism + execution backends

  • run.py rewrote the core loop from "expand one node per iteration" to a batched, multi-threaded expansion driven by an experiment budget (--batch_size, --n_threads, via ThreadPoolExecutor), with per-node locks, per-thread agents/work-dirs, and per-node failure isolation (a failed node is logged and skipped instead of aborting the run — supported by the new tiny future_utils.py).
  • agents.py added a --backend arg (local/process/modal; process is now the default). New ModalSandboxExecutor runs code in an ephemeral Modal/asta-sandbox with a GCS bucket mounted read-only at /data (--bucket_path). Datasets are now symlinked into the work dir instead of copied.
  • Vision model is configurable (--vision_model) instead of hardcoded gpt-4o.

4. Behavior/correctness changes worth flagging

These change results, not just plumbing:

  • beliefs.py: the Bayesian math is byte-identical after reformatting, BUT a break-on-success was added to the belief retry loop — previously a good result could be overwritten by continued retries (correctness + cost fix). Also --n_belief_samples default dropped 30 → 5, and belief calls got a separate --belief_reasoning_effort.
  • deduplication.py: dedup now keys off the raw hypothesis string instead of the structured decomposition, and the LLM merge prompt was rewritten to be stricter/more conservative (different clause ⇒ not a duplicate) → expect fewer merges. Output contract changed: dedupe() returns (deduped, duplicates) and writes duplicate_nodes.json. Nodes now use real IDs.
  • run.py: k_parents default 3 → 10; added a new normalized_surprisal metric per node (with a boolean_cat theoretical-max calc). MCTSNode gained created_at and normalized_surprisal fields.
  • dataset.py: new asta metadata format is now the default (was dbench); dataset paths are now metadata-relative rather than basename.

5. New user-facing surface: easy CLI + HTML reports

  • New easy.py: a simplified flat CLI (auto-discovery) — point it at raw CSV/TSV files, it auto-sniffs headers, generates the metadata JSON, runs, and produces a report. No hand-written metadata file needed.
  • New report.py (~850 lines): generates a self-contained static HTML report (report/index.html + data.js + figures/) with a D3 tree viz and markdown rendering. This replaces the old mcts_viz.html (the only file deleted).

New/removed files at a glance

Added in asta Purpose
llm_retry.py centralized backoff + Vertex token refresh
llm_usage.py token/cost accounting
vertex_client.py / vertex_config.py Vertex/Gemini plumbing
easy.py simplified "point at CSVs" CLI
report.py HTML report generator
future_utils.py concurrent-future helper for parallel expansion

Removed: mcts_viz.html (superseded by report.py).

Essentially unchanged (cosmetic only): structured_outputs.py, transitions.py, nodes_to_csv.py, log_utils.py. logger.py lost its logprob/choice-logging (the old autogen.ChatCompletion dependency was dropped).

Clone this wiki locally