diff --git a/docs/get-started/example-agent.mdx b/docs/get-started/example-agent.mdx index 226f334cb9..edffb4eb27 100644 --- a/docs/get-started/example-agent.mdx +++ b/docs/get-started/example-agent.mdx @@ -34,33 +34,34 @@ uv run nemo services start --config packages/nmp_platform/config/local.yaml You should now be able to navigate to `http://localhost:8080` and see the NeMo Platform web UI. -## 3. Prepare Tau2 and run the airline agent +## 3. Prepare τ-Bench and run the airline agent -Now that we have a running NeMo Platform, it's time to load up some data! This example uses the τ-Bench from Sierra. It's a great representation of a simplified agent that performs a business-critical task. Our first step will be to download the tau2 repository and install the dependencies to run it. +Now that we have a running NeMo Platform, it's time to load up some data! This example uses the τ-Bench from Sierra. It's a great representation of a simplified agent that performs a business-critical task. Our first step will be to set up our environment variables so that we can run models properly. -Then, we'll run a script to go through 30 of the tasks from the Tau 2 Airline dataset. In order to make our example more realistic and representative of the kind of agent data we typically see in production, we won't record whether the tasks passed verification – we want to see if we can gain useful insights even without ground truth data. This could take up to 20 minutes, and will cost about a dollar. +```bash +cp plugins/nemo-experimentalist/examples/tau3-nooa-agent/.env.example plugins/nemo-experimentalist/examples/tau3-nooa-agent/.env +``` + +Then, update the `.env` file to have an appropriate API key for your inference endpoint, and source the `.env` file. + +```bash +source plugins/nemo-experimentalist/examples/tau3-nooa-agent/.env +``` + +Then, we'll run a script to download the benchmark datasets and put it into the appropriate directories: ```bash -git init tmp/tau2-bench -git -C tmp/tau2-bench remote add origin https://github.com/sierra-research/tau2-bench.git -git -C tmp/tau2-bench fetch --depth 1 origin 8ebb7499622fc2be9b9d510d6f7a7653461f4f29 -git -C tmp/tau2-bench checkout --detach FETCH_HEAD -uv --directory tmp/tau2-bench sync --frozen -uv --directory tmp/tau2-bench run --frozen tau2 check-data - -export NMP_BASE_URL=http://localhost:8080 -export INFERENCE_API_KEY=sk-... -export OPENAI_API_KEY="$INFERENCE_API_KEY" -export OPENAI_API_BASE=https://inference-api.nvidia.com/v1 - -uv run --directory plugins/nemo-insights --frozen \ - python -m testbed run tau2-airline \ - --base "$NMP_BASE_URL" \ - --set include_rewards=false \ - --set tau2_repo="$PWD/tmp/tau2-bench" +plugins/nemo-experimentalist/examples/tau3-nooa-agent/prepare-airline-datasets.sh ``` -The testbed runs the 30-task airline train split and ingests realistic traces into the `tau2-airline` workspace. At this point you can navigate to the traces view and see the traces from the agent. +Finally, we'll run our agent using the dataset tasks and record the data to NeMo Platform: + +```bash +uv run --frozen plugins/nemo-experimentalist/examples/tau3-nooa-agent/record_tau_airline_traces.py +``` + +At this point you can navigate to the [traces tab](http://localhost:8080/studio/workspaces/tau3-airline/intake/traces) and see the traces from the agent. + ## 4. Run the Analyst @@ -72,38 +73,28 @@ Let's run it: ```bash uv run --frozen nemo insights analyze \ - --agent tau2-airline \ - --workspace tau2-airline \ + --agent nemo-experimentalist-tau3-nooa \ + --workspace tau3-airline \ --base-url "$NMP_BASE_URL" ``` -This will take a few minutes to run. Once it's done, you can navigate to the insights page to see the issues the analyst discovered in the τ-Bench Airline Agent. +This will take a few minutes to run. Once it's done, you can navigate to the [insights tab](http://localhost:8080/studio/workspaces/tau3-airline/optimizer) to see the issues the analyst discovered in the τ-Bench Airline Agent. -## 5. Prepare a Tau3 Airline smoke dataset +## 5. Optimize performance with the experimentalist The next step is to improve the τ-Bench agent using the experimentalist. The experimentalist will run your evals, debug failures using trace data, understand the root cause of the failure and attempt to fix it. After it makes the change, it will run the evals again to validate whether the change improved the -performance of your agent on the evalaution. +performance of your agent on the evaluation. -This is a shortened example that only uses a few tasks, but it can still take up to an hour to finish. First, set your environment variables and set up the task dataset: +This is a shortened example that only uses a few tasks, but it can still take up +to an hour to finish. First we will set up a new workspace for the optimization process: ```bash -export OPENAI_API_KEY="$INFERENCE_API_KEY" -export OPENAI_BASE_URL=https://inference-api.nvidia.com/v1 -export TAU2_USER_MODEL=openai/openai/openai/gpt-5-mini -export TAU2_NL_ASSERTIONS_MODEL=openai/openai/openai/gpt-5-mini -export AUT_MODEL_NAME=openai/openai/openai/gpt-5-mini -export EXPERIMENTALIST_SMART_MODEL_NAME=openai/openai/openai/gpt-5-mini -export EXPERIMENTALIST_MID_MODEL_NAME=openai/openai/openai/gpt-5-mini -export EXPERIMENTALIST_FAST_MODEL_NAME=openai/openai/openai/gpt-5-mini - uv run --frozen nemo workspaces create canonical-tau3-airline \ --description "Tau3 Airline Experimentalist runs" \ --exist-ok - -plugins/nemo-experimentalist/examples/tau3-nooa-agent/prepare-airline-smoke.sh ``` Now we're ready to run the experimentalist! @@ -113,8 +104,8 @@ uv run --frozen nemo experimentalist run \ --no-insight \ --agent plugins/nemo-experimentalist/examples/tau3-nooa-agent \ --agent-spec plugins/nemo-experimentalist/examples/tau3-nooa-agent/AGENT-SPEC.md \ - --train-dataset plugins/nemo-experimentalist/tmp/tau3-airline-smoke/train \ - --validation-dataset plugins/nemo-experimentalist/tmp/tau3-airline-smoke/validation \ + --train-dataset plugins/nemo-experimentalist/tmp/tau3-airline/experimentalist/train \ + --validation-dataset plugins/nemo-experimentalist/tmp/tau3-airline/experimentalist/validation \ --workspace canonical-tau3-airline \ --framework-skills plugins/nemo-experimentalist/framework-skills/nooa \ --config plugins/nemo-experimentalist/examples/tau3-nooa-agent/experimentalist-smoke.yaml \ @@ -123,6 +114,7 @@ uv run --frozen nemo experimentalist run \ ``` The trace records include the Experimentalist evaluation ID and Tau3 task ID. -After the run completes, inspect `eval-and-optimize/run.json` for the selected -winner and compare the `agent-0` and `agent-1` directories to review the code -change that was evaluated. +After the run completes, inspect +`plugins/nemo-experimentalist/tmp/tau3-airline-experimentalist/eval-and-optimize/run.json` +for the selected winner and compare the `agent-0` and `agent-1` directories to +review the code change that was evaluated. diff --git a/plugins/nemo-experimentalist/examples/tau3-nooa-agent/.env.example b/plugins/nemo-experimentalist/examples/tau3-nooa-agent/.env.example new file mode 100644 index 0000000000..185f9c661a --- /dev/null +++ b/plugins/nemo-experimentalist/examples/tau3-nooa-agent/.env.example @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +export NMP_BASE_URL=http://localhost:8080 +export INFERENCE_API_BASE=https://inference-api.nvidia.com/v1 +export INFERENCE_API_KEY=sk-... + +export OPENAI_API_KEY="$INFERENCE_API_KEY" +export OPENAI_BASE_URL="$INFERENCE_API_BASE" +export EXPERIMENTALIST_API_KEY="$INFERENCE_API_KEY" +export EXPERIMENTALIST_API_BASE="$INFERENCE_API_BASE" + +# Upstream Tau3 currently reads these compatibility variable names. +export TAU2_USER_MODEL=openai/openai/openai/gpt-5.6-luna +export TAU2_NL_ASSERTIONS_MODEL=openai/openai/openai/gpt-5.6-luna +export AUT_MODEL_NAME=openai/openai/openai/gpt-5.6-luna + +export EXPERIMENTALIST_SMART_MODEL_NAME=openai/openai/openai/gpt-5.6-sol +export EXPERIMENTALIST_MID_MODEL_NAME=openai/openai/openai/gpt-5.6-terra +export EXPERIMENTALIST_FAST_MODEL_NAME=openai/openai/openai/gpt-5.6-luna diff --git a/plugins/nemo-experimentalist/examples/tau3-nooa-agent/prepare-airline-datasets.sh b/plugins/nemo-experimentalist/examples/tau3-nooa-agent/prepare-airline-datasets.sh new file mode 100755 index 0000000000..96e02d2aee --- /dev/null +++ b/plugins/nemo-experimentalist/examples/tau3-nooa-agent/prepare-airline-datasets.sh @@ -0,0 +1,112 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +set -euo pipefail +shopt -s nullglob + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +PLUGIN_ROOT="$(cd -- "${SCRIPT_DIR}/../.." && pwd)" +OUTPUT_ROOT="${1:-${PLUGIN_ROOT}/tmp/tau3-airline}" +SOURCE_ROOT="${OUTPUT_ROOT}/source" +DATASET_ROOT="${SOURCE_ROOT}/tau3-bench" +INSIGHTS_ROOT="${OUTPUT_ROOT}/insights" +EXPERIMENTALIST_ROOT="${OUTPUT_ROOT}/experimentalist" +TRAIN_ROOT="${EXPERIMENTALIST_ROOT}/train" +VALIDATION_ROOT="${EXPERIMENTALIST_ROOT}/validation" +DATASET_REF="sierra-research/tau3-bench@1" + +INSIGHTS_TASKS=( + "tau3-bench__tau3-airline-0" + "tau3-bench__tau3-airline-1" + "tau3-bench__tau3-airline-4" + "tau3-bench__tau3-airline-5" + "tau3-bench__tau3-airline-9" + "tau3-bench__tau3-airline-10" + "tau3-bench__tau3-airline-12" + "tau3-bench__tau3-airline-14" + "tau3-bench__tau3-airline-17" + "tau3-bench__tau3-airline-20" + "tau3-bench__tau3-airline-23" + "tau3-bench__tau3-airline-27" + "tau3-bench__tau3-airline-33" + "tau3-bench__tau3-airline-34" + "tau3-bench__tau3-airline-38" + "tau3-bench__tau3-airline-39" + "tau3-bench__tau3-airline-41" + "tau3-bench__tau3-airline-42" + "tau3-bench__tau3-airline-46" + "tau3-bench__tau3-airline-47" +) +EXPERIMENTALIST_TRAIN_TASKS=( + "tau3-bench__tau3-airline-0" + "tau3-bench__tau3-airline-20" + "tau3-bench__tau3-airline-39" +) +EXPERIMENTALIST_VALIDATION_TASKS=( + "tau3-bench__tau3-airline-3" + "tau3-bench__tau3-airline-36" +) + +if [[ ! -f "${DATASET_ROOT}/${INSIGHTS_TASKS[0]}/task.toml" ]]; then + mkdir -p "${SOURCE_ROOT}" + ( + cd "${PLUGIN_ROOT}" + uv run --frozen harbor download "${DATASET_REF}" \ + --output-dir "${SOURCE_ROOT}" \ + --export \ + --overwrite + ) +fi + +validate_source() { + local task_names=("$@") + for task_name in "${task_names[@]}"; do + if [[ ! -f "${DATASET_ROOT}/${task_name}/task.toml" ]]; then + echo "Downloaded dataset is missing ${task_name}" >&2 + exit 1 + fi + done +} + +prepare_split() { + local split_root="$1" + shift + local task_names=("$@") + + if [[ -d "${split_root}" ]]; then + local existing_tasks=("${split_root}"/*/task.toml) + if [[ ${#existing_tasks[@]} -eq ${#task_names[@]} ]]; then + for task_name in "${task_names[@]}"; do + if [[ ! -f "${split_root}/${task_name}/task.toml" ]]; then + echo "Existing split is not the expected dataset: ${split_root}" >&2 + exit 1 + fi + done + echo "Reusing ${split_root}" + return + fi + echo "Existing split is incomplete: ${split_root}" >&2 + exit 1 + fi + + mkdir -p "${split_root}" + for task_name in "${task_names[@]}"; do + local source_task="${DATASET_ROOT}/${task_name}" + cp -R "${source_task}" "${split_root}/${task_name}" + done +} + +validate_source \ + "${INSIGHTS_TASKS[@]}" \ + "${EXPERIMENTALIST_TRAIN_TASKS[@]}" \ + "${EXPERIMENTALIST_VALIDATION_TASKS[@]}" + +prepare_split "${INSIGHTS_ROOT}" "${INSIGHTS_TASKS[@]}" +prepare_split "${TRAIN_ROOT}" "${EXPERIMENTALIST_TRAIN_TASKS[@]}" +prepare_split "${VALIDATION_ROOT}" "${EXPERIMENTALIST_VALIDATION_TASKS[@]}" + +echo "Tau3 Airline datasets are ready:" +echo " insights: ${INSIGHTS_ROOT}" +echo " experimentalist train: ${TRAIN_ROOT}" +echo " experimentalist validation: ${VALIDATION_ROOT}" diff --git a/plugins/nemo-experimentalist/examples/tau3-nooa-agent/prepare-airline-smoke.sh b/plugins/nemo-experimentalist/examples/tau3-nooa-agent/prepare-airline-smoke.sh deleted file mode 100755 index f57f437e6a..0000000000 --- a/plugins/nemo-experimentalist/examples/tau3-nooa-agent/prepare-airline-smoke.sh +++ /dev/null @@ -1,75 +0,0 @@ -#!/usr/bin/env bash -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -set -euo pipefail -shopt -s nullglob - -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" -PLUGIN_ROOT="$(cd -- "${SCRIPT_DIR}/../.." && pwd)" -OUTPUT_ROOT="${1:-${PLUGIN_ROOT}/tmp/tau3-airline-smoke}" -SOURCE_ROOT="${OUTPUT_ROOT}/source" -DATASET_ROOT="${SOURCE_ROOT}/tau3-bench" -TRAIN_ROOT="${OUTPUT_ROOT}/train" -VALIDATION_ROOT="${OUTPUT_ROOT}/validation" -DATASET_REF="sierra-research/tau3-bench@1" - -TRAIN_TASKS=( - "tau3-bench__tau3-airline-0" - "tau3-bench__tau3-airline-20" - "tau3-bench__tau3-airline-39" -) -VALIDATION_TASKS=( - "tau3-bench__tau3-airline-3" - "tau3-bench__tau3-airline-36" -) - -if [[ ! -f "${DATASET_ROOT}/${TRAIN_TASKS[0]}/task.toml" ]]; then - mkdir -p "${SOURCE_ROOT}" - ( - cd "${PLUGIN_ROOT}" - uv run --frozen harbor download "${DATASET_REF}" \ - --output-dir "${SOURCE_ROOT}" \ - --export \ - --overwrite - ) -fi - -prepare_split() { - local split_root="$1" - shift - local task_names=("$@") - - if [[ -d "${split_root}" ]]; then - local existing_tasks=("${split_root}"/*/task.toml) - if [[ ${#existing_tasks[@]} -eq ${#task_names[@]} ]]; then - for task_name in "${task_names[@]}"; do - if [[ ! -f "${split_root}/${task_name}/task.toml" ]]; then - echo "Existing split is not the expected dataset: ${split_root}" >&2 - exit 1 - fi - done - echo "Reusing ${split_root}" - return - fi - echo "Existing split is incomplete: ${split_root}" >&2 - exit 1 - fi - - mkdir -p "${split_root}" - for task_name in "${task_names[@]}"; do - local source_task="${DATASET_ROOT}/${task_name}" - if [[ ! -f "${source_task}/task.toml" ]]; then - echo "Downloaded dataset is missing ${task_name}" >&2 - exit 1 - fi - cp -R "${source_task}" "${split_root}/${task_name}" - done -} - -prepare_split "${TRAIN_ROOT}" "${TRAIN_TASKS[@]}" -prepare_split "${VALIDATION_ROOT}" "${VALIDATION_TASKS[@]}" - -echo "Tau3 Airline smoke datasets are ready:" -echo " train: ${TRAIN_ROOT}" -echo " validation: ${VALIDATION_ROOT}" diff --git a/plugins/nemo-experimentalist/examples/tau3-nooa-agent/record_tau_airline_traces.py b/plugins/nemo-experimentalist/examples/tau3-nooa-agent/record_tau_airline_traces.py new file mode 100644 index 0000000000..086a5e57aa --- /dev/null +++ b/plugins/nemo-experimentalist/examples/tau3-nooa-agent/record_tau_airline_traces.py @@ -0,0 +1,252 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Run the Tau3 Airline agent and upload only its execution traces to Intake.""" + +import argparse +import asyncio +import json +import os +import secrets +from datetime import UTC, datetime +from pathlib import Path + +from nemo_experimentalist_plugin.client import make_client +from nemo_experimentalist_plugin.experimentalist.components.evaluator.harbor import ( + HarborDataset, + HarborEvaluator, + HarborEvaluatorConfig, +) +from nemo_experimentalist_plugin.experimentalist.components.evaluator.models import TrialResult, local_path_from_uri +from nemo_experimentalist_plugin.experimentalist.otlp import jsonl_to_protobuf, read_trace_id +from nemo_platform import AsyncNeMoPlatform, NotFoundError + +SCRIPT_DIR = Path(__file__).resolve().parent +PLUGIN_ROOT = SCRIPT_DIR.parents[1] +DEFAULT_DATASET = PLUGIN_ROOT / "tmp" / "tau3-airline" / "insights" +DEFAULT_OUTPUT = PLUGIN_ROOT / "tmp" / "tau3-airline-insights" +DEFAULT_MODEL = "openai/openai/openai/gpt-5-mini" +DEFAULT_WORKSPACE = "tau3-airline" +DEFAULT_AGENT_NAME = "nemo-experimentalist-tau3-nooa" +DEFAULT_AGENT_VERSION = "1.0.0" + + +def _positive_int(value: str) -> int: + parsed = int(value) + if parsed < 1: + raise argparse.ArgumentTypeError("must be at least 1") + return parsed + + +def _experiment_id() -> str: + stamp = datetime.now(UTC).strftime("%Y%m%d-%H%M%S") + return f"tau3-airline-{stamp}-{secrets.token_hex(2)}" + + +def _configure_models(*, model: str, user_model: str, api_base: str) -> None: + api_key = os.environ.get("INFERENCE_API_KEY") or os.environ.get("EXPERIMENTALIST_API_KEY") + if not api_key: + raise RuntimeError("INFERENCE_API_KEY or EXPERIMENTALIST_API_KEY is required") + + normalized_base = api_base.rstrip("/") + openai_base = normalized_base if normalized_base.endswith("/v1") else f"{normalized_base}/v1" + os.environ["OPENAI_API_KEY"] = api_key + os.environ["OPENAI_BASE_URL"] = openai_base + os.environ["AUT_MODEL_NAME"] = model + os.environ["TAU2_USER_MODEL"] = user_model + os.environ["TAU2_NL_ASSERTIONS_MODEL"] = user_model + + +async def _upload_trials( + client: AsyncNeMoPlatform, + trials: list[TrialResult], + *, + workspace: str, + experiment_id: str, + agent_name: str, + agent_version: str, + model: str, +) -> dict[str, str]: + trace_ids: dict[str, str] = {} + url = f"/apis/intake/v2/workspaces/{workspace}/ingest/otlp/v1/traces" + + for trial in trials: + if trial.trace is None: + raise RuntimeError(f"Trial {trial.id} did not produce an agent execution trace") + trace_id = read_trace_id(trial.trace) + if trace_id in trace_ids.values(): + raise RuntimeError(f"Trace {trace_id} was produced by more than one trial") + + attrs = { + "nemo.experiment.id": experiment_id, + "nemo.test_case.id": trial.task_id, + "nemo.trial.id": trial.id, + "gen_ai.agent.name": agent_name, + "gen_ai.agent.version": agent_version, + "gen_ai.request.model": model, + } + path = local_path_from_uri(trial.trace.uri, context="Agent execution trace") + payloads = jsonl_to_protobuf(path, extra_resource_attrs=attrs) + if not payloads: + raise RuntimeError(f"Trial {trial.id} produced an empty agent execution trace") + for payload in payloads: + await client.post( + url, + cast_to=object, + content=payload, + options={"headers": {"Content-Type": "application/x-protobuf"}}, + ) + trace_ids[trial.id] = trace_id + + return trace_ids + + +async def _wait_for_traces( + client: AsyncNeMoPlatform, + trace_ids: set[str], + *, + workspace: str, + retries: int = 6, +) -> None: + pending = set(trace_ids) + delay = 1.0 + for attempt in range(retries): + for trace_id in tuple(pending): + try: + await client.intake.traces.retrieve(trace_id, workspace=workspace) + except NotFoundError: + continue + pending.remove(trace_id) + if not pending: + return + if attempt < retries - 1: + await asyncio.sleep(delay) + delay *= 2 + raise RuntimeError(f"Uploaded traces did not become readable from Intake: {sorted(pending)}") + + +async def run(args: argparse.Namespace) -> Path: + dataset_path = args.dataset.expanduser().resolve() + agent_path = args.agent.expanduser().resolve() + if not dataset_path.is_dir(): + raise FileNotFoundError(f"Tau3 dataset not found: {dataset_path}") + if not agent_path.is_dir(): + raise FileNotFoundError(f"Tau3 agent not found: {agent_path}") + + dataset = HarborDataset.from_path(dataset_path) + if args.task_ids: + dataset = dataset.subset(args.task_ids) + if len(dataset.tasks) != args.expected_task_count: + raise RuntimeError( + f"Expected {args.expected_task_count} Tau3 tasks in {dataset_path}, found {len(dataset.tasks)}" + ) + + experiment_id = args.experiment_id or _experiment_id() + run_dir = args.output.expanduser().resolve() / experiment_id + if run_dir.exists(): + raise FileExistsError(f"Output directory already exists: {run_dir}") + + _configure_models(model=args.model, user_model=args.user_model, api_base=args.api_base) + client = make_client(args.base_url) + try: + await client.workspaces.create( + name=args.workspace, + description="Tau3 Airline agent traces for Insights", + exist_ok=True, + ) + run_dir.mkdir(parents=True) + + options = HarborEvaluatorConfig( + job_name="tau3-airline-insights", + jobs_dir=Path("results"), + n_attempts=1, + n_concurrent_trials=args.concurrency, + quiet=not args.verbose, + agent_setup_timeout_multiplier=2.0, + environment_build_timeout_multiplier=3.0, + ) + result = await HarborEvaluator(experiment_dir=run_dir).run( + agent=agent_path, + dataset=dataset, + options=options, + ) + trials = list(result.trials) + if len(trials) != args.expected_task_count: + raise RuntimeError( + f"Expected {args.expected_task_count} Tau3 trial results, found {len(trials)} in {run_dir}" + ) + + trace_ids = await _upload_trials( + client, + trials, + workspace=args.workspace, + experiment_id=experiment_id, + agent_name=args.agent_name, + agent_version=args.agent_version, + model=args.model, + ) + await _wait_for_traces(client, set(trace_ids.values()), workspace=args.workspace) + finally: + await client.close() + + summary_path = run_dir / "uploaded-traces.json" + summary_path.write_text( + json.dumps( + { + "experiment_id": experiment_id, + "workspace": args.workspace, + "agent_name": args.agent_name, + "agent_version": args.agent_version, + "model": args.model, + "trace_count": len(trace_ids), + "traces": [ + { + "trial_id": trial.id, + "task_id": trial.task_id, + "trace_id": trace_ids[trial.id], + } + for trial in trials + ], + }, + indent=2, + sort_keys=True, + ) + + "\n", + encoding="utf-8", + ) + print(f"Uploaded and verified {len(trace_ids)} agent traces in workspace {args.workspace!r}.") + print(summary_path) + return summary_path + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--dataset", type=Path, default=DEFAULT_DATASET) + parser.add_argument("--agent", type=Path, default=SCRIPT_DIR) + parser.add_argument("--workspace", default=DEFAULT_WORKSPACE) + parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT) + parser.add_argument("--base-url", default=os.environ.get("NMP_BASE_URL", "http://localhost:8080")) + parser.add_argument( + "--api-base", + default=os.environ.get("INFERENCE_API_BASE", "https://inference-api.nvidia.com/v1"), + ) + parser.add_argument("--model", default=os.environ.get("AUT_MODEL_NAME", DEFAULT_MODEL)) + parser.add_argument("--user-model", default=os.environ.get("TAU2_USER_MODEL", DEFAULT_MODEL)) + parser.add_argument("--agent-name", default=DEFAULT_AGENT_NAME) + parser.add_argument("--agent-version", default=DEFAULT_AGENT_VERSION) + parser.add_argument("--experiment-id") + parser.add_argument("--expected-task-count", type=_positive_int, default=20) + parser.add_argument( + "--task-id", + dest="task_ids", + action="append", + help="Run only this task ID; repeat to select multiple tasks.", + ) + parser.add_argument("--concurrency", type=_positive_int, default=4) + parser.add_argument("--verbose", action="store_true") + return parser.parse_args() + + +if __name__ == "__main__": + asyncio.run(run(parse_args()))