diff --git a/.github/actions/insights-intake-stack/action.yml b/.github/actions/insights-intake-stack/action.yml index 815a885f3a..50e02dd808 100644 --- a/.github/actions/insights-intake-stack/action.yml +++ b/.github/actions/insights-intake-stack/action.yml @@ -1,5 +1,5 @@ name: insights-intake-stack -description: Self-contained Intake (ClickHouse + auth,entities,intake) on :8080 over $RUNNER_TEMP/state +description: Self-contained Analyst stack (ClickHouse + Intake and Platform model routing) on :8080 runs: using: composite steps: diff --git a/.github/workflows/insights-testbed.yml b/.github/workflows/insights-testbed.yml index 223e1f4e30..af538600b2 100644 --- a/.github/workflows/insights-testbed.yml +++ b/.github/workflows/insights-testbed.yml @@ -114,6 +114,10 @@ jobs: # Same gateway key; litellm/tau2 read it under the OpenAI-conventional name. OPENAI_API_KEY: ${{ secrets.NVIDIA_INFERENCE_KEY }} OPENAI_API_BASE: ${{ secrets.NVIDIA_INFERENCE_URL }} + # nemo setup registers this provider and persists its credential in Platform Secrets. + NEMO_DEFAULT_INFERENCE_KEY: ${{ secrets.NVIDIA_INFERENCE_KEY }} + NEMO_DEFAULT_INFERENCE_BASE_URL: ${{ secrets.NVIDIA_INFERENCE_URL }} + NEMO_DEFAULT_MODEL: ${{ vars.INSIGHTS_TESTBED_MODEL || 'default/openai-openai-gpt-5-5' }} steps: - name: Require secrets run: | @@ -129,6 +133,11 @@ jobs: persist-credentials: false - uses: ./nemo-platform/.github/actions/insights-testbed-prep - uses: ./nemo-platform/.github/actions/insights-intake-stack + - name: Configure Analyst models + working-directory: nemo-platform + run: >- + uv run nemo --base-url http://localhost:8080 setup --auto --workspace default + --no-start-services --no-install-skills --no-deploy-agent - name: Run subjects (tau2 -> ingest -> insights) working-directory: nemo-platform/plugins/nemo-insights env: @@ -188,14 +197,17 @@ jobs: matrix: subject: ${{ fromJSON(needs.plan.outputs.subjects) }} env: - INFERENCE_API_KEY: ${{ secrets.NVIDIA_INFERENCE_KEY }} + NEMO_DEFAULT_INFERENCE_KEY: ${{ secrets.NVIDIA_INFERENCE_KEY }} + NEMO_DEFAULT_INFERENCE_BASE_URL: ${{ secrets.NVIDIA_INFERENCE_URL }} + NEMO_DEFAULT_MODEL: ${{ vars.INSIGHTS_TESTBED_MODEL || 'default/openai-openai-gpt-5-5' }} steps: - name: Require secrets env: GH_TOKEN: ${{ secrets.TESTBED_STATE_GH_READ_TOKEN }} run: | [ -n "$GH_TOKEN" ] || { echo "secret TESTBED_STATE_GH_READ_TOKEN is not set"; exit 1; } - [ -n "$INFERENCE_API_KEY" ] || { echo "secret NVIDIA_INFERENCE_KEY is not set"; exit 1; } + [ -n "$NEMO_DEFAULT_INFERENCE_KEY" ] || { echo "secret NVIDIA_INFERENCE_KEY is not set"; exit 1; } + [ -n "$NEMO_DEFAULT_INFERENCE_BASE_URL" ] || { echo "secret NVIDIA_INFERENCE_URL is not set"; exit 1; } - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3 with: { path: nemo-platform, persist-credentials: false } - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3 @@ -207,6 +219,11 @@ jobs: - uses: ./nemo-platform/.github/actions/insights-testbed-prep with: { install-tau2: "false" } - uses: ./nemo-platform/.github/actions/insights-intake-stack + - name: Configure Analyst models + working-directory: nemo-platform + run: >- + uv run nemo --base-url http://localhost:8080 setup --auto --workspace default + --no-start-services --no-install-skills --no-deploy-agent - name: Generate insights working-directory: nemo-platform/plugins/nemo-insights env: diff --git a/README.md b/README.md index 57145848d3..f9e61778d6 100644 --- a/README.md +++ b/README.md @@ -41,7 +41,7 @@ nemo setup Source development still needs Git and GNU Make. `make bootstrap` supplies the rest of the toolchain, installing the uv, Node.js and pnpm versions pinned in `mise.toml`. Run it before other `make` targets — they call uv through mise but don't install it. See [SETUP.md](SETUP.md#toolchain-uv-nodejs-pnpm). -`nemo setup` starts local services, registers your LLM provider, discovers available models, installs agent skills, and deploys a sample agent (see more below). +`nemo setup` starts local services, registers your LLM provider, discovers available models, selects default and fast agent models, installs agent skills, and deploys a sample agent (see more below). Review [Telemetry and Privacy](docs/telemetry-and-privacy.mdx) for the omnibus disclosure covering anonymous telemetry, bundled library telemetry, third-party endpoint notes, and opt-out controls. @@ -95,6 +95,7 @@ If `make bootstrap` reports that Studio asset bootstrap did not complete, the AP ```bash export NVIDIA_API_KEY=nvapi... export NEMO_DEFAULT_MODEL=nvidia-nemotron-3-super-120b-a12b +export NEMO_FAST_MODEL="$NEMO_DEFAULT_MODEL" nemo setup --auto --start-services --install-skills --deploy-agent ``` diff --git a/SETUP.md b/SETUP.md index bd54e04508..f39e3a0df8 100644 --- a/SETUP.md +++ b/SETUP.md @@ -1,6 +1,6 @@ --- name: nemo-setup -description: Set up a local NeMo Platform (`make bootstrap` + `nemo setup`) — services, providers, plugins, default model, and an optional demo agent. Use when the user asks to install, bootstrap, set up, run, or start a local NeMo Platform. +description: Set up a local NeMo Platform (`make bootstrap` + `nemo setup`) — services, providers, plugins, default/fast models, and an optional demo agent. Use when the user asks to install, bootstrap, set up, run, or start a local NeMo Platform. version: "0.1" --- @@ -86,14 +86,14 @@ is unavailable, start Docker and retry—do not proceed to `rm -rf`. ## Bootstrap and start -The README documents the streamlined path. Prefer this over the manual steps below whenever the task fits — it covers prerequisites install, service startup, provider registration, default-model selection, and demo agent deployment in one shot: +The README documents the streamlined path. Prefer this over the manual steps below whenever the task fits — it covers prerequisites install, service startup, provider registration, default/fast model selection, and demo agent deployment in one shot: === "Interactive" ```bash make bootstrap # installs Python deps, Studio assets, and plugins (including demo calculator agent) source .venv/bin/activate -nemo setup # interactive: prompts for provider, picks default model, optionally deploys calculator-agent +nemo setup # interactive: prompts for provider, picks default/fast models, optionally deploys calculator-agent ``` === "Non-interactive (CI)" @@ -149,18 +149,20 @@ Without one of those, a system or nvm-managed Node.js takes precedence and may n If `nemo setup` is too high-level for the task (e.g. debugging startup, custom service set, custom plugin install after bootstrap), use the manual sections below. -### Default model selection (under `--auto`) +### Default and fast model selection -`$NEMO_DEFAULT_MODEL` **must be a hyphenated entity ID** from `nemo models list` (e.g. `nvidia-llama-3-3-nemotron-super-49b-v1-5` or `default/nvidia-llama-3-3-nemotron-super-49b-v1-5`). The slash-with-dots form (`nvidia/llama-3.3-nemotron-super-49b-v1-5`) is the upstream catalog's `served_model_name` — it's shown for human display but the gateway rejects it as a request input. +`$NEMO_DEFAULT_MODEL` and `$NEMO_FAST_MODEL` **must be hyphenated entity IDs** from `nemo models list` (e.g. `nvidia-llama-3-3-nemotron-super-49b-v1-5` or `default/nvidia-llama-3-3-nemotron-super-49b-v1-5`). The slash-with-dots form (`nvidia/llama-3.3-nemotron-super-49b-v1-5`) is the upstream catalog's `served_model_name` — it's shown for human display but the gateway rejects it as a request input. `nemo setup --auto` picks the default model in this order: 1. `$NEMO_DEFAULT_MODEL` (if set) — used as-is. The user may have exported this from a previous session; it takes precedence over anything discovered from the registered provider. 2. Otherwise, the first model entity returned by provider discovery. -If the user is surprised by which model got picked, check `echo $NEMO_DEFAULT_MODEL` first — that's the most common cause. +The fast model uses `$NEMO_FAST_MODEL` when set and otherwise reuses the effective default. Interactive setup prompts for both and defaults the fast choice to the selected default. Existing contexts without `fast_model` also reuse `default_model`. -The first-discovered fallback is intentionally simple — providers like NVIDIA Build expose dozens of models and "first one" rarely matches the user's intent. If the user wants a specific model as the default (or wants to compare options before committing), don't rely on the `--auto` fallback. After setup finishes, the `inference` skill's "Step 2 — Discover available models" enumerates entity IDs and shows the jq filters for picking one out by vendor or family. The user can then pin their choice via `export NEMO_DEFAULT_MODEL=/` or by overriding `body["model"]` per request. +If the user is surprised by which model got picked, check `NEMO_DEFAULT_MODEL` and `NEMO_FAST_MODEL` first — that's the most common cause. + +The first-discovered fallback is intentionally simple — providers like NVIDIA Build expose dozens of models and "first one" rarely matches the user's intent. If the user wants specific models (or wants to compare options before committing), don't rely on the `--auto` fallback. After setup finishes, the `inference` skill's "Step 2 — Discover available models" enumerates entity IDs and shows the jq filters for picking one out by vendor or family. The user can then pin the pair via `NEMO_DEFAULT_MODEL` and `NEMO_FAST_MODEL`, or override `body["model"]` per request. If `nemo agents invoke …` fails with HTTP 422 in under a second on the first call, the cause is almost always a slash-with-dots model name reaching the gateway (e.g. via a stale `NEMO_DEFAULT_MODEL` or an agent config that hardcoded the upstream catalog form). The `inference` skill's "Common failure: HTTP 422 from chat completion" subsection has the diagnose-and-recover steps — don't conclude the platform is broken. @@ -238,7 +240,7 @@ The platform is running. Don't leave the user with "you're good to go" — offer If the user's prompt doesn't already pin one down, ask: *"The platform is up. What would you like to do next — optimize an agent, deploy one, run inference, evaluate, generate data, or something else?"* -If the user wants to **pick or swap the default model** (e.g. they didn't like the one `--auto` selected, or they want to compare options), don't guess — hand off to the `inference` skill. Step 2 there enumerates `served_models[].model_entity_id` and shows jq filters for picking by vendor / family. To pin the choice for subsequent runs, export `NEMO_DEFAULT_MODEL=/` before the next `nemo setup --auto`. For one-off commands, pass the entity ID positionally: `nemo chat `. +If the user wants to **pick or swap the default/fast model pair**, don't guess — hand off to the `inference` skill. Step 2 there enumerates `served_models[].model_entity_id` and shows jq filters for picking by vendor / family. To pin the choices for subsequent runs, export `NEMO_DEFAULT_MODEL=/` and `NEMO_FAST_MODEL=/` before the next `nemo setup --auto`; fast falls back to default when omitted. For one-off commands, pass the entity ID positionally: `nemo chat `. ### Available skills diff --git a/docs/agents/insight-driven-optimization.mdx b/docs/agents/insight-driven-optimization.mdx index ee2de7130b..7955dac2ed 100644 --- a/docs/agents/insight-driven-optimization.mdx +++ b/docs/agents/insight-driven-optimization.mdx @@ -49,10 +49,9 @@ want to automate portions of the agent improvement loop. runtime configuration) and a way to run or evaluate the agent reproducibly. - Traces in Intake for the agent under test. The Analyst diagnoses only what it can observe, so an agent with no telemetry cannot be analyzed. -- Model access for the optimization agents themselves. The Analyst uses Claude Opus 4.8 through the - NVIDIA Inference Gateway. The Experimenter and Eval Author use configurable large language models - (LLMs) through an OpenAI-compatible endpoint. This access is separate from the model access that - your AUT needs at runtime. +- A supported inference provider and default/fast model pair configured by + `nemo setup` for the optimization agents. This access is separate from the + model access that your agent under test needs at runtime. - Train and validation datasets in a Harbor-compatible layout, plus a task template and an output directory for artifacts. This is what the Eval Author works from and what the Experimenter uses to validate changes. @@ -292,8 +291,8 @@ their implementations are available. ### Set Up -Point the CLI at a platform and set model access for the optimization agents. -The examples use a local platform; replace the base URL for a remote +Point the CLI at a platform and configure model access for the optimization +agents. The examples use a local platform; replace the base URL for a remote deployment. ```bash @@ -303,17 +302,16 @@ export NMP_STUDIO_URL="$NMP_BASE_URL/studio" export NMP_ACCESS_TOKEN="$(nemo auth token)" export AGENT=your-agent-name # This agent must have traces in Intake. -# The Analyst uses Claude Opus 4.8 through NVIDIA Inference Gateway. -export INFERENCE_API_KEY=sk-... - -# The Experimenter and Eval Author use an OpenAI-compatible endpoint. -export NEMO_EXPERIMENTALIST_API_BASE=https://inference-api.nvidia.com/v1 -export NEMO_EXPERIMENTALIST_API_KEY=sk-... -export NEMO_EXPERIMENTALIST_MODELS_SMART=openai/openai/openai/gpt-5.6-sol -export NEMO_EXPERIMENTALIST_MODELS_MID=openai/openai/openai/gpt-5.6-terra -export NEMO_EXPERIMENTALIST_MODELS_FAST=openai/openai/openai/gpt-5.6-luna +nemo setup ``` +During setup, register a supported provider and select the default and fast +Platform Model Entities. The default model handles quality-critical work and +the fast model handles latency-sensitive work; press Enter at the fast-model +prompt to reuse the default. The Analyst, Eval Author, and Experimentalist all +use this pair. Provider credentials are stored in Platform Secrets rather than +plugin-specific environment variables. + Set `NMP_STUDIO_URL` explicitly if Studio uses a different origin or path. For an authenticated remote deployment, log in against that base URL before requesting the token. For local development with authentication disabled, omit @@ -457,8 +455,8 @@ come from the profile: nemo agents experimentalist run ``` -The Experimenter's own models are configured through environment variables -(refer to [Models](#models)). Train and validation datasets are required, and they must be local +The Experimenter's own models come from the default/fast pair selected by +`nemo setup` (refer to [Models](#models)). Train and validation datasets are required, and they must be local paths for the Harbor evaluator. The loop runs a baseline evaluation, performs root cause analysis, implements changes, then validates and picks a winner. @@ -499,6 +497,10 @@ per agent that is *due* and has enough new telemetry. Runs are incremental: each successful run records a cursor (`last_successful_run_at`), and subsequent runs only consider traces newer than that cursor. +`enable` captures the effective default/fast pair in the server-side analysis +config. Re-run it after changing the pair with `nemo setup`; enabled records +created before model-pair persistence must also be re-enabled. + ### Configure the Analysis Schedule The global schedule lives in the Insights plugin configuration or environment. The `timezone` @@ -547,7 +549,8 @@ Run the Analyst for one analysis of an agent's traces. Manage per-agent opt-in for periodic analysis. `enable` and `disable` require `--agent`; `status` takes an optional `--agent` (omit it to list all configs in -the workspace). All three accept `--workspace` and `--base-url`. +the workspace). `enable` stores the active default/fast model pair for future +scheduled jobs. All three accept `--workspace` and `--base-url`. ### `nemo agents analyst doctor` @@ -577,7 +580,7 @@ Run the local Experimenter loop. ### `nemo agents experimentalist doctor` -Diagnose the Experimenter setup: profile, credentials, Insight resolution, +Diagnose the Experimenter setup: profile, configured models, Insight resolution, datasets, and the experiment plan. ### `nemo agents eval-author` @@ -588,21 +591,26 @@ exit with a nonzero status until their implementations are available. ### Models -The optimization agents use the following model configuration: +The optimization agents use the active Platform CLI context: | Variable | Used by | Default and notes | |----------|---------|-------------------| -| `INFERENCE_API_KEY` | Analyst | Required. API key for Claude Opus 4.8 through the NVIDIA Inference Gateway. The Analyst model and endpoint are fixed. | -| `NEMO_EXPERIMENTALIST_API_BASE` | Experimenter, Eval Author | OpenAI-compatible model API base URL. `run` and `doctor` fall back to the NVIDIA Inference Gateway when it is unset, so it is only required against another endpoint. Also settable as `api_base` under the `experimentalist:` config section; the environment wins. | -| `NEMO_EXPERIMENTALIST_API_KEY` | Experimenter, Eval Author | Required. Model API key. On the gateway, `INFERENCE_API_KEY` fills this. | -| `NEMO_EXPERIMENTALIST_MODELS_SMART` | Experimenter, Eval Author | Required. High-capability model for analysis, proposing, coding, and curation. No default: a model name is only meaningful against a specific endpoint. | -| `NEMO_EXPERIMENTALIST_MODELS_MID` | Experimenter, Eval Author | Required. Mid-tier model used for trajectory scoring and architecture documentation. No default. | -| `NEMO_EXPERIMENTALIST_MODELS_FAST` | Experimenter, Eval Author | Required. Low-latency model for lightweight steps such as termination checks and summarization. No default. | -| `AUTHOR_API_BASE` | Eval Author | Optional override for the OpenAI-compatible model API base URL. Falls back to `NEMO_EXPERIMENTALIST_API_BASE`. | -| `AUTHOR_API_KEY` | Eval Author | Optional override for the model API key. Falls back to `NEMO_EXPERIMENTALIST_API_KEY`, or to `INFERENCE_API_KEY` for the NVIDIA Inference Gateway. | -| `AUTHOR_SMART_MODEL_NAME` | Eval Author | Optional high-capability model override. | -| `AUTHOR_MID_MODEL_NAME` | Eval Author | Optional mid-tier model override for Experimentalist helpers used by Eval Author. | -| `AUTHOR_FAST_MODEL_NAME` | Eval Author | Optional low-latency model override. | +| stored `default_model` | Analyst, Experimenter, Eval Author | Required workspace-qualified Model Entity selected by `nemo setup`. Used for quality-critical analysis, authoring, proposing, and coding. | +| stored `fast_model` | Analyst, Experimenter, Eval Author | Selected by `nemo setup`; defaults to the default model. Used for high-volume scoring, summarization, and control steps. | +| `NEMO_DEFAULT_MODEL` | Analyst, Experimenter, Eval Author | Optional non-interactive override for `default_model`, in `workspace/model-name` form. | +| `NEMO_FAST_MODEL` | Analyst, Experimenter, Eval Author | Optional non-interactive override for `fast_model`, in `workspace/model-name` form. Falls back to the effective default model. | + +Each Model Entity owns its backend format and provider route. OpenAI-compatible +and Anthropic entities are adapted to Nooa's unified completion client through the +Platform route; the agents never inspect provider model-name conventions or +read provider credentials directly. + +For OpenAI Chat Completions entities, the shared client disables reasoning while +function tools are active and disables LiteLLM's automatic Responses bridge. +This keeps frontier models on the Model Entity's declared Chat Completions +contract without a model-name allowlist; LiteLLM drops the reasoning parameter +for models that do not support it. Anthropic entities use their native Messages +route. ## Troubleshooting @@ -616,7 +624,8 @@ flowing and that you are filtering by the right `agent` name. Check that the agent is enabled (`nemo insights analysis status`), that `insights.analyst.enabled` is true, that the scheduled window has passed in the configured `timezone`, and that at -least 10 new sessions have landed since the last run. +least 10 new sessions have landed since the last run. If status shows no model +selection, re-run `nemo insights analysis enable` after `nemo setup`. ### Insights Are Not in the Platform diff --git a/docs/cli/reference.mdx b/docs/cli/reference.mdx index d7c47002de..8883ac16b8 100644 --- a/docs/cli/reference.mdx +++ b/docs/cli/reference.mdx @@ -47,8 +47,8 @@ Set up NeMo Platform: connect or start services, configure a provider, install s Uses an already-running platform, starts local services, or connects the CLI to an existing remote deployment. Then selects and registers an -inference provider, picks a default model, installs coding agent skills, -and optionally deploys a demo agent. +inference provider, picks default and fast agent models, installs coding +agent skills, and optionally deploys a demo agent. The active config context remembers the Platform URL. When a remote deployment is already reachable, setup asks whether to continue with it, @@ -66,7 +66,7 @@ Requires an interactive terminal (TTY). In non-interactive contexts Use --auto for non-interactive setup from environment variables (NEMO_DEFAULT_INFERENCE_KEY, NVIDIA_API_KEY, OPENAI_API_KEY, ANTHROPIC_API_KEY, GEMINI_API_KEY). -Override the default model with NEMO_DEFAULT_MODEL. +Override the selected pair with NEMO_DEFAULT_MODEL and NEMO_FAST_MODEL. **Examples:** diff --git a/docs/get-started/example-agent.mdx b/docs/get-started/example-agent.mdx index 8dd44c1daf..03a66392b9 100644 --- a/docs/get-started/example-agent.mdx +++ b/docs/get-started/example-agent.mdx @@ -10,7 +10,7 @@ This provides a guide for how you can quickly load agent traces into your NeMo P - Git, GNU Make, uv, and pnpm - Docker Engine 29.6.2 or later - [Docker Sandboxes](https://docs.docker.com/ai/sandboxes/) (`sbx`) 0.37.1 or later -- An NVIDIA Inference Gateway virtual key (`sk-...`) from `inference.nvidia.com`. +- An API key for the inference provider used by the Tau3 agent under test. ## 1. Clone and bootstrap NeMo Platform @@ -32,6 +32,12 @@ platform services in the background and walks you through the configuration requ uv run nemo setup ``` +Register the provider you want the Platform agents to use, then select a +default model for quality-critical work and a fast model for latency-sensitive +work. Press Enter at the fast-model prompt to reuse the default. These +workspace-qualified Model Entities are used by the Analyst, Eval Author, and +Experimentalist; their provider credentials stay in Platform Secrets. + To use an external ClickHouse instead, set `NMP_INTAKE_CLICKHOUSE_URL` before running setup. You should now be able to navigate to `http://localhost:8080` and see the NeMo Platform web UI. @@ -44,7 +50,10 @@ Now that we have a running NeMo Platform, it's time to load up some data! This e 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. +Then, update the `.env` file with the inference endpoint and key for the Tau3 +agent under test. If you will run the Experimentalist in a Docker Sandbox, +also uncomment `NEMO_DEFAULT_MODEL` and `NEMO_FAST_MODEL` and copy the +workspace-qualified IDs reported by `nemo setup`. Source the file afterward. ```bash source plugins/nemo-experimentalist/examples/tau3-nooa-agent/.env @@ -140,6 +149,12 @@ pushes a validated winner to a new branch, and opens a draft PR or MR against configuration. Authenticate GitHub inside the sandbox so it can clone and push the private repository. + +Clone mode is an integrity boundary, not a confidentiality boundary. The sandbox +can read the cloned host repository and every credential passed with `sbx exec --env`. +Use dedicated, revocable, spending-limited credentials. + + ```bash repo="$(git rev-parse --show-toplevel)" sbx create --clone --name nemo-experimentalist-git shell "$repo" @@ -154,11 +169,8 @@ sbx exec --workdir "$repo" \ --env INFERENCE_API_BASE \ --env OPENAI_API_KEY \ --env OPENAI_BASE_URL \ - --env NEMO_EXPERIMENTALIST_API_KEY \ - --env NEMO_EXPERIMENTALIST_API_BASE \ - --env NEMO_EXPERIMENTALIST_MODELS_SMART \ - --env NEMO_EXPERIMENTALIST_MODELS_MID \ - --env NEMO_EXPERIMENTALIST_MODELS_FAST \ + --env NEMO_DEFAULT_MODEL \ + --env NEMO_FAST_MODEL \ --env TAU2_USER_MODEL \ --env TAU2_NL_ASSERTIONS_MODEL \ --env AUT_MODEL_NAME \ @@ -209,11 +221,8 @@ sbx exec --workdir "$repo" \ --env INFERENCE_API_BASE \ --env OPENAI_API_KEY \ --env OPENAI_BASE_URL \ - --env NEMO_EXPERIMENTALIST_API_KEY \ - --env NEMO_EXPERIMENTALIST_API_BASE \ - --env NEMO_EXPERIMENTALIST_MODELS_SMART \ - --env NEMO_EXPERIMENTALIST_MODELS_MID \ - --env NEMO_EXPERIMENTALIST_MODELS_FAST \ + --env NEMO_DEFAULT_MODEL \ + --env NEMO_FAST_MODEL \ --env TAU2_USER_MODEL \ --env TAU2_NL_ASSERTIONS_MODEL \ --env AUT_MODEL_NAME \ diff --git a/docs/get-started/setup.mdx b/docs/get-started/setup.mdx index fb99f940ba..e34b3b7429 100644 --- a/docs/get-started/setup.mdx +++ b/docs/get-started/setup.mdx @@ -70,7 +70,7 @@ The wizard walks through each stage: 2. **Choose a provider** — select your model provider and enter your API key 3. **Register the provider** — connects the provider to the platform 4. **Discover models** — finds available models from your provider -5. **Set a default model** — pick the model to use by default +5. **Select agent models** — pick a default model for quality-critical work and a fast model for latency-sensitive work; reuse the default if one model is enough 6. **Install skills** — adds NeMo skills to detected coding agents (Cursor, Claude Code, Codex) 7. **Deploy agent** — optionally deploys a demo calculator agent on the platform @@ -171,11 +171,13 @@ nemo setup --auto --start-services | `--install-skills` / `--no-install-skills` | prompt (interactive) or skip (auto) | Install NeMo skills for detected coding agents | | `--deploy-agent` / `--no-deploy-agent` | prompt (interactive) or skip (auto) | Deploy the demo calculator agent | -Override the default model with `NEMO_DEFAULT_MODEL`: +Override the default and fast models with `NEMO_DEFAULT_MODEL` and +`NEMO_FAST_MODEL`. If the fast override is omitted, it reuses the default. ```bash export OPENAI_API_KEY=sk-... export NEMO_DEFAULT_MODEL=default/openai-gpt-4.1 +export NEMO_FAST_MODEL="$NEMO_DEFAULT_MODEL" nemo setup --auto --start-services ``` diff --git a/docs/set-up/config-reference.mdx b/docs/set-up/config-reference.mdx index c840696658..195eaca783 100644 --- a/docs/set-up/config-reference.mdx +++ b/docs/set-up/config-reference.mdx @@ -888,26 +888,6 @@ safe_synthesizer: default_job_resource_cpu_limit: '4' ``` -### `experimentalist` - -LLM endpoint and model tiers the Experimentalist's own agents run on. - -```yaml wordWrap -experimentalist: - # Base URL of the OpenAI-compatible endpoint the optimizer's agents call. - api_base: - # Credential for that endpoint. Prefer NEMO_EXPERIMENTALIST_API_KEY over writing it here. - api_key: - # Model per tier, named as your endpoint names it. Every tier is required. - models: - # Model for the smart tier, as your endpoint names it. - smart: - # Model for the mid tier, as your endpoint names it. - mid: - # Model for the fast tier, as your endpoint names it. - fast: -``` - ### `studio` Configuration for the Studio service. diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/manifest_registry.py b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/manifest_registry.py index 962452d81d..8f61eaee98 100644 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/manifest_registry.py +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/manifest_registry.py @@ -37,8 +37,8 @@ Uses an already-running platform, starts local services, or connects the CLI to an existing remote deployment. Then selects and registers an -inference provider, picks a default model, installs coding agent skills, -and optionally deploys a demo agent. +inference provider, picks default and fast agent models, installs coding +agent skills, and optionally deploys a demo agent. The active config context remembers the Platform URL. When a remote deployment is already reachable, setup asks whether to continue with it, @@ -56,7 +56,7 @@ Use --auto for non-interactive setup from environment variables (NEMO_DEFAULT_INFERENCE_KEY, NVIDIA_API_KEY, OPENAI_API_KEY, ANTHROPIC_API_KEY, GEMINI_API_KEY). -Override the default model with NEMO_DEFAULT_MODEL. +Override the selected pair with NEMO_DEFAULT_MODEL and NEMO_FAST_MODEL. Examples: nemo setup diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/setup.py b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/setup.py index e6c3f2bbd3..ffe650a312 100644 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/setup.py +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/setup.py @@ -47,7 +47,7 @@ from nemo_platform_ext.cli.telemetry.events import OnboardingStepEvent, TaskStatusEnum from nemo_platform_ext.client.tls import client_verify_from_env from nemo_platform_ext.config.config import Config -from nemo_platform_ext.config.models import DEFAULT_BASE_URL, ConfigFile, ConfigParams, LocalServicesConfig +from nemo_platform_ext.config.models import DEFAULT_BASE_URL, ConfigFile, ConfigParams, LocalServicesConfig, NoAuthUser from nemo_platform_ext.local.process import ( check_port_available_for_start, compute_scope, @@ -218,6 +218,14 @@ class KeyValidationResult: message: str +@dataclass(frozen=True) +class ModelPair: + """Model entities selected for quality-critical and low-latency agent work.""" + + default: str + fast: str + + # Env vars probed during --auto mode, in priority order. _AUTO_ENV_VARS: tuple[tuple[str, str], ...] = ( ("NEMO_DEFAULT_INFERENCE_KEY", "NEMO_DEFAULT_INFERENCE_BASE_URL"), @@ -246,6 +254,7 @@ class KeyValidationResult: _POST_START_REACHABLE_DELAY = 2.0 _DEMO_AGENT_NAME = "calculator-agent" +_LOCAL_CONTEXT_NAME = "local" def _pause(seconds: float) -> None: @@ -272,8 +281,8 @@ def _bootstrap_config_if_missing(base_url: str, workspace: str) -> None: ``nemo setup`` can run before any config is on disk (first-time install) *or* with a partial config containing only ``local_services.data_dir`` written earlier in the same setup invocation by ``_save_data_dir``. - Later steps — in particular ``_save_default_model`` — call - ``Config.write`` with *only* a ``default_model`` param. If there's no + Later steps — in particular ``_save_model_pair`` — call + ``Config.write`` with only model-selection params. If there's no cluster on disk at that point, ``ensure_context`` will fail with ``Cluster '' does not exist and no base_url provided to create it`` because it has no ``base_url`` to attach to a new cluster. @@ -378,6 +387,49 @@ def _configure_remote_connection(cli_context: CLIContext, base_url: str, workspa cli_context.reset_sdk_context() +def _local_context_name(config_file: ConfigFile) -> str: + """Reuse a compatible local context or choose a name that cannot overwrite one.""" + for context in config_file.contexts: + cluster = next(cluster for cluster in config_file.clusters if cluster.name == context.cluster) + user = next(user for user in config_file.users if user.name == context.user) + if str(cluster.base_url).rstrip("/") == DEFAULT_BASE_URL.rstrip("/") and isinstance(user, NoAuthUser): + return context.name + + existing_names = {context.name for context in config_file.contexts} + if _LOCAL_CONTEXT_NAME not in existing_names: + return _LOCAL_CONTEXT_NAME + suffix = 2 + while f"{_LOCAL_CONTEXT_NAME}-{suffix}" in existing_names: + suffix += 1 + return f"{_LOCAL_CONTEXT_NAME}-{suffix}" + + +def _configure_local_connection(cli_context: CLIContext, workspace: str) -> None: + """Activate an isolated no-auth context for local services. + + Reusing an authenticated remote context after changing only its URL leaves + remote OAuth credentials attached to a local cluster. The SDK then tries to + refresh those credentials using the local auth discovery response. Keep the + remote context intact and use a dedicated local context instead. + """ + context_name = _local_context_name(Config.load().get_config_file()) + params: ConfigParams = { + "base_url": DEFAULT_BASE_URL, + "workspace": workspace, + "access_token": None, + "refresh_token": None, + "current_context": context_name, + } + Config.write( + params, + context_name=context_name, + set_current_on_create=True, + ) + cli_context.overrides["base_url"] = DEFAULT_BASE_URL + cli_context.overrides["current_context"] = context_name + cli_context.reset_sdk_context() + + def _ensure_platform_auth(cli_context: CLIContext) -> None: """Authenticate the active context when it lacks usable credentials.""" from nemo_platform_ext.cli.commands.auth import _login_with_oidc, _runtime_token_source_label @@ -714,12 +766,19 @@ def _wait_for_models_impl( return [] -def _get_all_model_entity_ids(client: NeMoPlatform, workspace: str) -> list[str]: - """Return all model entity IDs across all providers.""" +def _get_all_model_entity_ids( + client: NeMoPlatform, + workspace: str, + *, + provider_name: str | None = None, +) -> list[str]: + """Return model entity IDs, optionally scoped to one provider.""" entity_ids: list[str] = [] try: page = client.inference.providers.list(workspace=workspace) for provider in page.data: + if provider_name is not None and getattr(provider, "name", None) != provider_name: + continue for model in getattr(provider, "served_models", None) or []: if hasattr(model, "model_entity_id") and model.model_entity_id: entity_ids.append(model.model_entity_id) @@ -728,17 +787,24 @@ def _get_all_model_entity_ids(client: NeMoPlatform, workspace: str) -> list[str] return sorted(set(entity_ids)) -def _get_all_model_choices(client: NeMoPlatform, workspace: str) -> list[tuple[str, str]]: - """Return picker choices as (entity_id, label) across all providers.""" +def _get_all_model_choices( + client: NeMoPlatform, + workspace: str, + *, + provider_name: str | None = None, +) -> list[tuple[str, str]]: + """Return picker choices, optionally scoped to one provider.""" choices: list[tuple[str, str]] = [] try: page = client.inference.providers.list(workspace=workspace) for provider in page.data: - provider_name = getattr(provider, "name", "unknown-provider") + current_provider_name = getattr(provider, "name", "unknown-provider") + if provider_name is not None and current_provider_name != provider_name: + continue for model in getattr(provider, "served_models", None) or []: model_entity_id = getattr(model, "model_entity_id", None) if model_entity_id: - label = f"{_display_model_name(model_entity_id)} ({provider_name})" + label = f"{_display_model_name(model_entity_id)} ({current_provider_name})" choices.append((model_entity_id, label)) except Exception: logger.debug("Failed to list model choices", exc_info=True) @@ -1815,24 +1881,42 @@ def _validate_api_key( return KeyValidationResult(passed=True, message=f"Could not validate API key ({exc}).") -def _select_default_model(client: NeMoPlatform, workspace: str) -> str | None: - """Let the user pick a default model from discovered models.""" - display_models = _get_all_model_choices(client, workspace) +def _select_model_pair( + client: NeMoPlatform, + workspace: str, + *, + provider_name: str | None = None, +) -> ModelPair | None: + """Let the user pick default and fast models from one provider.""" + display_models = _get_all_model_choices(client, workspace, provider_name=provider_name) if not display_models: - console.print(f" {WARN} No models discovered yet. You can set a default later.") + console.print(f" {WARN} No models discovered yet. You can select models later.") return None - result = prompt_select( - "Choose your default model:", + first_model = display_models[0][0] + default_model = prompt_select( + "Choose your default model (used for quality-critical agent work):", choices=display_models, + default=first_model, + hint="Press Enter to accept the default.", ) - return result + fast = prompt_select( + "Choose your fast model (used for latency-sensitive agent work):", + choices=display_models, + default=default_model, + hint="Press Enter to reuse the default model.", + ) + return ModelPair(default=default_model, fast=fast) -def _save_default_model(cli_context: CLIContext, model_entity_id: str) -> None: - """Persist the default model to the CLI config file.""" +def _save_model_pair(cli_context: CLIContext, model_pair: ModelPair) -> None: + """Persist the default quality model and the low-latency model.""" context = cli_context.get_sdk_context() - Config.write({"default_model": model_entity_id}, context_name=context.context_name) + params: ConfigParams = { + "default_model": model_pair.default, + "fast_model": model_pair.fast, + } + Config.write(params, context_name=context.context_name) def _check_ollama_running(host_url: str) -> bool: @@ -1849,8 +1933,8 @@ def _check_ollama_running(host_url: str) -> bool: # --------------------------------------------------------------------------- -def _auto_setup(client: NeMoPlatform, workspace: str) -> bool: - """Register a provider from environment variables. Returns True on success.""" +def _auto_setup(client: NeMoPlatform, workspace: str) -> str | None: + """Register a provider from environment variables and return its name.""" for key_var, url_var in _AUTO_ENV_VARS: api_key = os.environ.get(key_var) if not api_key: @@ -1923,9 +2007,9 @@ def _auto_setup(client: NeMoPlatform, workspace: str) -> bool: ) console.print(f" {CHECK} Registered provider '{provider_name}' ({host_url})") - return True + return provider_name - return False + return None # --------------------------------------------------------------------------- @@ -2001,8 +2085,8 @@ def setup_command( Uses an already-running platform, starts local services, or connects the CLI to an existing remote deployment. Then selects and registers an - inference provider, picks a default model, installs coding agent skills, - and optionally deploys a demo agent. + inference provider, picks default and fast agent models, installs coding + agent skills, and optionally deploys a demo agent. The active config context remembers the Platform URL. When a remote deployment is already reachable, setup asks whether to continue with it, @@ -2020,7 +2104,7 @@ def setup_command( Use --auto for non-interactive setup from environment variables (NEMO_DEFAULT_INFERENCE_KEY, NVIDIA_API_KEY, OPENAI_API_KEY, ANTHROPIC_API_KEY, GEMINI_API_KEY). - Override the default model with NEMO_DEFAULT_MODEL. + Override the selected pair with NEMO_DEFAULT_MODEL and NEMO_FAST_MODEL. Examples: nemo setup @@ -2048,11 +2132,7 @@ def setup_command( configured_base_url = base_url service_result = _maybe_start_services(base_url, auto, start_services, timeout=effective_timeout) if service_result == "start_local": - context_name = cli_context.get_sdk_context().context_name - _bootstrap_config_if_missing(DEFAULT_BASE_URL, workspace) - Config.write({"base_url": DEFAULT_BASE_URL}, context_name=context_name) - cli_context.overrides["base_url"] = DEFAULT_BASE_URL - cli_context.reset_sdk_context() + _configure_local_connection(cli_context, workspace) base_url = DEFAULT_BASE_URL service_result = _maybe_start_services( base_url, @@ -2080,7 +2160,7 @@ def setup_command( # Ensure the config file exists on disk so later Config.write() calls # (e.g. saving the default model) can find the cluster and context. # Without this, a fresh install (no config.yaml) hits "Cluster - # 'default-cluster' does not exist" when _save_default_model runs. + # 'default-cluster' does not exist" when _save_model_pair runs. _bootstrap_config_if_missing(base_url, workspace) cli_context.reset_sdk_context() @@ -2156,7 +2236,8 @@ def _run_auto_mode( ) -> None: """Non-interactive provider registration from environment variables.""" console.print("[bold]Auto-detecting provider from environment...[/bold]\n") - if not _auto_setup(client, workspace): + provider_name = _auto_setup(client, workspace) + if provider_name is None: console.print(f"{CROSS} No provider credentials found in environment.") env_var_names = ", ".join(key for key, _ in _AUTO_ENV_VARS) console.print(f" Set one of: {env_var_names}") @@ -2167,7 +2248,7 @@ def _run_auto_mode( for attempt in range(_MODEL_DISCOVERY_MAX_ROUNDS): deadline = time.time() + _MODEL_DISCOVERY_ROUND_SECONDS while time.time() < deadline: - entity_ids = _get_all_model_entity_ids(client, workspace) + entity_ids = _get_all_model_entity_ids(client, workspace, provider_name=provider_name) if entity_ids: break _pause(_MODEL_DISCOVERY_POLL_INTERVAL) @@ -2189,14 +2270,20 @@ def _run_auto_mode( default_model = os.environ.get("NEMO_DEFAULT_MODEL", "").strip() if not default_model and entity_ids: default_model = entity_ids[0] + fast_model = os.environ.get("NEMO_FAST_MODEL", "").strip() or default_model if default_model: - _save_default_model(cli_context, default_model) - console.print(f" {CHECK} Default model: {default_model}") + model_pair = ModelPair(default=default_model, fast=fast_model) + _save_model_pair(cli_context, model_pair) + console.print(f" {CHECK} Default model: {model_pair.default}") + console.print(f" {CHECK} Fast model: {model_pair.fast}") else: + if fast_model: + console.print(f" {WARN} NEMO_FAST_MODEL is ignored until a default model is available") console.print(f" {WARN} No default model set (no models discovered yet)") - console.print(" Run [cyan]nemo setup[/cyan] again after models sync, or set via env var:") + console.print(" Run [cyan]nemo setup[/cyan] again after models sync, or set the models via env vars:") console.print(" [cyan]export NEMO_DEFAULT_MODEL=[/cyan]") + console.print(" [cyan]export NEMO_FAST_MODEL=[/cyan]") _maybe_install_skills( auto=True, @@ -2273,17 +2360,19 @@ def _run_interactive_mode( else: console.print(f" {WARN} No models discovered yet (provider may still be syncing)") - console.print("\n[bold]Step 5: Choose default model[/bold]\n") + console.print("\n[bold]Step 5: Choose agent models[/bold]\n") fallback_model_choices = _get_all_model_choices(client, workspace) if not models else [] if not models and fallback_model_choices: console.print(f" {WARN} Models from existing providers are available, but not from '{provider_name}' yet.") - default_model = _select_default_model(client, workspace) if models else None - if default_model: - _save_default_model(cli_context, default_model) - console.print(f" {CHECK} Default model set to {_display_model_name(default_model)}") + model_pair = _select_model_pair(client, workspace, provider_name=provider_name) if models else None + default_model = model_pair.default if model_pair else None + if model_pair: + _save_model_pair(cli_context, model_pair) + console.print(f" {CHECK} Default model set to {_display_model_name(model_pair.default)}") + console.print(f" {CHECK} Fast model set to {_display_model_name(model_pair.fast)}") else: - console.print(f" {WARN} No default model set for this provider yet.") + console.print(f" {WARN} No agent models set for this provider yet.") console.print(" Run [cyan]nemo setup[/cyan] again after models sync") console.print("\n[bold]Step 6: Install skills[/bold]\n") @@ -2305,7 +2394,13 @@ def _run_interactive_mode( headers=_platform_request_headers(cli_context), ) - _print_onboarding(base_url, provider_name, default_model, demo_deployed=demo_deployed) + _print_onboarding( + base_url, + provider_name, + default_model, + fast_model=model_pair.fast if model_pair else None, + demo_deployed=demo_deployed, + ) except UserCancelled: console.print(f"\n{WARN} Setup cancelled.") @@ -2372,6 +2467,7 @@ def _print_onboarding( provider_name: str, default_model: str | None, *, + fast_model: str | None = None, demo_deployed: bool = False, ) -> None: """Print setup summary, then present goal-oriented onboarding paths.""" @@ -2382,6 +2478,8 @@ def _print_onboarding( console.print(f" Provider: {provider_name}") if default_model: console.print(f" Default model: {_display_model_name(default_model)}") + if fast_model: + console.print(f" Fast model: {_display_model_name(fast_model)}") if demo_deployed: console.print(f" Demo agent: {_DEMO_AGENT_NAME}") diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/config/config.py b/packages/nemo_platform_ext/src/nemo_platform_ext/config/config.py index 341d41924d..144716f3c4 100644 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/config/config.py +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/config/config.py @@ -63,6 +63,13 @@ class _RuntimeAccessTokenSource: label: str +def _resolve_model_pair(default_model: str | None, fast_model: str | None) -> tuple[str | None, str | None]: + """Apply environment overrides and the fast-to-default fallback consistently.""" + effective_default = os.environ.get("NEMO_DEFAULT_MODEL") or default_model + effective_fast = os.environ.get("NEMO_FAST_MODEL") or fast_model or effective_default + return effective_default, effective_fast + + class Config(BaseModel): """ Configuration manager for nemo_platform. @@ -72,10 +79,10 @@ class Config(BaseModel): 2. Runtime overrides (from env vars, CLI, or code) 3. Resolution of effective configuration - Environment variables (prefix NMP_): + Environment variables: NMP_CURRENT_CONTEXT, NMP_WORKSPACE, NMP_OUTPUT_FORMAT, NMP_TIMESTAMP_FORMAT, NMP_PAGE_SIZE, NMP_COLOR_OUTPUT, - NMP_BASE_URL, NMP_ACCESS_TOKEN + NMP_BASE_URL, NMP_ACCESS_TOKEN, NEMO_DEFAULT_MODEL, NEMO_FAST_MODEL """ # Runtime overrides - can be set via env vars @@ -486,8 +493,12 @@ def resolve(self) -> Context: if effective_workspace is None: effective_workspace = DEFAULT_WORKSPACE - # Resolve default model: env var > config file - effective_default_model = os.environ.get("NEMO_DEFAULT_MODEL") or context.default_model + # The default model is the quality-oriented model for agent workloads. + # Existing single-model contexts also supply the fast role. + effective_default_model, effective_fast_model = _resolve_model_pair( + context.default_model, + context.fast_model, + ) # Create the resolved Context (runtime model) from the ContextDefinition (config file model) return Context( @@ -496,6 +507,7 @@ def resolve(self) -> Context: user=user, # Fully resolved User with authentication credentials workspace=effective_workspace, default_model=effective_default_model, + fast_model=effective_fast_model, preferences=prefs, ) @@ -545,12 +557,17 @@ def _create_default_config(self) -> Context: prefs.color_output = self.color_output # Return resolved Context (runtime model) + effective_default_model, effective_fast_model = _resolve_model_pair( + context_def.default_model, + context_def.fast_model, + ) return Context( context_name=context_name, cluster=cluster, user=user, workspace=context_def.workspace or DEFAULT_WORKSPACE, - default_model=os.environ.get("NEMO_DEFAULT_MODEL") or context_def.default_model, + default_model=effective_default_model, + fast_model=effective_fast_model, preferences=prefs, ) diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/config/models.py b/packages/nemo_platform_ext/src/nemo_platform_ext/config/models.py index d114ef2d4f..689cebc8ae 100644 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/config/models.py +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/config/models.py @@ -140,6 +140,7 @@ class ContextDefinition(BaseModel): user: str = Field(..., min_length=1, description="Reference to user name (string, not resolved)") workspace: str | None = Field(default=None, description="Default workspace") default_model: str | None = Field(default=None, description="Default model entity ID for inference") + fast_model: str | None = Field(default=None, description="Low-latency model entity ID for agent workloads") preferences: Preferences = Field(default_factory=Preferences, description="Context-specific preferences") metadata: dict[str, Any] = Field(default_factory=dict, description="Additional context metadata") @@ -174,6 +175,7 @@ class ConfigParams(TypedDict, total=False): workspace: str default_model: str + fast_model: str output_format: OutputFormat timestamp_format: TimestampFormat @@ -316,6 +318,8 @@ def ensure_context( context.workspace = params["workspace"] if "default_model" in params: context.default_model = params["default_model"] + if "fast_model" in params: + context.fast_model = params["fast_model"] if "output_format" in params: context.preferences.output_format = params["output_format"] if "timestamp_format" in params: @@ -383,4 +387,5 @@ class Context(BaseModel): user: User | None = Field(default=None, description="Resolved user with authentication credentials") workspace: str = Field(..., description="Active workspace (required)") default_model: str | None = Field(default=None, description="Default model entity ID for inference") + fast_model: str | None = Field(default=None, description="Low-latency model entity ID for agent workloads") preferences: Preferences = Field(..., description="Effective preferences") diff --git a/packages/nemo_platform_ext/tests/cli/commands/test_setup.py b/packages/nemo_platform_ext/tests/cli/commands/test_setup.py index 0c2ec42a6f..a41f3dd8f8 100644 --- a/packages/nemo_platform_ext/tests/cli/commands/test_setup.py +++ b/packages/nemo_platform_ext/tests/cli/commands/test_setup.py @@ -34,6 +34,7 @@ KNOWN_PROVIDERS, ONBOARDING_PATHS, KeyValidationResult, + ModelPair, _agent_config_path, _agents_plugin_available, _auto_setup, @@ -42,6 +43,7 @@ _check_ollama_running, _check_platform_reachable, _check_platform_reachable_with_retries, + _configure_local_connection, _create_provider, _deploy_demo_agent, _detect_coding_agents, @@ -62,9 +64,10 @@ _render_onboarding_card, _resolve_provider_for_url, _resolve_setup_workspace, + _run_auto_mode, _run_interactive_mode, _save_data_dir, - _select_default_model, + _select_model_pair, _start_services_background, _validate_api_key, _verify_platform_health, @@ -419,8 +422,8 @@ def test_no_auth_header_format_skips_auth_fields(self): workspace="default", ) call_kwargs = client.inference.providers.create.call_args.kwargs - assert "auth_header_format" not in call_kwargs assert "required_extra_headers" not in call_kwargs + assert "auth_header_format" not in call_kwargs def test_provider_without_secret(self): """Providers without secrets (e.g. Ollama) should omit api_key_secret_name.""" @@ -455,16 +458,16 @@ def _bypass_key_validation(): @pytest.mark.usefixtures("_bypass_key_validation") class TestAutoSetup: - def test_no_env_vars_returns_false(self): + def test_no_env_vars_returns_none(self): client = _make_mock_client() with patch.dict("os.environ", {}, clear=True): - assert _auto_setup(client, "default") is False + assert _auto_setup(client, "default") is None def test_openai_key_creates_provider(self): client = _make_mock_client() with patch.dict("os.environ", {"OPENAI_API_KEY": "sk-test123"}, clear=True): result = _auto_setup(client, "default") - assert result is True + assert result == "openai" client.mock_secrets.create_secret.assert_called_once() client.inference.providers.create.assert_called_once() create_kwargs = client.inference.providers.create.call_args @@ -475,7 +478,7 @@ def test_nvidia_key_creates_build_provider(self): client = _make_mock_client() with patch.dict("os.environ", {"NVIDIA_API_KEY": "nvapi-test"}, clear=True): result = _auto_setup(client, "default") - assert result is True + assert result == "nvidia-build" create_kwargs = client.inference.providers.create.call_args assert create_kwargs.kwargs["name"] == "nvidia-build" @@ -487,7 +490,7 @@ def test_nemo_default_key_with_url(self): } with patch.dict("os.environ", env, clear=True): result = _auto_setup(client, "default") - assert result is True + assert result == "openai" create_kwargs = client.inference.providers.create.call_args assert create_kwargs.kwargs["name"] == "openai" @@ -499,7 +502,7 @@ def test_nemo_default_key_with_custom_url(self): } with patch.dict("os.environ", env, clear=True): result = _auto_setup(client, "default") - assert result is True + assert result == "my-custom-llm-example-com" create_kwargs = client.inference.providers.create.call_args assert create_kwargs.kwargs["name"] == "my-custom-llm-example-com" assert create_kwargs.kwargs["host_url"] == "https://my-custom-llm.example.com/v1" @@ -508,7 +511,7 @@ def test_existing_provider_updated_not_recreated(self): client = _make_mock_client(provider_exists=True, secret_exists=True) with patch.dict("os.environ", {"OPENAI_API_KEY": "sk-test123"}, clear=True): result = _auto_setup(client, "default") - assert result is True + assert result == "openai" client.inference.providers.create.assert_not_called() client.inference.providers.update.assert_called_once() @@ -522,7 +525,7 @@ def test_priority_order(self): } with patch.dict("os.environ", env, clear=True): result = _auto_setup(client, "default") - assert result is True + assert result == "anthropic" create_kwargs = client.inference.providers.create.call_args assert create_kwargs.kwargs["name"] == "anthropic" @@ -532,7 +535,7 @@ def test_anthropic_auto_setup_maps_auth_header(self): api_key = "sk-ant-test" with patch.dict("os.environ", {"ANTHROPIC_API_KEY": api_key}, clear=True): result = _auto_setup(client, "default") - assert result is True + assert result == "anthropic" call_kwargs = client.inference.providers.create.call_args.kwargs assert call_kwargs["name"] == "anthropic" assert call_kwargs["api_key_secret_name"] == "anthropic-api-key" @@ -554,9 +557,25 @@ def test_context_definition_has_default_model(self): ctx2 = ContextDefinition(name="test", cluster="c", user="u", default_model="my-model") assert ctx2.default_model == "my-model" + def test_context_definition_has_optional_fast_model(self): + ctx = ContextDefinition(name="test", cluster="c", user="u") + assert ctx.fast_model is None + + configured = ContextDefinition( + name="test", + cluster="c", + user="u", + fast_model="workspace/fast", + ) + assert configured.fast_model == "workspace/fast" + def test_config_params_accepts_default_model(self): - params: ConfigParams = {"default_model": "workspace/openai-gpt-4o"} + params: ConfigParams = { + "default_model": "workspace/openai-gpt-4o", + "fast_model": "workspace/openai-gpt-4o-mini", + } assert params["default_model"] == "workspace/openai-gpt-4o" + assert params["fast_model"] == "workspace/openai-gpt-4o-mini" def test_ensure_context_applies_default_model(self): config_file = ConfigFile( @@ -565,9 +584,13 @@ def test_ensure_context_applies_default_model(self): users=[{"name": "test-user", "type": "no-auth"}], contexts=[{"name": "test", "cluster": "test-cluster", "user": "test-user"}], ) - params: ConfigParams = {"default_model": "workspace/my-model"} + params: ConfigParams = { + "default_model": "workspace/my-model", + "fast_model": "workspace/fast", + } _, _, ctx_def = config_file.ensure_context("test", params) assert ctx_def.default_model == "workspace/my-model" + assert ctx_def.fast_model == "workspace/fast" def test_context_has_default_model(self): ctx = Context( @@ -575,9 +598,79 @@ def test_context_has_default_model(self): cluster=Cluster(name="c", base_url="http://localhost:8080"), workspace="default", default_model="workspace/gpt-4o", + fast_model="workspace/gpt-4o-mini", preferences={}, ) assert ctx.default_model == "workspace/gpt-4o" + assert ctx.fast_model == "workspace/gpt-4o-mini" + + def test_default_model_fills_missing_fast_role(self, tmp_path, monkeypatch): + monkeypatch.delenv("NEMO_DEFAULT_MODEL", raising=False) + monkeypatch.delenv("NEMO_FAST_MODEL", raising=False) + config_file = ConfigFile( + current_context="test", + clusters=[{"name": "test-cluster", "base_url": "http://localhost:8080"}], + users=[{"name": "test-user", "type": "no-auth"}], + contexts=[ + { + "name": "test", + "cluster": "test-cluster", + "user": "test-user", + "default_model": "workspace/legacy-model", + } + ], + ) + + context = Config.create(tmp_path / "config.yaml", config_file).resolve() + + assert context.default_model == "workspace/legacy-model" + assert context.fast_model == "workspace/legacy-model" + + def test_explicit_model_pair_resolves_independently(self, tmp_path, monkeypatch): + monkeypatch.delenv("NEMO_DEFAULT_MODEL", raising=False) + monkeypatch.delenv("NEMO_FAST_MODEL", raising=False) + config_file = ConfigFile( + current_context="test", + clusters=[{"name": "test-cluster", "base_url": "http://localhost:8080"}], + users=[{"name": "test-user", "type": "no-auth"}], + contexts=[ + { + "name": "test", + "cluster": "test-cluster", + "user": "test-user", + "default_model": "workspace/default", + "fast_model": "workspace/fast", + } + ], + ) + + context = Config.create(tmp_path / "config.yaml", config_file).resolve() + + assert context.default_model == "workspace/default" + assert context.fast_model == "workspace/fast" + + def test_model_environment_overrides_resolve_independently(self, tmp_path, monkeypatch): + monkeypatch.setenv("NEMO_DEFAULT_MODEL", "workspace/env-default") + monkeypatch.setenv("NEMO_FAST_MODEL", "workspace/env-fast") + config_file = ConfigFile( + current_context="test", + clusters=[{"name": "test-cluster", "base_url": "http://localhost:8080"}], + users=[{"name": "test-user", "type": "no-auth"}], + contexts=[ + { + "name": "test", + "cluster": "test-cluster", + "user": "test-user", + "default_model": "workspace/config-default", + "fast_model": "workspace/config-fast", + } + ], + ) + + context = Config.create(tmp_path / "config.yaml", config_file).resolve() + + assert context.default_model == "workspace/env-default" + assert context.fast_model == "workspace/env-fast" # --------------------------------------------------------------------------- @@ -918,6 +1011,72 @@ def test_updates_selected_context_and_runtime_url_override(self, tmp_path, monke assert str(dev_cluster.base_url) == "https://new-dev.example.com/" assert cli_context.overrides["base_url"] == "https://new-dev.example.com" + def test_local_connection_preserves_remote_context_and_uses_no_auth(self, tmp_path, monkeypatch): + config_path = tmp_path / "config.yaml" + monkeypatch.setenv("NMP_CONFIG_FILE", str(config_path)) + Config.write( + { + "base_url": "https://remote.example.com", + "access_token": "remote-token", + "refresh_token": "remote-refresh-token", + "current_context": "dev", + }, + context_name="dev", + ) + cli_context = MagicMock() + cli_context.overrides = { + "current_context": "dev", + "base_url": "https://remote.example.com", + } + + _configure_local_connection(cli_context, "local-workspace") + + config = Config.load(config_path=config_path) + config_file = config.get_config_file() + assert config_file.current_context == "local" + remote = next(context for context in config_file.contexts if context.name == "dev") + remote_cluster = next(cluster for cluster in config_file.clusters if cluster.name == remote.cluster) + remote_user = next(user for user in config_file.users if user.name == remote.user) + assert str(remote_cluster.base_url) == "https://remote.example.com/" + assert isinstance(remote_user, OAuthUser) + + local = config.resolve() + assert local.context_name == "local" + assert str(local.cluster.base_url) == "http://localhost:8080/" + assert local.workspace == "local-workspace" + assert isinstance(local.user, NoAuthUser) + assert cli_context.overrides == { + "current_context": "local", + "base_url": "http://localhost:8080", + } + cli_context.reset_sdk_context.assert_called_once_with() + + def test_local_connection_does_not_overwrite_existing_local_context(self, tmp_path, monkeypatch): + config_path = tmp_path / "config.yaml" + monkeypatch.setenv("NMP_CONFIG_FILE", str(config_path)) + Config.write( + { + "base_url": "https://remote.example.com", + "access_token": "remote-token", + "current_context": "local", + }, + context_name="local", + ) + cli_context = MagicMock() + cli_context.overrides = {} + + _configure_local_connection(cli_context, "local-workspace") + + config_file = Config.load(config_path=config_path).get_config_file() + original = next(context for context in config_file.contexts if context.name == "local") + original_cluster = next(cluster for cluster in config_file.clusters if cluster.name == original.cluster) + assert str(original_cluster.base_url) == "https://remote.example.com/" + assert config_file.current_context == "local-2" + assert cli_context.overrides == { + "current_context": "local-2", + "base_url": "http://localhost:8080", + } + def test_preserves_active_workspace_when_flag_not_explicit(self): ctx = MagicMock(spec=typer.Context) ctx.get_parameter_source.return_value = ParameterSource.DEFAULT @@ -1080,7 +1239,7 @@ def test_bootstrap_seeds_cluster_when_only_data_dir_persisted(self, tmp_path, mo ``_bootstrap_config_if_missing`` used to short-circuit on ``config_path.exists()``, leaving the file without any cluster. Later - steps (`_save_default_model` → ``Config.write`` → ``ensure_context``) + steps (`_save_model_pair` → ``Config.write`` → ``ensure_context``) then raised ``Cluster 'default-cluster' does not exist and no base_url provided to create it!`` when the user picked a default model. @@ -1719,7 +1878,7 @@ def test_auto_setup_updates_existing_anthropic_auth(self): api_key = "sk-ant-updated" with patch.dict("os.environ", {"ANTHROPIC_API_KEY": api_key}, clear=True): result = _auto_setup(client, "default") - assert result is True + assert result == "anthropic" call_kwargs = client.inference.providers.update.call_args.kwargs assert call_kwargs["api_key_secret_name"] == "anthropic-api-key" assert call_kwargs["auth_header_format"] == "X-Api-Key: {{ auth_secret }}" @@ -1733,7 +1892,7 @@ def test_auto_setup_updates_existing_provider_secret_binding(self): client = _make_mock_client(provider_exists=True, secret_exists=False) with patch.dict("os.environ", {"OPENAI_API_KEY": "sk-new-key"}, clear=True): result = _auto_setup(client, "default") - assert result is True + assert result == "openai" client.mock_secrets.create_secret.assert_called_once() client.inference.providers.update.assert_called_once() call_kwargs = client.inference.providers.update.call_args.kwargs @@ -1744,7 +1903,7 @@ def test_auto_setup_updates_existing_secret_value(self): client = _make_mock_client(provider_exists=True, secret_exists=True) with patch.dict("os.environ", {"NVIDIA_API_KEY": "nvapi-new"}, clear=True): result = _auto_setup(client, "default") - assert result is True + assert result == "nvidia-build" client.mock_secrets.update_secret.assert_called_once() update_call = client.mock_secrets.update_secret.call_args assert update_call.kwargs["name"] == "nvidia-build-api-key" @@ -1759,9 +1918,47 @@ def test_auto_setup_updates_existing_secret_value(self): # --------------------------------------------------------------------------- -class TestInteractiveDefaultModelSelection: +class TestInteractiveModelPairSelection: _MOD = "nemo_platform_ext.cli.commands.setup" + def test_run_scopes_picker_to_registered_provider(self): + client = MagicMock() + cli_context = MagicMock() + model_pair = ModelPair( + default="default/claude-sonnet-4-6", + fast="default/claude-haiku-4-5-20251001", + ) + + with ( + patch( + f"{self._MOD}._interactive_collect_provider", + return_value=("anthropic", "https://api.anthropic.com", None, None, None), + ), + patch(f"{self._MOD}._register_provider_interactive"), + patch( + f"{self._MOD}._wait_for_models", + return_value=[ + "default/claude-sonnet-4-6", + "default/claude-haiku-4-5-20251001", + ], + ), + patch(f"{self._MOD}._select_model_pair", return_value=model_pair) as select_model_pair, + patch(f"{self._MOD}._save_model_pair"), + patch(f"{self._MOD}._maybe_install_skills"), + patch(f"{self._MOD}._maybe_deploy_agent", return_value=False), + patch(f"{self._MOD}._print_onboarding"), + ): + _run_interactive_mode( + cli_context, + client, + "default", + "http://localhost:8080", + install_skills=False, + deploy_agent=False, + ) + + select_model_pair.assert_called_once_with(client, "default", provider_name="anthropic") + def test_skips_default_model_picker_when_new_provider_has_no_models(self): """When the new provider is still syncing, setup should not show a misleading picker.""" client = MagicMock() @@ -1774,8 +1971,8 @@ def test_skips_default_model_picker_when_new_provider_has_no_models(self): ), patch(f"{self._MOD}._register_provider_interactive"), patch(f"{self._MOD}._wait_for_models", return_value=[]), - patch(f"{self._MOD}._select_default_model") as mock_select_default_model, - patch(f"{self._MOD}._save_default_model"), + patch(f"{self._MOD}._select_model_pair") as mock_select_model_pair, + patch(f"{self._MOD}._save_model_pair"), patch(f"{self._MOD}._maybe_install_skills"), patch(f"{self._MOD}._maybe_deploy_agent"), patch(f"{self._MOD}._print_onboarding"), @@ -1790,11 +1987,11 @@ def test_skips_default_model_picker_when_new_provider_has_no_models(self): deploy_agent=False, ) - mock_select_default_model.assert_not_called() + mock_select_model_pair.assert_not_called() printed_lines = [call.args[0] for call in mock_console.print.call_args_list if call.args] assert any("No models discovered yet (provider may still be syncing)" in line for line in printed_lines) - assert any("Step 5: Choose default model" in line for line in printed_lines) - assert any("No default model set for this provider yet." in line for line in printed_lines) + assert any("Step 5: Choose agent models" in line for line in printed_lines) + assert any("No agent models set for this provider yet." in line for line in printed_lines) assert any("Run [cyan]nemo setup[/cyan] again after models sync" in line for line in printed_lines) def test_warns_when_only_existing_provider_models_are_available(self): @@ -1818,8 +2015,8 @@ def test_warns_when_only_existing_provider_models_are_available(self): ) ], ), - patch(f"{self._MOD}._select_default_model") as mock_select_default_model, - patch(f"{self._MOD}._save_default_model"), + patch(f"{self._MOD}._select_model_pair") as mock_select_model_pair, + patch(f"{self._MOD}._save_model_pair"), patch(f"{self._MOD}._maybe_install_skills"), patch(f"{self._MOD}._maybe_deploy_agent"), patch(f"{self._MOD}._print_onboarding"), @@ -1834,15 +2031,15 @@ def test_warns_when_only_existing_provider_models_are_available(self): deploy_agent=False, ) - mock_select_default_model.assert_not_called() + mock_select_model_pair.assert_not_called() printed_lines = [call.args[0] for call in mock_console.print.call_args_list if call.args] assert any( "Models from existing providers are available, but not from 'my-ollama-custom' yet." in line for line in printed_lines ) - def test_picker_labels_include_provider_names(self): - """Default model choices should show which provider each model comes from.""" + def test_picker_only_includes_models_from_selected_provider(self): + """Models from another provider in the workspace must not enter the picker.""" client = MagicMock() provider_a = MagicMock() provider_a.name = "nvidia-build" @@ -1856,20 +2053,167 @@ def test_picker_labels_include_provider_names(self): ] client.inference.providers.list.return_value = MagicMock(data=[provider_a, provider_b]) - with patch(f"{self._MOD}.prompt_select", return_value="default/qwen2.5:1.5b") as mock_prompt_select: - result = _select_default_model(client, "default") + with patch( + f"{self._MOD}.prompt_select", + side_effect=["default/qwen2.5:1.5b", "default/qwen2.5:1.5b"], + ) as mock_prompt_select: + result = _select_model_pair(client, "default", provider_name="my-ollama-custom") - assert result == "default/qwen2.5:1.5b" - assert mock_prompt_select.call_args.kwargs["choices"] == [ - ( - "default/meta-llama-3-1-8b-instruct", - "meta-llama-3-1-8b-instruct (nvidia-build)", - ), + assert result == ModelPair( + default="default/qwen2.5:1.5b", + fast="default/qwen2.5:1.5b", + ) + expected_choices = [ ( "default/qwen2.5:1.5b", "qwen2.5:1.5b (my-ollama-custom)", ), ] + assert mock_prompt_select.call_count == 2 + assert all(call.kwargs["choices"] == expected_choices for call in mock_prompt_select.call_args_list) + default_call, fast_call = mock_prompt_select.call_args_list + assert default_call.args[0] == "Choose your default model (used for quality-critical agent work):" + assert default_call.kwargs["default"] == "default/qwen2.5:1.5b" + assert default_call.kwargs["hint"] == "Press Enter to accept the default." + assert fast_call.args[0] == "Choose your fast model (used for latency-sensitive agent work):" + assert fast_call.kwargs["default"] == "default/qwen2.5:1.5b" + assert fast_call.kwargs["hint"] == "Press Enter to reuse the default model." + + +class TestAutoModelPairSelection: + def test_entity_discovery_filters_other_workspace_providers(self): + client = MagicMock() + openai = MagicMock() + openai.name = "openai" + openai.served_models = [MagicMock(model_entity_id="default/gpt-4.1")] + anthropic = MagicMock() + anthropic.name = "anthropic" + anthropic.served_models = [MagicMock(model_entity_id="default/claude-sonnet-4-6")] + client.inference.providers.list.return_value = MagicMock(data=[openai, anthropic]) + + result = setup_commands._get_all_model_entity_ids( + client, + "default", + provider_name="anthropic", + ) + + assert result == ["default/claude-sonnet-4-6"] + + def test_explicit_default_and_fast_models_are_saved_independently(self): + cli_context = MagicMock() + with ( + patch.dict( + "os.environ", + { + "NEMO_DEFAULT_MODEL": "default/quality", + "NEMO_FAST_MODEL": "default/fast", + }, + clear=True, + ), + patch(f"{SETUP_MOD}._auto_setup", return_value="anthropic"), + patch(f"{SETUP_MOD}._get_all_model_entity_ids", return_value=["default/discovered"]), + patch(f"{SETUP_MOD}._save_model_pair") as save_pair, + patch(f"{SETUP_MOD}._maybe_install_skills"), + patch(f"{SETUP_MOD}._maybe_deploy_agent"), + patch(f"{SETUP_MOD}._verify_platform_health", return_value=True), + ): + _run_auto_mode( + cli_context, + MagicMock(), + "default", + "http://localhost:8080", + install_skills=False, + deploy_agent=False, + ) + + save_pair.assert_called_once_with(cli_context, ModelPair(default="default/quality", fast="default/fast")) + + def test_discovered_default_is_scoped_to_auto_configured_provider(self): + cli_context = MagicMock() + client = MagicMock() + with ( + patch.dict("os.environ", {}, clear=True), + patch(f"{SETUP_MOD}._auto_setup", return_value="anthropic"), + patch( + f"{SETUP_MOD}._get_all_model_entity_ids", + return_value=["default/claude-sonnet-4-6"], + ) as get_model_ids, + patch(f"{SETUP_MOD}._save_model_pair") as save_pair, + patch(f"{SETUP_MOD}._maybe_install_skills"), + patch(f"{SETUP_MOD}._maybe_deploy_agent"), + patch(f"{SETUP_MOD}._verify_platform_health", return_value=True), + ): + _run_auto_mode( + cli_context, + client, + "default", + "http://localhost:8080", + install_skills=False, + deploy_agent=False, + ) + + get_model_ids.assert_called_once_with(client, "default", provider_name="anthropic") + save_pair.assert_called_once_with( + cli_context, + ModelPair( + default="default/claude-sonnet-4-6", + fast="default/claude-sonnet-4-6", + ), + ) + + def test_first_discovered_model_is_default_and_fast_fallback(self): + cli_context = MagicMock() + with ( + patch.dict("os.environ", {}, clear=True), + patch(f"{SETUP_MOD}._auto_setup", return_value="openai"), + patch( + f"{SETUP_MOD}._get_all_model_entity_ids", + return_value=["default/a-model", "default/z-model"], + ), + patch(f"{SETUP_MOD}._save_model_pair") as save_pair, + patch(f"{SETUP_MOD}._maybe_install_skills"), + patch(f"{SETUP_MOD}._maybe_deploy_agent"), + patch(f"{SETUP_MOD}._verify_platform_health", return_value=True), + ): + _run_auto_mode( + cli_context, + MagicMock(), + "default", + "http://localhost:8080", + install_skills=False, + deploy_agent=False, + ) + + save_pair.assert_called_once_with( + cli_context, + ModelPair(default="default/a-model", fast="default/a-model"), + ) + + def test_fast_override_without_default_warns_and_is_not_saved(self): + cli_context = MagicMock() + with ( + patch.dict("os.environ", {"NEMO_FAST_MODEL": "default/fast"}, clear=True), + patch(f"{SETUP_MOD}._MODEL_DISCOVERY_MAX_ROUNDS", 1), + patch(f"{SETUP_MOD}._MODEL_DISCOVERY_ROUND_SECONDS", 0), + patch(f"{SETUP_MOD}._auto_setup", return_value="openai"), + patch(f"{SETUP_MOD}._get_all_model_entity_ids", return_value=[]), + patch(f"{SETUP_MOD}._save_model_pair") as save_pair, + patch(f"{SETUP_MOD}._maybe_install_skills"), + patch(f"{SETUP_MOD}._maybe_deploy_agent"), + patch(f"{SETUP_MOD}._verify_platform_health", return_value=True), + patch(f"{SETUP_MOD}.console.print") as print_message, + ): + _run_auto_mode( + cli_context, + MagicMock(), + "default", + "http://localhost:8080", + install_skills=False, + deploy_agent=False, + ) + + save_pair.assert_not_called() + assert any("NEMO_FAST_MODEL is ignored" in call.args[0] for call in print_message.call_args_list) # --------------------------------------------------------------------------- @@ -2756,6 +3100,7 @@ def _patch_setup_command( run_interactive=stack.enter_context(patch(f"{SETUP_MOD}._run_interactive_mode")), prompt_remote=None, configure_remote=None, + configure_local=stack.enter_context(patch(f"{SETUP_MOD}._configure_local_connection")), ensure_auth=None, config_write=None, run_auto=None, @@ -2850,7 +3195,6 @@ def test_start_local_persists_default_url_and_starts_services(self): ctx, cli_context = _make_setup_command_ctx(base_url="https://remote.example.com") with _patch_setup_command( maybe_start_services=["start_local", "ready"], - include_config_write=True, ) as mocks: setup_command(ctx) @@ -2861,10 +3205,7 @@ def test_start_local_persists_default_url_and_starts_services(self): start_services=True, timeout=_SERVICE_STARTUP_TIMEOUT_SECONDS, ) - mocks.bootstrap.assert_any_call(DEFAULT_BASE_URL, "default") - mocks.config_write.assert_called_once_with({"base_url": DEFAULT_BASE_URL}, context_name="default") - assert cli_context.overrides["base_url"] == DEFAULT_BASE_URL - cli_context.reset_sdk_context.assert_called() + mocks.configure_local.assert_called_once_with(cli_context, "default") assert mocks.run_interactive.call_args.args[3] == DEFAULT_BASE_URL diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/nooa_model_client.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/nooa_model_client.py new file mode 100644 index 0000000000..1ed2f55a26 --- /dev/null +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/nooa_model_client.py @@ -0,0 +1,242 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Run-scoped Nooa clients backed by configured Platform Model Entities. + +This module is an optional integration: consumers must depend on ``nooa``. +Keeping the integration here lets plugins share Platform config, routing, auth, +and client lifetime behavior without making Nooa a required dependency of the +public plugin contract. +""" + +from collections.abc import Iterator +from contextlib import contextmanager +from contextvars import ContextVar +from dataclasses import dataclass + +from litellm import get_model_info +from nemo_platform import AsyncNeMoPlatform +from nemo_platform.config import get_context +from nemo_platform.types.inference import ModelProvider +from nemo_platform.types.models import ModelEntity +from nooa.unifiedllm import CompletionClient, UnifiedLLM + +_PLACEHOLDER_API_KEY = "not-needed" +_OPENAI_FORMAT = "OPENAI_CHAT" +_ANTHROPIC_FORMAT = "ANTHROPIC_MESSAGES" +_OPENAI_CHAT_REASONING_EFFORT = "none" +_ACCEPT_ENCODING_HEADER = "accept-encoding" +_IDENTITY_ENCODING = "identity" + + +def _openai_chat_reasoning_effort(model: str) -> str | None: + """Disable reasoning unless LiteLLM explicitly says this value is unsupported.""" + try: + model_info = get_model_info(model) + except Exception as exc: + # Custom provider registrations need not appear in LiteLLM's model map. + # Preserve the value that makes frontier Chat Completions tool calls work. + if "This model isn't mapped yet" not in str(exc): + raise + return _OPENAI_CHAT_REASONING_EFFORT + if model_info.get("supports_none_reasoning_effort") is False: + return None + return _OPENAI_CHAT_REASONING_EFFORT + + +@dataclass(frozen=True) +class ConfiguredModelRefs: + """Workspace-qualified Model Entity IDs for default and fast agent work.""" + + default: str + fast: str + + +@dataclass(frozen=True) +class ConfiguredModelClients: + """Resolved Nooa clients for default and low-latency agent work.""" + + default: UnifiedLLM + fast: UnifiedLLM + refs: ConfiguredModelRefs | None = None + + async def aclose(self) -> None: + """Close each distinct client owned by this pair.""" + close_error: Exception | None = None + clients = (self.default,) if self.fast is self.default else (self.default, self.fast) + for client in clients: + try: + await client.aclose() + except Exception as exc: + if close_error is None: + close_error = exc + else: + close_error.add_note(f"Another model client also failed to close: {exc!r}") + if close_error is not None: + raise close_error + + +_ACTIVE_MODEL_CLIENTS: ContextVar[ConfiguredModelClients | None] = ContextVar( + "nemo_platform_active_model_clients", + default=None, +) + + +def configured_model_refs() -> ConfiguredModelRefs: + """Read the default/fast Model Entity pair from the active CLI context.""" + context = get_context() + default = context.default_model + fast = context.fast_model or default + if not default: + raise ValueError("No default model is configured. Run `nemo setup` and select agent models.") + if not fast: + raise ValueError("No fast model is configured. Run `nemo setup` and select agent models.") + return ConfiguredModelRefs(default=default, fast=fast) + + +def _parse_model_ref(model_ref: str) -> tuple[str, str]: + workspace, separator, name = model_ref.partition("/") + if separator != "/" or not workspace or not name: + raise ValueError(f"Model reference {model_ref!r} must use workspace/name format") + return workspace, name + + +def _completion_client( + client: AsyncNeMoPlatform, + model_entity: ModelEntity, + served_model_name: str, +) -> CompletionClient: + """Build a Nooa client that routes through the Model Entity's Platform URL.""" + api_base = client.models.get_model_entity_route_openai_url(model_entity) + extra_headers = dict(client.models.get_client_default_headers()) + # Inference Gateway's direct passthrough session preserves compressed + # response bytes. Nooa needs decoded JSON/SSE, so make that requirement + # explicit at this adapter boundary for every configured agent client. + extra_headers[_ACCEPT_ENCODING_HEADER] = _IDENTITY_ENCODING + # Backend format is the Platform-facing wire contract, not the upstream + # provider identity. The LiteLLM prefix selects the adapter for that shape. + if model_entity.backend_format == _OPENAI_FORMAT: + litellm_model = f"openai/{served_model_name}" + return CompletionClient( + litellm_model, + api_base=api_base, + api_key=_PLACEHOLDER_API_KEY, + # Platform rewrites the response model to the Model Entity ID. Keep + # the provider's authoritative served name available to LiteLLM for + # post-response capability and cost lookup. + base_model=litellm_model, + extra_headers=extra_headers, + drop_params=True, + # The configured backend contract is explicitly OPENAI_CHAT. Newer + # LiteLLM versions otherwise bridge GPT-5.4+ tool calls with any + # reasoning_effort value (including "none") to /responses. + _skip_responses_api_bridge=True, + # Chat Completions tool calls on current OpenAI frontier models + # require reasoning to be disabled. Use LiteLLM's capability map to + # omit that value only when the served model explicitly rejects it. + reasoning_effort=_openai_chat_reasoning_effort(litellm_model), + ) + elif model_entity.backend_format == _ANTHROPIC_FORMAT: + api_base = api_base.removesuffix("/v1") + else: + raise ValueError( + f"Model '{model_entity.workspace}/{model_entity.name}' has unsupported backend format " + f"{model_entity.backend_format!r}; expected {_OPENAI_FORMAT} or {_ANTHROPIC_FORMAT}" + ) + + litellm_model = f"anthropic/{served_model_name}" + return CompletionClient( + litellm_model, + api_base=api_base, + api_key=_PLACEHOLDER_API_KEY, + base_model=litellm_model, + extra_headers=extra_headers, + drop_params=True, + ) + + +async def _served_model_name( + client: AsyncNeMoPlatform, + model_entity: ModelEntity, + provider_cache: dict[str, ModelProvider], +) -> str: + """Return the provider-recorded model name, without inferring from the entity ID.""" + entity_ref = f"{model_entity.workspace}/{model_entity.name}" + for provider_ref in model_entity.model_providers or (): + provider = provider_cache.get(provider_ref) + if provider is None: + workspace, name = _parse_model_ref(provider_ref) + provider = await client.inference.providers.retrieve(name, workspace=workspace) + provider_cache[provider_ref] = provider + for mapping in provider.served_models or (): + if mapping.model_entity_id == entity_ref: + return mapping.served_model_name + + if model_entity.api_endpoint is not None and model_entity.api_endpoint.model_id: + return model_entity.api_endpoint.model_id + return model_entity.name + + +async def resolve_model_clients( + client: AsyncNeMoPlatform, + refs: ConfiguredModelRefs | None = None, +) -> ConfiguredModelClients: + """Resolve configured Model Entities and construct each distinct client once.""" + selected = refs or configured_model_refs() + resolved: dict[str, UnifiedLLM] = {} + provider_cache: dict[str, ModelProvider] = {} + try: + for model_ref in (selected.default, selected.fast): + if model_ref in resolved: + continue + workspace, name = _parse_model_ref(model_ref) + entity = await client.models.retrieve(name, workspace=workspace) + served_model_name = await _served_model_name(client, entity, provider_cache) + resolved[model_ref] = _completion_client(client, entity, served_model_name) + except Exception as resolution_error: + for model_client in resolved.values(): + try: + await model_client.aclose() + except Exception as cleanup_error: + resolution_error.add_note(f"A constructed model client also failed to close: {cleanup_error!r}") + raise + return ConfiguredModelClients( + default=resolved[selected.default], + fast=resolved[selected.fast], + refs=selected, + ) + + +@contextmanager +def activate_model_clients(model_clients: ConfiguredModelClients) -> Iterator[None]: + """Expose a resolved pair to agents constructed within this run context.""" + token = _ACTIVE_MODEL_CLIENTS.set(model_clients) + try: + yield + finally: + _ACTIVE_MODEL_CLIENTS.reset(token) + + +def _active_model_clients() -> ConfiguredModelClients: + model_clients = _ACTIVE_MODEL_CLIENTS.get() + if model_clients is None: + raise RuntimeError("Configured model clients were not activated for this agent run") + return model_clients + + +def get_default_model() -> UnifiedLLM: + """Return the active default model client.""" + return _active_model_clients().default + + +def get_fast_model() -> UnifiedLLM: + """Return the active low-latency model client.""" + return _active_model_clients().fast + + +def get_configured_model_refs() -> ConfiguredModelRefs: + """Return the model references resolved for the active agent run.""" + refs = _active_model_clients().refs + if refs is None: + raise RuntimeError("Configured model references are unavailable for this agent run") + return refs diff --git a/packages/nemo_platform_plugin/tests/test_nooa_model_client.py b/packages/nemo_platform_plugin/tests/test_nooa_model_client.py new file mode 100644 index 0000000000..8a92eec514 --- /dev/null +++ b/packages/nemo_platform_plugin/tests/test_nooa_model_client.py @@ -0,0 +1,282 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest +from nemo_platform_plugin import nooa_model_client +from nemo_platform_plugin.nooa_model_client import ( + ConfiguredModelClients, + ConfiguredModelRefs, + _openai_chat_reasoning_effort, + activate_model_clients, + configured_model_refs, + get_configured_model_refs, + get_default_model, + get_fast_model, + resolve_model_clients, +) + + +def test_configured_model_refs_uses_default_for_missing_fast(monkeypatch): + monkeypatch.setattr( + nooa_model_client, + "get_context", + lambda: SimpleNamespace(default_model="default/quality", fast_model=None), + ) + + assert configured_model_refs() == ConfiguredModelRefs( + default="default/quality", + fast="default/quality", + ) + + +def test_configured_model_refs_requires_default(monkeypatch): + monkeypatch.setattr( + nooa_model_client, + "get_context", + lambda: SimpleNamespace(default_model=None, fast_model="default/fast"), + ) + + with pytest.raises(ValueError, match="No default model"): + configured_model_refs() + + +async def test_resolve_model_clients_deduplicates_same_model(monkeypatch): + model_entity = SimpleNamespace( + workspace="default", + name="gpt-4-1", + backend_format="OPENAI_CHAT", + model_providers=[], + api_endpoint=None, + ) + client = MagicMock() + client.models.retrieve = AsyncMock(return_value=model_entity) + client.models.get_model_entity_route_openai_url.return_value = "http://platform/model/gpt-4-1/-/v1" + default_headers = {"x-test": "value"} + client.models.get_client_default_headers.return_value = default_headers + completion_client = MagicMock() + factory = MagicMock(return_value=completion_client) + monkeypatch.setattr(nooa_model_client, "CompletionClient", factory) + + result = await resolve_model_clients( + client, + ConfiguredModelRefs(default="default/gpt-4-1", fast="default/gpt-4-1"), + ) + + assert result.default is completion_client + assert result.fast is completion_client + client.models.retrieve.assert_awaited_once_with("gpt-4-1", workspace="default") + factory.assert_called_once_with( + "openai/gpt-4-1", + api_base="http://platform/model/gpt-4-1/-/v1", + api_key="not-needed", + base_model="openai/gpt-4-1", + extra_headers={"x-test": "value", "accept-encoding": "identity"}, + drop_params=True, + _skip_responses_api_bridge=True, + reasoning_effort="none", + ) + assert default_headers == {"x-test": "value"} + + +async def test_resolve_model_clients_uses_anthropic_route_shape(monkeypatch): + model_entity = SimpleNamespace( + workspace="default", + name="claude-sonnet-4", + backend_format="ANTHROPIC_MESSAGES", + model_providers=[], + api_endpoint=None, + ) + client = MagicMock() + client.models.retrieve = AsyncMock(return_value=model_entity) + client.models.get_model_entity_route_openai_url.return_value = "http://platform/model/claude-sonnet-4/-/v1" + client.models.get_client_default_headers.return_value = {} + factory = MagicMock(return_value=MagicMock()) + monkeypatch.setattr(nooa_model_client, "CompletionClient", factory) + + await resolve_model_clients( + client, + ConfiguredModelRefs(default="default/claude-sonnet-4", fast="default/claude-sonnet-4"), + ) + + factory.assert_called_once_with( + "anthropic/claude-sonnet-4", + api_base="http://platform/model/claude-sonnet-4/-", + api_key="not-needed", + base_model="anthropic/claude-sonnet-4", + extra_headers={"accept-encoding": "identity"}, + drop_params=True, + ) + + +async def test_resolve_model_clients_rejects_unsupported_backend_format(): + model_entity = SimpleNamespace( + workspace="default", + name="responses-only", + backend_format="OPENAI_RESPONSES", + model_providers=[], + api_endpoint=None, + ) + client = MagicMock() + client.models.retrieve = AsyncMock(return_value=model_entity) + + with pytest.raises(ValueError, match="unsupported backend format 'OPENAI_RESPONSES'"): + await resolve_model_clients( + client, + ConfiguredModelRefs(default="default/responses-only", fast="default/responses-only"), + ) + + +async def test_resolve_model_clients_requires_workspace_qualified_refs(): + client = MagicMock() + + with pytest.raises(ValueError, match="workspace/name"): + await resolve_model_clients( + client, + ConfiguredModelRefs(default="unqualified", fast="unqualified"), + ) + + +async def test_resolve_model_clients_uses_provider_served_name(monkeypatch): + model_entity = SimpleNamespace( + workspace="default", + name="gpt-5-6-sol", + backend_format="OPENAI_CHAT", + model_providers=["default/openai"], + api_endpoint=SimpleNamespace(model_id="gpt-5-6-sol"), + ) + provider = SimpleNamespace( + served_models=[ + SimpleNamespace( + model_entity_id="default/gpt-5-6-sol", + served_model_name="gpt-5.6-sol", + ) + ] + ) + client = MagicMock() + client.models.retrieve = AsyncMock(return_value=model_entity) + client.inference.providers.retrieve = AsyncMock(return_value=provider) + client.models.get_model_entity_route_openai_url.return_value = "http://platform/model/gpt-5-6-sol/-/v1" + client.models.get_client_default_headers.return_value = {} + factory = MagicMock(return_value=MagicMock()) + monkeypatch.setattr(nooa_model_client, "CompletionClient", factory) + + await resolve_model_clients( + client, + ConfiguredModelRefs(default="default/gpt-5-6-sol", fast="default/gpt-5-6-sol"), + ) + + client.inference.providers.retrieve.assert_awaited_once_with("openai", workspace="default") + factory.assert_called_once_with( + "openai/gpt-5.6-sol", + api_base="http://platform/model/gpt-5-6-sol/-/v1", + api_key="not-needed", + base_model="openai/gpt-5.6-sol", + extra_headers={"accept-encoding": "identity"}, + drop_params=True, + _skip_responses_api_bridge=True, + reasoning_effort="none", + ) + + +def test_active_model_clients_are_scoped(): + default = MagicMock() + fast = MagicMock() + refs = ConfiguredModelRefs(default="default/quality", fast="default/fast") + pair = ConfiguredModelClients(default=default, fast=fast, refs=refs) + + with activate_model_clients(pair): + assert get_default_model() is default + assert get_fast_model() is fast + assert get_configured_model_refs() == refs + + with pytest.raises(RuntimeError, match="were not activated"): + get_default_model() + + +async def test_model_clients_close_each_distinct_client_once(): + default = MagicMock() + default.aclose = AsyncMock() + fast = MagicMock() + fast.aclose = AsyncMock() + + await ConfiguredModelClients(default=default, fast=fast).aclose() + + default.aclose.assert_awaited_once() + fast.aclose.assert_awaited_once() + + +async def test_model_clients_close_shared_client_once(): + shared = MagicMock() + shared.aclose = AsyncMock() + + await ConfiguredModelClients(default=shared, fast=shared).aclose() + + shared.aclose.assert_awaited_once() + + +async def test_model_clients_close_fast_after_default_close_fails(): + default = MagicMock() + default.aclose = AsyncMock(side_effect=RuntimeError("default close failed")) + fast = MagicMock() + fast.aclose = AsyncMock() + + with pytest.raises(RuntimeError, match="default close failed"): + await ConfiguredModelClients(default=default, fast=fast).aclose() + + default.aclose.assert_awaited_once() + fast.aclose.assert_awaited_once() + + +async def test_resolve_model_clients_closes_constructed_client_after_failure(monkeypatch): + default_entity = SimpleNamespace( + workspace="default", + name="quality", + backend_format="OPENAI_CHAT", + model_providers=[], + api_endpoint=SimpleNamespace(model_id="quality"), + ) + client = MagicMock() + client.models.retrieve = AsyncMock(side_effect=[default_entity, RuntimeError("fast resolution failed")]) + constructed = MagicMock() + constructed.aclose = AsyncMock() + monkeypatch.setattr(nooa_model_client, "_completion_client", MagicMock(return_value=constructed)) + + with pytest.raises(RuntimeError, match="fast resolution failed"): + await resolve_model_clients( + client, + ConfiguredModelRefs(default="default/quality", fast="default/fast"), + ) + + constructed.aclose.assert_awaited_once() + + +def test_openai_chat_reasoning_effort_uses_litellm_capability(monkeypatch): + monkeypatch.setattr( + nooa_model_client, + "get_model_info", + lambda model: {"supports_none_reasoning_effort": False}, + ) + + assert _openai_chat_reasoning_effort("openai/gpt-5-mini") is None + + +def test_openai_chat_reasoning_effort_preserves_none_for_unmapped_models(monkeypatch): + def unmapped(model: str): + raise Exception("This model isn't mapped yet") + + monkeypatch.setattr(nooa_model_client, "get_model_info", unmapped) + + assert _openai_chat_reasoning_effort("openai/custom-model") == "none" + + +def test_openai_chat_reasoning_effort_does_not_hide_lookup_failures(monkeypatch): + def failed(model: str): + raise RuntimeError("model registry failed") + + monkeypatch.setattr(nooa_model_client, "get_model_info", failed) + + with pytest.raises(RuntimeError, match="model registry failed"): + _openai_chat_reasoning_effort("openai/gpt-5-mini") diff --git a/plugins/nemo-eval-author/.env.example b/plugins/nemo-eval-author/.env.example index d45305de9a..c363e4ef80 100644 --- a/plugins/nemo-eval-author/.env.example +++ b/plugins/nemo-eval-author/.env.example @@ -1,43 +1,9 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # -# Example environment for standalone Eval Author runs. -# Copy to `.env` next to this file (or export in your shell): -# -# cp .env.example .env -# -# AUTHOR_* is Eval Author's credential contract. The EXPERIMENTALIST_* interplay -# below is transitional and goes away once Eval Author no longer reuses any -# Experimentalist code: -# -# - an unset AUTHOR_* falls back to its EXPERIMENTALIST_* equivalent, so insight -# mode works from one Experimentalist profile .env -# - a set AUTHOR_* is copied into any unset EXPERIMENTALIST_* slot, because the -# Experimentalist helpers Eval Author still borrows read only EXPERIMENTALIST_* -# -# TODO(eval-author-standalone): remove both, at which point AUTHOR_* becomes required -# here rather than optional. Set AUTHOR_* explicitly now to avoid that break. -# -# TODO(cli): wire `nemo eval-author` to load this `.env` the way Experimentalist -# loads the profile-dir `.env`. - -# Required for LLM-backed Eval Author / TraceAnalyzer (NOOA CompletionClient). -AUTHOR_API_BASE=https://inference-api.nvidia.com/v1 -AUTHOR_API_KEY= - -# Optional model overrides (defaults shown). Eval Author itself uses the smart and fast -# tiers; the mid tier is listed only because it is bridged to Experimentalist helpers, -# and it goes away with the bridge. -# AUTHOR_SMART_MODEL_NAME=openai/openai/openai/gpt-5.5 -# AUTHOR_MID_MODEL_NAME=openai/gcp/google/gemini-3.5-flash -# AUTHOR_FAST_MODEL_NAME=openai/openai/openai/gpt-5-mini - -# Optional: NVIDIA Inference Gateway virtual key. When AUTHOR_API_KEY is unset and -# AUTHOR_API_BASE is the gateway over HTTPS, standalone runners fall back to -# INFERENCE_API_KEY the same way Experimentalist does today. The HTTPS requirement -# is deliberate: the key is never forwarded to a plain-http base. -# INFERENCE_API_KEY= - -# Optional: NeMo Platform URL for Insight / Fileset / agent-code access. -# Leave unset to use the active `nemo auth` context, or set for a local stack. +# Run `nemo setup` to register a provider and select default and fast models. +# The stored active context is used by default. These environment variables are +# only needed to override it in a non-interactive or isolated environment. # NMP_BASE_URL=http://localhost:8080 +# NEMO_DEFAULT_MODEL=default/gpt-4-1 +# NEMO_FAST_MODEL=default/gpt-4-1-mini diff --git a/plugins/nemo-eval-author/README.md b/plugins/nemo-eval-author/README.md index 17819e18aa..6e105239c5 100644 --- a/plugins/nemo-eval-author/README.md +++ b/plugins/nemo-eval-author/README.md @@ -38,31 +38,23 @@ from nemo_experimentalist_plugin.experimentalist.components.trace_analyzer impor from nemo_experimentalist_plugin.experimentalist.components.trace_explorer import TraceExplorer ``` -## Credentials (standalone) +## Agent models -Copy [`.env.example`](.env.example) to `.env` and set `AUTHOR_API_KEY` (and optionally model names / `NMP_BASE_URL`): +Run `nemo setup` and select the default and fast models for the active Platform +context. Eval Author uses the default model for authoring and the fast model for +summarization. Press Enter at the fast-model prompt to reuse the default model. -```bash -cp plugins/nemo-eval-author/.env.example plugins/nemo-eval-author/.env -``` - -`AUTHOR_*` is Eval Author's credential contract, and `model_config` imports nothing -from Experimentalist. When the API base is the NVIDIA Inference Gateway over HTTPS, -`INFERENCE_API_KEY` is also accepted. - -Two pieces of that module are transitional and disappear with the last -Experimentalist import, both tagged `TODO(eval-author-standalone)`: +The selections are workspace-qualified Platform Model Entity IDs. The Platform +routes each request to the provider registered for that entity and reads its +credential from Platform Secrets; Eval Author does not accept separate provider, +endpoint, key, or model environment variables. -- unset `AUTHOR_*` variables fall back to `NEMO_EXPERIMENTALIST_*`, so insight mode works - from a single Experimentalist profile `.env`. Setting `AUTHOR_*` explicitly today - avoids the break when the fallback is removed. -- `EvalAuthor.__init__` calls `bridge_author_env_to_experimentalist()`, copying - `AUTHOR_*` into unset `NEMO_EXPERIMENTALIST_*` slots so the Experimentalist helpers - Eval Author still borrows see credentials during a standalone run. +For non-interactive and isolated environments, `NEMO_DEFAULT_MODEL` and +`NEMO_FAST_MODEL` can override the stored selections. Values must still use +`workspace/model-name` and refer to Model Entities on the target Platform. A `nemo agents eval-author` CLI is registered under `nemo.cli.agents` and mounted by the agents plugin. Verb scaffolding is in place (`discover`, `audit`, `propose`, `run`, `doctor`); bodies are still -placeholders until ASE-673–678 land. The CLI does not auto-load this `.env` -yet — set credentials in the environment (or rely on the transitional -`EXPERIMENTALIST_*` fallback) before invoking it. +placeholders until ASE-673–678 land. The library runner already uses the +configured Platform model pair. diff --git a/plugins/nemo-eval-author/src/nemo_eval_author_plugin/cli.py b/plugins/nemo-eval-author/src/nemo_eval_author_plugin/cli.py index 8e9b2d3f79..3f35bcb7a3 100644 --- a/plugins/nemo-eval-author/src/nemo_eval_author_plugin/cli.py +++ b/plugins/nemo-eval-author/src/nemo_eval_author_plugin/cli.py @@ -11,11 +11,8 @@ until its own ticket lands. Flags belong to those tickets, so nothing here declares options yet. -This module imports nothing from ``eval_author``: those agents build their LLM client -while the class body executes, so importing one here would make the whole CLI require -``AUTHOR_*`` credentials — including the ``doctor`` verb whose job is to report that they -are missing. Experimentalist's CLI keeps its runner behind a lazily-assigned module -global for the same reason. +This module imports nothing from ``eval_author`` while the command bodies remain +placeholders. The completed runner resolves the active Platform model pair at run time. """ from typing import ClassVar, NoReturn diff --git a/plugins/nemo-eval-author/src/nemo_eval_author_plugin/eval_author/agent.py b/plugins/nemo-eval-author/src/nemo_eval_author_plugin/eval_author/agent.py index 0144a61c54..2b1b292a60 100644 --- a/plugins/nemo-eval-author/src/nemo_eval_author_plugin/eval_author/agent.py +++ b/plugins/nemo-eval-author/src/nemo_eval_author_plugin/eval_author/agent.py @@ -14,11 +14,6 @@ from nemo_eval_author_plugin.eval_author.materialization import InsightSuite, validate_metric_contracts from nemo_eval_author_plugin.eval_author.models import EvalAuthorConfig, EvalAuthorResult, MetricAuthoringResult -from nemo_eval_author_plugin.model_config import ( - bridge_author_env_to_experimentalist, - get_fast_model, - get_smart_model, -) from nemo_experimentalist_plugin.entities import ( Dataset, DatasetValidationError, @@ -43,6 +38,7 @@ from nemo_experimentalist_plugin.experimentalist.reporting import RunReporter from nemo_insights_plugin.entities import Insight from nemo_platform import AsyncNeMoPlatform +from nemo_platform_plugin.nooa_model_client import get_default_model, get_fast_model from nooa import Agent, CodeActStrategy, strategy from nooa.agentdoc import doc from nooa.agents import TokenBudgetSummarizer @@ -76,11 +72,7 @@ def __init__( When set, emits mid-run progress lines; never owns header/footer. **kwargs: Forwarded to ``Agent.__init__``. """ - # Eval Author still reuses a few Experimentalist agents (TraceAnalyzer, - # TraceExplorer). They read NEMO_EXPERIMENTALIST_* when constructed, so bridge the - # AUTHOR_* credentials before any of them is built. - bridge_author_env_to_experimentalist() - super().__init__(llm=kwargs.pop("llm", None) or get_smart_model(), **kwargs) + super().__init__(llm=kwargs.pop("llm", None) or get_default_model(), **kwargs) self._config = config or EvalAuthorConfig() self._reporter = reporter self.experiment_dir = experiment_dir diff --git a/plugins/nemo-eval-author/src/nemo_eval_author_plugin/eval_author/run.py b/plugins/nemo-eval-author/src/nemo_eval_author_plugin/eval_author/run.py index 179bd45870..e43e840c8d 100644 --- a/plugins/nemo-eval-author/src/nemo_eval_author_plugin/eval_author/run.py +++ b/plugins/nemo-eval-author/src/nemo_eval_author_plugin/eval_author/run.py @@ -3,9 +3,8 @@ """Reusable Eval Author run orchestration.""" -import importlib from pathlib import Path -from typing import Protocol, cast +from typing import Protocol from nemo_eval_author_plugin.eval_author.models import EvalAuthorConfig, EvalAuthorResult from nemo_experimentalist_plugin.client import make_client @@ -19,10 +18,13 @@ from nemo_experimentalist_plugin.experimentalist.reporting import RunReporter from nemo_insights_plugin.entities import Insight from nemo_platform import AsyncNeMoPlatform - - -class _LiteLLMModule(Protocol): - drop_params: bool +from nemo_platform_plugin.nooa_model_client import ( + ConfiguredModelClients, + ConfiguredModelRefs, + activate_model_clients, + configured_model_refs, + resolve_model_clients, +) class _EvalAuthorAgent(Protocol): @@ -50,6 +52,7 @@ async def run_eval_author( config: EvalAuthorConfig, agent: Path | str | None = None, evaluator_type: EvaluatorType = "harbor", + model_refs: ConfiguredModelRefs | None = None, ) -> EvalAuthorResult: """Stage evaluation inputs, resolve one Insight, then run Eval Author. @@ -62,20 +65,24 @@ async def run_eval_author( workspace: Platform workspace. base_url: Platform base URL. ``None`` uses the active platform context. config: Eval Author tuning parameters. - agent: Optional agent source override. The Insight's agent is the default. - evaluator_type: Evaluator adapter used to parse the task template. + agent: Optional agent source override. When absent, the Insight's agent is used. + evaluator_type: Evaluator adapter used to parse datasets and task template. + model_refs: Optional explicit default/fast Model Entity IDs. Unset uses + the active Platform CLI context. Returns: EvalAuthorResult: containing the modified and newly created datasets, additional metrics and summary. """ - _enable_litellm_drop_params() + selected_model_refs = model_refs if model_refs is not None else configured_model_refs() experiment_dir.mkdir(parents=True, exist_ok=True) experiment_dir = experiment_dir.resolve() insight_locator = str(insight.resolve()) if isinstance(insight, Path) else insight client = make_client(base_url) + model_clients: ConfiguredModelClients | None = None try: + model_clients = await resolve_model_clients(client, selected_model_refs) backend = make_experimentalist_backend( client=client, experiments_output=str(experiment_dir), @@ -113,20 +120,25 @@ async def run_eval_author( evaluator_type, staged_inputs.task_template, ) - eval_author = build_eval_author_agent( - experiment_dir=experiment_dir, - config=config, - ) - return await eval_author.run( - insight=resolved_insight, - agent_path=agent_path, - task_template=parsed_template, - train_dataset=parsed_train, - validation_dataset=parsed_validation, - client=client, - ) + with activate_model_clients(model_clients): + eval_author = build_eval_author_agent( + experiment_dir=experiment_dir, + config=config, + ) + return await eval_author.run( + insight=resolved_insight, + agent_path=agent_path, + task_template=parsed_template, + train_dataset=parsed_train, + validation_dataset=parsed_validation, + client=client, + ) finally: - await client.close() + try: + if model_clients is not None: + await model_clients.aclose() + finally: + await client.close() def build_eval_author_agent( @@ -143,12 +155,3 @@ def build_eval_author_agent( config=config, reporter=reporter, ) - - -def _enable_litellm_drop_params() -> None: - """Let LiteLLM omit unsupported model parameters when it is installed.""" - try: - litellm = cast(_LiteLLMModule, importlib.import_module("litellm")) - except ModuleNotFoundError: - return - litellm.drop_params = True diff --git a/plugins/nemo-eval-author/src/nemo_eval_author_plugin/model_config.py b/plugins/nemo-eval-author/src/nemo_eval_author_plugin/model_config.py deleted file mode 100644 index 0eb9eb2f8b..0000000000 --- a/plugins/nemo-eval-author/src/nemo_eval_author_plugin/model_config.py +++ /dev/null @@ -1,211 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""LLM clients for Eval Author. - -This module deliberately imports nothing from Experimentalist. Eval Author is meant to -become a standalone plugin, so its credential handling duplicates Experimentalist's rather -than sharing it. See the plugin README for the wider split, and -``tests/test_plugin_boundary.py`` for the check that holds this module independent. - -``AUTHOR_*`` is the real contract. The ``NEMO_EXPERIMENTALIST_*`` fallback and -:func:`bridge_author_env_to_experimentalist` are both transitional, and exist only while -Eval Author still reuses Experimentalist agents. - -TODO(eval-author-standalone): drop the ``NEMO_EXPERIMENTALIST_*`` fallback and the bridge once -this plugin imports nothing from ``nemo_experimentalist_plugin``. Concretely: delete -``_BRIDGED_ENV_PAIRS`` and :func:`bridge_author_env_to_experimentalist`, delete -the ``bridge_author_env_to_experimentalist`` call in ``eval_author/agent.py``, and -narrow every ``_env`` / -``_env_or_default`` call here to its ``AUTHOR_*`` name. That is a breaking change for -anyone running insight mode from an Experimentalist-only profile ``.env``, so it needs a -release note. ``rg 'eval-author-standalone'`` finds the other sites. -""" - -import functools -import logging -import os -from urllib.parse import urlsplit - -from nemo_platform_plugin.config import Configuration -from nooa.unifiedllm import CompletionClient - -logger = logging.getLogger(__name__) - -# Duplicated from `nemo_experimentalist_plugin.cli` on purpose: sharing it would mean -# importing Experimentalist from a module that is otherwise already standalone. -_GATEWAY_HOST = "inference-api.nvidia.com" - -_SMART_MODEL_DEFAULT = "openai/openai/openai/gpt-5.5" -_FAST_MODEL_DEFAULT = "openai/openai/openai/gpt-5-mini" - -# Transitional; see TODO(eval-author-standalone) above. Every pair targets an Experimentalist -# helper rather than a client Eval Author owns, which is why the mid tier appears here -# despite Eval Author having no mid-tier client of its own. -_BRIDGED_ENV_PAIRS = ( - ("AUTHOR_API_BASE", "NEMO_EXPERIMENTALIST_API_BASE"), - ("AUTHOR_API_KEY", "NEMO_EXPERIMENTALIST_API_KEY"), - ("AUTHOR_SMART_MODEL_NAME", "NEMO_EXPERIMENTALIST_MODELS_SMART"), - ("AUTHOR_MID_MODEL_NAME", "NEMO_EXPERIMENTALIST_MODELS_MID"), - ("AUTHOR_FAST_MODEL_NAME", "NEMO_EXPERIMENTALIST_MODELS_FAST"), -) - - -def _env(*names: str) -> str | None: - """Return the first non-empty value among ``names``, or ``None`` when all are empty.""" - for name in names: - value = os.environ.get(name, "").strip() - if value: - return value - return None - - -def _env_or_default(*names: str, default: str) -> str: - """Return the first non-empty value among ``names``, falling back to ``default``.""" - return _env(*names) or default - - -def _api_base() -> str: - base = _env("AUTHOR_API_BASE", "NEMO_EXPERIMENTALIST_API_BASE") - if not base: - raise ValueError("AUTHOR_API_BASE (or NEMO_EXPERIMENTALIST_API_BASE) must be set") - return base - - -def _is_gateway_base(value: str) -> bool: - """Report whether ``value`` is the NVIDIA Inference Gateway reached over HTTPS. - - Mirrors ``nemo_experimentalist_plugin.cli._is_gateway_base``. The scheme check is not - cosmetic: a host-only match would forward ``INFERENCE_API_KEY`` to a plain-HTTP base. - """ - try: - parsed = urlsplit(value) - except ValueError: - return False - return parsed.scheme == "https" and parsed.hostname == _GATEWAY_HOST - - -def _api_key() -> str: - """Resolve the LLM API key for Eval Author. - - Preference order: - 1. ``AUTHOR_API_KEY`` - 2. ``NEMO_EXPERIMENTALIST_API_KEY`` - 3. ``INFERENCE_API_KEY``, but only when the resolved base is the HTTPS NVIDIA - Inference Gateway, matching how the Experimentalist CLI forwards that key. - - Returns: - str: the resolved API key. - - Raises: - ValueError: if no key is available for the resolved base. - - """ - key = _env("AUTHOR_API_KEY", "NEMO_EXPERIMENTALIST_API_KEY") - if key: - return key - if _is_gateway_base(_api_base()): - inference = _env("INFERENCE_API_KEY") - if inference: - return inference - raise ValueError( - "AUTHOR_API_KEY (or NEMO_EXPERIMENTALIST_API_KEY) must be set; " - "when using the NVIDIA Inference Gateway, INFERENCE_API_KEY is also accepted" - ) - - -@functools.cache -def _completion_client(name: str, api_base: str, api_key: str) -> CompletionClient: - """Return Eval Author's client for an already-resolved model name, base URL, and key. - - Those three arguments are the cache key, so any tier resolving to the same model gets the - same client. The tier getters below are cached as well, but that alone only promises one - client per tier: two tiers configured with the same model would each build their own, and - every ``CompletionClient`` opens its own pool of keepalive connections. The base URL and - key belong in the key too, since the same model reached through a different endpoint, or - with different credentials, is a different client and must not be reused. - - Experimentalist caches its clients separately, so an insight-mode run holds two clients - for what is usually the same model. Sharing one factory across the plugins would fix - that, but only by importing Experimentalist into the one Eval Author module that has no - such import (see the module docstring). A duplicate connection pool is the cheaper cost. - - Args: - name: Fully qualified model name, as ``CompletionClient`` expects it. - api_base: Resolved API base URL. - api_key: Resolved API key. - - Returns: - CompletionClient: the client for this name, base URL, and key. - - """ - return CompletionClient(name, api_base=api_base, api_key=api_key) - - -@functools.cache -def get_smart_model() -> CompletionClient: - """Return the cached smart (high-capability) LLM client configured from environment variables. - - Returns: - CompletionClient: the singleton smart-model client. - - Raises: - ValueError: if Eval Author API credentials are unset or empty. - - """ - name = _env_or_default( - "AUTHOR_SMART_MODEL_NAME", - "NEMO_EXPERIMENTALIST_MODELS_SMART", - default=_SMART_MODEL_DEFAULT, - ) - return _completion_client(name, _api_base(), _api_key()) - - -@functools.cache -def get_fast_model() -> CompletionClient: - """Return the cached fast (low-latency) LLM client configured from environment variables. - - Returns: - CompletionClient: the singleton fast-model client. - - Raises: - ValueError: if Eval Author API credentials are unset or empty. - - """ - name = _env_or_default( - "AUTHOR_FAST_MODEL_NAME", - "NEMO_EXPERIMENTALIST_MODELS_FAST", - default=_FAST_MODEL_DEFAULT, - ) - return _completion_client(name, _api_base(), _api_key()) - - -def bridge_author_env_to_experimentalist() -> list[str]: - """Copy ``AUTHOR_*`` into unset ``NEMO_EXPERIMENTALIST_*`` slots. - - Eval Author owns ``AUTHOR_*`` for its own clients, but the Experimentalist helpers it - reuses (``TraceAnalyzer`` and the summarizers it pulls in) resolve their LLM from - Experimentalist's own settings. Bridging lets a standalone ``AUTHOR_*``-only run supply - those helpers without the caller duplicating credentials. Values already present in - ``NEMO_EXPERIMENTALIST_*`` always win, so insight-mode runs are left alone. - - Experimentalist reads those names through a cached settings object, so a write here is - only visible if the cache has not been populated yet. Clearing it keeps the bridge - order-independent -- the alternative is a silent no-op whenever Eval Author is - constructed after the first model lookup. - - Returns: - list[str]: the ``NEMO_EXPERIMENTALIST_*`` names this call populated, in declaration - order. Empty once everything is in place, which makes repeat calls no-ops. - - """ - applied: list[str] = [] - for src, dst in _BRIDGED_ENV_PAIRS: - value = os.environ.get(src, "").strip() - if value and not os.environ.get(dst, "").strip(): - os.environ[dst] = value - applied.append(dst) - if applied: - Configuration.clear_cache() - logger.info("Bridged Eval Author credentials into %s", ", ".join(applied)) - return applied diff --git a/plugins/nemo-eval-author/tests/conftest.py b/plugins/nemo-eval-author/tests/conftest.py index f9179848c0..029bbf49b1 100644 --- a/plugins/nemo-eval-author/tests/conftest.py +++ b/plugins/nemo-eval-author/tests/conftest.py @@ -1,30 +1,13 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Top-level conftest: runs before any test module is imported. - -Eval Author and Experimentalist agents call get_smart_model() at class-definition -time. Placeholder credentials keep collection working when real keys are absent. -setdefault does not overwrite CI or developer-provided values. - -Both stand-in values contain "placeholder" on purpose: real-model canaries such as -test_eval_author_repair_e2e cannot use "is it set?" to decide whether to skip, because -these assignments make every credential look present. Keep the substring if you change -the values. -""" +"""Eval Author test-wide state isolation.""" import os import litellm import pytest -# An unroutable HTTPS host: if a real key is present but the base is not, an -# unmocked call fails to connect rather than sending the key over plaintext. -os.environ.setdefault("AUTHOR_API_BASE", "https://placeholder.invalid") -os.environ.setdefault("AUTHOR_API_KEY", "placeholder-for-import") -os.environ.setdefault("NEMO_EXPERIMENTALIST_API_BASE", "https://placeholder.invalid") -os.environ.setdefault("NEMO_EXPERIMENTALIST_API_KEY", "placeholder-for-import") - litellm.drop_params = True diff --git a/plugins/nemo-eval-author/tests/test_eval_author_repair_e2e.py b/plugins/nemo-eval-author/tests/test_eval_author_repair_e2e.py index cb67f8b258..b506fa0a99 100644 --- a/plugins/nemo-eval-author/tests/test_eval_author_repair_e2e.py +++ b/plugins/nemo-eval-author/tests/test_eval_author_repair_e2e.py @@ -10,9 +10,10 @@ from pathlib import Path import pytest +import pytest_asyncio from nemo_eval_author_plugin.eval_author.agent import EvalAuthor from nemo_eval_author_plugin.eval_author.models import EvalAuthorConfig -from nemo_eval_author_plugin.model_config import get_fast_model +from nemo_experimentalist_plugin.client import make_client from nemo_experimentalist_plugin.entities import DatasetValidationError, local_path_from_uri from nemo_experimentalist_plugin.experimentalist.components.evaluator.harbor import ( HarborDataset, @@ -21,22 +22,24 @@ ) from nemo_experimentalist_plugin.experimentalist.components.trace_analyzer import Diagnostic from nemo_insights_plugin.entities import Insight +from nemo_platform_plugin.nooa_model_client import ( + activate_model_clients, + configured_model_refs, + get_fast_model, + resolve_model_clients, +) -def _has_real_llm_credentials() -> bool: - """Report whether real credentials are present, rather than the conftest stand-ins. - - conftest fills placeholder credentials so agent classes can be imported at collection - time, which means a plain "is it set?" check always passes. These canaries call a live - model, so they have to look past the stand-ins. - """ - base = os.environ.get("AUTHOR_API_BASE") or os.environ.get("NEMO_EXPERIMENTALIST_API_BASE") or "" - key = os.environ.get("AUTHOR_API_KEY") or os.environ.get("NEMO_EXPERIMENTALIST_API_KEY") or "" - return bool(base and key) and "placeholder" not in base and "placeholder" not in key +def _has_configured_models() -> bool: + try: + configured_model_refs() + except (FileNotFoundError, RuntimeError, ValueError): + return False + return True -_HAS_LLM = _has_real_llm_credentials() -_CREDENTIALS_HINT = "AUTHOR_API_BASE and AUTHOR_API_KEY (or their NEMO_EXPERIMENTALIST_* equivalents)" +_HAS_MODELS = _has_configured_models() +_MODELS_HINT = "run `nemo setup` and select default and fast agent models" _RUN_EVAL_AUTHOR_REPAIR_E2E = os.environ.get("RUN_EVAL_AUTHOR_REPAIR_E2E") == "1" _RUN_EVAL_AUTHOR_HARBOR_E2E = os.environ.get("RUN_EVAL_AUTHOR_HARBOR_E2E") == "1" _MALFORMED_VERIFIER = """\ @@ -114,6 +117,18 @@ def check_tool_hallucination() -> float: } +@pytest_asyncio.fixture +async def configured_models(): + client = make_client(None) + model_clients = await resolve_model_clients(client) + try: + with activate_model_clients(model_clients): + yield + finally: + await model_clients.aclose() + await client.close() + + def _write_malformed_task(dataset_dir: Path, task_id: str) -> None: task_dir = dataset_dir / task_id verifier_dir = task_dir / "tests" @@ -251,14 +266,14 @@ async def run( @pytest.mark.skipif( - not (_RUN_EVAL_AUTHOR_REPAIR_E2E and _HAS_LLM), - reason=f"Set RUN_EVAL_AUTHOR_REPAIR_E2E=1 with {_CREDENTIALS_HINT} to run the Eval Author repair canary.", + not (_RUN_EVAL_AUTHOR_REPAIR_E2E and _HAS_MODELS), + reason=f"Set RUN_EVAL_AUTHOR_REPAIR_E2E=1 and {_MODELS_HINT} to run the Eval Author repair canary.", ) -async def test_gpt5_mini_repairs_malformed_harbor_verifiers( +async def test_configured_fast_model_repairs_malformed_harbor_verifiers( tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, + configured_models: None, ) -> None: - """GPT-5 mini repairs try-without-except failures across an Insight suite.""" + """The configured fast model repairs try-without-except failures across an Insight suite.""" insight_suite_dir = tmp_path / "insight-suite" _write_malformed_task(insight_suite_dir, "task-a") _write_malformed_task(insight_suite_dir, "task-b") @@ -273,9 +288,6 @@ async def test_gpt5_mini_repairs_malformed_harbor_verifiers( assert "SyntaxError: expected 'except' or 'finally' block" in validation_feedback llm = get_fast_model() - # Construction must not need NEMO_EXPERIMENTALIST_API_KEY: Eval Author owns AUTHOR_* and only - # bridges into the Experimentalist slots, which are already resolved by this point. - monkeypatch.delenv("NEMO_EXPERIMENTALIST_API_KEY", raising=False) eval_author = EvalAuthor( experiment_dir=tmp_path, config=EvalAuthorConfig(), @@ -320,12 +332,12 @@ async def test_gpt5_mini_repairs_malformed_harbor_verifiers( @pytest.mark.skipif( - not (_RUN_EVAL_AUTHOR_HARBOR_E2E and _HAS_LLM), - reason=f"Set RUN_EVAL_AUTHOR_HARBOR_E2E=1 with {_CREDENTIALS_HINT} to run the live Harbor metric canary.", + not (_RUN_EVAL_AUTHOR_HARBOR_E2E and _HAS_MODELS), + reason=f"Set RUN_EVAL_AUTHOR_HARBOR_E2E=1 and {_MODELS_HINT} to run the live Harbor metric canary.", ) async def test_eval_author_metric_scores_known_failing_harbor_baseline_low( tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, + configured_models: None, ) -> None: """An authored root-cause metric scores a known-failing Harbor baseline low.""" insight_suite_dir = tmp_path / "insight-suite" @@ -337,8 +349,6 @@ async def test_eval_author_metric_scores_known_failing_harbor_baseline_low( llm = get_fast_model() llm.config["temperature"] = 0.0 - # See the repair canary above: construction must not need NEMO_EXPERIMENTALIST_API_KEY. - monkeypatch.delenv("NEMO_EXPERIMENTALIST_API_KEY", raising=False) eval_author = EvalAuthor( experiment_dir=tmp_path, config=EvalAuthorConfig(), @@ -437,8 +447,8 @@ async def test_eval_author_metric_scores_known_failing_harbor_baseline_low( @pytest.mark.skipif( - not (_RUN_EVAL_AUTHOR_HARBOR_E2E and _HAS_LLM), - reason=f"Set RUN_EVAL_AUTHOR_HARBOR_E2E=1 with {_CREDENTIALS_HINT} to run the live Harbor metric canary.", + not (_RUN_EVAL_AUTHOR_HARBOR_E2E and _HAS_MODELS), + reason=f"Set RUN_EVAL_AUTHOR_HARBOR_E2E=1 and {_MODELS_HINT} to run the live Harbor metric canary.", ) async def test_eval_author_metric_discriminates_controlled_harbor_tool_evidence( tmp_path: Path, diff --git a/plugins/nemo-eval-author/tests/test_eval_author_run.py b/plugins/nemo-eval-author/tests/test_eval_author_run.py index 722af79f7c..81ad26135b 100644 --- a/plugins/nemo-eval-author/tests/test_eval_author_run.py +++ b/plugins/nemo-eval-author/tests/test_eval_author_run.py @@ -7,6 +7,7 @@ from dataclasses import dataclass from pathlib import Path from typing import Any +from unittest.mock import MagicMock import pytest from nemo_eval_author_plugin.eval_author import run as eval_author_run @@ -25,6 +26,32 @@ async def close(self) -> None: self.closed = True +@dataclass +class ClosingModelClients: + default: object = None + fast: object = None + closed: bool = False + + async def aclose(self) -> None: + self.closed = True + + +@pytest.fixture +def model_clients(monkeypatch: pytest.MonkeyPatch) -> ClosingModelClients: + clients = ClosingModelClients() + refs = eval_author_run.ConfiguredModelRefs( + default="workspace-a/default-model", + fast="workspace-a/fast-model", + ) + + async def resolve(*_: object) -> ClosingModelClients: + return clients + + monkeypatch.setattr(eval_author_run, "configured_model_refs", lambda: refs) + monkeypatch.setattr(eval_author_run, "resolve_model_clients", resolve) + return clients + + class FakeBackend: def __init__(self, insight: Insight) -> None: self.insight = insight @@ -83,9 +110,40 @@ async def run( ) +@pytest.mark.asyncio +async def test_run_eval_author_fails_before_side_effects_when_model_configuration_is_missing( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + experiment_dir = tmp_path / "eval_author" + make_client = MagicMock() + + def missing_model_refs() -> eval_author_run.ConfiguredModelRefs: + raise ValueError("No default model is configured. Run `nemo setup` and select agent models.") + + monkeypatch.setattr(eval_author_run, "configured_model_refs", missing_model_refs) + monkeypatch.setattr(eval_author_run, "make_client", make_client) + + with pytest.raises(ValueError, match="No default model is configured"): + await eval_author_run.run_eval_author( + insight="insight-remote-123", + train_dataset=DatasetRef(uri="train"), + validation_dataset=DatasetRef(uri="validation"), + task_template=DatasetRef(uri="template"), + experiment_dir=experiment_dir, + workspace="workspace-a", + base_url="http://platform.test", + config=EvalAuthorConfig(), + ) + + assert not experiment_dir.exists() + make_client.assert_not_called() + + @pytest.mark.asyncio async def test_run_eval_author_resolves_inputs_and_returns_datasets( monkeypatch: pytest.MonkeyPatch, + model_clients: ClosingModelClients, tmp_path: Path, ) -> None: client = ClosingClient() @@ -103,7 +161,6 @@ async def test_run_eval_author_resolves_inputs_and_returns_datasets( monkeypatch.setattr(eval_author_run, "make_experimentalist_backend", lambda **_: backend) monkeypatch.setattr(eval_author_run, "DatasetFactory", lambda: dataset_factory) monkeypatch.setattr(eval_author_run, "build_eval_author_agent", lambda **_: eval_author) - monkeypatch.setattr(eval_author_run, "_enable_litellm_drop_params", lambda: None) template = tmp_path / "template" template.mkdir() (template / "task.toml").write_text("template\n", encoding="utf-8") @@ -144,6 +201,7 @@ async def test_run_eval_author_resolves_inputs_and_returns_datasets( client, ) assert client.closed + assert model_clients.closed def test_public_apis_accept_train_validation_and_generated_task_inputs() -> None: @@ -169,6 +227,7 @@ def test_public_apis_accept_train_validation_and_generated_task_inputs() -> None @pytest.mark.asyncio async def test_run_eval_author_hydrates_fileset_task_template( monkeypatch: pytest.MonkeyPatch, + model_clients: ClosingModelClients, tmp_path: Path, ) -> None: downloads: list[tuple[str, str, str]] = [] @@ -189,7 +248,6 @@ async def download(self, *, remote_path: str, local_path: str, workspace: str) - monkeypatch.setattr(eval_author_run, "make_experimentalist_backend", lambda **_: backend) monkeypatch.setattr(eval_author_run, "DatasetFactory", lambda: dataset_factory) monkeypatch.setattr(eval_author_run, "build_eval_author_agent", lambda **_: FakeEvalAuthor()) - monkeypatch.setattr(eval_author_run, "_enable_litellm_drop_params", lambda: None) template_ref = DatasetRef(uri="fileset://workspace-a/template") train = tmp_path / "train" validation = tmp_path / "validation" @@ -212,11 +270,52 @@ async def download(self, *, remote_path: str, local_path: str, workspace: str) - assert dataset_factory.template_calls == [("harbor", template_ref.model_copy(update={"uri": str(staged)}))] assert [Path(ref.uri).name for _, ref in dataset_factory.dataset_calls] == ["train", "validation"] assert client.closed + assert model_clients.closed + + +@pytest.mark.asyncio +async def test_run_eval_author_uses_agent_override( + monkeypatch: pytest.MonkeyPatch, + model_clients: ClosingModelClients, + tmp_path: Path, +) -> None: + client = ClosingClient() + backend = FakeBackend( + Insight(workspace="workspace-a", title="failure", description="description", agent="insight-agent") + ) + monkeypatch.setattr(eval_author_run, "make_client", lambda _: client) + monkeypatch.setattr(eval_author_run, "make_experimentalist_backend", lambda **_: backend) + monkeypatch.setattr(eval_author_run, "DatasetFactory", FakeDatasetFactory) + monkeypatch.setattr(eval_author_run, "build_eval_author_agent", lambda **_: FakeEvalAuthor()) + template = tmp_path / "template" + train = tmp_path / "train" + validation = tmp_path / "validation" + template.mkdir() + train.mkdir() + validation.mkdir() + override = tmp_path / "override-agent" + + await eval_author_run.run_eval_author( + insight="insight-123", + agent=override, + train_dataset=DatasetRef(uri=str(train)), + validation_dataset=DatasetRef(uri=str(validation)), + task_template=DatasetRef(uri=str(template)), + experiment_dir=tmp_path / "experiment", + workspace="workspace-a", + base_url="http://platform.test", + config=EvalAuthorConfig(), + ) + + assert backend.agent_calls[0][1] == override + assert client.closed + assert model_clients.closed @pytest.mark.asyncio -async def test_run_eval_author_closes_client_on_failure( +async def test_run_eval_author_closes_clients_on_failure( monkeypatch: pytest.MonkeyPatch, + model_clients: ClosingModelClients, tmp_path: Path, ) -> None: client = ClosingClient() @@ -240,3 +339,4 @@ async def test_run_eval_author_closes_client_on_failure( ) assert client.closed + assert model_clients.closed diff --git a/plugins/nemo-eval-author/tests/test_model_config.py b/plugins/nemo-eval-author/tests/test_model_config.py deleted file mode 100644 index 8f0a92ce5e..0000000000 --- a/plugins/nemo-eval-author/tests/test_model_config.py +++ /dev/null @@ -1,275 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Credential and model-client resolution for Eval Author.""" - -import os -import subprocess -import sys -from collections.abc import Iterator - -import pytest -from nemo_eval_author_plugin import model_config -from nemo_platform_plugin.config import Configuration - - -@pytest.fixture(autouse=True) -def _clear_model_cache() -> Iterator[None]: - """Drop every cached client so each test resolves credentials from scratch. - - ``Configuration`` goes with them: the bridge writes ``NEMO_EXPERIMENTALIST_*`` and - Experimentalist resolves those through a process-global settings cache that - ``monkeypatch`` does not touch. Without this, a test that bridges an endpoint leaves - it visible to every test that runs after it. - """ - - def clear() -> None: - model_config.get_smart_model.cache_clear() - model_config.get_fast_model.cache_clear() - model_config._completion_client.cache_clear() - Configuration.clear_cache() - - clear() - yield - clear() - - -def test_api_base_prefers_eval_author_over_experimentalist( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setenv("AUTHOR_API_BASE", "https://eval-author.example/v1") - monkeypatch.setenv("NEMO_EXPERIMENTALIST_API_BASE", "https://experimentalist.example/v1") - monkeypatch.setenv("AUTHOR_API_KEY", "eval-key") - - assert model_config._api_base() == "https://eval-author.example/v1" - - -def test_api_base_falls_back_to_experimentalist( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.delenv("AUTHOR_API_BASE", raising=False) - monkeypatch.setenv("NEMO_EXPERIMENTALIST_API_BASE", "https://experimentalist.example/v1") - monkeypatch.setenv("NEMO_EXPERIMENTALIST_API_KEY", "exp-key") - - assert model_config._api_base() == "https://experimentalist.example/v1" - - -def test_api_base_requires_configuration(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.delenv("AUTHOR_API_BASE", raising=False) - monkeypatch.delenv("NEMO_EXPERIMENTALIST_API_BASE", raising=False) - - with pytest.raises(ValueError, match="AUTHOR_API_BASE"): - model_config._api_base() - - -def test_whitespace_only_values_count_as_unset(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("AUTHOR_API_BASE", " ") - monkeypatch.setenv("NEMO_EXPERIMENTALIST_API_BASE", " https://experimentalist.example/v1 ") - - assert model_config._api_base() == "https://experimentalist.example/v1" - - -def test_api_key_accepts_inference_key_on_gateway( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setenv("AUTHOR_API_BASE", "https://inference-api.nvidia.com/v1") - monkeypatch.delenv("AUTHOR_API_KEY", raising=False) - monkeypatch.delenv("NEMO_EXPERIMENTALIST_API_KEY", raising=False) - monkeypatch.setenv("INFERENCE_API_KEY", "sk-gateway") - - assert model_config._api_key() == "sk-gateway" - - -def test_api_key_refuses_to_forward_inference_key_over_plain_http( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setenv("AUTHOR_API_BASE", "http://inference-api.nvidia.com/v1") - monkeypatch.delenv("AUTHOR_API_KEY", raising=False) - monkeypatch.delenv("NEMO_EXPERIMENTALIST_API_KEY", raising=False) - monkeypatch.setenv("INFERENCE_API_KEY", "sk-gateway") - - with pytest.raises(ValueError, match="AUTHOR_API_KEY"): - model_config._api_key() - - -def test_api_key_refuses_to_forward_inference_key_to_a_lookalike_host( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setenv("AUTHOR_API_BASE", "https://inference-api.nvidia.com.attacker.example/v1") - monkeypatch.delenv("AUTHOR_API_KEY", raising=False) - monkeypatch.delenv("NEMO_EXPERIMENTALIST_API_KEY", raising=False) - monkeypatch.setenv("INFERENCE_API_KEY", "sk-gateway") - - with pytest.raises(ValueError, match="AUTHOR_API_KEY"): - model_config._api_key() - - -def test_api_key_requires_credentials_when_missing( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setenv("AUTHOR_API_BASE", "https://custom.example/v1") - monkeypatch.delenv("AUTHOR_API_KEY", raising=False) - monkeypatch.delenv("NEMO_EXPERIMENTALIST_API_KEY", raising=False) - monkeypatch.delenv("INFERENCE_API_KEY", raising=False) - - with pytest.raises(ValueError, match="AUTHOR_API_KEY"): - model_config._api_key() - - -@pytest.mark.parametrize("tier", ["SMART", "FAST"]) -def test_model_name_prefers_author_then_experimentalist_then_default( - monkeypatch: pytest.MonkeyPatch, - tier: str, -) -> None: - getter = {"SMART": model_config.get_smart_model, "FAST": model_config.get_fast_model}[tier] - default = {"SMART": model_config._SMART_MODEL_DEFAULT, "FAST": model_config._FAST_MODEL_DEFAULT}[tier] - author_var = f"AUTHOR_{tier}_MODEL_NAME" - experimentalist_var = f"NEMO_EXPERIMENTALIST_MODELS_{tier}" - - monkeypatch.setenv("AUTHOR_API_BASE", "https://eval-author.example/v1") - monkeypatch.setenv("AUTHOR_API_KEY", "eval-key") - - monkeypatch.setenv(author_var, "vendor/from-author") - monkeypatch.setenv(experimentalist_var, "vendor/from-experimentalist") - assert getter().model == "vendor/from-author" - - getter.cache_clear() - monkeypatch.delenv(author_var) - assert getter().model == "vendor/from-experimentalist" - - getter.cache_clear() - monkeypatch.delenv(experimentalist_var) - assert getter().model == default - - -def test_tiers_pointing_at_one_model_share_a_client(monkeypatch: pytest.MonkeyPatch) -> None: - # The @functools.cache on each getter only promises one client per tier. Keying the - # factory on the resolved name, base URL, and key is what collapses identical tiers. - monkeypatch.setenv("AUTHOR_API_BASE", "https://eval-author.example/v1") - monkeypatch.setenv("AUTHOR_API_KEY", "eval-key") - monkeypatch.setenv("AUTHOR_SMART_MODEL_NAME", "vendor/one-model") - monkeypatch.setenv("AUTHOR_FAST_MODEL_NAME", "vendor/one-model") - - assert model_config.get_smart_model() is model_config.get_fast_model() - - -def test_differing_tiers_get_separate_clients(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("AUTHOR_API_BASE", "https://eval-author.example/v1") - monkeypatch.setenv("AUTHOR_API_KEY", "eval-key") - monkeypatch.setenv("AUTHOR_SMART_MODEL_NAME", "vendor/smart") - monkeypatch.setenv("AUTHOR_FAST_MODEL_NAME", "vendor/fast") - - assert model_config.get_smart_model() is not model_config.get_fast_model() - - -def test_bridge_author_env_fills_unset_experimentalist_slots( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setenv("AUTHOR_API_BASE", "https://eval-author.example/v1") - monkeypatch.setenv("AUTHOR_API_KEY", "eval-key") - monkeypatch.delenv("NEMO_EXPERIMENTALIST_API_BASE", raising=False) - monkeypatch.delenv("NEMO_EXPERIMENTALIST_API_KEY", raising=False) - - applied = model_config.bridge_author_env_to_experimentalist() - - assert applied == ["NEMO_EXPERIMENTALIST_API_BASE", "NEMO_EXPERIMENTALIST_API_KEY"] - assert os.environ["NEMO_EXPERIMENTALIST_API_BASE"] == "https://eval-author.example/v1" - assert os.environ["NEMO_EXPERIMENTALIST_API_KEY"] == "eval-key" - - -def test_bridge_author_env_covers_every_declared_pair( - monkeypatch: pytest.MonkeyPatch, -) -> None: - for src, dst in model_config._BRIDGED_ENV_PAIRS: - monkeypatch.setenv(src, f"value-from-{src}") - monkeypatch.delenv(dst, raising=False) - - applied = model_config.bridge_author_env_to_experimentalist() - - assert applied == [dst for _, dst in model_config._BRIDGED_ENV_PAIRS] - for src, dst in model_config._BRIDGED_ENV_PAIRS: - assert os.environ[dst] == f"value-from-{src}" - - -def test_bridge_author_env_does_not_overwrite_existing_experimentalist( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setenv("AUTHOR_API_BASE", "https://eval-author.example/v1") - monkeypatch.setenv("NEMO_EXPERIMENTALIST_API_BASE", "https://experimentalist.example/v1") - - applied = model_config.bridge_author_env_to_experimentalist() - - assert applied == [] - assert os.environ["NEMO_EXPERIMENTALIST_API_BASE"] == "https://experimentalist.example/v1" - - -def test_bridge_author_env_is_a_no_op_once_applied( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setenv("AUTHOR_API_BASE", "https://eval-author.example/v1") - monkeypatch.delenv("NEMO_EXPERIMENTALIST_API_BASE", raising=False) - - assert model_config.bridge_author_env_to_experimentalist() == ["NEMO_EXPERIMENTALIST_API_BASE"] - assert model_config.bridge_author_env_to_experimentalist() == [] - - -def test_bridge_is_visible_even_after_settings_were_already_resolved( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """Experimentalist memoizes its settings, so bridging afterwards must invalidate them. - - Without the cache clear this passes silently in isolation and fails in a real run, - where something resolves a model before ``EvalAuthor.__init__`` gets to bridge. - """ - from nemo_experimentalist_plugin.experimentalist.components import model_config as exp_model_config - - for name in ("NEMO_EXPERIMENTALIST_API_BASE", "NEMO_EXPERIMENTALIST_API_KEY"): - monkeypatch.delenv(name, raising=False) - Configuration.clear_cache() - with pytest.raises(ValueError): - exp_model_config.api_base() # populates the cache with "nothing configured" - - monkeypatch.setenv("AUTHOR_API_BASE", "https://eval-author.example/v1") - monkeypatch.setenv("AUTHOR_API_KEY", "author-key") - model_config.bridge_author_env_to_experimentalist() - - assert exp_model_config.api_base() == "https://eval-author.example/v1" - - -def test_agent_constructs_under_author_credentials_alone() -> None: - """An AUTHOR_*-only run must still be able to build the Experimentalist agents. - - ``TraceAnalyzer`` and friends resolve ``NEMO_EXPERIMENTALIST_*`` when they are - constructed, so ``EvalAuthor.__init__`` bridges the credentials first. Checked in a - subprocess with a pruned environment because this suite's conftest already populates - ``NEMO_EXPERIMENTALIST_*`` and modules are imported once per session -- it fails with - ``ValueError`` if the bridge stops running early enough. - """ - env = { - key: value - for key, value in os.environ.items() - if not key.startswith(("AUTHOR_", "NEMO_EXPERIMENTALIST_", "INFERENCE_API_KEY")) - } - env["AUTHOR_API_BASE"] = "https://placeholder.invalid" - env["AUTHOR_API_KEY"] = "placeholder-for-import" - - # Importing nooa loads a .env, which can pre-populate NEMO_EXPERIMENTALIST_* and mask the - # bridge. Clear those slots after the imports so construction is the only thing that - # can fill them. - probe = ( - "import os, tempfile\n" - "from pathlib import Path\n" - "from nemo_eval_author_plugin.eval_author.agent import EvalAuthor\n" - "for k in [k for k in os.environ if k.startswith('NEMO_EXPERIMENTALIST_')]: del os.environ[k]\n" - "EvalAuthor(experiment_dir=Path(tempfile.mkdtemp()))\n" - "assert os.environ['NEMO_EXPERIMENTALIST_API_KEY'] == 'placeholder-for-import', os.environ.get('NEMO_EXPERIMENTALIST_API_KEY')\n" - ) - result = subprocess.run( - [sys.executable, "-c", probe], - env=env, - capture_output=True, - text=True, - check=False, - ) - - assert result.returncode == 0, result.stderr diff --git a/plugins/nemo-eval-author/tests/test_plugin_boundary.py b/plugins/nemo-eval-author/tests/test_plugin_boundary.py index 50c07a70b1..ce64320958 100644 --- a/plugins/nemo-eval-author/tests/test_plugin_boundary.py +++ b/plugins/nemo-eval-author/tests/test_plugin_boundary.py @@ -12,8 +12,7 @@ ``_BORROWED_BEHAVIOUR`` is debt: Harbor, tools, trace analysis, the backend factory. Eval Author is meant to end up standalone, so this list can only shrink, and duplicating a helper beats adding a row. Emptying it is what unblocks the remaining -TODO(eval-author-standalone) cleanup — the ``NEMO_EXPERIMENTALIST_*`` credential fallback and -the ``bridge_author_env_to_experimentalist`` call in ``EvalAuthor.__init__``. +TODO(eval-author-standalone) cleanup. """ import ast @@ -93,14 +92,3 @@ def test_experimentalist_imports_only_shrink() -> None: + ", ".join(sorted(removed)) + ". Drop them from _BORROWED_BEHAVIOUR so the list keeps meaning something." ) - - -def test_model_config_is_already_standalone() -> None: - """Credential resolution is the one part of the plugin that owes Experimentalist nothing. - - It was briefly refactored to share Experimentalist's client cache, which read as a - cleanup but moved the plugin away from standalone. Keep it independent. - """ - importers = {file for files in _imported_experimentalist_modules().values() for file in files} - - assert "model_config.py" not in importers diff --git a/plugins/nemo-experimentalist/AGENTS.md b/plugins/nemo-experimentalist/AGENTS.md index 9e3eff5217..94d57f6d77 100644 --- a/plugins/nemo-experimentalist/AGENTS.md +++ b/plugins/nemo-experimentalist/AGENTS.md @@ -28,7 +28,7 @@ to the user. ### 2026-07-28: Eval Author extracted to its own plugin, heading for standalone `plugins/nemo-eval-author/` (`nemo-eval-author-plugin`) owns the Eval Author agent package -(`eval_author/`) plus its own `AUTHOR_*` `model_config`. +(`eval_author/`). **The target is one arrow: Experimentalist → Eval Author.** Eval Author is meant to stop depending on Experimentalist entirely, even where that means duplicating code. Today the @@ -42,23 +42,11 @@ arrow points both ways: and the backend factory. `plugins/nemo-eval-author/tests/test_plugin_boundary.py` pins that second list so it can -only shrink. **When you are tempted to share a helper between the two plugins, duplicate it -into Eval Author instead.** Sharing reads like a cleanup and is a regression here; the -boundary test will reject it, which is the intended answer, not an obstacle to route -around. `uv` resolves the current cycle fine — install both with -`uv sync --group experimentalist`. - -Two transitional mechanisms exist only because of the second arrow. Both should be deleted -with the last `nemo_experimentalist_plugin` import, and both are tagged -`TODO(eval-author-standalone)` — `rg 'eval-author-standalone'` lists every site: - -- `EvalAuthor.__init__` calls `bridge_author_env_to_experimentalist()`, copying `AUTHOR_*` - into unset `NEMO_EXPERIMENTALIST_*` slots so the Experimentalist agents Eval Author borrows - resolve their models. This was a `_env_bridge.py` side-effect import until agents stopped - binding their LLM in the class body; an ordinary call in `__init__` now suffices because - those agents resolve when constructed, not when imported. -- `AUTHOR_*` falls back to `NEMO_EXPERIMENTALIST_*` in `model_config`. `AUTHOR_*` is the real - contract; the fallback only keeps one credential set working in insight mode. +only shrink. **Do not share Eval Author implementation helpers with Experimentalist.** +The Platform-owned model client integration is intentionally different: both plugins use +`nemo_platform_plugin.nooa_model_client` so provider routing, Platform authentication, +and configured model selection have one owner outside either plugin. `uv` resolves the +remaining package cycle; install both with `uv sync --group experimentalist`. ### 2026-07-24: Optimizer renamed to Experimentalist diff --git a/plugins/nemo-experimentalist/README.md b/plugins/nemo-experimentalist/README.md index c2fd71c7ca..61bb6ef697 100644 --- a/plugins/nemo-experimentalist/README.md +++ b/plugins/nemo-experimentalist/README.md @@ -111,12 +111,8 @@ repo="$(git rev-parse --show-toplevel)" sbx create --clone --name nemo-experimentalist shell "$repo" sbx exec --workdir "$repo" \ --env UV_PROJECT_ENVIRONMENT=/home/agent/.venvs/nemo-platform \ - --env INFERENCE_API_KEY \ - --env NEMO_EXPERIMENTALIST_API_BASE \ - --env NEMO_EXPERIMENTALIST_API_KEY \ - --env NEMO_EXPERIMENTALIST_MODELS_SMART \ - --env NEMO_EXPERIMENTALIST_MODELS_MID \ - --env NEMO_EXPERIMENTALIST_MODELS_FAST \ + --env NEMO_DEFAULT_MODEL \ + --env NEMO_FAST_MODEL \ nemo-experimentalist \ uv run --frozen --python 3.13 --package nemo-experimentalist-plugin --with ./plugins/nemo-agents \ nemo agents experimentalist run @@ -142,41 +138,27 @@ Harbor tasks that publish only `linux/amd64` images do not run in the `linux/arm64` sandbox. This currently includes the Terminal-Bench `fix-git` task; use an x86_64 machine or VM for that suite. -### Endpoint and model settings +### Platform models -Which endpoint the Experimentalist talks to, and with which models, is a *deployment* -setting: one per install, not per experiment. Like every other NeMo plugin it is a -`NemoConfig`, so it can be set either in the `experimentalist:` section of the platform -config file or through the environment, and **the environment wins**. -`nemo experimentalist doctor` reports what is unset. +Run `nemo setup` once for the active Platform context. Setup registers an inference +provider, stores its credential as a Platform Secret, and asks for two +workspace-qualified Model Entities: -| Variable | Config key | Default | Used by | -|---|---|---|---| -| `NEMO_EXPERIMENTALIST_API_BASE` | `api_base` | `https://inference-api.nvidia.com/v1` | all optimizer agents | -| `NEMO_EXPERIMENTALIST_API_KEY` | `api_key` | — (required) | all optimizer agents | -| `NEMO_EXPERIMENTALIST_MODELS_SMART` | `models.smart` | — (required) | Coder, Analyzer, Proposer, Rationalizer, TraceAnalyzer | -| `NEMO_EXPERIMENTALIST_MODELS_MID` | `models.mid` | — (required) | trajectory scorer, architecture doc | -| `NEMO_EXPERIMENTALIST_MODELS_FAST` | `models.fast` | — (required) | Terminator, goal tree, summarizers | +- the default model for quality-critical analysis, proposing, and coding; +- the fast model for latency-sensitive scoring, summarization, and control steps. -The tiers exist to buy capability where it changes the result and speed where it does not, -so give them different models — smart writes the code, fast runs the high-volume judging: +The default model is also the compatibility value for existing single-model +contexts. Press Enter at the fast-model prompt to reuse it. The Experimentalist +resolves both entities through Platform and lets the entity's backend format route +the Nooa completion request through any registered provider exposed as OpenAI Chat +Completions or Anthropic Messages. It does not read a separate optimizer endpoint, +provider key, or provider model name. -```bash -export NEMO_EXPERIMENTALIST_API_KEY=sk-... -export NEMO_EXPERIMENTALIST_MODELS_SMART=openai/openai/openai/gpt-5.6-sol -export NEMO_EXPERIMENTALIST_MODELS_MID=openai/openai/openai/gpt-5.6-terra -export NEMO_EXPERIMENTALIST_MODELS_FAST=openai/openai/openai/gpt-5.6-luna -``` - -Model names have no default: a name is only meaningful against a specific endpoint, so -an unset tier fails before the run starts rather than at the first LLM call. Name them -as *your* endpoint does — `openai/openai/openai/gpt-5.6-sol` on the NVIDIA gateway, -`gpt-5.6-sol` against OpenAI directly. The -[example agent's `.env.example`](examples/tau3-nooa-agent/.env.example) is a working set. - -Credentials are the one place a default applies: when the API base is the NVIDIA gateway, -a set `INFERENCE_API_KEY` fills `NEMO_EXPERIMENTALIST_API_KEY`. A custom base never -inherits it, so a key scoped to the gateway is not forwarded elsewhere. +`nemo agents experimentalist doctor` reports the effective pair. For +non-interactive or isolated environments, `NEMO_DEFAULT_MODEL` and +`NEMO_FAST_MODEL` override the stored selections; both values remain Platform +Model Entity IDs in `workspace/model-name` form. The sandbox examples pass these +overrides because the host's `~/.config/nmp/config.yaml` is not part of the clone. The `--config` YAML is the other kind of configuration: it holds what *one experiment* does (`max_rounds`, `max_survivors`, per-component tuning) and takes no environment @@ -216,9 +198,9 @@ run-level objective metrics. The configured objective targets move to `regression_metrics`, alongside the existing guardrails, so the insight is improved without giving up the run's original priorities. -The agent under test is separate: it reads `AUT_MODEL_NAME` -plus `OPENAI_API_KEY` / `OPENAI_BASE_URL`, which are the only variables forwarded -into the evaluation container. +The agent under test remains separate. The Tau3 example reads `AUT_MODEL_NAME` +plus `OPENAI_API_KEY` / `OPENAI_BASE_URL`; those are the only variables forwarded +into its evaluation container and are not used by the Experimentalist agents. ### Insight-driven optimization diff --git a/plugins/nemo-experimentalist/benchmarks/README.md b/plugins/nemo-experimentalist/benchmarks/README.md index 6d8b60e03d..4e12ad452c 100644 --- a/plugins/nemo-experimentalist/benchmarks/README.md +++ b/plugins/nemo-experimentalist/benchmarks/README.md @@ -117,9 +117,10 @@ Required credentials: export INFERENCE_API_KEY=... ``` -`NEMO_EXPERIMENTALIST_API_KEY` may be provided instead. The runner maps either credential -to both the AUT and optimizer, using -`https://inference-api.nvidia.com/v1` unless an API base is explicitly set. +This credential is for the agent under test. Run `nemo setup` separately to +register the provider and select the default and fast Platform Model Entities +used by the optimization agents. Set `INFERENCE_API_BASE` to the provider +endpoint used by the agent under test. Validate provenance and task IDs without Docker or model calls: diff --git a/plugins/nemo-experimentalist/benchmarks/configs/tau3-quality.yaml b/plugins/nemo-experimentalist/benchmarks/configs/tau3-quality.yaml index ae953389a3..08d6d09aae 100644 --- a/plugins/nemo-experimentalist/benchmarks/configs/tau3-quality.yaml +++ b/plugins/nemo-experimentalist/benchmarks/configs/tau3-quality.yaml @@ -5,9 +5,6 @@ suite_partition: quality test_attempts: 3 models: aut: openai/openai/openai/gpt-5-mini - experimentalist_smart: openai/openai/openai/gpt-5-mini - experimentalist_mid: openai/openai/openai/gpt-5-mini - experimentalist_fast: openai/openai/openai/gpt-5-mini user_simulator: openai/openai/openai/gpt-5-mini optimizer: max_rounds: 3 diff --git a/plugins/nemo-experimentalist/benchmarks/configs/tau3-smoke.yaml b/plugins/nemo-experimentalist/benchmarks/configs/tau3-smoke.yaml index c4be604fb3..9d2359f909 100644 --- a/plugins/nemo-experimentalist/benchmarks/configs/tau3-smoke.yaml +++ b/plugins/nemo-experimentalist/benchmarks/configs/tau3-smoke.yaml @@ -5,9 +5,6 @@ suite_partition: fast test_attempts: 1 models: aut: openai/openai/openai/gpt-5-mini - experimentalist_smart: openai/openai/openai/gpt-5-mini - experimentalist_mid: openai/openai/openai/gpt-5-mini - experimentalist_fast: openai/openai/openai/gpt-5-mini user_simulator: openai/openai/openai/gpt-5-mini optimizer: max_rounds: 1 diff --git a/plugins/nemo-experimentalist/benchmarks/configs/terminal-bench-quality.yaml b/plugins/nemo-experimentalist/benchmarks/configs/terminal-bench-quality.yaml index 2937d4c874..08e1270887 100644 --- a/plugins/nemo-experimentalist/benchmarks/configs/terminal-bench-quality.yaml +++ b/plugins/nemo-experimentalist/benchmarks/configs/terminal-bench-quality.yaml @@ -5,9 +5,6 @@ suite_partition: quality test_attempts: 2 models: aut: aws/anthropic/claude-haiku-4-5-v1 - experimentalist_smart: openai/openai/openai/gpt-5.5 - experimentalist_mid: openai/openai/openai/gpt-5-mini - experimentalist_fast: openai/openai/openai/gpt-5-mini optimizer: max_rounds: 2 min_rounds_before_stopping: 2 diff --git a/plugins/nemo-experimentalist/benchmarks/configs/terminal-bench-smoke.yaml b/plugins/nemo-experimentalist/benchmarks/configs/terminal-bench-smoke.yaml index af97ec02ce..1dec23bd5c 100644 --- a/plugins/nemo-experimentalist/benchmarks/configs/terminal-bench-smoke.yaml +++ b/plugins/nemo-experimentalist/benchmarks/configs/terminal-bench-smoke.yaml @@ -5,9 +5,6 @@ suite_partition: fast test_attempts: 1 models: aut: aws/anthropic/claude-haiku-4-5-v1 - experimentalist_smart: openai/openai/openai/gpt-5-mini - experimentalist_mid: openai/openai/openai/gpt-5-mini - experimentalist_fast: openai/openai/openai/gpt-5-mini optimizer: max_rounds: 1 min_rounds_before_stopping: 1 diff --git a/plugins/nemo-experimentalist/benchmarks/run.py b/plugins/nemo-experimentalist/benchmarks/run.py index 19c4331ad2..a85dcefeab 100644 --- a/plugins/nemo-experimentalist/benchmarks/run.py +++ b/plugins/nemo-experimentalist/benchmarks/run.py @@ -124,9 +124,6 @@ def framework_skills_dirs(self, plugin_root: Path) -> list[Path]: class ModelSpec(BaseModel): aut: str - experimentalist_smart: str - experimentalist_mid: str - experimentalist_fast: str user_simulator: str | None = None @@ -194,19 +191,14 @@ def validate_canonical_suite( def _configure_models(models: ModelSpec) -> None: - api_key = os.environ.get("INFERENCE_API_KEY") or os.environ.get("NEMO_EXPERIMENTALIST_API_KEY") + api_key = os.environ.get("INFERENCE_API_KEY") if not api_key: - raise RuntimeError("INFERENCE_API_KEY or NEMO_EXPERIMENTALIST_API_KEY is required") - api_base = os.environ.get("INFERENCE_API_BASE") or os.environ.get("NEMO_EXPERIMENTALIST_API_BASE") + raise RuntimeError("INFERENCE_API_KEY is required for the agent under test") + api_base = os.environ.get("INFERENCE_API_BASE") api_base = api_base or "https://inference-api.nvidia.com/v1" os.environ.setdefault("INFERENCE_API_KEY", api_key) - os.environ.setdefault("NEMO_EXPERIMENTALIST_API_KEY", api_key) os.environ.setdefault("INFERENCE_API_BASE", api_base) - os.environ.setdefault("NEMO_EXPERIMENTALIST_API_BASE", api_base) os.environ["AUT_MODEL_NAME"] = models.aut - os.environ["NEMO_EXPERIMENTALIST_MODELS_SMART"] = models.experimentalist_smart - os.environ["NEMO_EXPERIMENTALIST_MODELS_MID"] = models.experimentalist_mid - os.environ["NEMO_EXPERIMENTALIST_MODELS_FAST"] = models.experimentalist_fast os.environ.setdefault("NEMO_EXPERIMENTALIST_RUNTIME_CACHE", str(DEFAULT_RUNTIME_CACHE)) if models.user_simulator is not None: # tau-style tasks run a user simulator and NL-assertion judge inside the task @@ -329,14 +321,22 @@ def summarize_experimentalist_jobs(results_dir: Path) -> dict[str, Any]: if not isinstance(payload, dict) or not isinstance(payload.get("stats"), dict): continue stats = payload["stats"] + integer_fields = { + "trials": payload.get("n_total_trials"), + "completed_trials": stats.get("n_completed_trials"), + "errored_trials": stats.get("n_errored_trials"), + "input_tokens": stats.get("n_input_tokens"), + "cache_tokens": stats.get("n_cache_tokens"), + "output_tokens": stats.get("n_output_tokens"), + } + for field, value in integer_fields.items(): + if value is None: + integer_fields[field] = 0 + elif isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"Expected integer job summary field {field!r}, got {type(value).__name__}") job = { "name": result_path.parent.name, - "trials": int(payload.get("n_total_trials") or 0), - "completed_trials": int(stats.get("n_completed_trials") or 0), - "errored_trials": int(stats.get("n_errored_trials") or 0), - "input_tokens": int(stats.get("n_input_tokens") or 0), - "cache_tokens": int(stats.get("n_cache_tokens") or 0), - "output_tokens": int(stats.get("n_output_tokens") or 0), + **integer_fields, "cost_usd": stats.get("cost_usd") if isinstance(stats.get("cost_usd"), int | float) else None, "evals": stats.get("evals") if isinstance(stats.get("evals"), dict) else {}, } @@ -350,7 +350,7 @@ def summarize_experimentalist_jobs(results_dir: Path) -> dict[str, Any]: "cache_tokens", "output_tokens", ): - totals[key] += job[key] + totals[key] += integer_fields[key] if job["cost_usd"] is not None: costs.append(float(job["cost_usd"])) return {**totals, "cost_usd": sum(costs) if costs else None, "job_results": jobs} @@ -395,7 +395,9 @@ async def run_benchmark(args: argparse.Namespace) -> Path: package_client = PackageDatasetClient() metadata = await package_client.get_dataset_metadata(suite.dataset.requested_reference) - canonical_task_ids = {task.name for task in metadata.task_ids} + canonical_task_ids = {task.get_name() for task in metadata.task_ids} + if metadata.version is None: + raise ValueError(f"Dataset {suite.dataset.requested_reference} did not resolve to an immutable version") canonical_task_ids = validate_canonical_suite( suite, canonical_task_ids=canonical_task_ids, resolved_ref=metadata.version ) @@ -413,6 +415,9 @@ async def run_benchmark(args: argparse.Namespace) -> Path: _configure_models(benchmark_config.models) from nemo_experimentalist_plugin.experimentalist.run import run_experimentalist # noqa: PLC0415 + from nemo_platform_plugin.nooa_model_client import configured_model_refs # noqa: PLC0415 + + optimizer_model_refs = configured_model_refs() run_dir = args.output.resolve() run_dir.mkdir(parents=True, exist_ok=True) @@ -460,6 +465,7 @@ async def run_benchmark(args: argparse.Namespace) -> Path: client=None, config=benchmark_config.optimizer, framework_skills_dirs=framework_skills_dirs, + model_refs=optimizer_model_refs, ) run_document = json.loads((experimentalist_dir / "eval-and-optimize" / "run.json").read_text(encoding="utf-8")) winner_label = run_document.get("winner_agent") @@ -499,6 +505,10 @@ async def run_benchmark(args: argparse.Namespace) -> Path: "config_path": str(args.config.resolve()), }, "models": benchmark_config.models.model_dump(), + "optimizer_models": { + "default": optimizer_model_refs.default, + "fast": optimizer_model_refs.fast, + }, "config": benchmark_config.model_dump(mode="json"), "baseline": baseline, "optimizer": { diff --git a/plugins/nemo-experimentalist/examples/tau3-nooa-agent/.env.example b/plugins/nemo-experimentalist/examples/tau3-nooa-agent/.env.example index 33ae2ebbda..0cc99b3faf 100644 --- a/plugins/nemo-experimentalist/examples/tau3-nooa-agent/.env.example +++ b/plugins/nemo-experimentalist/examples/tau3-nooa-agent/.env.example @@ -7,14 +7,14 @@ export INFERENCE_API_KEY=sk-... export OPENAI_API_KEY="$INFERENCE_API_KEY" export OPENAI_BASE_URL="$INFERENCE_API_BASE" -export NEMO_EXPERIMENTALIST_API_KEY="$INFERENCE_API_KEY" -export NEMO_EXPERIMENTALIST_API_BASE="$INFERENCE_API_BASE" + +# `nemo setup` stores the optimizer's default and fast Platform Model Entities. +# A Docker Sandbox cannot read the host's CLI config, so uncomment these with the +# workspace-qualified IDs shown by setup before following the sandbox example. +# export NEMO_DEFAULT_MODEL=default/gpt-4-1 +# export NEMO_FAST_MODEL=default/gpt-4-1-mini # 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-mini - -export NEMO_EXPERIMENTALIST_MODELS_SMART=openai/openai/openai/gpt-5.6-sol -export NEMO_EXPERIMENTALIST_MODELS_MID=openai/openai/openai/gpt-5.6-terra -export NEMO_EXPERIMENTALIST_MODELS_FAST=openai/openai/openai/gpt-5.6-luna 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 index 2a1b31c0d8..a8590c1907 100644 --- 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 @@ -49,9 +49,9 @@ def _experiment_id() -> str: def _configure_models(*, model: str, user_model: str, api_base: str) -> None: - api_key = os.environ.get("INFERENCE_API_KEY") or os.environ.get("NEMO_EXPERIMENTALIST_API_KEY") + api_key = os.environ.get("INFERENCE_API_KEY") if not api_key: - raise RuntimeError("INFERENCE_API_KEY or NEMO_EXPERIMENTALIST_API_KEY is required") + raise RuntimeError("INFERENCE_API_KEY is required for the agent under test") normalized_base = api_base.rstrip("/") openai_base = normalized_base if normalized_base.endswith("/v1") else f"{normalized_base}/v1" diff --git a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/cli.py b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/cli.py index 4e2f3f0af7..ff5d6a4c3f 100644 --- a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/cli.py +++ b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/cli.py @@ -8,13 +8,10 @@ """ import asyncio -import os import uuid -from collections.abc import MutableMapping from datetime import UTC, datetime from pathlib import Path from typing import ClassVar -from urllib.parse import urlsplit import httpx import typer @@ -228,11 +225,9 @@ async def _flow() -> str: f"Insight file: {effective_insight.ref} (default; pass --insight to override)", err=True, ) - _announce_credential_defaults() # Phase 1 — cheap environment checks BEFORE resolution: certain failures - # (missing creds, docker down, harbor absent) surface as the grouped - # report before any dataset download or the loop-chain import (which - # itself requires NEMO_EXPERIMENTALIST_API_* env) can preempt them. check_profile is + # (missing models, docker down, harbor absent) surface as the grouped + # report before any dataset download. check_profile is # doctor-only on purpose: flags may fully specify the run, and resolution # reports missing inputs with the skeleton. _preflight_or_exit( @@ -340,7 +335,7 @@ def doctor( envvar="NMP_BASE_URL", ), ) -> None: - """Diagnose Experimentalist setup: profile, artifacts, credentials, platform, runtime.""" + """Diagnose Experimentalist setup: profile, artifacts, models, platform, runtime.""" found = profile_path or discover_profile() profile_obj, profile_error = None, None env_results: list[CheckResult] = [] @@ -363,7 +358,6 @@ def doctor( ) ) base_url_resolved = resolve_base_url(base_url) - _announce_credential_defaults() insight_results: list[CheckResult] = [] try: effective_insight = resolve_effective_insight( @@ -495,49 +489,6 @@ def _announce_env_file(profile_dir: Path) -> None: typer.echo(f"Loaded .env from {env_path} ({len(loaded)} vars)", err=True) -_GATEWAY_BASE = "https://inference-api.nvidia.com/v1" -_GATEWAY_HOST = "inference-api.nvidia.com" - - -def _is_gateway_base(value: str) -> bool: - try: - parsed = urlsplit(value) - except ValueError: - return False - return parsed.scheme == "https" and parsed.hostname == _GATEWAY_HOST - - -def _apply_credential_defaults(env: MutableMapping[str, str] = os.environ) -> list[str]: - """Fill gateway-shaped credential gaps; returns what was applied. - - ``NEMO_EXPERIMENTALIST_API_BASE`` defaults to the NVIDIA Inference Gateway, and when - the effective base IS the gateway, ``INFERENCE_API_KEY`` can power the - experimentalist too. A custom base never inherits the gateway key. - - This still writes to the environment rather than to - :class:`~nemo_experimentalist_plugin.settings.ExperimentalistConfig`, because the - environment is the highest-precedence source and this runs before anything resolves - the config. Whether the copy should exist at all is a separate open question. - """ - applied: list[str] = [] - if not env.get("NEMO_EXPERIMENTALIST_API_BASE", "").strip(): - env["NEMO_EXPERIMENTALIST_API_BASE"] = _GATEWAY_BASE - applied.append(f"NEMO_EXPERIMENTALIST_API_BASE={_GATEWAY_BASE}") - if _is_gateway_base(env["NEMO_EXPERIMENTALIST_API_BASE"]): - inference = env.get("INFERENCE_API_KEY", "").strip() - optimizer = env.get("NEMO_EXPERIMENTALIST_API_KEY", "").strip() - if inference and not optimizer: - env["NEMO_EXPERIMENTALIST_API_KEY"] = inference - applied.append("NEMO_EXPERIMENTALIST_API_KEY=INFERENCE_API_KEY") - return applied - - -def _announce_credential_defaults() -> None: - applied = _apply_credential_defaults() - if applied: - typer.echo("Credential defaults: " + ", ".join(applied), err=True) - - def _load_config_payload(config: Path | None) -> dict | None: """Read a --config file; an explicitly passed empty file is an error, not a silent fallback to the profile's experiment_config.""" diff --git a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/config.py b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/config.py index 2df38fc1a3..440c175897 100644 --- a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/config.py +++ b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/config.py @@ -9,9 +9,8 @@ says otherwise would be a bad failure, and it would make ``config_snapshot`` a dishonest record of what ran. It is therefore a plain ``BaseModel`` with no environment binding. -The other half of the configuration -- which endpoint this install talks to and with which -models -- is a deployment setting, lives in ``settings.py`` as a :class:`NemoConfig` like -every other plugin's, and *does* let the environment win over the config file. +The optimizer's own default/fast model pair is selected by ``nemo setup`` and +stored in the active Platform CLI context. Component-owned slices (``CoderConfig``, ``AnalyzerConfig``, ...) are imported from the components that consume them rather than redeclared here -- ``resolve.py`` used to carry a @@ -86,8 +85,7 @@ class EvolutionaryOptimizerConfig(BaseModel): Deliberately not a :class:`NemoConfig`: these values describe a single experiment, are named explicitly on the command line, and are recorded in ``config_snapshot`` as the account of what ran. Letting an ambient environment variable override them would make - that account wrong. Endpoint and model settings live in - :class:`~nemo_experimentalist_plugin.settings.ExperimentalistConfig`. + that account wrong. Agent model settings live in the active Platform CLI context. """ @model_validator(mode="before") @@ -99,9 +97,7 @@ def reject_legacy_curator_config(cls, data: Any) -> Any: raise ValueError("'curator' was renamed to 'eval_author'; update the optimizer configuration") if "models" in data: raise ValueError( - "'models' is no longer a run-config key; model tiers are deployment settings. " - "Set them under the 'experimentalist:' config section or as " - "NEMO_EXPERIMENTALIST_MODELS_{SMART,MID,FAST}." + "'models' is no longer a run-config key. Run `nemo setup` to select the default and fast agent models." ) return data diff --git a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/analyzer.py b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/analyzer.py index 9728abe8a2..f8316e179a 100644 --- a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/analyzer.py +++ b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/analyzer.py @@ -19,6 +19,7 @@ ) from nemo_experimentalist_plugin.experimentalist.components.trace_explorer import TraceExplorer # noqa: F401 from nemo_platform import AsyncNeMoPlatform +from nemo_platform_plugin.nooa_model_client import get_default_model, get_fast_model from nooa import Agent, CodeActStrategy, strategy from nooa.agentdoc import doc, spec from nooa.agents import TokenBudgetSummarizer @@ -29,7 +30,6 @@ from pydantic import BaseModel, Field from . import cache -from .model_config import get_fast_model, get_smart_model from .rationalizer import Rationale, Rationalizer, RationalizerConfig # noqa: F401 from .tools import GuardedShellTools from .util import load_framework_skills @@ -244,7 +244,7 @@ def __init__( **kwargs: Forwarded to ``Agent.__init__``. """ - super().__init__(llm=kwargs.pop("llm", None) or get_smart_model(), **kwargs) + super().__init__(llm=kwargs.pop("llm", None) or get_default_model(), **kwargs) self._config = config or AnalyzerConfig() self._workspace_path = workspace self._framework_skills_dirs: list[Path] = framework_skills_dirs or [] diff --git a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/cards.py b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/cards.py index 44f559531a..60a5faebc9 100644 --- a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/cards.py +++ b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/cards.py @@ -1144,9 +1144,9 @@ class OptimizeModelCapability(Skill): default. Duplicate provider route aliases are intentionally omitted so tuning choices are not split across equivalent IDs. If editing an LLM model id, also call `available_model_ids = await self.list_available_models()` first and - choose a catalog model that is currently available. - - > **Tip**: The full live catalog can be queried at runtime — `requests.get("https://inference-api.nvidia.com/v1/models", headers={"Authorization": f"Bearer {key}"})`. `assets/models.yaml` is a curated subset. + choose one of those curated IDs. Optimizer model configuration is separate: + the default and fast Platform Model Entities selected by `nemo setup` run the + optimization agents and do not constrain model IDs used by the agent under test. **Code example** — upgrading from a fast model to a stronger one for classification: diff --git a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/coder.py b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/coder.py index d72e991456..cdcd3b3560 100644 --- a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/coder.py +++ b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/coder.py @@ -11,9 +11,9 @@ from pathlib import Path from typing import Any -import httpx from nemo_experimentalist_plugin.entities import Candidate, Dataset, EvaluationResult, Task, TrialResult from nemo_experimentalist_plugin.experimentalist.components.evaluator import Evaluator, EvaluatorConfig +from nemo_platform_plugin.nooa_model_client import get_default_model, get_fast_model from nooa import Agent, CodeActStrategy, strategy from nooa.agentdoc import doc, spec from nooa.agents import TokenBudgetSummarizer @@ -25,7 +25,6 @@ from pydantic import BaseModel, Field from .cards import Optimize -from .model_config import api_base, api_key, get_fast_model, get_mid_model, get_smart_model from .tools import GuardedShellTools from .util import load_framework_skills @@ -41,10 +40,6 @@ class CoderConfig(BaseModel): default=2, description="Max LLM repair iterations inside integration_check before giving up on a candidate.", ) - timeout_model_list_secs: float = Field( - default=10.0, - description="HTTP timeout in seconds when fetching the list of available LLM models.", - ) model_catalog_path: Path | None = Field( default=None, description="Optional YAML model catalog path overriding the packaged assets/models.yaml.", @@ -285,7 +280,7 @@ def summarize(hits: list[str]) -> str: # stochastic method (LLM) solid entrypoint from framework, all LLM-driven calls dashed. ```python - class MyAgent(Agent, llm=get_smart_model()): + class MyAgent(Agent, llm=get_default_model()): '''You are a research assistant.''' def __init__(self): @@ -383,7 +378,7 @@ def create_graph(self): with a different model and non-default temperature; deterministic orchestrator method. ```python - class WriterAgent(Agent, llm=get_smart_model()): + class WriterAgent(Agent, llm=get_default_model()): '''You are a writing assistant. Draft and refine content.''' def __init__(self): @@ -437,7 +432,7 @@ async def run(self, brief: str) -> str: sequence as a chain (not fan-out), every stochastic method gets dashed arrows to all tools. ```python - class Orchestrator(Agent, llm=get_smart_model()): + class Orchestrator(Agent, llm=get_default_model()): '''You are an optimization orchestrator.''' def __init__(self): @@ -583,10 +578,10 @@ def __init__( **kwargs: Any, ): """Initialize the coder for the given workspace.""" - super().__init__(llm=kwargs.pop("llm", None) or get_smart_model(), **kwargs) - # create_architecture_doc runs on the mid tier. Resolved here, like every other - # tier this component uses, and read off the instance by the decorator's callable. - self._mid_model = get_mid_model() + super().__init__(llm=kwargs.pop("llm", None) or get_default_model(), **kwargs) + # Architecture extraction requires the same quality-oriented model as + # the rest of the coding work. + self._architecture_model = get_default_model() self._config = config or CoderConfig() self._workspace_path = workspace.resolve() self.shell = GuardedShellTools(cwd=self._workspace_path) @@ -604,7 +599,6 @@ def __init__( self.optimize = Optimize(model_catalog_path=self._config.model_catalog_path) self.skills.register("ext.optimize", self.optimize) self.skills.activate(["cmd.*", "ext.*"]) - self._models_cache: list[str] | None = None TokenBudgetSummarizer.install( self, llm=get_fast_model(), @@ -612,44 +606,14 @@ def __init__( ) async def list_available_models(self) -> list[str]: - """Fetch available LLM model IDs from the configured inference API. - - Results are cached in memory for the lifetime of this instance. + """Return the curated model IDs allowed for agent-under-test mutations. Returns: - list[str]: model ID strings as returned by the API. - - Raises: - ValueError: if the endpoint or its credential is configured nowhere. - httpx.HTTPStatusError: if the API returned a non-2xx response. - httpx.RequestError: if there was a network or connection failure. + list[str]: Model IDs from the configured Experimentalist catalog. """ - if self._models_cache is not None: - return self._models_cache - - base, key = api_base(), api_key() - async with httpx.AsyncClient() as client: - resp = await client.get( - f"{base}/models", - headers={"Authorization": f"Bearer {key}"}, - timeout=self._config.timeout_model_list_secs, - ) - resp.raise_for_status() - data = resp.json() - - rows = data.get("data", []) - if not isinstance(rows, list): - raise ValueError("Invalid /models response: expected data to be a list") - models = [ - model_id - for row in rows - if isinstance(row, dict) and isinstance(model_id := row.get("id"), str) and model_id - ] - if not models: - raise ValueError("No model ids found in /models response") - self._models_cache = models - return models + catalog = self.optimize.optimize_model_capability.read_model_catalog() + return [model.model_id for model in catalog.models] async def run( self, @@ -1035,7 +999,8 @@ async def run_smoke_eval( ) @strategy( - CodeActStrategy(config=CodeActConfig(max_iterations=50, cell_timeout=3600.0)), llm=lambda self: self._mid_model + CodeActStrategy(config=CodeActConfig(max_iterations=50, cell_timeout=3600.0)), + llm=lambda self: self._architecture_model, ) async def create_architecture_doc( self, agent_id: str, source_path: str | None = None, entrypoint: str | None = None diff --git a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/goal_tree.py b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/goal_tree.py index c55f022fec..8f3f48ad84 100644 --- a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/goal_tree.py +++ b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/goal_tree.py @@ -9,6 +9,7 @@ from typing import Any from nemo_experimentalist_plugin.entities import Dataset +from nemo_platform_plugin.nooa_model_client import get_fast_model from nooa import Agent, CodeActStrategy, strategy from nooa.agentdoc import doc, spec from nooa.agents import TokenBudgetSummarizer @@ -18,7 +19,6 @@ from nooa.tools import Match, ShellTools, TodoManager from pydantic import BaseModel, Field, model_validator -from .model_config import get_fast_model from .tools import WorkspaceTool from .util import load_framework_skills @@ -464,7 +464,7 @@ async def _update( analysis: str, round_num: int, agent_spec: Path | None = None, - ) -> GoalTree: # pyright: ignore[reportReturnType] + ) -> GoalTree: # pyright: ignore[reportReturnType] # ty: ignore[invalid-return-type] """Propose a reweighted goal tree informed by a round of agent analysis. # What you receive diff --git a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/loop.py b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/loop.py index 3160188ce1..4c7c7eafb6 100644 --- a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/loop.py +++ b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/loop.py @@ -51,10 +51,6 @@ ensure_heldout_hidden, restore_heldout_splits, ) -from nemo_experimentalist_plugin.experimentalist.components.model_config import ( - get_fast_model, - get_smart_model, -) from nemo_experimentalist_plugin.experimentalist.components.models import ( EvolutionTree, OptimizationType, @@ -76,6 +72,7 @@ ) from nemo_experimentalist_plugin.experimentalist.result import ExperimentalistResult from nemo_platform import AsyncNeMoPlatform +from nemo_platform_plugin.nooa_model_client import get_default_model, get_fast_model from nooa import Agent, CodeActStrategy, strategy from nooa.agentdoc import doc, spec from nooa.agents import TokenBudgetSummarizer @@ -337,7 +334,7 @@ def __init__( framework_skills_dirs: list[Path] | None = None, **kwargs: Any, ) -> None: - super().__init__(llm=kwargs.pop("llm", None) or get_smart_model(), **kwargs) + super().__init__(llm=kwargs.pop("llm", None) or get_default_model(), **kwargs) self.working_dir = working_dir.resolve() self.config = config or EvolutionaryOptimizerConfig() self._config = self.config diff --git a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/model_config.py b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/model_config.py deleted file mode 100644 index 24846e68ef..0000000000 --- a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/model_config.py +++ /dev/null @@ -1,144 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -import functools -from collections.abc import Callable -from urllib.parse import urlparse - -from nemo_experimentalist_plugin.settings import TIERS, ExperimentalistConfig -from nooa.unifiedllm import CompletionClient - -_ENV_PREFIX = "NEMO_EXPERIMENTALIST_" - - -def _settings() -> ExperimentalistConfig: - """Return the resolved deployment settings. - - ``NemoConfig.get()`` caches, so this is cheap to call per lookup and picks up a - ``Configuration.clear_cache()`` in tests. - """ - return ExperimentalistConfig.get() - - -def _required(value: str | None, env_name: str) -> str: - if not value or not value.strip(): - raise ValueError(f"{env_name} must be set") - return value.strip() - - -def model_name(tier: str) -> str: - """Return the configured model name for *tier* (``smart``/``mid``/``fast``). - - No default: a model name is only meaningful against a specific endpoint, so there is - no portable value to fall back to. Failing here beats failing at the first LLM call, - minutes into a run. - - Raises: - ValueError: if the tier is unknown, or is configured nowhere. - - """ - if tier not in TIERS: - raise ValueError(f"unknown model tier {tier!r}") - configured = getattr(_settings().models, tier) - return _required(configured, f"{_ENV_PREFIX}MODELS_{tier.upper()}") - - -@functools.cache -def _client(name: str, api_base: str, api_key: str) -> CompletionClient: - """Cache on the full identity, not the tier: two tiers may name the same model, - and a caller may legitimately switch endpoint or key within one process.""" - return CompletionClient(name, api_base=api_base, api_key=api_key) - - -def api_base() -> str: - """Return the configured endpoint base URL. - - Raises: - ValueError: if it is configured nowhere. - - """ - return _required(_settings().api_base, f"{_ENV_PREFIX}API_BASE") - - -def api_key() -> str: - """Return the configured endpoint credential. - - Raises: - ValueError: if it is configured nowhere. - - """ - secret = _settings().api_key - return _required(secret.get_secret_value() if secret else None, f"{_ENV_PREFIX}API_KEY") - - -def get_model(tier: str) -> CompletionClient: - """Build (or reuse) the client for *tier*, resolving settings at call time. - - Raises: - ValueError: if the endpoint, the credential, or the tier's model name is - configured nowhere. - - """ - return _client(model_name(tier), api_base(), api_key()) - - -def get_smart_model() -> CompletionClient: - """Return the smart (high-capability) client, resolved from the environment.""" - return get_model("smart") - - -def get_mid_model() -> CompletionClient: - """Return the mid-tier client, resolved from the environment.""" - return get_model("mid") - - -def get_fast_model() -> CompletionClient: - """Return the fast (low-latency) client, resolved from the environment.""" - return get_model("fast") - - -def _mask_key(value: str) -> str: - if len(value) <= 4: - return "*" * len(value) - return f"*****{value[-4:]}" - - -def _sanitize_url(url: str) -> str: - """Strip userinfo, query, and fragment so credentials never leak via the log banner.""" - parsed = urlparse(url) - host = parsed.hostname or "" - port = f":{parsed.port}" if parsed.port else "" - return f"{parsed.scheme}://{host}{port}{parsed.path}" - - -def log_model_config() -> str: - """Return a formatted string summarising the current optimizer model configuration. - - Returns: - str: a multi-line summary of smart model, fast model, API base, and masked API key. - - """ - - def _or_unset(resolve: Callable[[], str]) -> str | None: - """Return the resolved value, or None when it is configured nowhere. - - Every field is optional here on purpose: this banner is most useful on exactly - the misconfigured install that cannot resolve one, so it must never be the thing - that raises. The callers that actually need a value still go through - ``api_base()`` / ``api_key()`` / ``model_name()`` and fail there. - """ - try: - return resolve() - except ValueError: - return None - - smart, mid, fast = (_or_unset(lambda t=tier: model_name(t)) or "(unset)" for tier in TIERS) - base = _or_unset(api_base) - key = _or_unset(api_key) - return ( - "Experimentalist model config:\n" - f" smart model: {smart}\n" - f" mid model: {mid}\n" - f" fast model: {fast}\n" - f" api base: {_sanitize_url(base) if base else '(unset)'}\n" - f" api key: {_mask_key(key) if key else '(unset)'}" - ) diff --git a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/proposer.py b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/proposer.py index ea14766328..6c1d647d5f 100644 --- a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/proposer.py +++ b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/proposer.py @@ -9,6 +9,7 @@ EvolutionTree, OptimizationType, ) +from nemo_platform_plugin.nooa_model_client import get_default_model, get_fast_model from nooa import Agent, CodeActStrategy, strategy from nooa.agentdoc import doc, spec from nooa.agents import TokenBudgetSummarizer @@ -18,7 +19,6 @@ from pydantic import BaseModel, Field from .cards import Optimize -from .model_config import get_fast_model, get_smart_model from .tools import WorkspaceTool from .util import load_framework_skills @@ -76,7 +76,7 @@ def __init__( framework_skills_dirs: list[Path] | None = None, **kwargs: Any, ): - super().__init__(llm=kwargs.pop("llm", None) or get_smart_model(), **kwargs) + super().__init__(llm=kwargs.pop("llm", None) or get_default_model(), **kwargs) self._config = config or ProposerConfig() self._workspace_path = workspace.resolve() self.workspace = WorkspaceTool(workspace=self._workspace_path) diff --git a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/rationalizer.py b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/rationalizer.py index 8a2ab24f57..b82cb1e4ec 100644 --- a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/rationalizer.py +++ b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/rationalizer.py @@ -10,6 +10,7 @@ from typing import Any from nemo_experimentalist_plugin.entities import DependencyRuntime, Task +from nemo_platform_plugin.nooa_model_client import get_default_model, get_fast_model from nooa import Agent, CodeActStrategy, strategy from nooa.agentdoc import doc, spec from nooa.agents import TokenBudgetSummarizer @@ -20,7 +21,6 @@ from pydantic import BaseModel, Field from . import cache -from .model_config import get_fast_model, get_smart_model from .rationale import Rationale, RationaleStep from .tools import GuardedShellTools from .util import load_framework_skills @@ -78,7 +78,7 @@ def __init__( **kwargs: Forwarded to ``Agent.__init__``. """ - super().__init__(llm=kwargs.pop("llm", None) or get_smart_model(), **kwargs) + super().__init__(llm=kwargs.pop("llm", None) or get_default_model(), **kwargs) self._config = config or RationalizerConfig() self._workspace_path = workspace self.shell = GuardedShellTools(cwd=workspace) diff --git a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/selector.py b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/selector.py index 511425cc39..2b904de9e5 100644 --- a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/selector.py +++ b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/selector.py @@ -10,8 +10,8 @@ from nemo_experimentalist_plugin.config import MetricTarget from nemo_experimentalist_plugin.entities import Candidate -from nemo_experimentalist_plugin.experimentalist.components.model_config import get_smart_model from nemo_experimentalist_plugin.experimentalist.components.tools import WorkspaceTool +from nemo_platform_plugin.nooa_model_client import get_default_model from nooa import Agent, CodeActStrategy, strategy from nooa.config import CodeActConfig @@ -20,7 +20,7 @@ class SurvivorSelector(Agent): """Select survivors from a Pareto-ranked population.""" def __init__(self, workspace: Path, **kwargs: Any) -> None: - super().__init__(llm=kwargs.pop("llm", None) or get_smart_model(), **kwargs) + super().__init__(llm=kwargs.pop("llm", None) or get_default_model(), **kwargs) self.workspace = WorkspaceTool(workspace=workspace) @strategy(CodeActStrategy(config=CodeActConfig(max_iterations=100, cell_timeout=3600.0))) diff --git a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/terminator.py b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/terminator.py index 623051a6d3..5e12ceb353 100644 --- a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/terminator.py +++ b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/terminator.py @@ -24,13 +24,12 @@ ) from nemo_experimentalist_plugin.experimentalist.components.models import EvolutionTree, pareto_front from nemo_experimentalist_plugin.skills import skills_dir +from nemo_platform_plugin.nooa_model_client import get_fast_model from nooa import Agent, CodeActStrategy, TextSkill, hidden, strategy from nooa.agentdoc import doc from nooa.config import CodeActConfig from pydantic import BaseModel -from .model_config import get_fast_model - logger = logging.getLogger(__name__) @@ -194,7 +193,7 @@ async def qualitative_stop_check( analysis: str, objective_metrics: list[MetricTarget], regression_metrics: list[MetricTarget], - ) -> bool: # pyright: ignore[reportReturnType] + ) -> bool: # pyright: ignore[reportReturnType] # ty: ignore[invalid-return-type] """Decide whether the optimization has qualitatively plateaued; return True to stop. Judge the round ``analysis`` text against the terminator skill's stop diff --git a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/trace_analyzer.py b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/trace_analyzer.py index 433a7091b7..fc2feb5e73 100644 --- a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/trace_analyzer.py +++ b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/trace_analyzer.py @@ -13,10 +13,10 @@ from nemo_experimentalist_plugin.entities import DependencyRuntime, MetricResult, Task, TrialResult from nemo_experimentalist_plugin.experimentalist.components import cache -from nemo_experimentalist_plugin.experimentalist.components.model_config import get_fast_model, get_smart_model from nemo_experimentalist_plugin.experimentalist.components.tools import GuardedShellTools, WorkspaceTool from nemo_insights_plugin.entities import Insight from nemo_platform import AsyncNeMoPlatform +from nemo_platform_plugin.nooa_model_client import get_default_model, get_fast_model from nooa import Agent, CodeActStrategy, strategy from nooa.agentdoc import doc, spec from nooa.agents import TokenBudgetSummarizer @@ -122,7 +122,7 @@ def __init__( **kwargs: Forwarded to ``Agent.__init__``. """ - super().__init__(llm=kwargs.pop("llm", None) or get_smart_model(), **kwargs) + super().__init__(llm=kwargs.pop("llm", None) or get_default_model(), **kwargs) self._config = config or TraceAnalyzerConfig() self._experiment_dir = experiment_dir self.shell = GuardedShellTools(cwd=experiment_dir) diff --git a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/trace_scorer.py b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/trace_scorer.py index 4ddc8908d7..0e0b2b72ab 100644 --- a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/trace_scorer.py +++ b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/trace_scorer.py @@ -8,13 +8,13 @@ from nemo_experimentalist_plugin.entities import Dataset, TrialResult from nemo_experimentalist_plugin.experimentalist.components.trace_explorer import TraceExplorer # noqa: F401 from nemo_platform import AsyncNeMoPlatform +from nemo_platform_plugin.nooa_model_client import get_fast_model from nooa import Agent, CodeActStrategy, strategy from nooa.agentdoc import doc from nooa.config import CodeActConfig from pydantic import BaseModel, Field from .goal_tree import GoalNode -from .model_config import get_mid_model logger = logging.getLogger(__name__) @@ -49,7 +49,7 @@ def __init__( Args: workspace: Path to the workspace directory. - llm: Language model instance; defaults to mid-tier model selection. + llm: Language model instance; defaults to the configured fast model. client: NeMo Platform client; required to load ``intake://`` traces. nmp_workspace: NeMo Platform workspace name; required to load ``intake://`` traces. **kwargs: Additional arguments passed to parent Agent class. @@ -58,7 +58,7 @@ def __init__( ValueError: if workspace does not exist or is invalid. """ - super().__init__(llm=llm or get_mid_model(), **kwargs) + super().__init__(llm=llm or get_fast_model(), **kwargs) self.workspace = workspace self._client = client self._nmp_workspace = nmp_workspace @@ -70,7 +70,7 @@ async def run( node: GoalNode, trials: dict[str, TrialResult], dataset: Dataset, - ) -> dict[str, GroupLeafScore]: # pyright: ignore[reportReturnType] + ) -> dict[str, GroupLeafScore]: # pyright: ignore[reportReturnType] # ty: ignore[invalid-return-type] """Compute the relative group advantage score for a group of traces coming from different agents for a given node. All scores must be strictly ordered: score_a < score_b < ... < score_n (no ties). diff --git a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/run.py b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/run.py index b122a5363a..c54055067e 100644 --- a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/run.py +++ b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/run.py @@ -3,10 +3,9 @@ """Reusable optimizer Experimentalist run orchestration.""" -import importlib import logging from pathlib import Path -from typing import Protocol, TextIO, cast +from typing import TextIO from nemo_experimentalist_plugin.entities import DatasetRef from nemo_experimentalist_plugin.experimentalist.agent import build_experimentalist_agent @@ -17,14 +16,16 @@ ) from nemo_experimentalist_plugin.experimentalist.reporting import RunReporter, Verbosity from nemo_platform import AsyncNeMoPlatform +from nemo_platform_plugin.nooa_model_client import ( + ConfiguredModelClients, + ConfiguredModelRefs, + activate_model_clients, + resolve_model_clients, +) logger = logging.getLogger(__name__) -class _LiteLLMModule(Protocol): - drop_params: bool - - def build_run_reporter( *, run_dir: Path, @@ -42,7 +43,7 @@ def build_run_reporter( async def run_experimentalist( *, - agent: str | None = None, + agent: Path | str | None = None, agent_spec: str | None = None, insight: Path | str | None, train_dataset: DatasetRef, @@ -53,6 +54,7 @@ async def run_experimentalist( config: EvolutionaryOptimizerConfig, task_template: DatasetRef | None = None, framework_skills_dirs: list[Path] | None = None, + model_refs: ConfiguredModelRefs | None = None, ) -> str: """Build and run the Experimentalist against an agent and dataset. @@ -70,10 +72,12 @@ async def run_experimentalist( task_template: Evaluator-specific task template used for production traces. experiment_dir: Working directory for optimization artifacts. workspace: Platform workspace. - client: Optional caller-owned Platform client. Local-only Mode 2 runs - may pass ``None``; Platform Insight access, mirroring, and Intake - persistence require a client. + client: Optional caller-owned Platform client. When omitted, the run + creates one from the active Platform context for model resolution; + the data backend remains local-only. config: Evolutionary optimizer configuration. + model_refs: Optional explicit default/fast Model Entity IDs. Unset uses + the active Platform CLI context. Returns: Terminal optimization summary. @@ -81,8 +85,6 @@ async def run_experimentalist( # Logging is configured at the entry-point boundary (the root ``nemo`` CLI # callback runs ``configure_logging`` before dispatching this subcommand), # so this library function leaves root logging untouched. - _enable_litellm_drop_params() - experiment_dir.mkdir(parents=True, exist_ok=True) experiment_dir = experiment_dir.resolve() # Leave ``agent`` unresolved: it may be a git ``url@ref`` (not a filesystem path). @@ -92,7 +94,7 @@ async def run_experimentalist( reporter = build_run_reporter( run_dir=experiment_dir, - agent=agent or "(from insight)", + agent=str(agent) if agent else "(from insight)", insight=str(insight) if insight is not None else None, ) @@ -101,37 +103,41 @@ async def run_experimentalist( experiments_output=str(experiment_dir), storage=config.storage, ) - deps = ExperimentalistDeps( - workspace=workspace, - agent=agent, - agent_spec=agent_spec, - insight=insight, - train_dataset=train_dataset, - validation_dataset=validation_dataset, - task_template=task_template, - backend=backend, - reporter=reporter, - config=config, - ) - experimentalist = build_experimentalist_agent( - working_dir=experiment_dir, - config=config, - framework_skills_dirs=framework_skills_dirs, - ) - result = await experimentalist.run(deps) - winner = result.winner - reporter.run_finished( - winner=winner.label if winner is not None else None, - scores=dict(winner.reward("validation").metrics) if winner and winner.reward("validation").metrics else {}, - report_path=(experiment_dir / "eval-and-optimize" / "OPTIMIZATION.md") if winner is not None else None, - ) - return result.summary - - -def _enable_litellm_drop_params() -> None: - """Let LiteLLM omit unsupported model parameters when it is installed.""" + model_platform_client = client or AsyncNeMoPlatform() + owns_model_platform_client = client is None + model_clients: ConfiguredModelClients | None = None try: - litellm = cast(_LiteLLMModule, importlib.import_module("litellm")) - except ModuleNotFoundError: - return - litellm.drop_params = True + model_clients = await resolve_model_clients(model_platform_client, model_refs) + deps = ExperimentalistDeps( + workspace=workspace, + agent=agent, + agent_spec=agent_spec, + insight=insight, + train_dataset=train_dataset, + validation_dataset=validation_dataset, + task_template=task_template, + backend=backend, + reporter=reporter, + config=config, + ) + with activate_model_clients(model_clients): + experimentalist = build_experimentalist_agent( + working_dir=experiment_dir, + config=config, + framework_skills_dirs=framework_skills_dirs, + ) + result = await experimentalist.run(deps) + winner = result.winner + reporter.run_finished( + winner=winner.label if winner is not None else None, + scores=dict(winner.reward("validation").metrics) if winner and winner.reward("validation").metrics else {}, + report_path=(experiment_dir / "eval-and-optimize" / "OPTIMIZATION.md") if winner is not None else None, + ) + return result.summary + finally: + try: + if model_clients is not None: + await model_clients.aclose() + finally: + if owns_model_platform_client: + await model_platform_client.close() diff --git a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/preflight.py b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/preflight.py index a0efb16548..29d636a66b 100644 --- a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/preflight.py +++ b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/preflight.py @@ -9,7 +9,6 @@ """ import importlib.util -import os import subprocess import tomllib from collections.abc import Callable, Mapping @@ -33,6 +32,7 @@ ) from nemo_insights_plugin.contracts.checks import CheckResult, make_check_result from nemo_insights_plugin.contracts.profile import resolve_profile_path +from nemo_platform_plugin.nooa_model_client import configured_model_refs from pydantic import BaseModel _SKELETON_HINT = ( @@ -62,13 +62,12 @@ def _default_http_ok(url: str) -> bool: class Probes(BaseModel): - """Injectable environment probes; defaults hit the real system.""" + """Injectable system probes; defaults hit the real system.""" model_config = {"arbitrary_types_allowed": True} run_cmd: Callable[[list[str]], tuple[int, str]] = _default_run_cmd http_ok: Callable[[str], bool] = _default_http_ok - env: Mapping[str, str] = os.environ def check_profile(profile: AgentProfile | None, profile_error: str | None) -> list[CheckResult]: @@ -119,34 +118,14 @@ def check_environment( ) -> list[CheckResult]: """Environment checks for Experiment and Doctor. - Checks the experimentalist credentials and model endpoint, platform - reachability, an optional insight file, Docker, and Harbor. + Checks configured agent models, platform reachability, an optional insight + file, Docker, and Harbor. ``enforce_insight_agent=False`` skips the local-insight agent match (an explicit ``--agent`` overrides the insight's agent). """ p = probes or Probes() - results: list[CheckResult] = [] - profile_dir = profile.profile_dir if profile is not None else None - results += _check_env( - p, "credentials-experiment", ("NEMO_EXPERIMENTALIST_API_BASE", "NEMO_EXPERIMENTALIST_API_KEY"), profile_dir - ) - base = p.env.get("NEMO_EXPERIMENTALIST_API_BASE") - if base: - model_url = f"{base.rstrip('/')}/models" - ok = p.http_ok(model_url) - display_model_url = f"{_redact_url(base).rstrip('/')}/models" - results.append( - make_check_result( - "model-endpoint", - "credentials-experiment", - ok, - "advisory", - f"{display_model_url} reachable", - "model endpoint unreachable", - hint="check NEMO_EXPERIMENTALIST_API_BASE and network", - ) - ) + results = check_models() ok = p.http_ok(f"{base_url.rstrip('/')}/health/ready") display_base_url = _redact_url(base_url) results.append( @@ -189,6 +168,32 @@ def check_environment( return results +def check_models() -> list[CheckResult]: + """Check that the active Platform context has a complete model pair.""" + try: + refs = configured_model_refs() + except (FileNotFoundError, RuntimeError, ValueError) as exc: + return [ + CheckResult( + name="agent-models", + group="models", + status="fail", + severity="required", + message=str(exc), + hint="run `nemo setup` and select default and fast agent models", + ) + ] + return [ + CheckResult( + name="agent-models", + group="models", + status="pass", + severity="required", + message=f"default={refs.default}; fast={refs.fast}", + ) + ] + + def check_artifacts( profile: AgentProfile | None, *, @@ -446,37 +451,6 @@ def _check_agent_source( return results -# Where each credential comes from — surfaced in the missing-var hint so the -# fix is actionable without hunting through the README. -_ENV_SOURCES = { - "NEMO_EXPERIMENTALIST_API_BASE": ( - "OpenAI-compatible LLM endpoint for the experimentalist (defaults to https://inference-api.nvidia.com/v1)" - ), - "NEMO_EXPERIMENTALIST_API_KEY": "API key for NEMO_EXPERIMENTALIST_API_BASE (on the gateway, INFERENCE_API_KEY fills this)", -} - -_ENV_EXAMPLE_POINTER = "see docs/get-started/example-agent.mdx" - - -def _check_env(p: Probes, group: str, names: tuple[str, ...], profile_dir: Path | None) -> list[CheckResult]: - env_location = str(profile_dir / ".env") if profile_dir else ".env next to your optimizer.yaml" - return [ - make_check_result( - name, - group, - bool(p.env.get(name, "").strip()), - "required", - f"{name} set", - f"{name} not set", - hint=( - f"{_ENV_SOURCES.get(name, 'required credential')}. " - f"Save it in {env_location} (auto-loaded; {_ENV_EXAMPLE_POINTER}) or export {name}=..." - ), - ) - for name in names - ] - - def _check_insight_file( insight: str | None, selector: str | None, diff --git a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/settings.py b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/settings.py deleted file mode 100644 index 65a9ec14da..0000000000 --- a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/settings.py +++ /dev/null @@ -1,66 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Deployment settings: which endpoint this install talks to, and with which models. - -Separate from ``config.py`` for layering, not taste. ``config.py`` imports each -component's config slice from the component that owns it, and those components resolve -their models through ``components/model_config.py`` -- so anything ``model_config`` needs -has to sit *below* the components, and ``config.py`` sits above them. This module imports -nothing from the plugin. - -For the run-parameter half of the configuration, and why the two do not share a -precedence rule, see ``config.py``. -""" - -from typing import ClassVar - -from nemo_platform_plugin.config import NemoConfig -from pydantic import BaseModel, Field, SecretStr - -TIERS: tuple[str, ...] = ("smart", "mid", "fast") - - -class ModelsConfig(BaseModel): - """Model name per tier. - - No tier has a default. A model name is only meaningful against a specific endpoint, so - there is no portable value to fall back to, and failing at startup beats failing at the - first LLM call minutes into a run. - """ - - smart: str | None = Field(default=None, description="Model for the smart tier, as your endpoint names it.") - mid: str | None = Field(default=None, description="Model for the mid tier, as your endpoint names it.") - fast: str | None = Field(default=None, description="Model for the fast tier, as your endpoint names it.") - - -class ExperimentalistConfig(NemoConfig): - """LLM endpoint and model tiers the Experimentalist's own agents run on. - - These are the models that *do* the optimizing -- proposing, coding, analyzing. The - agent being optimized is configured separately, by whatever the evaluator hands it. - - ``plugin_name`` gives this the ``NEMO_EXPERIMENTALIST_`` environment prefix and the - ``experimentalist:`` section of the platform config file. Precedence is the platform's: - environment, then config file, then these defaults. Nested fields use ``_``, so the - smart tier is ``NEMO_EXPERIMENTALIST_MODELS_SMART``. - - Contrast :class:`~nemo_experimentalist_plugin.config.EvolutionaryOptimizerConfig`, - which describes one experiment and takes no environment override at all. - """ - - plugin_name: ClassVar[str] = "experimentalist" - plugin_description: ClassVar[str] = "Endpoint and model settings for the NeMo Experimentalist." - - api_base: str | None = Field( - default=None, - description="Base URL of the OpenAI-compatible endpoint the optimizer's agents call.", - ) - api_key: SecretStr | None = Field( - default=None, - description="Credential for that endpoint. Prefer NEMO_EXPERIMENTALIST_API_KEY over writing it here.", - ) - models: ModelsConfig = Field( - default_factory=ModelsConfig, - description="Model per tier, named as your endpoint names it. Every tier is required.", - ) diff --git a/plugins/nemo-experimentalist/tests/conftest.py b/plugins/nemo-experimentalist/tests/conftest.py index eab95cd28c..292f4c3f34 100644 --- a/plugins/nemo-experimentalist/tests/conftest.py +++ b/plugins/nemo-experimentalist/tests/conftest.py @@ -1,26 +1,15 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Top-level conftest: runs before any test module is imported. - -Coder and Proposer call get_smart_model() at class-definition time (class keyword -argument). That call needs an endpoint and a credential, so we satisfy it before pytest -imports any test module, here at module level using setdefault — real CI values are not -overwritten. -""" +"""Experimentalist test-wide state isolation.""" import os +from typing import cast import litellm import pytest -from nemo_platform_plugin.config import Configuration - -os.environ.setdefault("NEMO_EXPERIMENTALIST_API_BASE", "http://placeholder-for-import") -os.environ.setdefault("NEMO_EXPERIMENTALIST_API_KEY", "placeholder-for-import") -# Tiers too, so constructing an agent does not depend on a developer's .env being loaded -# by the nooa import. Tests that care about tier resolution set their own values. -for _tier in ("SMART", "MID", "FAST"): - os.environ.setdefault(f"NEMO_EXPERIMENTALIST_MODELS_{_tier}", "placeholder/for-import") +from nemo_platform_plugin.nooa_model_client import ConfiguredModelClients, ConfiguredModelRefs, activate_model_clients +from nooa.unifiedllm import CompletionClient, FakeLLMClient # Some NVIDIA inference endpoint models reject the tool_choice parameter. # Drop unsupported params silently so the CodeAct strategy can call tools. @@ -29,23 +18,23 @@ @pytest.fixture(autouse=True) def _restore_environ(): - """Undo environment changes that monkeypatch cannot, and unshare the settings cache. + """Undo environment changes and activate hermetic agent models. The CLI loads a profile's .env straight into os.environ. When the variable was previously unset, monkeypatch has nothing recorded to restore, so the value survives the test and leaks into whatever else the xdist worker runs next. - ``ExperimentalistConfig.get()`` memoizes, so a test that sets endpoint or model - variables would otherwise be read back by the next test — or, worse, would silently - read the previous test's values itself. Overrides installed with - ``Configuration.set_override`` outlive the test that set them for the same reason. - Clearing both on either side of the yield keeps each test resolving its own settings. """ snapshot = os.environ.copy() - Configuration.clear_cache() - Configuration.clear_overrides() - yield + os.environ["NEMO_DEFAULT_MODEL"] = "default/fake" + os.environ["NEMO_FAST_MODEL"] = "default/fake" + fake = FakeLLMClient() + clients = ConfiguredModelClients( + default=cast(CompletionClient, fake), + fast=cast(CompletionClient, fake), + refs=ConfiguredModelRefs(default="default/fake", fast="default/fake"), + ) + with activate_model_clients(clients): + yield os.environ.clear() os.environ.update(snapshot) - Configuration.clear_cache() - Configuration.clear_overrides() diff --git a/plugins/nemo-experimentalist/tests/experimentalist/test_model_config.py b/plugins/nemo-experimentalist/tests/experimentalist/test_model_config.py deleted file mode 100644 index 9f20928b3f..0000000000 --- a/plugins/nemo-experimentalist/tests/experimentalist/test_model_config.py +++ /dev/null @@ -1,244 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -import os -import subprocess -import sys -import textwrap - -import pytest -from nemo_experimentalist_plugin.experimentalist.components.model_config import ( - _client, - api_base, - api_key, - log_model_config, - model_name, -) -from nemo_experimentalist_plugin.settings import ExperimentalistConfig, ModelsConfig -from nemo_platform_plugin.config import NMP_CONFIG_FILE_PATH_ENV_VAR, Configuration - -# Only the tests that assert on *where* a value comes from touch the environment. Everything -# else injects the settings object, which is both closer to what the code reads and immune -# to an ``/etc/nmp/config.yaml`` existing on the machine running the suite. -TIER_ENV = { - "smart": "NEMO_EXPERIMENTALIST_MODELS_SMART", - "mid": "NEMO_EXPERIMENTALIST_MODELS_MID", - "fast": "NEMO_EXPERIMENTALIST_MODELS_FAST", -} - - -def _configure(**kwargs) -> None: - """Install settings directly, bypassing both the environment and any config file. - - The environment has to be cleared first, and that is not belt-and-braces: this plugin's - settings are env-first by design, so ``ExperimentalistConfig(models=...)`` would read - conftest's import-time placeholders over the values passed here. The autouse fixture in - ``conftest`` snapshots and restores ``os.environ``, so removing them is safe. - """ - for name in [n for n in os.environ if n.startswith("NEMO_EXPERIMENTALIST_")]: - del os.environ[name] - Configuration.set_override(ExperimentalistConfig(**kwargs)) - - -def _configured(base: str = "https://llm.example/v1", key: str = "k", **tiers) -> None: - """The common case: a usable endpoint plus whichever tiers the test names.""" - _configure(api_base=base, api_key=key, models=ModelsConfig(**tiers)) - - -def test_model_name_is_required_per_tier() -> None: - """There is no portable default for a model name, so an unset tier must fail here. - - Any built-in default would name a model on exactly one endpoint and be wrong on every - other, surfacing as an opaque provider error at the first LLM call instead of a - configuration error before the run starts. - """ - _configured(base="https://inference-api.nvidia.com/v1") - - with pytest.raises(ValueError, match=TIER_ENV["smart"]): - model_name("smart") - - -def test_model_name_reads_its_tier() -> None: - _configured(smart="vendor/smart-1", fast="vendor/fast-1") - - assert model_name("smart") == "vendor/smart-1" - assert model_name("fast") == "vendor/fast-1" - - -def test_unknown_tier_is_rejected() -> None: - with pytest.raises(ValueError, match="unknown model tier"): - model_name("enormous") - - -def test_log_model_config_tolerates_unset_tiers() -> None: - """A display helper must not be the thing that fails.""" - _configured() - - assert "(unset)" in log_model_config() - - -def test_log_model_config_survives_a_wholly_unconfigured_install() -> None: - """The banner must not raise on the install it exists to diagnose. - - ``_shown`` always guarded the tiers, but the endpoint and key were resolved - unguarded, so the one case where you most want to see the banner was the one case - that raised. - """ - _configure() - - rendered = log_model_config() - - assert rendered.count("(unset)") == 5 - - -def test_log_model_config_masks_the_key() -> None: - """The banner is printed to logs, so it must never carry the whole credential.""" - _configured(key="sk-abcdefghijkl") - - rendered = log_model_config() - - assert "sk-abcdefghijkl" not in rendered - assert "ijkl" in rendered - - -def test_model_tiers_cache_on_full_identity() -> None: - """Repeat identities share a client; a changed key or base does not.""" - _client.cache_clear() - try: - first = _client("m", "https://base", "key-1") - assert _client("m", "https://base", "key-1") is first - assert _client("m", "https://base", "key-2") is not first - assert _client("m", "https://other", "key-1") is not first - finally: - _client.cache_clear() - - -def test_importing_components_resolves_no_model() -> None: - """Importing a component must not build an LLM client. - - Tiers resolve when an agent is constructed, not when its module is imported. - If this regresses, component modules cannot be imported without credentials -- - which is what registry discovery and ``nemo experimentalist doctor`` require. - - Runs in a subprocess because the client cache is process-global: any earlier - test that built a client would mask the regression. - """ - probe = ( - "from nemo_experimentalist_plugin.experimentalist.components import model_config\n" - "import nemo_experimentalist_plugin.experimentalist.components.loop\n" - "import nemo_experimentalist_plugin.experimentalist.components.coder\n" - "import nemo_eval_author_plugin.eval_author.agent\n" - "info = model_config._client.cache_info()\n" - "raise SystemExit(0 if info.hits + info.misses == 0 else 1)\n" - ) - completed = subprocess.run([sys.executable, "-c", probe], capture_output=True, text=True, check=False) - assert completed.returncode == 0, f"a model was resolved at import time\n{completed.stderr}" - - -# --------------------------------------------------------------------------- -# Where settings come from -# --------------------------------------------------------------------------- - - -def _write_platform_config(tmp_path, monkeypatch, body: str): - path = tmp_path / "config.yaml" - path.write_text(textwrap.dedent(body)) - monkeypatch.setenv(NMP_CONFIG_FILE_PATH_ENV_VAR, str(path)) - Configuration.clear_cache() - return path - - -def test_settings_come_from_the_platform_config_file(tmp_path, monkeypatch) -> None: - """Model tiers and endpoint are deployment settings, so the config file can supply them.""" - for name in (*TIER_ENV.values(), "NEMO_EXPERIMENTALIST_API_BASE", "NEMO_EXPERIMENTALIST_API_KEY"): - monkeypatch.delenv(name, raising=False) - _write_platform_config( - tmp_path, - monkeypatch, - """ - experimentalist: - api_base: https://from-file.example/v1 - api_key: file-key - models: - smart: vendor/from-file - """, - ) - - assert api_base() == "https://from-file.example/v1" - assert api_key() == "file-key" - assert model_name("smart") == "vendor/from-file" - - -def test_environment_overrides_the_config_file(tmp_path, monkeypatch) -> None: - """The platform's precedence, which this plugin now follows: env beats file beats default. - - This is the inverse of the old behaviour, where a config block was written into the - environment and so silently won over anything the operator had exported. - """ - _write_platform_config( - tmp_path, - monkeypatch, - """ - experimentalist: - api_base: https://from-file.example/v1 - models: - smart: vendor/from-file - """, - ) - monkeypatch.setenv("NEMO_EXPERIMENTALIST_API_BASE", "https://from-env.example/v1") - monkeypatch.setenv(TIER_ENV["smart"], "vendor/from-env") - Configuration.clear_cache() - - assert api_base() == "https://from-env.example/v1" - assert model_name("smart") == "vendor/from-env" - - -def test_api_key_is_not_exposed_by_repr() -> None: - """A settings object gets logged; the credential inside it must not come along.""" - _configured(key="sk-not-in-the-logs") - - assert "sk-not-in-the-logs" not in repr(ExperimentalistConfig.get()) - - -def test_reward_summary_is_separate_from_the_dimensions() -> None: - """A scalar rollup must not sit in metrics, where a selector would treat it as a - dimension that dominates every real one.""" - from nemo_experimentalist_plugin.entities import RewardRecord - - record = RewardRecord(metrics={"node-1": 0.6, "node-2": 0.8}, summary=0.7) - - assert "summary" not in record.metrics - assert record.summary == 0.7 - - -def test_record_reward_merges_rather_than_replaces() -> None: - """A second writer must not drop what the first measured. - - No production path writes a channel twice today, so nothing else would catch a - change to replace semantics — but the channel set is open by design, and the loss - would be silent. - """ - from nemo_experimentalist_plugin.entities import Candidate - - candidate = Candidate(run_id="run-1", label="agent-0", round=0, optimization="baseline") - candidate.record_reward("custom", metrics={"reward": 0.5}) - candidate.record_reward("custom", metadata={"source": "external"}) - - record = candidate.reward("custom") - assert record.metrics == {"reward": 0.5} - assert record.metadata == {"source": "external"} - - -def test_an_unmeasured_channel_is_distinguishable_from_an_empty_one() -> None: - """``reward()`` erases the difference on purpose; ``in`` is how you ask. - - Eight call sites gate evaluation on ``channel not in candidate.rewards``, so a - read must never make a channel appear measured. This is why the mapping cannot - become a ``defaultdict``, whose ``__missing__`` inserts. - """ - from nemo_experimentalist_plugin.entities import Candidate - - candidate = Candidate(run_id="run-1", label="agent-0", round=0, optimization="baseline") - - assert candidate.reward("train").metrics == {} - assert "train" not in candidate.rewards diff --git a/plugins/nemo-experimentalist/tests/experimentalist/test_terminator.py b/plugins/nemo-experimentalist/tests/experimentalist/test_terminator.py index 4335c37c1b..f6f24e3837 100644 --- a/plugins/nemo-experimentalist/tests/experimentalist/test_terminator.py +++ b/plugins/nemo-experimentalist/tests/experimentalist/test_terminator.py @@ -6,7 +6,7 @@ These guard the behavior-preserving extraction of the termination logic (``_has_converged`` + ``qualitative_stop_check`` + ``max_rounds``) out of the loop and into :mod:`terminator`. Every Terminator is constructed with an injected -``FakeLLMClient`` so the tests need no ``NEMO_EXPERIMENTALIST_API_*`` env vars and make no +``FakeLLMClient`` so the tests need no configured provider credentials and make no network calls. """ diff --git a/plugins/nemo-experimentalist/tests/experimentalist/test_tools.py b/plugins/nemo-experimentalist/tests/experimentalist/test_tools.py index 7a0ea7af33..99984b95b4 100644 --- a/plugins/nemo-experimentalist/tests/experimentalist/test_tools.py +++ b/plugins/nemo-experimentalist/tests/experimentalist/test_tools.py @@ -2,12 +2,15 @@ # SPDX-License-Identifier: Apache-2.0 from pathlib import Path +from typing import cast -from nemo_experimentalist_plugin.experimentalist.components.coder import Coder +from nemo_experimentalist_plugin.experimentalist.components.coder import Coder, CoderConfig from nemo_experimentalist_plugin.experimentalist.components.holdout_utils import BLOCKED_MESSAGE from nemo_experimentalist_plugin.experimentalist.components.tools import GuardedShellTools +from nemo_platform_plugin.nooa_model_client import ConfiguredModelClients, ConfiguredModelRefs, activate_model_clients from nooa.agentdoc import pformat from nooa.tools import ShellResult +from nooa.unifiedllm import CompletionClient, FakeLLMClient async def test_guarded_shell_tools_runs_allowed_commands(tmp_path): @@ -45,3 +48,42 @@ def test_coder_hides_skill_registry_that_can_replace_guarded_shell(tmp_path): assert "skills=SkillRegistry" not in rendered assert "shell=ShellTools" in rendered assert "nooa=Nooa" in rendered + + +def test_coder_uses_default_model_for_architecture_docs(tmp_path: Path) -> None: + default = cast(CompletionClient, FakeLLMClient()) + fast = cast(CompletionClient, FakeLLMClient()) + clients = ConfiguredModelClients( + default=default, + fast=fast, + refs=ConfiguredModelRefs(default="default/quality", fast="default/fast"), + ) + + with activate_model_clients(clients): + coder = Coder(workspace=tmp_path) + + assert coder._architecture_model is default + + +async def test_coder_lists_agent_mutation_models_from_catalog(tmp_path: Path) -> None: + catalog = tmp_path / "models.yaml" + catalog.write_text( + """\ +default_endpoint: https://models.example/v1 +default_api_key_env: EXAMPLE_API_KEY +id_field: model_id +models: + - model_id: provider/quality + provider: Example + range: quality + notes: Quality model. + - model_id: provider/fast + provider: Example + range: fast + notes: Fast model. +""", + encoding="utf-8", + ) + coder = Coder(workspace=tmp_path, config=CoderConfig(model_catalog_path=catalog)) + + assert await coder.list_available_models() == ["provider/quality", "provider/fast"] diff --git a/plugins/nemo-experimentalist/tests/test_cli_profile.py b/plugins/nemo-experimentalist/tests/test_cli_profile.py index eb69b6c8c9..8c093780d8 100644 --- a/plugins/nemo-experimentalist/tests/test_cli_profile.py +++ b/plugins/nemo-experimentalist/tests/test_cli_profile.py @@ -7,10 +7,11 @@ import os from dataclasses import dataclass from pathlib import Path +from typing import Any import httpx import pytest -from nemo_experimentalist_plugin import cli +from nemo_experimentalist_plugin import cli, preflight from nemo_experimentalist_plugin.preflight import Probes from nemo_insights_plugin.contracts.profile import DEFAULT_BASE_URL from typer.testing import CliRunner @@ -18,16 +19,10 @@ runner = CliRunner() -def quiet_probes(env: dict | None = None) -> Probes: +def quiet_probes() -> Probes: return Probes( run_cmd=lambda argv: (0, "ok"), http_ok=lambda url: True, - env=env - or { - "NEMO_EXPERIMENTALIST_API_BASE": "http://llm", - "NEMO_EXPERIMENTALIST_API_KEY": "k", - "INFERENCE_API_KEY": "k", - }, ) @@ -76,7 +71,12 @@ def profile_tree(tmp_path: Path) -> Path: class RunRecorder: def __init__(self) -> None: - self.kwargs: dict | None = None + self.kwargs: dict[str, Any] | None = None + + @property + def call_kwargs(self) -> dict[str, Any]: + assert self.kwargs is not None + return self.kwargs async def __call__(self, **kwargs) -> str: self.kwargs = kwargs @@ -91,9 +91,9 @@ def test_experiment_insight_only_with_profile(app, profile_tree: Path, monkeypat monkeypatch.chdir(profile_tree) result = runner.invoke(app, ["run", "--insight", "ins-1", "-o", str(profile_tree / "out")]) assert result.exit_code == 0, result.output - assert recorder.kwargs["insight"] == "ins-1" - assert recorder.kwargs["train_dataset"].uri == str((profile_tree / "evals" / "train").resolve()) - assert recorder.kwargs["agent"] == str(profile_tree.resolve()) + assert recorder.call_kwargs["insight"] == "ins-1" + assert recorder.call_kwargs["train_dataset"].uri == str((profile_tree / "evals" / "train").resolve()) + assert recorder.call_kwargs["agent"] == str(profile_tree.resolve()) def test_experiment_explicit_profile_flag(app, profile_tree: Path, monkeypatch, tmp_path: Path) -> None: @@ -117,7 +117,7 @@ def test_experiment_explicit_profile_flag(app, profile_tree: Path, monkeypatch, ], ) assert result.exit_code == 0, result.output - assert recorder.kwargs["workspace"] == "default" + assert recorder.call_kwargs["workspace"] == "default" def test_experiment_no_profile_missing_flags_errors_with_skeleton(app, tmp_path: Path, monkeypatch) -> None: @@ -150,7 +150,7 @@ def test_experiment_all_flags_no_profile_still_works(app, tmp_path: Path, monkey ], ) assert result.exit_code == 0, result.output - assert recorder.kwargs["insight"] is None + assert recorder.call_kwargs["insight"] is None def test_profileless_experiment_blocks_missing_effective_agent_path(app, tmp_path: Path, monkeypatch) -> None: @@ -258,9 +258,9 @@ def now(*, tz): ] first_result = runner.invoke(app, args) - first_dir = Path(recorder.kwargs["experiment_dir"]) + first_dir = Path(recorder.call_kwargs["experiment_dir"]) second_result = runner.invoke(app, args) - second_dir = Path(recorder.kwargs["experiment_dir"]) + second_dir = Path(recorder.call_kwargs["experiment_dir"]) expected_first = Path("tmp") / f"20260714-123456-{first_hex}" expected_second = Path("tmp") / f"20260714-123456-{second_hex}" @@ -292,7 +292,7 @@ def test_experiment_insight_id_selects_from_analyst_file(app, profile_tree: Path ["run", "--insight", str(insights), "--insight-id", "i-b", "-o", str(profile_tree / "o")], ) assert result.exit_code == 0, result.output - selected = json.loads(Path(recorder.kwargs["insight"]).read_text(encoding="utf-8")) + selected = json.loads(Path(recorder.call_kwargs["insight"]).read_text(encoding="utf-8")) assert selected["id"] == "i-b" @@ -407,25 +407,16 @@ def test_doctor_invalid_git_agent_path_is_structured_without_traceback( assert "Traceback" not in result.output -def test_doctor_redacts_model_and_platform_display_urls(app, profile_tree: Path, monkeypatch) -> None: +def test_doctor_redacts_platform_display_url(app, profile_tree: Path, monkeypatch) -> None: write_task_toml(profile_tree) - model_base = "https://model-user:model-secret@models.example:8443/v1" # trufflehog:ignore platform_base = "https://platform-user:platform-secret@platform.example:9443/api" # trufflehog:ignore - monkeypatch.setattr( - cli, - "_PREFLIGHT_PROBES", - quiet_probes(env={"NEMO_EXPERIMENTALIST_API_BASE": model_base, "NEMO_EXPERIMENTALIST_API_KEY": "k"}), - ) monkeypatch.chdir(profile_tree) result = runner.invoke(app, ["doctor", "--base-url", platform_base]) assert result.exit_code == 0, result.output - assert "https://***@models.example:8443/v1/models reachable" in result.output assert "https://***@platform.example:9443/api reachable" in result.output - assert not any( - secret in result.output for secret in ("model-user", "model-secret", "platform-user", "platform-secret") - ) + assert not any(secret in result.output for secret in ("platform-user", "platform-secret")) def test_experiment_hard_fails_on_required_check(app, profile_tree: Path, monkeypatch) -> None: @@ -437,7 +428,6 @@ def test_experiment_hard_fails_on_required_check(app, profile_tree: Path, monkey Probes( run_cmd=lambda argv: (1, "docker down"), http_ok=lambda url: True, - env={"NEMO_EXPERIMENTALIST_API_BASE": "b", "NEMO_EXPERIMENTALIST_API_KEY": "k"}, ), ) monkeypatch.chdir(profile_tree) @@ -457,7 +447,6 @@ def test_experiment_advisory_warns_but_runs(app, profile_tree: Path, monkeypatch Probes( run_cmd=lambda argv: (0, "ok"), http_ok=lambda url: False, # platform unreachable → advisory - env={"NEMO_EXPERIMENTALIST_API_BASE": "b", "NEMO_EXPERIMENTALIST_API_KEY": "k"}, ), ) monkeypatch.chdir(profile_tree) @@ -485,7 +474,6 @@ def missing_git(argv: list[str]) -> tuple[int, str]: Probes( run_cmd=missing_git, http_ok=lambda url: True, - env={"NEMO_EXPERIMENTALIST_API_BASE": "b", "NEMO_EXPERIMENTALIST_API_KEY": "k"}, ), ) config = profile_tree / "experiment.yaml" @@ -536,23 +524,27 @@ def test_explicit_profile_is_not_announced(app, profile_tree: Path, monkeypatch, assert "Using profile:" not in result.output -def test_missing_creds_reported_before_any_resolution(app, profile_tree: Path, monkeypatch) -> None: - # Phase-1 ordering: the grouped credentials report must surface even though - # resolution (whose loop import would raise a bare ValueError) never runs. +def test_missing_models_reported_before_any_resolution(app, profile_tree: Path, monkeypatch) -> None: + # Phase-1 ordering: the grouped model report must surface before resolution. write_task_toml(profile_tree) recorder = RunRecorder() monkeypatch.setattr(cli, "run_experimentalist", recorder) monkeypatch.setattr( cli, "_PREFLIGHT_PROBES", - Probes(run_cmd=lambda argv: (0, "ok"), http_ok=lambda url: True, env={}), + Probes(run_cmd=lambda argv: (0, "ok"), http_ok=lambda url: True), + ) + monkeypatch.setattr( + preflight, + "configured_model_refs", + lambda: (_ for _ in ()).throw(ValueError("No default model is configured")), ) monkeypatch.chdir(profile_tree) result = runner.invoke(app, ["run", "--insight", "ins-1", "-o", str(profile_tree / "o")]) assert result.exit_code == 1 assert recorder.kwargs is None - assert "NEMO_EXPERIMENTALIST_API_BASE" in result.output - assert "export NEMO_EXPERIMENTALIST_API_BASE" in result.output # the hint, not a bare ValueError + assert "No default model is configured" in result.output + assert "nemo setup" in result.output def test_insightless_run_skips_template_checks(app, profile_tree: Path, monkeypatch) -> None: @@ -564,7 +556,7 @@ def test_insightless_run_skips_template_checks(app, profile_tree: Path, monkeypa result = runner.invoke(app, ["run", "-o", str(profile_tree / "o")]) assert result.exit_code == 0, result.output assert recorder.kwargs is not None - assert recorder.kwargs["insight"] is None + assert recorder.call_kwargs["insight"] is None def test_empty_config_flag_is_an_error_not_silent_fallback(app, profile_tree: Path, monkeypatch) -> None: @@ -637,7 +629,7 @@ def test_experiment_loads_nmp_base_url_from_profile_env(app, profile_tree: Path, ) assert result.exit_code == 0, result.output - client = recorder.kwargs["client"] + client = recorder.call_kwargs["client"] assert client.base_url == "https://platform.example" assert client.closed @@ -656,7 +648,6 @@ def record_http(url: str) -> bool: Probes( run_cmd=lambda argv: (0, "ok"), http_ok=record_http, - env={"NEMO_EXPERIMENTALIST_API_BASE": "http://llm", "NEMO_EXPERIMENTALIST_API_KEY": "k"}, ), ) monkeypatch.delenv("NMP_BASE_URL", raising=False) @@ -680,7 +671,7 @@ def test_default_experiment_dir_is_profile_scoped_and_announced(app, profile_tre result = runner.invoke(app, ["run", "--insight", "ins-1"]) assert result.exit_code == 0, result.output assert "Experiment dir:" in result.output - experiment_dir = Path(recorder.kwargs["experiment_dir"]) + experiment_dir = Path(recorder.call_kwargs["experiment_dir"]) assert experiment_dir.is_relative_to(profile_tree / ".nemo-optimizer" / "experiments") @@ -717,7 +708,7 @@ def now(*, tz): expected = experiments_root / f"20260714-123456-{unique_hex}" assert result.exit_code == 0, result.output - assert Path(recorder.kwargs["experiment_dir"]) == expected + assert Path(recorder.call_kwargs["experiment_dir"]) == expected assert expected.is_dir() assert colliding_dir.is_dir() @@ -731,7 +722,6 @@ def test_default_experiment_dir_is_not_reserved_before_required_preflight(app, p Probes( run_cmd=lambda argv: (1, "docker down"), http_ok=lambda url: True, - env={"NEMO_EXPERIMENTALIST_API_BASE": "b", "NEMO_EXPERIMENTALIST_API_KEY": "k"}, ), ) monkeypatch.chdir(profile_tree) @@ -751,48 +741,7 @@ def test_explicit_experiment_dir_still_wins(app, profile_tree: Path, monkeypatch result = runner.invoke(app, ["run", "--insight", "ins-1", "-o", str(profile_tree / "custom")]) assert result.exit_code == 0, result.output assert "Experiment dir:" not in result.output - assert Path(recorder.kwargs["experiment_dir"]) == profile_tree / "custom" - - -def test_credential_defaults_inference_key_powers_gateway_experiment() -> None: - env = {"INFERENCE_API_KEY": "sk-gateway"} - applied = cli._apply_credential_defaults(env) - assert env["NEMO_EXPERIMENTALIST_API_BASE"] == "https://inference-api.nvidia.com/v1" - assert env["NEMO_EXPERIMENTALIST_API_KEY"] == "sk-gateway" - assert len(applied) == 2 - - -def test_credential_defaults_never_fill_analyst_key() -> None: - env = {"NEMO_EXPERIMENTALIST_API_KEY": "sk-gateway"} - cli._apply_credential_defaults(env) - assert "INFERENCE_API_KEY" not in env - - -def test_credential_defaults_custom_base_never_inherits_gateway_key() -> None: - custom = {"NEMO_EXPERIMENTALIST_API_BASE": "https://api.openai.com/v1", "INFERENCE_API_KEY": "sk-gateway"} - cli._apply_credential_defaults(custom) - assert "NEMO_EXPERIMENTALIST_API_KEY" not in custom - - -def test_credential_defaults_reject_gateway_lookalike_hosts() -> None: - for base in ( - "https://inference-api.nvidia.com.attacker.example/v1", - "https://evil-inference-api.nvidia.com/v1", - "http://inference-api.nvidia.com/v1", - ): - env = {"NEMO_EXPERIMENTALIST_API_BASE": base, "INFERENCE_API_KEY": "sk-secret"} - cli._apply_credential_defaults(env) - assert "NEMO_EXPERIMENTALIST_API_KEY" not in env - - -def test_credential_defaults_never_override() -> None: - env = { - "INFERENCE_API_KEY": "sk-a", - "NEMO_EXPERIMENTALIST_API_BASE": "https://inference-api.nvidia.com/v1", - "NEMO_EXPERIMENTALIST_API_KEY": "sk-b", - } - assert cli._apply_credential_defaults(env) == [] - assert env["NEMO_EXPERIMENTALIST_API_KEY"] == "sk-b" + assert Path(recorder.call_kwargs["experiment_dir"]) == profile_tree / "custom" def test_experiment_defaults_to_shared_insights_file(app, profile_tree: Path, monkeypatch) -> None: @@ -809,7 +758,7 @@ def test_experiment_defaults_to_shared_insights_file(app, profile_tree: Path, mo result = runner.invoke(app, ["run", "-o", str(profile_tree / "o")]) assert result.exit_code == 0, result.output assert "Insight file:" in result.output - selected = json.loads(Path(recorder.kwargs["insight"]).read_text(encoding="utf-8")) + selected = json.loads(Path(recorder.call_kwargs["insight"]).read_text(encoding="utf-8")) assert selected["id"] == "i-1" @@ -838,7 +787,7 @@ def test_experiment_selected_insight_uses_shared_default_without_stale_warning( ["run", "--insight-id", "1", "-o", str(profile_tree / "o")], ) assert result.exit_code == 0, result.output - selected = json.loads(Path(recorder.kwargs["insight"]).read_text(encoding="utf-8")) + selected = json.loads(Path(recorder.call_kwargs["insight"]).read_text(encoding="utf-8")) assert selected["id"] == "i-b" assert "will require --insight-id" not in result.output @@ -902,7 +851,7 @@ def test_experiment_insight_id_matching_single_local_insight_passes(app, profile ], ) assert result.exit_code == 0, result.output - selected = json.loads(Path(recorder.kwargs["insight"]).read_text(encoding="utf-8")) + selected = json.loads(Path(recorder.call_kwargs["insight"]).read_text(encoding="utf-8")) assert selected["id"] == "i-1" @@ -921,7 +870,7 @@ def test_experiment_no_insight_skips_existing_shared_default(app, profile_tree: ["run", "--no-insight", "-o", str(profile_tree / "o")], ) assert result.exit_code == 0, result.output - assert recorder.kwargs["insight"] is None + assert recorder.call_kwargs["insight"] is None assert "Insight disabled" in result.output assert "Insight file:" not in result.output @@ -1128,10 +1077,10 @@ def test_doctor_env_file_permission_failure_is_required_and_actionable( env_file.write_text("NMP_BASE_URL=https://profile.example\n", encoding="utf-8") original_read_text = Path.read_text - def deny_env_file(path: Path, *args: object, **kwargs: object) -> str: + def deny_env_file(path: Path, encoding: str | None = None, errors: str | None = None) -> str: if path == env_file: raise PermissionError("permission denied by test") - return original_read_text(path, *args, **kwargs) + return original_read_text(path, encoding=encoding, errors=errors) monkeypatch.setattr(Path, "read_text", deny_env_file) monkeypatch.chdir(profile_tree) @@ -1181,7 +1130,7 @@ def test_experiment_nmp_base_url_precedence_and_ignores_nemo_base_url( result = runner.invoke(app, args) assert result.exit_code == 0, result.output - client = recorder.kwargs["client"] + client = recorder.call_kwargs["client"] assert client.base_url == expected assert client.closed @@ -1213,7 +1162,6 @@ def test_doctor_nmp_base_url_precedence_and_ignores_nemo_base_url( Probes( run_cmd=lambda argv: (0, "ok"), http_ok=lambda url: not urls.append(url), - env={"NEMO_EXPERIMENTALIST_API_BASE": "http://llm", "NEMO_EXPERIMENTALIST_API_KEY": "k"}, ), ) monkeypatch.setenv("NEMO_BASE_URL", "https://legacy-must-be-ignored.example") @@ -1231,7 +1179,7 @@ def test_doctor_nmp_base_url_precedence_and_ignores_nemo_base_url( result = runner.invoke(app, args) assert result.exit_code == 0, result.output - assert urls == ["http://llm/models", f"{expected}/health/ready"] + assert urls == [f"{expected}/health/ready"] def test_experiment_help_names_insights_writer(app) -> None: diff --git a/plugins/nemo-experimentalist/tests/test_experiment_cli.py b/plugins/nemo-experimentalist/tests/test_experiment_cli.py index af537f7753..6a37ad0ebb 100644 --- a/plugins/nemo-experimentalist/tests/test_experiment_cli.py +++ b/plugins/nemo-experimentalist/tests/test_experiment_cli.py @@ -28,7 +28,6 @@ def quiet_preflight(monkeypatch: pytest.MonkeyPatch) -> None: Probes( run_cmd=lambda argv: (0, "ok"), http_ok=lambda url: True, - env={"NEMO_EXPERIMENTALIST_API_BASE": "http://llm", "NEMO_EXPERIMENTALIST_API_KEY": "k"}, ), ) @@ -226,7 +225,9 @@ def test_experiment_cli_passes_dataset_driven_contract_to_runner( # reads insight files with json.loads, so YAML-authored files must not # reach it raw); platform ids still pass through verbatim. assert runner.captured.insight == str(paths.experiment / "resolved" / "insight.json") - assert Path(runner.captured.insight).read_text(encoding="utf-8").strip() == "{}" + captured_insight = runner.captured.insight + assert captured_insight is not None + assert Path(captured_insight).read_text(encoding="utf-8").strip() == "{}" # Datasets and the task template are forwarded as DatasetRef URI handles, # tagged with a stable id the evaluator adapter uses when building datasets. assert runner.captured.train_dataset == DatasetRef(uri=str(paths.train), metadata={"id": "train"}) diff --git a/plugins/nemo-experimentalist/tests/test_experimentalist_benchmark.py b/plugins/nemo-experimentalist/tests/test_experimentalist_benchmark.py index d1cf96f80b..4da38b4e23 100644 --- a/plugins/nemo-experimentalist/tests/test_experimentalist_benchmark.py +++ b/plugins/nemo-experimentalist/tests/test_experimentalist_benchmark.py @@ -169,9 +169,6 @@ def configure( runner._configure_models( runner.ModelSpec( aut="a", - experimentalist_smart="b", - experimentalist_mid="c", - experimentalist_fast="d", user_simulator=user_simulator, ) ) @@ -297,3 +294,28 @@ def test_optimizer_job_summary_reports_full_trial_and_usage_totals(tmp_path: Pat assert summary["input_tokens"] == 150 assert summary["cost_usd"] == pytest.approx(0.2) assert len(summary["job_results"]) == 2 + + +@pytest.mark.parametrize("invalid_value", [True, 1.5, "2"]) +def test_optimizer_job_summary_rejects_non_integer_fields(tmp_path: Path, invalid_value: object) -> None: + runner = _load_runner() + job_dir = tmp_path / "agent-0-train" + job_dir.mkdir() + (job_dir / "result.json").write_text( + json.dumps( + { + "n_total_trials": 2, + "stats": { + "n_completed_trials": invalid_value, + "n_errored_trials": 0, + "n_input_tokens": 10, + "n_cache_tokens": 0, + "n_output_tokens": 5, + }, + } + ), + encoding="utf-8", + ) + + with pytest.raises(TypeError, match="completed_trials"): + runner.summarize_experimentalist_jobs(tmp_path) diff --git a/plugins/nemo-experimentalist/tests/test_experimentalist_run.py b/plugins/nemo-experimentalist/tests/test_experimentalist_run.py index 2b88ee807d..682a95035c 100644 --- a/plugins/nemo-experimentalist/tests/test_experimentalist_run.py +++ b/plugins/nemo-experimentalist/tests/test_experimentalist_run.py @@ -15,6 +15,7 @@ from nemo_experimentalist_plugin.experimentalist.experimentalist_backend import LocalExperimentalistBackend from nemo_experimentalist_plugin.experimentalist.result import ExperimentalistResult from nemo_platform import AsyncNeMoPlatform +from nemo_platform_plugin.nooa_model_client import ConfiguredModelRefs @dataclass @@ -25,6 +26,28 @@ async def close(self) -> None: self.closed = True +@dataclass +class ClosingModelClients: + default: object = "default" + fast: object = "fast" + refs: ConfiguredModelRefs = ConfiguredModelRefs(default="default/quality", fast="default/fast") + closed: bool = False + + async def aclose(self) -> None: + self.closed = True + + +@pytest.fixture +def model_clients(monkeypatch: pytest.MonkeyPatch) -> ClosingModelClients: + clients = ClosingModelClients() + + async def resolve(*_: object) -> ClosingModelClients: + return clients + + monkeypatch.setattr(experimentalist_run, "resolve_model_clients", resolve) + return clients + + @dataclass class ExperimentRunPaths: agent: Path @@ -83,6 +106,7 @@ def _make_run_paths(tmp_path: Path) -> ExperimentRunPaths: @pytest.mark.asyncio async def test_run_experimentalist_builds_and_runs_complete_local_contract( monkeypatch: pytest.MonkeyPatch, + model_clients: ClosingModelClients, tmp_path: Path, ) -> None: paths = _make_run_paths(tmp_path) @@ -92,7 +116,6 @@ async def test_run_experimentalist_builds_and_runs_complete_local_contract( experimentalist = FakeExperimentalist() backend_calls: list[BackendFactoryCall] = [] agent_calls: list[AgentFactoryCall] = [] - litellm_calls: list[bool] = [] def make_backend( *, @@ -117,7 +140,6 @@ def build_agent( monkeypatch.setattr(experimentalist_run, "make_experimentalist_backend", make_backend) monkeypatch.setattr(experimentalist_run, "build_experimentalist_agent", build_agent) - monkeypatch.setattr(experimentalist_run, "_enable_litellm_drop_params", lambda: litellm_calls.append(True)) train_dataset = DatasetRef(uri=str(paths.train)) validation_dataset = DatasetRef(uri=str(paths.validation)) @@ -142,8 +164,8 @@ def build_agent( ) ] assert agent_calls == [AgentFactoryCall(working_dir=paths.experiment.resolve(), config=optimizer_config)] - assert litellm_calls == [True] assert not client.closed + assert model_clients.closed assert experimentalist.deps is not None assert experimentalist.deps.workspace == "workspace-a" # ``agent`` is forwarded verbatim (it may be a git url@ref); the loop resolves it. @@ -159,6 +181,7 @@ def build_agent( @pytest.mark.asyncio async def test_run_experimentalist_forwards_platform_insight_id_verbatim( monkeypatch: pytest.MonkeyPatch, + model_clients: ClosingModelClients, tmp_path: Path, ) -> None: paths = _make_run_paths(tmp_path) @@ -176,7 +199,6 @@ async def test_run_experimentalist_forwards_platform_insight_id_verbatim( "build_experimentalist_agent", lambda **_: experimentalist, ) - monkeypatch.setattr(experimentalist_run, "_enable_litellm_drop_params", lambda: None) await experimentalist_run.run_experimentalist( insight="insight-remote-123", @@ -192,20 +214,22 @@ async def test_run_experimentalist_forwards_platform_insight_id_verbatim( assert experimentalist.deps is not None # A str id is not resolved to a Path — it flows through untouched to the backend. assert experimentalist.deps.insight == "insight-remote-123" + assert model_clients.closed @pytest.mark.asyncio async def test_run_experimentalist_forwards_agent_spec_uri_to_deps( monkeypatch: pytest.MonkeyPatch, + model_clients: ClosingModelClients, tmp_path: Path, ) -> None: paths = _make_run_paths(tmp_path) backend = LocalExperimentalistBackend(path=tmp_path / "backend") experimentalist = FakeExperimentalist() + client = ClosingClient() monkeypatch.setattr(experimentalist_run, "make_experimentalist_backend", lambda **_: backend) monkeypatch.setattr(experimentalist_run, "build_experimentalist_agent", lambda **_: experimentalist) - monkeypatch.setattr(experimentalist_run, "_enable_litellm_drop_params", lambda: None) spec_uri = "/path/to/AGENT-SPEC.md" await experimentalist_run.run_experimentalist( @@ -216,17 +240,19 @@ async def test_run_experimentalist_forwards_agent_spec_uri_to_deps( validation_dataset=DatasetRef(uri=str(paths.validation)), experiment_dir=paths.experiment, workspace="default", - client=None, + client=cast(AsyncNeMoPlatform, client), config=EvolutionaryOptimizerConfig(), ) assert experimentalist.deps is not None assert experimentalist.deps.agent_spec == spec_uri + assert model_clients.closed @pytest.mark.asyncio async def test_run_experimentalist_does_not_close_caller_client_when_backend_creation_fails( monkeypatch: pytest.MonkeyPatch, + model_clients: ClosingModelClients, tmp_path: Path, ) -> None: paths = _make_run_paths(tmp_path) @@ -250,3 +276,4 @@ def fail_backend_creation(**_: object) -> object: ) assert not client.closed + assert not model_clients.closed diff --git a/plugins/nemo-experimentalist/tests/test_preflight.py b/plugins/nemo-experimentalist/tests/test_preflight.py index 98009de5be..8835edbdde 100644 --- a/plugins/nemo-experimentalist/tests/test_preflight.py +++ b/plugins/nemo-experimentalist/tests/test_preflight.py @@ -19,11 +19,10 @@ from nemo_insights_plugin.contracts.checks import CheckResult, required_failures -def make_probes(*, cmd_ok: bool = True, http: bool = True, env: dict | None = None) -> Probes: +def make_probes(*, cmd_ok: bool = True, http: bool = True) -> Probes: return Probes( run_cmd=lambda argv: (0, "ok") if cmd_ok else (1, "boom"), http_ok=lambda url: http, - env=env or {}, ) @@ -50,7 +49,7 @@ def test_preflight_results_use_shared_check_result() -> None: def test_healthy_experiment_setup_all_pass(tmp_path: Path) -> None: - probes = make_probes(env={"NEMO_EXPERIMENTALIST_API_BASE": "http://llm", "NEMO_EXPERIMENTALIST_API_KEY": "k"}) + probes = make_probes() profile = full_profile(tmp_path) results = ( check_profile(profile, None) @@ -112,23 +111,19 @@ def test_unreachable_platform_is_advisory(tmp_path: Path) -> None: profile=full_profile(tmp_path), insight=None, base_url="http://x", - probes=make_probes( - http=False, env={"NEMO_EXPERIMENTALIST_API_BASE": "http://llm", "NEMO_EXPERIMENTALIST_API_KEY": "k"} - ), + probes=make_probes(http=False), ) assert any(r.name == "platform-reachable" and r.status == "warn" for r in results) assert all(r.severity == "advisory" for r in results if r.status != "pass") -def test_environment_display_urls_are_sanitized_without_changing_probes(tmp_path: Path) -> None: - model_base = "https://model-user:model-secret@models.example:8443/v1" # trufflehog:ignore - model_probe = f"{model_base}/models" +def test_environment_display_url_is_sanitized_without_changing_probe(tmp_path: Path) -> None: platform_base = "https://platform-user:platform-secret@platform.example:9443/api" # trufflehog:ignore probed_urls: list[str] = [] def http_ok(url: str) -> bool: probed_urls.append(url) - return url == model_probe + return False results = check_environment( profile=full_profile(tmp_path), @@ -137,31 +132,31 @@ def http_ok(url: str) -> bool: probes=Probes( run_cmd=lambda argv: (0, "ok"), http_ok=http_ok, - env={"NEMO_EXPERIMENTALIST_API_BASE": model_base, "NEMO_EXPERIMENTALIST_API_KEY": "k"}, ), ) - model = next(result for result in results if result.name == "model-endpoint") platform = next(result for result in results if result.name == "platform-reachable") - assert probed_urls == [model_probe, f"{platform_base}/health/ready"] - assert model.message == "https://***@models.example:8443/v1/models reachable" + assert probed_urls == [f"{platform_base}/health/ready"] assert platform.message == "https://***@platform.example:9443/api unreachable" - assert not any( - secret in f"{model.message}\n{platform.message}" - for secret in ("model-user", "model-secret", "platform-user", "platform-secret") - ) + assert not any(secret in platform.message for secret in ("platform-user", "platform-secret")) -def test_missing_experiment_credentials_are_required(tmp_path: Path) -> None: +def test_missing_agent_models_are_required(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + preflight, + "configured_model_refs", + lambda: (_ for _ in ()).throw(ValueError("No default model is configured")), + ) results = check_environment( profile=full_profile(tmp_path), insight=None, base_url="http://x", - probes=make_probes(env={"NEMO_EXPERIMENTALIST_API_BASE": " "}), # whitespace-only = unset + probes=make_probes(), ) - assert any(r.name == "NEMO_EXPERIMENTALIST_API_BASE" and r.status == "fail" for r in results) - assert any(r.name == "NEMO_EXPERIMENTALIST_API_KEY" and r.status == "fail" for r in results) - assert not any(r.name == "INFERENCE_API_KEY" for r in results) + models = next(r for r in results if r.name == "agent-models") + assert models.status == "fail" + assert models.severity == "required" + assert "nemo setup" in (models.hint or "") def test_git_source_probe_failure_is_advisory(tmp_path: Path) -> None: @@ -170,7 +165,7 @@ def remote_unreachable(argv: list[str]) -> tuple[int, str]: results = check_artifacts( full_profile(tmp_path, agent_source="https://host/g/repo.git@main"), - probes=Probes(run_cmd=remote_unreachable, http_ok=lambda url: True, env={}), + probes=Probes(run_cmd=remote_unreachable, http_ok=lambda url: True), ) git = next(r for r in results if r.name == "agent-source-git") assert git.severity == "advisory" @@ -200,7 +195,7 @@ def missing_git(argv: list[str]) -> tuple[int, str]: results = check_artifacts( full_profile(tmp_path, agent_source=agent_source), - probes=Probes(run_cmd=missing_git, http_ok=lambda url: True, env={}), + probes=Probes(run_cmd=missing_git, http_ok=lambda url: True), storage=storage, ) @@ -235,7 +230,7 @@ def record(argv: list[str]) -> tuple[int, str]: results = check_artifacts( full_profile(tmp_path, agent_source=source), - probes=Probes(run_cmd=record, http_ok=lambda url: True, env={}), + probes=Probes(run_cmd=record, http_ok=lambda url: True), ) assert calls == [["git", "--version"], ["git", "ls-remote", expected_repo]] assert all(r.status == "pass" for r in results) @@ -258,7 +253,7 @@ def record(argv: list[str]) -> tuple[int, str]: results = check_artifacts( full_profile(tmp_path, agent_source=f"https://host/g/repo.git#{fragment}"), - probes=Probes(run_cmd=record, http_ok=lambda _url: True, env={}), + probes=Probes(run_cmd=record, http_ok=lambda _url: True), ) failure = next(result for result in results if result.name == "agent-source-path") @@ -315,7 +310,7 @@ def run_cmd(argv: list[str]) -> tuple[int, str]: results = check_artifacts( full_profile(tmp_path, agent_source=source), - probes=Probes(run_cmd=run_cmd, http_ok=lambda url: True, env={}), + probes=Probes(run_cmd=run_cmd, http_ok=lambda url: True), ) failure = next(result for result in results if result.name == "agent-source-git") @@ -363,9 +358,7 @@ def test_unreadable_insight_file_is_required_failure(tmp_path: Path) -> None: profile=full_profile(tmp_path), insight=str(insight), base_url="http://x", - probes=make_probes( - env={"NEMO_EXPERIMENTALIST_API_BASE": "http://llm", "NEMO_EXPERIMENTALIST_API_KEY": "k"} - ), + probes=make_probes(), ) finally: insight.chmod(0o644) @@ -388,7 +381,7 @@ def test_non_json_insight_content_is_required_failure(tmp_path: Path) -> None: insight=str(insight), insight_id="0", base_url="http://x", - probes=make_probes(env={"NEMO_EXPERIMENTALIST_API_BASE": "http://llm", "NEMO_EXPERIMENTALIST_API_KEY": "k"}), + probes=make_probes(), ) failure = next(r for r in required_failures(results) if r.name == "insight-file") assert "JSON-serializable" in failure.message @@ -469,7 +462,7 @@ def run_cmd(argv: list[str]) -> tuple[int, str]: results = check_artifacts( full_profile(tmp_path, agent_source=source), - probes=Probes(run_cmd=run_cmd, http_ok=lambda url: True, env={}), + probes=Probes(run_cmd=run_cmd, http_ok=lambda url: True), storage={"publish_winner": True}, ) @@ -489,7 +482,7 @@ def run_cmd(argv: list[str]) -> tuple[int, str]: results = check_artifacts( full_profile(tmp_path, agent_source="https://bitbucket.org/org/repo.git"), - probes=Probes(run_cmd=run_cmd, http_ok=lambda url: True, env={}), + probes=Probes(run_cmd=run_cmd, http_ok=lambda url: True), storage={"publish_winner": True}, ) @@ -517,7 +510,7 @@ def run_cmd(argv: list[str]) -> tuple[int, str]: results = check_artifacts( full_profile(tmp_path), - probes=Probes(run_cmd=run_cmd, http_ok=lambda url: True, env={}), + probes=Probes(run_cmd=run_cmd, http_ok=lambda url: True), storage=storage, ) @@ -538,7 +531,7 @@ def run_cmd(argv: list[str]) -> tuple[int, str]: results = check_artifacts( full_profile(tmp_path, agent_source="https://github.com/org/repo.git"), - probes=Probes(run_cmd=run_cmd, http_ok=lambda url: True, env={}), + probes=Probes(run_cmd=run_cmd, http_ok=lambda url: True), storage={"archive_candidates": True}, ) @@ -563,7 +556,6 @@ def run_cmd(argv: list[str]) -> tuple[int, str]: probes=Probes( run_cmd=run_cmd, http_ok=lambda url: True, - env={}, ), ) @@ -597,7 +589,7 @@ def test_local_multi_insight_requires_selector(tmp_path: Path) -> None: profile=full_profile(tmp_path), insight=str(insights), base_url="http://x", - probes=make_probes(env={"NEMO_EXPERIMENTALIST_API_BASE": "http://llm", "NEMO_EXPERIMENTALIST_API_KEY": "k"}), + probes=make_probes(), ) failure = next(r for r in results if r.name == "insight-file") assert failure.status == "fail" and failure.severity == "required" @@ -613,7 +605,7 @@ def test_single_insight_file_with_matching_selector_passes(tmp_path: Path) -> No insight=str(insights), insight_id="i-a", base_url="http://x", - probes=make_probes(env={"NEMO_EXPERIMENTALIST_API_BASE": "http://llm", "NEMO_EXPERIMENTALIST_API_KEY": "k"}), + probes=make_probes(), ) parse = next(r for r in results if r.name == "insight-file") assert parse.status == "pass" @@ -631,7 +623,7 @@ def test_selected_insight_checks_only_selected_agent_without_selector_warning(tm insight=str(insights), insight_id="i-a", base_url="http://x", - probes=make_probes(env={"NEMO_EXPERIMENTALIST_API_BASE": "http://llm", "NEMO_EXPERIMENTALIST_API_KEY": "k"}), + probes=make_probes(), ) agent = next(r for r in results if r.name == "insight-agent") assert agent.status == "pass" @@ -650,7 +642,7 @@ def test_ambiguous_selected_insight_is_required_failure(tmp_path: Path) -> None: insight=str(insights), insight_id="same", base_url="http://x", - probes=make_probes(env={"NEMO_EXPERIMENTALIST_API_BASE": "http://llm", "NEMO_EXPERIMENTALIST_API_KEY": "k"}), + probes=make_probes(), ) failure = next(r for r in required_failures(results) if r.name == "insight-file") assert "ambiguous" in failure.message @@ -694,16 +686,17 @@ def test_check_artifacts_has_no_dataset_results(tmp_path: Path) -> None: assert not any(r.name.startswith("dataset-") for r in results) -def test_missing_env_hint_names_source_and_env_file(tmp_path: Path) -> None: - profile = full_profile(tmp_path) +def test_missing_model_hint_names_setup(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + preflight, + "configured_model_refs", + lambda: (_ for _ in ()).throw(ValueError("No default model is configured")), + ) results = check_environment( - profile=profile, + profile=full_profile(tmp_path), insight=None, base_url="http://x", - probes=make_probes(env={}), + probes=make_probes(), ) - base = next(r for r in results if r.name == "NEMO_EXPERIMENTALIST_API_BASE") - assert "inference-api.nvidia.com" in (base.hint or "") # names the real endpoint - assert str(tmp_path / ".env") in (base.hint or "") # says exactly where to save it - assert "example-agent.mdx" in (base.hint or "") # points at a doc that exists - assert "export NEMO_EXPERIMENTALIST_API_BASE" in (base.hint or "") + models = next(r for r in results if r.name == "agent-models") + assert models.hint == "run `nemo setup` and select default and fast agent models" diff --git a/plugins/nemo-experimentalist/tests/test_resolve.py b/plugins/nemo-experimentalist/tests/test_resolve.py index fbacdc6e7e..f216d90cd3 100644 --- a/plugins/nemo-experimentalist/tests/test_resolve.py +++ b/plugins/nemo-experimentalist/tests/test_resolve.py @@ -703,6 +703,7 @@ def test_full_resolution_from_profile(tmp_path: Path) -> None: assert inputs.agent_spec == str((tmp_path / "AGENT-SPEC.md").resolve()) # auto-pick assert inputs.train_dataset.uri == str((tmp_path / "evals" / "train").resolve()) assert inputs.train_dataset.metadata == {"id": "train"} + assert inputs.task_template is not None assert inputs.task_template.uri == str((tmp_path / "evals" / "task_template").resolve()) assert inputs.workspace == "default" assert inputs.config.storage.publish_winner is True @@ -913,7 +914,7 @@ def test_run_config_rejects_a_models_block() -> None: import pytest from nemo_experimentalist_plugin.config import EvolutionaryOptimizerConfig - with pytest.raises(ValueError, match="NEMO_EXPERIMENTALIST_MODELS"): + with pytest.raises(ValueError, match="nemo setup"): EvolutionaryOptimizerConfig.model_validate({"models": {"smart": "a/b"}}) diff --git a/plugins/nemo-insights/README.md b/plugins/nemo-insights/README.md index 799a1b3d42..2ee4dcee60 100644 --- a/plugins/nemo-insights/README.md +++ b/plugins/nemo-insights/README.md @@ -22,6 +22,11 @@ uv run nemo agents analyst doctor uv run nemo agents analyst run ``` +Run `nemo setup` first to select the default and fast Platform Model Entities. +The Analyst uses the default model for analysis and the fast model for context +summarization; an existing context without `fast_model` reuses `default_model`. +Provider credentials remain in Platform Secrets. + The profile contract consumed by Insights is deliberately small: ```yaml @@ -79,6 +84,11 @@ uv run nemo insights analysis status uv run nemo insights analysis disable --agent research-agent ``` +`analysis enable` stores the effective default/fast pair in the server-side +analysis config so scheduled jobs do not depend on the operator's local CLI +file. Re-run `enable` after changing the pair with `nemo setup`. Existing +enabled records created before model-pair persistence must also be re-enabled. + `--base-url` defaults to `NMP_BASE_URL`, then `http://localhost:8080`. ## API and SDK diff --git a/plugins/nemo-insights/examples/research-agent/README.md b/plugins/nemo-insights/examples/research-agent/README.md index 24a8bc43be..a3628f4be8 100644 --- a/plugins/nemo-insights/examples/research-agent/README.md +++ b/plugins/nemo-insights/examples/research-agent/README.md @@ -67,11 +67,12 @@ Desktop on macOS/Windows or the Docker service on Linux, then restart `nemo serv ### Configure a provider -In a second shell, register the NVIDIA Build provider and pick a default model: +In a second shell, register the NVIDIA Build provider and pick default and fast +models (press Enter at the fast prompt to reuse the default): ```bash set -a && source .env && set +a -uv run nemo setup --auto +uv run nemo setup ``` ### Register the workflow as an agent diff --git a/plugins/nemo-insights/examples/research-agent/tests/test_analyst_e2e.py b/plugins/nemo-insights/examples/research-agent/tests/test_analyst_e2e.py index a57e4d11d0..12cc287532 100644 --- a/plugins/nemo-insights/examples/research-agent/tests/test_analyst_e2e.py +++ b/plugins/nemo-insights/examples/research-agent/tests/test_analyst_e2e.py @@ -19,10 +19,11 @@ - ``NMP_INSIGHTS_E2E=1`` — opt in; otherwise the test skips. - ``NVIDIA_API_KEY`` and ``TAVILY_API_KEY`` — for the research agent's NIM model + Tavily search. Read from the example's ``.env`` (or the shell). -- ``INFERENCE_API_KEY`` — the ``sk-...`` NVIDIA Inference Gateway virtual key - for the analyst's Claude Opus. Nooa routes it through LiteLLM using the - gateway's OpenAI-compatible API. The model and base URL are pinned in - :mod:`nemo_insights_plugin.analyst.model_config`, so no override is required. + The test also runs ``nemo setup --auto`` against its isolated Platform, so + the NVIDIA provider supplies the persisted default and fast model entities + used by the analyst. ``NEMO_DEFAULT_INFERENCE_KEY`` and its base URL can be + supplied in addition when the optimizer models use another provider; they do + not replace the NVIDIA credential required by the research agent under test. - Docker — required to auto-start ClickHouse if one isn't already at ``NMP_INTAKE_CLICKHOUSE_URL`` (default ``http://localhost:8123``). A missing Docker daemon fails the test. @@ -120,6 +121,8 @@ def _subprocess_env() -> dict[str, str]: env.update(_load_dotenv(EXAMPLE_DIR / ".env")) env.update( { + "NMP_BASE_URL": BASE_URL, + "NMP_CONFIG_FILE": str(TMP_DIR / "e2e-nmp-config.yaml"), "NMP_INTAKE_CLICKHOUSE_URL": CLICKHOUSE_URL, "NMP_INTAKE_CLICKHOUSE_USER": CLICKHOUSE_USER, "NMP_INTAKE_CLICKHOUSE_PASSWORD": CLICKHOUSE_PASSWORD, @@ -147,11 +150,17 @@ def _cli_cmd(name: str, *args: str) -> list[str]: def _require_keys() -> None: env = _subprocess_env() - missing = [key for key in ("NVIDIA_API_KEY", "TAVILY_API_KEY", "INFERENCE_API_KEY") if not env.get(key)] + missing = [key for key in ("NVIDIA_API_KEY", "TAVILY_API_KEY") if not env.get(key)] if missing: pytest.skip(f"missing required API keys: {', '.join(missing)}") +@pytest.fixture(scope="module", autouse=True) +def required_keys() -> None: + """Skip before provisioning services when the opt-in credentials are absent.""" + _require_keys() + + # --------------------------------------------------------------------------- # # ClickHouse # # --------------------------------------------------------------------------- # @@ -281,6 +290,28 @@ def platform_server(clickhouse: None) -> Iterator[str]: # noqa: ARG001 - orderi time.sleep(2.0) else: pytest.fail(f"platform server not ready after {SERVER_START_TIMEOUT_S}s; see {log_path}") + + setup_result = subprocess.run( + _cli_cmd( + "nemo", + "--base-url", + BASE_URL, + "setup", + "--auto", + "--workspace", + WORKSPACE, + "--no-start-services", + "--no-install-skills", + "--no-deploy-agent", + ), + cwd=str(EXAMPLE_DIR), + env=env, + capture_output=True, + text=True, + timeout=300, + ) + if setup_result.returncode != 0: + pytest.fail(f"isolated `nemo setup --auto` failed:\n{setup_result.stdout}\n{setup_result.stderr}") yield BASE_URL finally: try: @@ -416,8 +447,6 @@ def _wait_for_analyst_spans(timeout_s: int) -> list[dict]: # The test # # --------------------------------------------------------------------------- # def test_analyst_creates_insight_end_to_end(platform_server: str) -> None: # noqa: ARG001 - _require_keys() - # 1. Clean slate: drop any spans + Insights left by a previous run. _clear_spans_for_project() _delete_insights(_list_insight_ids()) diff --git a/plugins/nemo-insights/src/nemo_insights_plugin/analyst/agent.py b/plugins/nemo-insights/src/nemo_insights_plugin/analyst/agent.py index 882d97ee1e..cc0d0e797f 100644 --- a/plugins/nemo-insights/src/nemo_insights_plugin/analyst/agent.py +++ b/plugins/nemo-insights/src/nemo_insights_plugin/analyst/agent.py @@ -30,8 +30,8 @@ from nemo_insights_plugin.analyst.deps import AnalystDeps from nemo_insights_plugin.analyst.functions import annotations, insights, spans -from nemo_insights_plugin.analyst.model_config import get_fast_model, get_smart_model from nemo_insights_plugin.analyst.result import AnalystResult +from nemo_platform_plugin.nooa_model_client import get_default_model, get_fast_model from nooa import Agent, CodeActStrategy, hidden, strategy from nooa.agents import TokenBudgetSummarizer from nooa.config import CodeActConfig @@ -173,7 +173,7 @@ def __init__( agent_spec: str | None = None, **kwargs: Any, ) -> None: - super().__init__(llm=kwargs.pop("llm", None) or get_smart_model(), **kwargs) + super().__init__(llm=kwargs.pop("llm", None) or get_default_model(), **kwargs) self._deps = deps self.todos = TodoManager() TokenBudgetSummarizer.install( diff --git a/plugins/nemo-insights/src/nemo_insights_plugin/analyst/model_config.py b/plugins/nemo-insights/src/nemo_insights_plugin/analyst/model_config.py deleted file mode 100644 index 2edeec48c1..0000000000 --- a/plugins/nemo-insights/src/nemo_insights_plugin/analyst/model_config.py +++ /dev/null @@ -1,33 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""LLM clients for the Insights Analyst.""" - -import functools -import os - -from nooa.unifiedllm import CompletionClient - -_API_BASE = "https://inference-api.nvidia.com/v1" -# The leading ``openai/`` selects LiteLLM's OpenAI-compatible transport. The -# remaining value is the model alias sent unchanged to NVIDIA's gateway. -_SMART_MODEL = "openai/aws/anthropic/bedrock-claude-opus-4-8" -_FAST_MODEL = "openai/openai/openai/gpt-5-mini" - - -@functools.cache -def _completion_client(name: str, api_base: str, api_key: str) -> CompletionClient: - """Reuse clients by their complete inference identity.""" - return CompletionClient(name, api_base=api_base, api_key=api_key) - - -@functools.cache -def get_smart_model() -> CompletionClient: - """Return the cached high-capability model used for analysis.""" - return _completion_client(_SMART_MODEL, _API_BASE, os.environ["INFERENCE_API_KEY"]) - - -@functools.cache -def get_fast_model() -> CompletionClient: - """Return the cached low-latency model used for context summarization.""" - return _completion_client(_FAST_MODEL, _API_BASE, os.environ["INFERENCE_API_KEY"]) diff --git a/plugins/nemo-insights/src/nemo_insights_plugin/analyst/run.py b/plugins/nemo-insights/src/nemo_insights_plugin/analyst/run.py index 4a1e5c6026..54958ff9d1 100644 --- a/plugins/nemo-insights/src/nemo_insights_plugin/analyst/run.py +++ b/plugins/nemo-insights/src/nemo_insights_plugin/analyst/run.py @@ -3,12 +3,10 @@ """Reusable insights analyst run orchestration.""" -import importlib import os import sys from datetime import datetime from pathlib import Path -from typing import Protocol, cast from nemo_insights_plugin.analyst.agent import ( KICKOFF, @@ -23,6 +21,12 @@ ) from nemo_insights_plugin.analyst.result import AnalystResult from nemo_platform import AsyncNeMoPlatform +from nemo_platform_plugin.nooa_model_client import ( + ConfiguredModelClients, + ConfiguredModelRefs, + activate_model_clients, + resolve_model_clients, +) from nooa.context_blocks import EventBase from nooa.events import LLMComplete, PythonOutput @@ -31,10 +35,6 @@ _VERBOSE_TRUNCATE = 2000 -class _LiteLLMModule(Protocol): - drop_params: bool - - class ClientConstructionError(Exception): """The analyst's NeMo Platform client could not be constructed.""" @@ -51,6 +51,7 @@ async def run_analyst( verbose: bool = False, since: datetime | None = None, evaluation_id: str | None = None, + model_refs: ConfiguredModelRefs | None = None, ) -> str: """Build and run the analyst agent against an agent's telemetry. @@ -71,11 +72,14 @@ async def run_analyst( verbose: Whether to stream model/tool events to stderr. since: Optional incremental lower bound enforced on trace/span reads. evaluation_id: Optional run scope; AND-pinned onto every span read. + model_refs: Optional explicit default/fast Model Entity IDs. Unset uses + the active Platform CLI context. """ observability = None + model_clients: ConfiguredModelClients | None = None insights_output_path = str(insights_output) if insights_output else None try: - _enable_litellm_drop_params() + model_clients = await resolve_model_clients(client, model_refs) backend = make_analyst_backend( client=client, insights_output=insights_output_path, @@ -96,19 +100,24 @@ async def run_analyst( workspace=workspace, target_agent=agent, ) - analyst = build_analyst_agent( - deps=deps, - agent=agent, - agent_spec=agent_spec, - ) - result = await _run_agent(analyst, verbose=verbose) + with activate_model_clients(model_clients): + analyst = build_analyst_agent( + deps=deps, + agent=agent, + agent_spec=agent_spec, + ) + result = await _run_agent(analyst, verbose=verbose) return await backend.persist_result(workspace=workspace, agent=agent, result=result) finally: try: if observability is not None: observability.shutdown() finally: - await client.close() + try: + if model_clients is not None: + await model_clients.aclose() + finally: + await client.close() def _analyst_observability_enabled() -> bool: @@ -117,12 +126,6 @@ def _analyst_observability_enabled() -> bool: return value.strip().lower() in {"1", "true", "yes", "on"} -def _enable_litellm_drop_params() -> None: - """Let LiteLLM omit parameters unsupported by the configured model.""" - litellm = cast(_LiteLLMModule, importlib.import_module("litellm")) - litellm.drop_params = True - - async def _run_agent( analyst: Analyst, *, diff --git a/plugins/nemo-insights/src/nemo_insights_plugin/cli.py b/plugins/nemo-insights/src/nemo_insights_plugin/cli.py index 6fde946391..16af1bcdbf 100644 --- a/plugins/nemo-insights/src/nemo_insights_plugin/cli.py +++ b/plugins/nemo-insights/src/nemo_insights_plugin/cli.py @@ -33,14 +33,15 @@ ) from nemo_insights_plugin.preflight import ( AnalysisProbes, - check_credentials, check_environment, + check_models, check_profile, read_agent_spec, ) from nemo_insights_plugin.profile import AnalysisProfile, load_profile, pick_agent_spec from nemo_platform import NeMoPlatformError from nemo_platform_plugin.cli import NemoCLI +from nemo_platform_plugin.nooa_model_client import configured_model_refs from nooa import GenerationError DEFAULT_WORKSPACE = "default" @@ -168,7 +169,7 @@ def _prepare_mirror(insights_output: Path | None) -> Path | None: async def _run_analysis(analysis: _ResolvedAnalysis, *, verbose: bool) -> str: checks = list(analysis.spec_checks) - checks.extend(check_credentials(analysis.profile_dir, probes=_PREFLIGHT_PROBES)) + checks.extend(check_models()) _preflight_or_exit(checks) insights_output = _prepare_mirror(analysis.insights_output) @@ -458,12 +459,23 @@ async def _analysis_config_command( base_url: str, ) -> str: """Run one analysis-config CLI action and return JSON for stdout.""" + model_refs = None + if action == "enable": + if agent is None: + raise ValueError("agent is required for enable") + model_refs = configured_model_refs() + client = make_client(base_url) try: if action == "enable": - if agent is None: - raise ValueError("agent is required for enable") - result = await client.insights.analysis_configs.enable(workspace=workspace, agent=agent) + assert agent is not None + assert model_refs is not None + result = await client.insights.analysis_configs.enable( + workspace=workspace, + agent=agent, + default_model=model_refs.default, + fast_model=model_refs.fast, + ) return _json(result.model_dump(mode="json")) if action == "disable": if agent is None: diff --git a/plugins/nemo-insights/src/nemo_insights_plugin/config.py b/plugins/nemo-insights/src/nemo_insights_plugin/config.py index 260c24551a..8a1e087932 100644 --- a/plugins/nemo-insights/src/nemo_insights_plugin/config.py +++ b/plugins/nemo-insights/src/nemo_insights_plugin/config.py @@ -88,14 +88,6 @@ def _validate_timezone(cls, value: str) -> str: "Optional platform base URL passed to analyst jobs. When unset, jobs use their active platform context." ), ) - inference_api_key_secret_name: str | None = Field( - default=None, - description=( - "Optional platform secret name whose value is exposed to analyst " - "jobs as INFERENCE_API_KEY. Temporary until FP-202 moves analyst " - "model execution to platform-registered models." - ), - ) class InsightsConfig(NemoConfig): diff --git a/plugins/nemo-insights/src/nemo_insights_plugin/controller.py b/plugins/nemo-insights/src/nemo_insights_plugin/controller.py index 1541033457..e784041839 100644 --- a/plugins/nemo-insights/src/nemo_insights_plugin/controller.py +++ b/plugins/nemo-insights/src/nemo_insights_plugin/controller.py @@ -125,6 +125,16 @@ async def reconcile_one(self, obj: object) -> None: async def _reconcile_config(self, config: AnalysisConfig) -> None: if not config.enabled: return + if not config.default_model: + logger.error( + "Analysis config for agent '%s' in workspace '%s' has no model selection; " + "run `nemo insights analysis enable --agent %s --workspace %s` again", + config.agent, + config.workspace, + config.agent, + config.workspace, + ) + return if await self._has_active_job(config): return status = await self._get_run_status(config) @@ -232,7 +242,8 @@ async def _submit_analysis_job( agent=config.agent, base_url=self.insights_config.analyst.base_url, since=status.last_successful_run_at if status is not None else None, - inference_api_key_secret_name=(self.insights_config.analyst.inference_api_key_secret_name), + default_model=config.default_model, + fast_model=config.fast_model or config.default_model, ) job_name = _job_name(config, submitted_at) platform_spec = await self._compile_job_spec( diff --git a/plugins/nemo-insights/src/nemo_insights_plugin/entities.py b/plugins/nemo-insights/src/nemo_insights_plugin/entities.py index aa4765f6a6..5865bfa851 100644 --- a/plugins/nemo-insights/src/nemo_insights_plugin/entities.py +++ b/plugins/nemo-insights/src/nemo_insights_plugin/entities.py @@ -72,9 +72,11 @@ class Insight(NemoEntity, entity_type="insights_insight"): class AnalysisConfig(NemoEntity, entity_type="insights_analysis_config"): """Per-agent opt-in state for framework-managed periodic analysis. - The cadence is intentionally global insights configuration. This entity is - only the per-agent switch. Machine-written run state lives on - :class:`AnalysisRunStatus` so the controller never races the running job. + The cadence is intentionally global insights configuration. The model pair + is captured when analysis is enabled because the controller runs in the + Platform process and cannot read the operator's local CLI configuration. + Machine-written run state lives on :class:`AnalysisRunStatus` so the + controller never races the running job. """ agent: str = Field(description="Name of the agent this analysis config targets.") @@ -82,6 +84,14 @@ class AnalysisConfig(NemoEntity, entity_type="insights_analysis_config"): default=True, description="Whether the periodic insights controller should analyze this agent.", ) + default_model: str = Field( + default="", + description="Workspace-qualified Model Entity used for quality-critical analysis work.", + ) + fast_model: str = Field( + default="", + description="Workspace-qualified Model Entity used for latency-sensitive analysis work.", + ) class AnalysisRunStatus(NemoEntity, entity_type="insights_analysis_run_status"): diff --git a/plugins/nemo-insights/src/nemo_insights_plugin/jobs/analyze.py b/plugins/nemo-insights/src/nemo_insights_plugin/jobs/analyze.py index 3bef2ddadc..294cb65a59 100644 --- a/plugins/nemo-insights/src/nemo_insights_plugin/jobs/analyze.py +++ b/plugins/nemo-insights/src/nemo_insights_plugin/jobs/analyze.py @@ -7,12 +7,10 @@ import asyncio import logging -import os from datetime import datetime, timezone from typing import ClassVar from nemo_insights_plugin.analyst.run import run_analyst -from nemo_insights_plugin.client import make_client from nemo_insights_plugin.entities import AnalysisConfigStatus from nemo_platform import NeMoPlatform from nemo_platform_plugin.job import NemoJob @@ -21,7 +19,6 @@ ContainerSpec, CPUExecutionProviderSpec, EnvironmentVariable, - EnvironmentVariableFromSecret, PlatformJobSpec, PlatformJobStep, ) @@ -30,6 +27,8 @@ PERSISTENT_JOB_STORAGE_PATH_ENVVAR, ) from nemo_platform_plugin.jobs.image import get_qualified_image +from nemo_platform_plugin.nooa_model_client import ConfiguredModelRefs +from nemo_platform_plugin.sdk_provider import get_async_task_sdk from pydantic import BaseModel, ConfigDict, Field logger = logging.getLogger(__name__) @@ -67,10 +66,8 @@ class AnalyzeSpec(BaseModel): default=True, description="Update the matching AnalysisRunStatus with run metadata.", ) - inference_api_key_secret_name: str | None = Field( - default=None, - description=("Optional platform secret exposed as INFERENCE_API_KEY for the current analyst model path."), - ) + default_model: str = Field(description="Workspace-qualified default Model Entity ID selected during setup.") + fast_model: str = Field(description="Workspace-qualified fast Model Entity ID selected during setup.") class AnalyzeJob(NemoJob): @@ -102,18 +99,6 @@ async def compile( value=DEFAULT_JOB_STORAGE_PATH, ) ] - if canonical.inference_api_key_secret_name: - environment.append( - EnvironmentVariable( - name="INFERENCE_API_KEY", - from_secret=EnvironmentVariableFromSecret(name=canonical.inference_api_key_secret_name), - ) - ) - elif inference_api_key := os.environ.get("INFERENCE_API_KEY"): - # Local smoke-test fallback until FP-202 moves analyst model - # execution onto platform-registered models/secrets. - environment.append(EnvironmentVariable(name="INFERENCE_API_KEY", value=inference_api_key)) - return PlatformJobSpec( steps=[ PlatformJobStep( @@ -153,15 +138,17 @@ def run( ) try: + async_client = get_async_task_sdk("insights") report = asyncio.run( run_analyst( agent=spec.agent, agent_spec=spec.agent_spec, workspace=ctx.workspace, base_url=spec.base_url, - client=make_client(spec.base_url), + client=async_client, insights_output=spec.insights_output, since=spec.since, + model_refs=ConfiguredModelRefs(default=spec.default_model, fast=spec.fast_model), ) ) except Exception as exc: # pragma: no cover - exercised by integration paths diff --git a/plugins/nemo-insights/src/nemo_insights_plugin/preflight.py b/plugins/nemo-insights/src/nemo_insights_plugin/preflight.py index edbcb06234..e73d9ab673 100644 --- a/plugins/nemo-insights/src/nemo_insights_plugin/preflight.py +++ b/plugins/nemo-insights/src/nemo_insights_plugin/preflight.py @@ -3,9 +3,8 @@ """Read-only readiness checks for Insights analysis.""" -import os -from collections.abc import Awaitable, Callable, Mapping -from dataclasses import dataclass, field +from collections.abc import Awaitable, Callable +from dataclasses import dataclass from pathlib import Path import httpx @@ -14,6 +13,7 @@ from nemo_insights_plugin.contracts.checks import CheckResult, make_check_result from nemo_insights_plugin.profile import AnalysisProfile from nemo_platform import AsyncNeMoPlatform, NeMoPlatformError +from nemo_platform_plugin.nooa_model_client import configured_model_refs _EXPECTED_PLATFORM_ERRORS = (NeMoPlatformError, httpx.HTTPError, OSError, RuntimeError, ValueError) @@ -53,7 +53,6 @@ async def _default_workspace_ok(base_url: str, workspace: str, agent: str) -> bo class AnalysisProbes: """Dependencies used by read-only environment checks.""" - env: Mapping[str, str] = field(default_factory=lambda: os.environ) http_ok: Callable[[str], bool] = _default_http_ok workspace_ok: Callable[[str, str, str], Awaitable[bool]] = _default_workspace_ok @@ -153,28 +152,28 @@ def read_agent_spec( ] -def check_credentials( - profile_dir: Path | None, - probes: AnalysisProbes | None = None, -) -> list[CheckResult]: - """Check that the analyst's required inference credential is present.""" - active = probes or AnalysisProbes() - env_path = profile_dir / ".env" if profile_dir is not None else None - credential_hint = ( - f"save it in {env_path} or export INFERENCE_API_KEY=" - if env_path is not None - else "export INFERENCE_API_KEY=" - ) - credential = bool(active.env.get("INFERENCE_API_KEY", "").strip()) +def check_models() -> list[CheckResult]: + """Check that the active Platform context has a complete model pair.""" + try: + refs = configured_model_refs() + except (FileNotFoundError, RuntimeError, ValueError) as exc: + return [ + CheckResult( + name="agent-models", + group="models", + status="fail", + severity="required", + message=str(exc), + hint="run `nemo setup` and select default and fast agent models", + ) + ] return [ - make_check_result( - "INFERENCE_API_KEY", - "credentials", - credential, - "required", - "INFERENCE_API_KEY set", - "INFERENCE_API_KEY not set", - hint=credential_hint, + CheckResult( + name="agent-models", + group="models", + status="pass", + severity="required", + message=f"default={refs.default}; fast={refs.fast}", ) ] @@ -187,13 +186,14 @@ async def check_environment( profile_dir: Path | None, probes: AnalysisProbes | None = None, ) -> list[CheckResult]: - """Run credential and advisory platform checks without persisting state. + """Run model-selection and advisory platform checks without persisting state. The workspace probe is profile-dependent and skipped when *agent* or - *workspace* is unknown; the credential and reachability checks always run. + *workspace* is unknown; model selection and reachability checks always run. """ active = probes or AnalysisProbes() - results = check_credentials(profile_dir, active) + del profile_dir + results = check_models() reachable = active.http_ok(base_url) results.append( make_check_result( diff --git a/plugins/nemo-insights/src/nemo_insights_plugin/schema.py b/plugins/nemo-insights/src/nemo_insights_plugin/schema.py index 23aadd7104..94ffd493a7 100644 --- a/plugins/nemo-insights/src/nemo_insights_plugin/schema.py +++ b/plugins/nemo-insights/src/nemo_insights_plugin/schema.py @@ -66,6 +66,13 @@ class InsightListItem(Insight, entity_type="insights_insight"): InsightPage = NemoListResponse[InsightListItem] +class EnableAnalysisConfigRequest(BaseModel): + """Model selection captured when periodic analysis is enabled.""" + + default_model: str = Field(min_length=1) + fast_model: str = Field(min_length=1) + + class UpdateAnalysisConfigRequest(BaseModel): """Body for ``PATCH /analysis-configs/{agent}``.""" diff --git a/plugins/nemo-insights/src/nemo_insights_plugin/sdk_resources/analysis_configs.py b/plugins/nemo-insights/src/nemo_insights_plugin/sdk_resources/analysis_configs.py index 9400d49a9e..6708e29367 100644 --- a/plugins/nemo-insights/src/nemo_insights_plugin/sdk_resources/analysis_configs.py +++ b/plugins/nemo-insights/src/nemo_insights_plugin/sdk_resources/analysis_configs.py @@ -14,6 +14,7 @@ from nemo_insights_plugin.schema import ( AnalysisConfigPage, AnalysisRunStatusPage, + EnableAnalysisConfigRequest, UpdateAnalysisConfigRequest, UpdateAnalysisRunStatusRequest, ) @@ -49,6 +50,11 @@ def _build_update_body( return body.model_dump(mode="json", exclude_none=True, exclude_unset=True) +def _build_enable_body(*, default_model: str, fast_model: str) -> dict[str, Any]: + body = EnableAnalysisConfigRequest(default_model=default_model, fast_model=fast_model) + return body.model_dump(mode="json") + + def _build_status_update_body( *, status: AnalysisConfigStatus | str | None, @@ -87,9 +93,10 @@ class _AnalysisConfigResource: def __init__(self, parent: _ResourceParent) -> None: self._parent = parent - def enable(self, *, workspace: str, agent: str) -> AnalysisConfig: + def enable(self, *, workspace: str, agent: str, default_model: str, fast_model: str) -> AnalysisConfig: response = self._parent._http_client.post( - self._parent._url(f"/v2/workspaces/{workspace}/analysis-configs/{agent}/enable") + self._parent._url(f"/v2/workspaces/{workspace}/analysis-configs/{agent}/enable"), + json=_build_enable_body(default_model=default_model, fast_model=fast_model), ) response.raise_for_status() return entity_from_response(AnalysisConfig, response.json()) @@ -145,9 +152,10 @@ class _AsyncAnalysisConfigResource: def __init__(self, parent: _ResourceParent) -> None: self._parent = parent - async def enable(self, *, workspace: str, agent: str) -> AnalysisConfig: + async def enable(self, *, workspace: str, agent: str, default_model: str, fast_model: str) -> AnalysisConfig: response = await self._parent._http_client.post( - self._parent._url(f"/v2/workspaces/{workspace}/analysis-configs/{agent}/enable") + self._parent._url(f"/v2/workspaces/{workspace}/analysis-configs/{agent}/enable"), + json=_build_enable_body(default_model=default_model, fast_model=fast_model), ) response.raise_for_status() return entity_from_response(AnalysisConfig, response.json()) diff --git a/plugins/nemo-insights/src/nemo_insights_plugin/service.py b/plugins/nemo-insights/src/nemo_insights_plugin/service.py index 6c70c77791..ab168e9981 100644 --- a/plugins/nemo-insights/src/nemo_insights_plugin/service.py +++ b/plugins/nemo-insights/src/nemo_insights_plugin/service.py @@ -24,6 +24,7 @@ AnalysisConfigPage, AnalysisRunStatusPage, CreateInsightRequest, + EnableAnalysisConfigRequest, InsightListItem, InsightPage, UpdateAnalysisConfigRequest, @@ -328,6 +329,8 @@ async def _get_or_create_analysis_config( workspace: str, agent: str, enabled_if_created: bool, + default_model: str = "", + fast_model: str = "", ) -> tuple[AnalysisConfig, bool]: """Fetch an analysis config, creating a default one if it is absent. @@ -343,7 +346,14 @@ async def _get_or_create_analysis_config( logger.exception("Failed to fetch analysis config") raise HTTPException(status_code=500, detail="Failed to fetch analysis config.") from exc - config = AnalysisConfig(name=agent, workspace=workspace, agent=agent, enabled=enabled_if_created) + config = AnalysisConfig( + name=agent, + workspace=workspace, + agent=agent, + enabled=enabled_if_created, + default_model=default_model, + fast_model=fast_model, + ) try: return await entity_client.create(config), True except NemoEntityConflictError: @@ -386,16 +396,24 @@ def _build_analysis_configs_router() -> APIRouter: async def enable_analysis_config( workspace: str, agent: str, + body: EnableAnalysisConfigRequest, entity_client: NemoEntitiesClient = Depends(get_entity_client), ) -> AnalysisConfig: """Enable periodic insights analysis for one agent.""" config, created = await _get_or_create_analysis_config( - entity_client, workspace=workspace, agent=agent, enabled_if_created=True + entity_client, + workspace=workspace, + agent=agent, + enabled_if_created=True, + default_model=body.default_model, + fast_model=body.fast_model, ) if created: return config config.enabled = True + config.default_model = body.default_model + config.fast_model = body.fast_model try: return await entity_client.update(config) except Exception as exc: diff --git a/plugins/nemo-insights/testbed/README.md b/plugins/nemo-insights/testbed/README.md index 07758f9aaa..eb18a6beaf 100644 --- a/plugins/nemo-insights/testbed/README.md +++ b/plugins/nemo-insights/testbed/README.md @@ -206,8 +206,10 @@ written to `testbed/tmp/insights_.yaml`. ## Config split: secrets in `.env`, everything else in `testbeds.toml` On startup the CLI auto-loads `testbed/.env` (gitignored) as `KEY=VALUE` lines. Keep -**only secrets/endpoints** there — `INFERENCE_API_KEY` (analyst) and -`OPENAI_API_KEY`/`OPENAI_API_BASE` (the proxy litellm uses for the benchmark sim LLMs). +**only secrets/endpoints** there — `OPENAI_API_KEY`/`OPENAI_API_BASE` for the +benchmark simulator LLMs. Run `nemo setup` against the analysis Platform to +select the Analyst's default and fast Model Entities; its provider credential +stays in Platform Secrets. GLAMR live analysis additionally reads `GLAMR_INTAKE_USER` and `GLAMR_INTAKE_PASSWORD` from `.env`; `testbeds.toml` stores only those environment-variable names, never their credential values. @@ -383,7 +385,7 @@ locally: `analyze --state state-vN`). Secrets: the workflow uses `TESTBED_STATE_GH_READ_TOKEN`, a least-privilege GitHub App/PAT credential with release-read access to `TESTBED_STATE_REPO`, -plus `NVIDIA_INFERENCE_KEY` and `NVIDIA_INFERENCE_URL`. The latter are exposed -under the `INFERENCE_API_KEY` and OpenAI-compatible environment names expected -by the analyst and litellm/tau2. The analyst and tau2 sim LLMs need no VPN and -work on public runners. +plus `NVIDIA_INFERENCE_KEY` and `NVIDIA_INFERENCE_URL`. The workflow registers +the latter with `nemo setup --auto` for the Analyst and also exposes it under +the OpenAI-compatible environment names expected by litellm/tau2. The Analyst +and tau2 simulator LLMs need no VPN and work on public runners. diff --git a/plugins/nemo-insights/testbed/cli.py b/plugins/nemo-insights/testbed/cli.py index fc3c6a9699..27c1b9b12f 100644 --- a/plugins/nemo-insights/testbed/cli.py +++ b/plugins/nemo-insights/testbed/cli.py @@ -104,8 +104,6 @@ def _doctor(subjects: dict[str, Subject], name: str | None) -> None: print(f"✗ {subject_name}: unknown subject") continue unmet: list[str] = [] - if not os.environ.get("INFERENCE_API_KEY"): - unmet.append("env INFERENCE_API_KEY (analyst key, in testbed/.env)") if shutil.which("gh") is None: unmet.append("gh CLI (needed for pinned/--state analyze; https://cli.github.com)") unmet += build_adapter(subject).check() @@ -1019,8 +1017,6 @@ def main() -> None: drop_auth=not args.live, ) subject = _apply_overrides(subject, args.sets) - if not os.environ.get("INFERENCE_API_KEY"): - sys.exit("Set INFERENCE_API_KEY (NVIDIA Inference Gateway sk-... key) and re-run.") # One backup destination per invocation: the restore's clobber-backup and the # fresh-insights move both park files here, so a single analyze yields at most # one tmp/backup-/ dir (created lazily, only if something moves). diff --git a/plugins/nemo-insights/testbed/eval/stack.py b/plugins/nemo-insights/testbed/eval/stack.py index def8a45770..b4f1c780e7 100644 --- a/plugins/nemo-insights/testbed/eval/stack.py +++ b/plugins/nemo-insights/testbed/eval/stack.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Self-contained Intake stack (ClickHouse + auth, entities, intake on :8080). +"""Self-contained Analyst stack (ClickHouse + Platform model routing on :8080). State lives under $RUNNER_TEMP/state. `--verify` only re-checks both health endpoints and writes the summary line. @@ -102,7 +102,9 @@ def main() -> None: "services", "run", "--services", - "auth,entities,intake", + "auth,entities,intake,models,inference-gateway,secrets", + "--controllers", + "models", "--host", "127.0.0.1", "--port", diff --git a/plugins/nemo-insights/tests/test_analysis_config_contract.py b/plugins/nemo-insights/tests/test_analysis_config_contract.py new file mode 100644 index 0000000000..5d52d9bba2 --- /dev/null +++ b/plugins/nemo-insights/tests/test_analysis_config_contract.py @@ -0,0 +1,143 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Contract tests for carrying scheduled-analysis models into Platform state.""" + +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient +from nemo_insights_plugin import cli +from nemo_insights_plugin.entities import AnalysisConfig +from nemo_insights_plugin.sdk_resources.analysis_configs import _build_enable_body +from nemo_insights_plugin.service import InsightsService +from nemo_platform_plugin.entity_client import NemoEntityNotFoundError, get_entity_client + + +def _app(entity_client: AsyncMock) -> FastAPI: + app = FastAPI() + for spec in InsightsService().get_routers(): + app.include_router(spec.router, prefix=spec.prefix) + app.dependency_overrides[get_entity_client] = lambda: entity_client + return app + + +def test_enable_request_requires_and_persists_model_pair() -> None: + entity_client = AsyncMock() + entity_client.get.side_effect = NemoEntityNotFoundError("missing") + entity_client.create.side_effect = lambda config: config + + response = TestClient(_app(entity_client)).post( + "/v2/workspaces/default/analysis-configs/calculator-agent/enable", + json={ + "default_model": "default/gpt-5", + "fast_model": "default/gpt-5-mini", + }, + ) + + assert response.status_code == 200 + created = entity_client.create.await_args.args[0] + assert created.default_model == "default/gpt-5" + assert created.fast_model == "default/gpt-5-mini" + assert response.json()["default_model"] == "default/gpt-5" + assert response.json()["fast_model"] == "default/gpt-5-mini" + + +def test_reenable_refreshes_persisted_model_pair() -> None: + entity_client = AsyncMock() + existing = AnalysisConfig( + name="calculator-agent", + workspace="default", + agent="calculator-agent", + enabled=False, + default_model="default/old-model", + fast_model="default/old-model", + ) + entity_client.get.return_value = existing + entity_client.update.side_effect = lambda config: config + + response = TestClient(_app(entity_client)).post( + "/v2/workspaces/default/analysis-configs/calculator-agent/enable", + json={ + "default_model": "default/gpt-5", + "fast_model": "default/gpt-5-mini", + }, + ) + + assert response.status_code == 200 + updated = entity_client.update.await_args.args[0] + assert updated.enabled is True + assert updated.default_model == "default/gpt-5" + assert updated.fast_model == "default/gpt-5-mini" + + +def test_enable_sdk_body_contains_both_model_refs() -> None: + assert _build_enable_body( + default_model="default/gpt-5", + fast_model="default/gpt-5-mini", + ) == { + "default_model": "default/gpt-5", + "fast_model": "default/gpt-5-mini", + } + + +@pytest.mark.asyncio +async def test_enable_cli_sends_locally_configured_models(monkeypatch: pytest.MonkeyPatch) -> None: + analysis_configs = SimpleNamespace( + enable=AsyncMock( + return_value=AnalysisConfig( + name="calculator-agent", + workspace="default", + agent="calculator-agent", + default_model="default/gpt-5", + fast_model="default/gpt-5-mini", + ) + ) + ) + client = SimpleNamespace( + insights=SimpleNamespace(analysis_configs=analysis_configs), + close=AsyncMock(), + ) + monkeypatch.setattr(cli, "make_client", lambda base_url: client) + monkeypatch.setattr( + cli, + "configured_model_refs", + lambda: SimpleNamespace(default="default/gpt-5", fast="default/gpt-5-mini"), + ) + + await cli._analysis_config_command( + action="enable", + agent="calculator-agent", + workspace="default", + base_url="http://localhost:8080", + ) + + analysis_configs.enable.assert_awaited_once_with( + workspace="default", + agent="calculator-agent", + default_model="default/gpt-5", + fast_model="default/gpt-5-mini", + ) + client.close.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_enable_cli_checks_local_models_before_constructing_client(monkeypatch: pytest.MonkeyPatch) -> None: + def missing_models(): + raise ValueError("No default model is configured. Run `nemo setup` and select agent models.") + + make_client = AsyncMock() + monkeypatch.setattr(cli, "configured_model_refs", missing_models) + monkeypatch.setattr(cli, "make_client", make_client) + + with pytest.raises(ValueError, match="No default model is configured"): + await cli._analysis_config_command( + action="enable", + agent="calculator-agent", + workspace="default", + base_url="http://localhost:8080", + ) + + make_client.assert_not_called() diff --git a/plugins/nemo-insights/tests/test_analyst_agent.py b/plugins/nemo-insights/tests/test_analyst_agent.py index b7157092c3..563956a033 100644 --- a/plugins/nemo-insights/tests/test_analyst_agent.py +++ b/plugins/nemo-insights/tests/test_analyst_agent.py @@ -8,7 +8,6 @@ from typing import Any, cast import pytest -from nemo_insights_plugin.analyst import model_config from nemo_insights_plugin.analyst.agent import KICKOFF, Analyst, build_analyst_agent from nemo_insights_plugin.analyst.analyst_backend import AnalystBackend from nemo_insights_plugin.analyst.deps import AnalystDeps @@ -39,24 +38,6 @@ def _exec_response(code: str) -> LLMResponse: ) -def test_smart_model_uses_opus_through_openai_compatible_gateway(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("INFERENCE_API_KEY", "test-key") - model_config.get_smart_model.cache_clear() - model_config._completion_client.cache_clear() - - client = model_config.get_smart_model() - try: - assert client.model == "openai/aws/anthropic/bedrock-claude-opus-4-8" - assert client.config == { - "api_base": "https://inference-api.nvidia.com/v1", - "api_key": "test-key", - } - finally: - client.close() - model_config.get_smart_model.cache_clear() - model_config._completion_client.cache_clear() - - async def test_nooa_codeact_returns_typed_analyst_result_and_receives_prompt() -> None: fake = FakeLLMClient( scripted_responses=[ diff --git a/plugins/nemo-insights/tests/test_analyst_run.py b/plugins/nemo-insights/tests/test_analyst_run.py index 66169a9dc7..8a347f66d3 100644 --- a/plugins/nemo-insights/tests/test_analyst_run.py +++ b/plugins/nemo-insights/tests/test_analyst_run.py @@ -3,12 +3,13 @@ """The ``run_analyst`` client injection contract.""" -from typing import cast +from typing import Any, cast import pytest from nemo_insights_plugin.analyst import run as run_module from nemo_insights_plugin.analyst.deps import AnalystDeps from nemo_platform import AsyncNeMoPlatform +from nemo_platform_plugin.nooa_model_client import ConfiguredModelClients from nooa.context_blocks import ResultStatus from nooa.events import LLMComplete, PythonOutput @@ -21,12 +22,33 @@ async def close(self) -> None: self.closed = True +class FakeModelClient: + def __init__(self) -> None: + self.closed = False + + async def aclose(self) -> None: + self.closed = True + + class FakeBackend: async def persist_result(self, *, workspace: str, agent: str, result: object) -> str: return "REPORT" def _stub_pipeline(monkeypatch: pytest.MonkeyPatch, seen: dict[str, object]) -> None: + default = FakeModelClient() + fast = FakeModelClient() + model_clients = ConfiguredModelClients( + default=cast(Any, default), + fast=cast(Any, fast), + ) + + async def fake_resolve_model_clients(client: object, refs: object) -> ConfiguredModelClients: + seen["model_client"] = client + seen["model_refs"] = refs + seen["model_clients"] = model_clients + return model_clients + def fake_make_backend(*, client: FakeClient, insights_output: str | None, local_only: bool) -> FakeBackend: seen["backend_client"] = client seen["local_only"] = local_only @@ -36,6 +58,7 @@ async def fake_run_agent(analyst: object, *, verbose: bool) -> object: return object() monkeypatch.setattr(run_module, "make_analyst_backend", fake_make_backend) + monkeypatch.setattr(run_module, "resolve_model_clients", fake_resolve_model_clients) def fake_build_agent(**kwargs: object) -> object: seen["build_kwargs"] = kwargs @@ -62,27 +85,28 @@ async def test_injected_client_is_used_and_closed(monkeypatch: pytest.MonkeyPatc assert seen["backend_client"] is client build_kwargs = cast(dict[str, object], seen["build_kwargs"]) assert cast(AnalystDeps, build_kwargs["deps"]).backend is not None + assert seen["model_client"] is client + model_clients = cast(ConfiguredModelClients, seen["model_clients"]) + assert cast(FakeModelClient, model_clients.default).closed + assert cast(FakeModelClient, model_clients.fast).closed assert client.closed -def test_litellm_compatibility_is_enabled(monkeypatch: pytest.MonkeyPatch) -> None: - class FakeLiteLLM: - drop_params = False - - fake_litellm = FakeLiteLLM() - monkeypatch.setattr(run_module.importlib, "import_module", lambda name: fake_litellm) - - run_module._enable_litellm_drop_params() - - assert fake_litellm.drop_params is True - - async def test_client_closed_when_backend_construction_raises(monkeypatch: pytest.MonkeyPatch) -> None: client = FakeClient() + model = FakeModelClient() + pair = ConfiguredModelClients( + default=cast(Any, model), + fast=cast(Any, model), + ) + + async def fake_resolve_model_clients(client: object, refs: object) -> ConfiguredModelClients: + return pair def raising_backend(*, client: FakeClient, insights_output: str | None, local_only: bool) -> FakeBackend: raise RuntimeError("backend failed") + monkeypatch.setattr(run_module, "resolve_model_clients", fake_resolve_model_clients) monkeypatch.setattr(run_module, "make_analyst_backend", raising_backend) with pytest.raises(RuntimeError, match="backend failed"): @@ -94,18 +118,19 @@ def raising_backend(*, client: FakeClient, insights_output: str | None, local_on client=cast(AsyncNeMoPlatform, client), ) + assert model.closed assert client.closed -async def test_client_closed_when_litellm_initialization_raises(monkeypatch: pytest.MonkeyPatch) -> None: +async def test_client_closed_when_model_resolution_raises(monkeypatch: pytest.MonkeyPatch) -> None: client = FakeClient() - def raising_litellm_initialization() -> None: - raise RuntimeError("litellm failed") + async def raising_model_resolution(client: object, refs: object) -> ConfiguredModelClients: + raise RuntimeError("model resolution failed") - monkeypatch.setattr(run_module, "_enable_litellm_drop_params", raising_litellm_initialization) + monkeypatch.setattr(run_module, "resolve_model_clients", raising_model_resolution) - with pytest.raises(RuntimeError, match="litellm failed"): + with pytest.raises(RuntimeError, match="model resolution failed"): await run_module.run_analyst( agent="agent", agent_spec=None, diff --git a/plugins/nemo-insights/tests/test_cli_profile.py b/plugins/nemo-insights/tests/test_cli_profile.py index c03fd0565e..75224b8134 100644 --- a/plugins/nemo-insights/tests/test_cli_profile.py +++ b/plugins/nemo-insights/tests/test_cli_profile.py @@ -4,12 +4,14 @@ import os from collections.abc import Iterator from pathlib import Path +from types import SimpleNamespace import httpx import pytest import typer from nemo_insights_plugin import cli from nemo_insights_plugin.analyst.cli import AnalystCLI +from nemo_insights_plugin.contracts.checks import CheckResult from nemo_insights_plugin.contracts.profile import DEFAULT_BASE_URL from nemo_insights_plugin.preflight import AnalysisProbes from nemo_platform import NeMoPlatformError @@ -17,7 +19,7 @@ from typer.testing import CliRunner runner = CliRunner() -_PROFILE_ENV_KEYS = ("NMP_BASE_URL", "INFERENCE_API_KEY") +_PROFILE_ENV_KEYS = ("NMP_BASE_URL", "TEST_PROFILE_ENV") class AnalystRecorder: @@ -57,12 +59,28 @@ async def queryable(base_url: str, workspace: str, agent: str) -> bool: cli, "_PREFLIGHT_PROBES", AnalysisProbes( - env={"INFERENCE_API_KEY": "k"}, http_ok=lambda base_url: True, workspace_ok=queryable, ), ) monkeypatch.setattr(cli, "make_client", lambda base_url: object()) + monkeypatch.setattr( + cli, + "check_models", + lambda: [ + CheckResult( + name="agent-models", + group="models", + status="pass", + severity="required", + message="default=default/gpt-5; fast=default/gpt-5-mini", + ) + ], + ) + monkeypatch.setattr( + "nemo_insights_plugin.preflight.configured_model_refs", + lambda: SimpleNamespace(default="default/gpt-5", fast="default/gpt-5-mini"), + ) @pytest.fixture @@ -128,10 +146,15 @@ def test_unusable_mirror_directory_drops_the_mirror_and_still_analyzes( output = tmp_path / "unwritable" / "insights.yaml" original_mkdir = Path.mkdir - def refuse_mirror_dir(self: Path, *args: object, **kwargs: object) -> None: + def refuse_mirror_dir( + self: Path, + mode: int = 0o777, + parents: bool = False, + exist_ok: bool = False, + ) -> None: if self == output.parent: raise PermissionError(13, "Permission denied") - original_mkdir(self, *args, **kwargs) # type: ignore[arg-type] + original_mkdir(self, mode=mode, parents=parents, exist_ok=exist_ok) monkeypatch.setattr(Path, "mkdir", refuse_mirror_dir) monkeypatch.setattr(cli, "run_analyst", recorder) @@ -190,7 +213,6 @@ async def record_workspace_probe(base_url: str, workspace: str, agent: str) -> b cli, "_PREFLIGHT_PROBES", AnalysisProbes( - env={"INFERENCE_API_KEY": "k"}, http_ok=lambda base_url: probe_calls.append("http") or True, workspace_ok=record_workspace_probe, ), @@ -228,7 +250,6 @@ async def record_workspace_probe(base_url: str, workspace: str, agent: str) -> b cli, "_PREFLIGHT_PROBES", AnalysisProbes( - env={"INFERENCE_API_KEY": "k"}, http_ok=lambda base_url: probe_calls.append("http") or True, workspace_ok=record_workspace_probe, ), @@ -289,7 +310,6 @@ async def record_workspace_probe(base_url: str, workspace: str, agent: str) -> b cli, "_PREFLIGHT_PROBES", AnalysisProbes( - env={"INFERENCE_API_KEY": "k"}, http_ok=lambda base_url: http_urls.append(base_url) or True, workspace_ok=record_workspace_probe, ), @@ -406,15 +426,15 @@ def test_malformed_discovered_profile_warns_with_agent_only(app: typer.Typer, tm def test_malformed_discovered_profile_still_loads_env_file(app: typer.Typer, tmp_path: Path, monkeypatch) -> None: recorder = AnalystRecorder() monkeypatch.setattr(cli, "run_analyst", recorder) - monkeypatch.delenv("INFERENCE_API_KEY", raising=False) + monkeypatch.delenv("TEST_PROFILE_ENV", raising=False) (tmp_path / "optimizer.yaml").write_text("agent: ''\n", encoding="utf-8") - (tmp_path / ".env").write_text("INFERENCE_API_KEY=from-env-file\n", encoding="utf-8") + (tmp_path / ".env").write_text("TEST_PROFILE_ENV=from-env-file\n", encoding="utf-8") monkeypatch.chdir(tmp_path) result = runner.invoke(app, ["run", "--agent", "other"]) assert result.exit_code == 0, result.output - assert os.environ["INFERENCE_API_KEY"] == "from-env-file" + assert os.environ["TEST_PROFILE_ENV"] == "from-env-file" assert recorder.kwargs is not None @@ -440,16 +460,28 @@ def test_missing_profile_and_agent_errors(app: typer.Typer, tmp_path: Path, monk assert "No --agent given and no optimizer.yaml profile found" in result.output -def test_analyze_blocks_before_runner_when_preflight_fails( +def test_analyze_blocks_before_runner_when_model_configuration_is_missing( app: typer.Typer, profile_tree: Path, monkeypatch: pytest.MonkeyPatch ) -> None: recorder = AnalystRecorder() monkeypatch.setattr(cli, "run_analyst", recorder) + monkeypatch.setattr( + cli, + "check_models", + lambda: [ + CheckResult( + name="agent-models", + group="models", + status="fail", + severity="required", + message="No default model is configured", + ) + ], + ) monkeypatch.setattr( cli, "_PREFLIGHT_PROBES", AnalysisProbes( - env={}, http_ok=lambda base_url: True, workspace_ok=lambda base_url, workspace, agent: _queryable(), ), @@ -459,11 +491,11 @@ def test_analyze_blocks_before_runner_when_preflight_fails( result = runner.invoke(app, ["run"]) assert result.exit_code == 1 - assert "INFERENCE_API_KEY not set" in result.output + assert "No default model is configured" in result.output assert recorder.kwargs is None -def test_analyze_runs_only_the_credential_check( +def test_analyze_runs_only_the_model_configuration_check( app: typer.Typer, profile_tree: Path, monkeypatch: pytest.MonkeyPatch ) -> None: recorder = AnalystRecorder() @@ -478,7 +510,6 @@ async def record_workspace_probe(base_url: str, workspace: str, agent: str) -> b cli, "_PREFLIGHT_PROBES", AnalysisProbes( - env={"INFERENCE_API_KEY": "k"}, http_ok=lambda base_url: probe_calls.append("http") or False, workspace_ok=record_workspace_probe, ), @@ -746,7 +777,6 @@ async def record_workspace_probe(base_url: str, workspace: str, agent: str) -> b cli, "_PREFLIGHT_PROBES", AnalysisProbes( - env={}, http_ok=lambda base_url: http_urls.append(base_url) or True, workspace_ok=record_workspace_probe, ), @@ -757,7 +787,7 @@ async def record_workspace_probe(base_url: str, workspace: str, agent: str) -> b assert result.exit_code == 1 assert "no optimizer.yaml found" in result.output - assert "INFERENCE_API_KEY not set" in result.output + assert "Models\n ✓ default=default/gpt-5; fast=default/gpt-5-mini" in result.output assert http_urls == ["https://flag.example"] assert workspace_urls == [] @@ -769,4 +799,4 @@ def test_doctor_reports_healthy_profile(app: typer.Typer, profile_tree: Path, mo assert result.exit_code == 0, result.output assert "Profile\n ✓ profile for agent 'flight-planner'" in result.output - assert "Credentials\n ✓ INFERENCE_API_KEY set" in result.output + assert "Models\n ✓ default=default/gpt-5; fast=default/gpt-5-mini" in result.output diff --git a/plugins/nemo-insights/tests/test_periodic_analysis.py b/plugins/nemo-insights/tests/test_periodic_analysis.py index f910ea8770..1e4b3c4872 100644 --- a/plugins/nemo-insights/tests/test_periodic_analysis.py +++ b/plugins/nemo-insights/tests/test_periodic_analysis.py @@ -6,6 +6,7 @@ from datetime import datetime, timezone from pathlib import Path from types import SimpleNamespace +from typing import cast from zoneinfo import ZoneInfo import httpx @@ -36,13 +37,16 @@ ) from nemo_insights_plugin.jobs.analyze import AnalyzeJob, AnalyzeSpec from nemo_insights_plugin.schedule import is_due, previous_scheduled +from nemo_platform import AsyncNeMoPlatform, NeMoPlatform from nemo_platform.types.intake.spans.span_group import SpanGroup -from nemo_platform_plugin.entity_client import NemoEntityNotFoundError +from nemo_platform_plugin.entity_client import NemoEntitiesClient, NemoEntityNotFoundError from nemo_platform_plugin.job_context import JobContext, StoragePaths +from nemo_platform_plugin.job_results import JobResults from nemo_platform_plugin.jobs.constants import ( DEFAULT_JOB_STORAGE_PATH, PERSISTENT_JOB_STORAGE_PATH_ENVVAR, ) +from nemo_platform_plugin.nooa_model_client import ConfiguredModelRefs from pydantic import ValidationError @@ -88,7 +92,7 @@ async def get(self, **kwargs: object) -> None: @pytest.mark.asyncio async def test_remote_persist_validates_updates_without_trace_refs() -> None: client = SimpleNamespace(insights=SimpleNamespace(insights=_MissingInsights())) - backend = RemoteAnalystBackend(client) # type: ignore[arg-type] + backend = RemoteAnalystBackend(cast(AsyncNeMoPlatform, client)) result = AnalystResult( summary="Nothing new.", updated_insights=[InsightUpdate(id="missing-insight")], @@ -149,7 +153,7 @@ async def update(self, *, workspace: str, insight_id: str, trace_refs: list[str] def _remote_backend_with_mirror(path: Path) -> tuple[RemoteAnalystBackend, _RecordingInsights]: insights = _RecordingInsights() client = SimpleNamespace(insights=SimpleNamespace(insights=insights)) - backend = RemoteAnalystBackend(client, mirror=InsightsFileStore(path)) # type: ignore[arg-type] + backend = RemoteAnalystBackend(cast(AsyncNeMoPlatform, client), mirror=InsightsFileStore(path)) return backend, insights @@ -219,8 +223,9 @@ def test_make_analyst_backend_always_writes_to_the_platform(tmp_path: Path) -> N """An output path adds a mirror; it never diverts writes off the platform.""" client = SimpleNamespace() - plain = make_analyst_backend(client=client, insights_output=None) # type: ignore[arg-type] - mirrored = make_analyst_backend(client=client, insights_output=str(tmp_path / "insights.yaml")) # type: ignore[arg-type] + platform_client = cast(AsyncNeMoPlatform, client) + plain = make_analyst_backend(client=platform_client, insights_output=None) + mirrored = make_analyst_backend(client=platform_client, insights_output=str(tmp_path / "insights.yaml")) assert isinstance(plain, RemoteAnalystBackend) and plain.mirror is None assert isinstance(mirrored, RemoteAnalystBackend) @@ -231,15 +236,16 @@ def test_local_only_is_testbed_plumbing_and_requires_a_path(tmp_path: Path) -> N """``local_only`` stays reachable for the testbed, which no CLI flag sets.""" client = SimpleNamespace() - local = make_analyst_backend( # type: ignore[arg-type] - client=client, # type: ignore[arg-type] + platform_client = cast(AsyncNeMoPlatform, client) + local = make_analyst_backend( + client=platform_client, insights_output=str(tmp_path / "insights.yaml"), local_only=True, ) assert isinstance(local, LocalAnalystBackend) with pytest.raises(ValueError, match="requires an insights output path"): - make_analyst_backend(client=client, insights_output=None, local_only=True) # type: ignore[arg-type] + make_analyst_backend(client=platform_client, insights_output=None, local_only=True) def test_local_backend_reads_and_writes_insights_file_with_explicit_utf8( @@ -250,18 +256,27 @@ def test_local_backend_reads_and_writes_insights_file_with_explicit_utf8( original_read_text = Path.read_text original_write_text = Path.write_text - def spy_read_text(self: Path, *args: object, **kwargs: object) -> str: - read_calls.append(kwargs) - return original_read_text(self, *args, **kwargs) + def spy_read_text(self: Path, encoding: str | None = None, errors: str | None = None) -> str: + read_calls.append({"encoding": encoding, "errors": errors}) + return original_read_text(self, encoding=encoding, errors=errors) - def spy_write_text(self: Path, *args: object, **kwargs: object) -> int: - write_calls.append(kwargs) - return original_write_text(self, *args, **kwargs) + def spy_write_text( + self: Path, + data: str, + encoding: str | None = None, + errors: str | None = None, + newline: str | None = None, + ) -> int: + write_calls.append({"encoding": encoding, "errors": errors, "newline": newline}) + return original_write_text(self, data, encoding=encoding, errors=errors, newline=newline) monkeypatch.setattr(Path, "read_text", spy_read_text) monkeypatch.setattr(Path, "write_text", spy_write_text) - backend = LocalAnalystBackend(client=SimpleNamespace(), path=tmp_path / "insights.yaml") # type: ignore[arg-type] + backend = LocalAnalystBackend( + client=cast(AsyncNeMoPlatform, SimpleNamespace()), + path=tmp_path / "insights.yaml", + ) backend.store.write_records([]) backend.store.read_records() @@ -272,7 +287,7 @@ def spy_write_text(self: Path, *args: object, **kwargs: object) -> int: def test_local_backend_write_preserves_other_top_level_keys(tmp_path: Path) -> None: path = tmp_path / "insights.yaml" path.write_text("metadata: retained\ninsights:\n- id: stale\n", encoding="utf-8") - backend = LocalAnalystBackend(client=SimpleNamespace(), path=path) # type: ignore[arg-type] + backend = LocalAnalystBackend(client=cast(AsyncNeMoPlatform, SimpleNamespace()), path=path) backend.store.write_records([{"id": "insight-1"}]) @@ -314,15 +329,17 @@ async def list(self, **kwargs: object) -> SimpleNamespace: class _SpanGroupClient: def __init__(self) -> None: + self.groups = _SpanGroups() self.intake = SimpleNamespace( - spans=SimpleNamespace(groups=_SpanGroups()), + spans=SimpleNamespace(groups=self.groups), ) @pytest.mark.asyncio async def test_count_agent_sessions_uses_server_side_session_groups() -> None: since = datetime(2026, 6, 4, 12, tzinfo=timezone.utc) - backend = RemoteAnalystBackend(_SpanGroupClient()) # type: ignore[arg-type] + client = _SpanGroupClient() + backend = RemoteAnalystBackend(cast(AsyncNeMoPlatform, client)) count = await backend.count_agent_sessions( agent="research-agent", @@ -331,7 +348,7 @@ async def test_count_agent_sessions_uses_server_side_session_groups() -> None: ) assert count == 7 - assert backend.client.intake.spans.groups.calls == [ + assert client.groups.calls == [ { "workspace": "default", "by": "session_id", @@ -370,7 +387,7 @@ def __init__(self, groups: _SpanGroupsPaged) -> None: async def test_list_span_groups_fans_out_over_sessions() -> None: since = datetime(2026, 6, 4, 12, tzinfo=timezone.utc) groups = _SpanGroupsPaged() - backend = RemoteAnalystBackend(_GroupsBackendClient(groups)) # type: ignore[arg-type] + backend = RemoteAnalystBackend(cast(AsyncNeMoPlatform, _GroupsBackendClient(groups))) result = await backend.list_span_groups( workspace="default", @@ -407,7 +424,8 @@ async def test_list_span_groups_fans_out_over_sessions() -> None: @pytest.mark.asyncio async def test_count_agent_sessions_pins_evaluation_id() -> None: - backend = RemoteAnalystBackend(_SpanGroupClient()) # type: ignore[arg-type] + client = _SpanGroupClient() + backend = RemoteAnalystBackend(cast(AsyncNeMoPlatform, client)) await backend.count_agent_sessions( agent="research-agent", @@ -415,7 +433,7 @@ async def test_count_agent_sessions_pins_evaluation_id() -> None: evaluation_id="run-1", ) - assert backend.client.intake.spans.groups.calls == [ + assert client.groups.calls == [ { "workspace": "default", "by": "session_id", @@ -430,7 +448,7 @@ async def test_count_agent_sessions_pins_evaluation_id() -> None: @pytest.mark.asyncio async def test_list_span_groups_pins_evaluation_id() -> None: groups = _SpanGroupsPaged() - backend = RemoteAnalystBackend(_GroupsBackendClient(groups)) # type: ignore[arg-type] + backend = RemoteAnalystBackend(cast(AsyncNeMoPlatform, _GroupsBackendClient(groups))) await backend.list_span_groups( workspace="default", @@ -568,7 +586,7 @@ def test_weekly_not_due_until_configured_weekday() -> None: def test_config_accepts_weekday_name() -> None: - config = AnalystSchedulerConfig(run_on_weekday="friday") + config = AnalystSchedulerConfig.model_validate({"run_on_weekday": "friday"}) assert config.run_on_weekday is Weekday.FRIDAY @@ -622,23 +640,38 @@ def _ctx(tmp_path: Path) -> JobContext: return JobContext( workspace="default", storage=StoragePaths(ephemeral=ephemeral, persistent=persistent), - results=_Results(), + results=cast(JobResults, _Results()), job_id="insights-job-1", ) +def _analyze_spec(agent: str = "research-agent") -> AnalyzeSpec: + return AnalyzeSpec( + agent=agent, + default_model="default/gpt-5", + fast_model="default/gpt-5-mini", + ) + + def test_analyze_job_records_success(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: - async def fake_run_analyst(**_: object) -> str: + calls: list[dict[str, object]] = [] + async_client = object() + + async def fake_run_analyst(**kwargs: object) -> str: + calls.append(kwargs) return "analysis report" monkeypatch.setattr("nemo_insights_plugin.jobs.analyze.run_analyst", fake_run_analyst) - monkeypatch.setattr("nemo_insights_plugin.jobs.analyze.make_client", lambda base_url: object()) + monkeypatch.setattr( + "nemo_insights_plugin.jobs.analyze.get_async_task_sdk", + lambda plugin: async_client, + ) sdk = _SyncSdk() result = AnalyzeJob().run( - AnalyzeSpec(agent="research-agent").model_dump(mode="json"), + _analyze_spec().model_dump(mode="json"), ctx=_ctx(tmp_path), - sdk=sdk, + sdk=cast(NeMoPlatform, sdk), ) assert result["status"] == "completed" @@ -653,77 +686,32 @@ async def fake_run_analyst(**_: object) -> str: ] assert updates[-1]["last_submitted_job"] == "insights-job-1" assert (tmp_path / "persistent" / "analysis-report.txt").read_text() == "analysis report" - - -@pytest.mark.asyncio -async def test_analyze_job_compile_requests_persistent_storage( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.delenv("INFERENCE_API_KEY", raising=False) - platform_spec = await AnalyzeJob.compile( - workspace="default", - spec=AnalyzeSpec(agent="research-agent"), - entity_client=object(), - job_name="opt-analyze-default-research-agent-20260608204901", - async_sdk=object(), - ) - - step = platform_spec["steps"][0] - assert step["environment"] == [ - { - "name": PERSISTENT_JOB_STORAGE_PATH_ENVVAR, - "value": DEFAULT_JOB_STORAGE_PATH, - }, - ] - - -@pytest.mark.asyncio -async def test_analyze_job_compile_can_reference_inference_secret() -> None: - platform_spec = await AnalyzeJob.compile( - workspace="default", - spec=AnalyzeSpec( - agent="calculator-agent", - inference_api_key_secret_name="insights-inference-api-key", - ), - entity_client=object(), - job_name="opt-analyze-default-calculator-agent-20260608210807", - async_sdk=object(), + assert calls[0]["client"] is async_client + assert calls[0]["model_refs"] == ConfiguredModelRefs( + default="default/gpt-5", + fast="default/gpt-5-mini", ) - step = platform_spec["steps"][0] - assert step["environment"] == [ - { - "name": PERSISTENT_JOB_STORAGE_PATH_ENVVAR, - "value": DEFAULT_JOB_STORAGE_PATH, - }, - { - "name": "INFERENCE_API_KEY", - "from_secret": {"name": "insights-inference-api-key"}, - }, - ] - @pytest.mark.asyncio -async def test_analyze_job_compile_can_forward_local_inference_env( +async def test_analyze_job_compile_requests_storage_without_provider_credentials( monkeypatch: pytest.MonkeyPatch, ) -> None: - monkeypatch.setenv("INFERENCE_API_KEY", "test-key") - + monkeypatch.setenv("INFERENCE_API_KEY", "must-not-be-forwarded") platform_spec = await AnalyzeJob.compile( workspace="default", - spec=AnalyzeSpec(agent="calculator-agent"), + spec=_analyze_spec(), entity_client=object(), - job_name="opt-analyze-default-calculator-agent-20260608210807", + job_name="opt-analyze-default-research-agent-20260608204901", async_sdk=object(), ) - step = platform_spec["steps"][0] + step = next(iter(platform_spec["steps"])) assert step["environment"] == [ { "name": PERSISTENT_JOB_STORAGE_PATH_ENVVAR, "value": DEFAULT_JOB_STORAGE_PATH, }, - {"name": "INFERENCE_API_KEY", "value": "test-key"}, ] @@ -777,7 +765,7 @@ def _controller( *, jobs: list[SimpleNamespace] | None = None, run_status: AnalysisRunStatus | None = None, -) -> InsightsAnalysisController: +) -> tuple[InsightsAnalysisController, _AsyncSdk, _Entities]: controller = InsightsAnalysisController() controller._config = InsightsConfig( analyst=AnalystSchedulerConfig( @@ -786,9 +774,11 @@ def _controller( job_profile="test-profile", ) ) - controller._sdk = _AsyncSdk(jobs=jobs) - controller._entities = _Entities(run_status=run_status) - return controller + sdk = _AsyncSdk(jobs=jobs) + entities = _Entities(run_status=run_status) + controller._sdk = cast(AsyncNeMoPlatform, sdk) + controller._entities = cast(NemoEntitiesClient, entities) + return controller, sdk, entities def test_generated_job_name_fits_derived_fileset_name_limit() -> None: @@ -807,8 +797,14 @@ def test_generated_job_name_fits_derived_fileset_name_limit() -> None: @pytest.mark.asyncio async def test_controller_submits_due_job(monkeypatch: pytest.MonkeyPatch) -> None: - controller = _controller() - config = AnalysisConfig(name="research-agent", workspace="default", agent="research-agent") + controller, sdk, entities = _controller() + config = AnalysisConfig( + name="research-agent", + workspace="default", + agent="research-agent", + default_model="default/gpt-5", + fast_model="default/gpt-5-mini", + ) async def fake_compile_job_spec(**_: object) -> dict[str, list[object]]: return {"steps": []} @@ -816,17 +812,35 @@ async def fake_compile_job_spec(**_: object) -> dict[str, list[object]]: monkeypatch.setattr(controller, "_compile_job_spec", fake_compile_job_spec) await controller._reconcile_config(config) - assert len(controller.sdk.jobs.created) == 1 - created = controller.sdk.jobs.created[0] + assert len(sdk.jobs.created) == 1 + created = sdk.jobs.created[0] + created_spec = cast(dict[str, object], created["spec"]) assert created["source"] == "insights" - assert created["spec"]["agent"] == "research-agent" - assert created["spec"]["since"] is None - assert controller.entities.updated == [] + assert created_spec["agent"] == "research-agent" + assert created_spec["since"] is None + assert created_spec["default_model"] == "default/gpt-5" + assert created_spec["fast_model"] == "default/gpt-5-mini" + assert entities.updated == [] + + +@pytest.mark.asyncio +async def test_controller_defers_legacy_config_without_persisted_models( + caplog: pytest.LogCaptureFixture, +) -> None: + controller, sdk, _ = _controller() + config = AnalysisConfig(name="research-agent", workspace="default", agent="research-agent") + + with caplog.at_level("ERROR", logger="nemo_insights_plugin.controller"): + await controller._reconcile_config(config) + + assert sdk.jobs.created == [] + assert "has no model selection" in caplog.text + assert "nemo insights analysis enable" in caplog.text @pytest.mark.asyncio async def test_controller_skips_active_job(monkeypatch: pytest.MonkeyPatch) -> None: - controller = _controller( + controller, sdk, _ = _controller( jobs=[ SimpleNamespace( status="active", @@ -846,4 +860,4 @@ async def fake_compile_job_spec(**_: object) -> dict[str, list[object]]: monkeypatch.setattr(controller, "_compile_job_spec", fake_compile_job_spec) await controller._reconcile_config(config) - assert controller.sdk.jobs.created == [] + assert sdk.jobs.created == [] diff --git a/plugins/nemo-insights/tests/test_preflight.py b/plugins/nemo-insights/tests/test_preflight.py index 3eb9d2d045..c847151fa1 100644 --- a/plugins/nemo-insights/tests/test_preflight.py +++ b/plugins/nemo-insights/tests/test_preflight.py @@ -3,6 +3,7 @@ import asyncio from pathlib import Path +from types import SimpleNamespace import httpx import pytest @@ -26,7 +27,20 @@ async def never_queryable(base_url: str, workspace: str, agent: str) -> bool: return False -def test_missing_inference_key_is_required_failure(tmp_path: Path) -> None: +@pytest.fixture(autouse=True) +def configured_models(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + preflight, + "configured_model_refs", + lambda: SimpleNamespace(default="default/gpt-5", fast="default/gpt-5-mini"), + ) + + +def test_missing_model_pair_is_required_failure(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + def missing_models() -> object: + raise ValueError("No default model is configured") + + monkeypatch.setattr(preflight, "configured_model_refs", missing_models) results = asyncio.run( check_environment( agent="a", @@ -34,14 +48,13 @@ def test_missing_inference_key_is_required_failure(tmp_path: Path) -> None: base_url="http://localhost:8080", profile_dir=tmp_path, probes=AnalysisProbes( - env={}, http_ok=lambda base_url: True, workspace_ok=always_queryable, ), ) ) - assert any(result.name == "INFERENCE_API_KEY" and result.status == "fail" for result in results) + assert any(result.name == "agent-models" and result.status == "fail" for result in results) assert required_failures(results) @@ -53,14 +66,13 @@ def test_check_environment_without_profile_skips_workspace_probe(tmp_path: Path) base_url="http://localhost:8080", profile_dir=None, probes=AnalysisProbes( - env={"INFERENCE_API_KEY": "k"}, http_ok=lambda base_url: True, workspace_ok=never_queryable, ), ) ) - assert [result.name for result in results] == ["INFERENCE_API_KEY", "platform-reachable"] + assert [result.name for result in results] == ["agent-models", "platform-reachable"] def test_default_http_probe_treats_invalid_base_url_as_unreachable() -> None: @@ -75,7 +87,6 @@ def test_workspace_query_failure_is_advisory(tmp_path: Path) -> None: base_url="http://localhost:8080", profile_dir=tmp_path, probes=AnalysisProbes( - env={"INFERENCE_API_KEY": "k"}, http_ok=lambda base_url: True, workspace_ok=never_queryable, ), @@ -160,7 +171,6 @@ def test_healthy_setup_formats_grouped_report(tmp_path: Path) -> None: base_url="http://localhost:8080", profile_dir=tmp_path, probes=AnalysisProbes( - env={"INFERENCE_API_KEY": "k"}, http_ok=lambda base_url: True, workspace_ok=always_queryable, ), @@ -170,6 +180,6 @@ def test_healthy_setup_formats_grouped_report(tmp_path: Path) -> None: report = format_report(results) assert "Profile\n ✓ profile for agent 'a'" in report - assert "Credentials\n ✓ INFERENCE_API_KEY set" in report + assert "Models\n ✓ default=default/gpt-5; fast=default/gpt-5-mini" in report assert "Platform\n ✓ http://localhost:8080 reachable" in report assert not required_failures(results) diff --git a/plugins/nemo-insights/tests/testbed/test_cli.py b/plugins/nemo-insights/tests/testbed/test_cli.py index 6caddc8507..24139fb910 100644 --- a/plugins/nemo-insights/tests/testbed/test_cli.py +++ b/plugins/nemo-insights/tests/testbed/test_cli.py @@ -191,16 +191,24 @@ def test_analyze_unknown_subject_exits(monkeypatch): cli.main() -def test_analyze_without_key_exits(monkeypatch): - # Neutralize .env loading so a developer's local testbed/.env can't supply the key. - # The key guard must fire before any state resolution/download (bare analyze = pinned mode). +def test_analyze_does_not_require_plugin_inference_key(monkeypatch, tmp_path, capsys): + # Analyst credentials belong to the configured Platform Model provider, not + # a plugin-owned environment variable. monkeypatch.setattr(cli, "_load_dotenv", lambda *a, **k: None) monkeypatch.delenv("INFERENCE_API_KEY", raising=False) - monkeypatch.setattr(sys, "argv", ["testbed", "analyze", "nvq"]) - with pytest.raises(SystemExit) as exc: - cli.main() - # Must exit for the key guard, not (e.g.) an "Unknown testbed" reason. - assert "INFERENCE_API_KEY" in str(exc.value) + monkeypatch.setattr(cli, "TMP", tmp_path) + monkeypatch.setattr( + sys, + "argv", + ["testbed", "analyze", "nvq", "--live", "--no-baseline-update"], + ) + + async def fake_analyze(self, *, record, since, verbose, out_path): + return "REPORT-OK" + + monkeypatch.setattr("testbed.adapters.IntakeAdapter.analyze", fake_analyze) + cli.main() + assert "REPORT-OK" in capsys.readouterr().out def test_analyze_live_happy_path(monkeypatch, tmp_path, capsys): @@ -866,7 +874,6 @@ def test_load_dotenv_missing_file_is_noop(tmp_path, monkeypatch): def test_doctor_ready(monkeypatch, capsys): monkeypatch.setattr(cli, "_load_dotenv", lambda *a, **k: None) - monkeypatch.setenv("INFERENCE_API_KEY", "sk") monkeypatch.setattr("testbed.adapters.IntakeAdapter.check", lambda self: []) monkeypatch.setattr(cli.shutil, "which", lambda _name: "/usr/bin/gh") # deterministic: gh "installed" monkeypatch.setattr(sys, "argv", ["testbed", "doctor", "nvq"]) @@ -881,7 +888,6 @@ def test_doctor_ready(monkeypatch, capsys): def test_doctor_flags_missing_gh(monkeypatch, capsys): """Pinned/--state analyze shells out to gh; doctor must flag its absence up front.""" monkeypatch.setattr(cli, "_load_dotenv", lambda *a, **k: None) - monkeypatch.setenv("INFERENCE_API_KEY", "sk") monkeypatch.setattr("testbed.adapters.IntakeAdapter.check", lambda self: []) monkeypatch.setattr(cli.shutil, "which", lambda _name: None) monkeypatch.setattr(sys, "argv", ["testbed", "doctor", "nvq"]) @@ -892,13 +898,14 @@ def test_doctor_flags_missing_gh(monkeypatch, capsys): def test_doctor_lists_unmet(monkeypatch, capsys): - monkeypatch.setattr(cli, "_load_dotenv", lambda *a, **k: None) # don't let .env supply the key + monkeypatch.setattr(cli, "_load_dotenv", lambda *a, **k: None) monkeypatch.delenv("INFERENCE_API_KEY", raising=False) monkeypatch.setattr("testbed.adapters.IntakeAdapter.check", lambda self: ["config key 'agent'"]) monkeypatch.setattr(sys, "argv", ["testbed", "doctor", "nvq"]) cli.main() out = capsys.readouterr().out - assert "✗" in out and "INFERENCE_API_KEY" in out and "config key 'agent'" in out + assert "✗" in out and "config key 'agent'" in out + assert "INFERENCE_API_KEY" not in out def test_run_records_and_prints(monkeypatch, tmp_path, capsys): @@ -923,7 +930,9 @@ async def fake_produce(self): assert "analyze tau2-airline --live" in out from testbed.runstore import load_run - assert load_run(tmp_path / "tau2-airline.run.json")["agent"] == "tau2-airline-xyz" + recorded_run = load_run(tmp_path / "tau2-airline.run.json") + assert recorded_run is not None + assert recorded_run["agent"] == "tau2-airline-xyz" def test_analyze_live_passes_record_to_analyze(monkeypatch, tmp_path): @@ -2327,7 +2336,9 @@ async def fake_analyze(self, *, record, since, verbose, out_path): (backup,) = tmp_path.glob("backup-*") # exactly one dir carries both parked files assert json.loads((backup / "tau2-airline.run.json").read_text(encoding="utf-8"))["experiment_id"] == "stale" assert (backup / "insights_tau2-airline.yaml").read_text(encoding="utf-8") == "prior" - assert load_run(tmp_path / "tau2-airline.run.json")["experiment_id"] == "run-1" # bundle record seeded + restored_run = load_run(tmp_path / "tau2-airline.run.json") + assert restored_run is not None + assert restored_run["experiment_id"] == "run-1" # bundle record seeded def test_analyze_bundle_without_run_record_exits(monkeypatch, tmp_path, stub_reingest): diff --git a/script/generate_config_docs.py b/script/generate_config_docs.py index 9158de05c4..9cdd50d4f7 100644 --- a/script/generate_config_docs.py +++ b/script/generate_config_docs.py @@ -39,7 +39,6 @@ from typing import Any, get_args, get_origin import yaml -from nemo_experimentalist_plugin.settings import ExperimentalistConfig from nemo_safe_synthesizer_plugin.config import SafeSynthesizerConfig from nmp.automodel.config import AutomodelConfig from nmp.common.config.base import CommonServiceConfig, PlatformConfig @@ -73,7 +72,6 @@ AutomodelConfig, UnslothConfig, SafeSynthesizerConfig, - ExperimentalistConfig, StudioConfig, ] diff --git a/sdk/python/nemo-platform/src/nemo_platform/cli/commands/manifest_registry.py b/sdk/python/nemo-platform/src/nemo_platform/cli/commands/manifest_registry.py index 88a718a0df..9db87eb049 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/cli/commands/manifest_registry.py +++ b/sdk/python/nemo-platform/src/nemo_platform/cli/commands/manifest_registry.py @@ -37,8 +37,8 @@ Uses an already-running platform, starts local services, or connects the CLI to an existing remote deployment. Then selects and registers an -inference provider, picks a default model, installs coding agent skills, -and optionally deploys a demo agent. +inference provider, picks default and fast agent models, installs coding +agent skills, and optionally deploys a demo agent. The active config context remembers the Platform URL. When a remote deployment is already reachable, setup asks whether to continue with it, @@ -56,7 +56,7 @@ Use --auto for non-interactive setup from environment variables (NEMO_DEFAULT_INFERENCE_KEY, NVIDIA_API_KEY, OPENAI_API_KEY, ANTHROPIC_API_KEY, GEMINI_API_KEY). -Override the default model with NEMO_DEFAULT_MODEL. +Override the selected pair with NEMO_DEFAULT_MODEL and NEMO_FAST_MODEL. Examples: nemo setup diff --git a/sdk/python/nemo-platform/src/nemo_platform/cli/commands/setup.py b/sdk/python/nemo-platform/src/nemo_platform/cli/commands/setup.py index 8883b93a00..7ab01befbb 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/cli/commands/setup.py +++ b/sdk/python/nemo-platform/src/nemo_platform/cli/commands/setup.py @@ -47,7 +47,7 @@ from nemo_platform.cli.telemetry.events import OnboardingStepEvent, TaskStatusEnum from nemo_platform.client.tls import client_verify_from_env from nemo_platform.config.config import Config -from nemo_platform.config.models import DEFAULT_BASE_URL, ConfigFile, ConfigParams, LocalServicesConfig +from nemo_platform.config.models import DEFAULT_BASE_URL, ConfigFile, ConfigParams, LocalServicesConfig, NoAuthUser from nemo_platform.local.process import ( check_port_available_for_start, compute_scope, @@ -218,6 +218,14 @@ class KeyValidationResult: message: str +@dataclass(frozen=True) +class ModelPair: + """Model entities selected for quality-critical and low-latency agent work.""" + + default: str + fast: str + + # Env vars probed during --auto mode, in priority order. _AUTO_ENV_VARS: tuple[tuple[str, str], ...] = ( ("NEMO_DEFAULT_INFERENCE_KEY", "NEMO_DEFAULT_INFERENCE_BASE_URL"), @@ -246,6 +254,7 @@ class KeyValidationResult: _POST_START_REACHABLE_DELAY = 2.0 _DEMO_AGENT_NAME = "calculator-agent" +_LOCAL_CONTEXT_NAME = "local" def _pause(seconds: float) -> None: @@ -272,8 +281,8 @@ def _bootstrap_config_if_missing(base_url: str, workspace: str) -> None: ``nemo setup`` can run before any config is on disk (first-time install) *or* with a partial config containing only ``local_services.data_dir`` written earlier in the same setup invocation by ``_save_data_dir``. - Later steps — in particular ``_save_default_model`` — call - ``Config.write`` with *only* a ``default_model`` param. If there's no + Later steps — in particular ``_save_model_pair`` — call + ``Config.write`` with only model-selection params. If there's no cluster on disk at that point, ``ensure_context`` will fail with ``Cluster '' does not exist and no base_url provided to create it`` because it has no ``base_url`` to attach to a new cluster. @@ -378,6 +387,49 @@ def _configure_remote_connection(cli_context: CLIContext, base_url: str, workspa cli_context.reset_sdk_context() +def _local_context_name(config_file: ConfigFile) -> str: + """Reuse a compatible local context or choose a name that cannot overwrite one.""" + for context in config_file.contexts: + cluster = next(cluster for cluster in config_file.clusters if cluster.name == context.cluster) + user = next(user for user in config_file.users if user.name == context.user) + if str(cluster.base_url).rstrip("/") == DEFAULT_BASE_URL.rstrip("/") and isinstance(user, NoAuthUser): + return context.name + + existing_names = {context.name for context in config_file.contexts} + if _LOCAL_CONTEXT_NAME not in existing_names: + return _LOCAL_CONTEXT_NAME + suffix = 2 + while f"{_LOCAL_CONTEXT_NAME}-{suffix}" in existing_names: + suffix += 1 + return f"{_LOCAL_CONTEXT_NAME}-{suffix}" + + +def _configure_local_connection(cli_context: CLIContext, workspace: str) -> None: + """Activate an isolated no-auth context for local services. + + Reusing an authenticated remote context after changing only its URL leaves + remote OAuth credentials attached to a local cluster. The SDK then tries to + refresh those credentials using the local auth discovery response. Keep the + remote context intact and use a dedicated local context instead. + """ + context_name = _local_context_name(Config.load().get_config_file()) + params: ConfigParams = { + "base_url": DEFAULT_BASE_URL, + "workspace": workspace, + "access_token": None, + "refresh_token": None, + "current_context": context_name, + } + Config.write( + params, + context_name=context_name, + set_current_on_create=True, + ) + cli_context.overrides["base_url"] = DEFAULT_BASE_URL + cli_context.overrides["current_context"] = context_name + cli_context.reset_sdk_context() + + def _ensure_platform_auth(cli_context: CLIContext) -> None: """Authenticate the active context when it lacks usable credentials.""" from nemo_platform.cli.commands.auth import _login_with_oidc, _runtime_token_source_label @@ -714,12 +766,19 @@ def _wait_for_models_impl( return [] -def _get_all_model_entity_ids(client: NeMoPlatform, workspace: str) -> list[str]: - """Return all model entity IDs across all providers.""" +def _get_all_model_entity_ids( + client: NeMoPlatform, + workspace: str, + *, + provider_name: str | None = None, +) -> list[str]: + """Return model entity IDs, optionally scoped to one provider.""" entity_ids: list[str] = [] try: page = client.inference.providers.list(workspace=workspace) for provider in page.data: + if provider_name is not None and getattr(provider, "name", None) != provider_name: + continue for model in getattr(provider, "served_models", None) or []: if hasattr(model, "model_entity_id") and model.model_entity_id: entity_ids.append(model.model_entity_id) @@ -728,17 +787,24 @@ def _get_all_model_entity_ids(client: NeMoPlatform, workspace: str) -> list[str] return sorted(set(entity_ids)) -def _get_all_model_choices(client: NeMoPlatform, workspace: str) -> list[tuple[str, str]]: - """Return picker choices as (entity_id, label) across all providers.""" +def _get_all_model_choices( + client: NeMoPlatform, + workspace: str, + *, + provider_name: str | None = None, +) -> list[tuple[str, str]]: + """Return picker choices, optionally scoped to one provider.""" choices: list[tuple[str, str]] = [] try: page = client.inference.providers.list(workspace=workspace) for provider in page.data: - provider_name = getattr(provider, "name", "unknown-provider") + current_provider_name = getattr(provider, "name", "unknown-provider") + if provider_name is not None and current_provider_name != provider_name: + continue for model in getattr(provider, "served_models", None) or []: model_entity_id = getattr(model, "model_entity_id", None) if model_entity_id: - label = f"{_display_model_name(model_entity_id)} ({provider_name})" + label = f"{_display_model_name(model_entity_id)} ({current_provider_name})" choices.append((model_entity_id, label)) except Exception: logger.debug("Failed to list model choices", exc_info=True) @@ -1815,24 +1881,42 @@ def _validate_api_key( return KeyValidationResult(passed=True, message=f"Could not validate API key ({exc}).") -def _select_default_model(client: NeMoPlatform, workspace: str) -> str | None: - """Let the user pick a default model from discovered models.""" - display_models = _get_all_model_choices(client, workspace) +def _select_model_pair( + client: NeMoPlatform, + workspace: str, + *, + provider_name: str | None = None, +) -> ModelPair | None: + """Let the user pick default and fast models from one provider.""" + display_models = _get_all_model_choices(client, workspace, provider_name=provider_name) if not display_models: - console.print(f" {WARN} No models discovered yet. You can set a default later.") + console.print(f" {WARN} No models discovered yet. You can select models later.") return None - result = prompt_select( - "Choose your default model:", + first_model = display_models[0][0] + default_model = prompt_select( + "Choose your default model (used for quality-critical agent work):", choices=display_models, + default=first_model, + hint="Press Enter to accept the default.", ) - return result + fast = prompt_select( + "Choose your fast model (used for latency-sensitive agent work):", + choices=display_models, + default=default_model, + hint="Press Enter to reuse the default model.", + ) + return ModelPair(default=default_model, fast=fast) -def _save_default_model(cli_context: CLIContext, model_entity_id: str) -> None: - """Persist the default model to the CLI config file.""" +def _save_model_pair(cli_context: CLIContext, model_pair: ModelPair) -> None: + """Persist the default quality model and the low-latency model.""" context = cli_context.get_sdk_context() - Config.write({"default_model": model_entity_id}, context_name=context.context_name) + params: ConfigParams = { + "default_model": model_pair.default, + "fast_model": model_pair.fast, + } + Config.write(params, context_name=context.context_name) def _check_ollama_running(host_url: str) -> bool: @@ -1849,8 +1933,8 @@ def _check_ollama_running(host_url: str) -> bool: # --------------------------------------------------------------------------- -def _auto_setup(client: NeMoPlatform, workspace: str) -> bool: - """Register a provider from environment variables. Returns True on success.""" +def _auto_setup(client: NeMoPlatform, workspace: str) -> str | None: + """Register a provider from environment variables and return its name.""" for key_var, url_var in _AUTO_ENV_VARS: api_key = os.environ.get(key_var) if not api_key: @@ -1923,9 +2007,9 @@ def _auto_setup(client: NeMoPlatform, workspace: str) -> bool: ) console.print(f" {CHECK} Registered provider '{provider_name}' ({host_url})") - return True + return provider_name - return False + return None # --------------------------------------------------------------------------- @@ -2001,8 +2085,8 @@ def setup_command( Uses an already-running platform, starts local services, or connects the CLI to an existing remote deployment. Then selects and registers an - inference provider, picks a default model, installs coding agent skills, - and optionally deploys a demo agent. + inference provider, picks default and fast agent models, installs coding + agent skills, and optionally deploys a demo agent. The active config context remembers the Platform URL. When a remote deployment is already reachable, setup asks whether to continue with it, @@ -2020,7 +2104,7 @@ def setup_command( Use --auto for non-interactive setup from environment variables (NEMO_DEFAULT_INFERENCE_KEY, NVIDIA_API_KEY, OPENAI_API_KEY, ANTHROPIC_API_KEY, GEMINI_API_KEY). - Override the default model with NEMO_DEFAULT_MODEL. + Override the selected pair with NEMO_DEFAULT_MODEL and NEMO_FAST_MODEL. Examples: nemo setup @@ -2048,11 +2132,7 @@ def setup_command( configured_base_url = base_url service_result = _maybe_start_services(base_url, auto, start_services, timeout=effective_timeout) if service_result == "start_local": - context_name = cli_context.get_sdk_context().context_name - _bootstrap_config_if_missing(DEFAULT_BASE_URL, workspace) - Config.write({"base_url": DEFAULT_BASE_URL}, context_name=context_name) - cli_context.overrides["base_url"] = DEFAULT_BASE_URL - cli_context.reset_sdk_context() + _configure_local_connection(cli_context, workspace) base_url = DEFAULT_BASE_URL service_result = _maybe_start_services( base_url, @@ -2080,7 +2160,7 @@ def setup_command( # Ensure the config file exists on disk so later Config.write() calls # (e.g. saving the default model) can find the cluster and context. # Without this, a fresh install (no config.yaml) hits "Cluster - # 'default-cluster' does not exist" when _save_default_model runs. + # 'default-cluster' does not exist" when _save_model_pair runs. _bootstrap_config_if_missing(base_url, workspace) cli_context.reset_sdk_context() @@ -2156,7 +2236,8 @@ def _run_auto_mode( ) -> None: """Non-interactive provider registration from environment variables.""" console.print("[bold]Auto-detecting provider from environment...[/bold]\n") - if not _auto_setup(client, workspace): + provider_name = _auto_setup(client, workspace) + if provider_name is None: console.print(f"{CROSS} No provider credentials found in environment.") env_var_names = ", ".join(key for key, _ in _AUTO_ENV_VARS) console.print(f" Set one of: {env_var_names}") @@ -2167,7 +2248,7 @@ def _run_auto_mode( for attempt in range(_MODEL_DISCOVERY_MAX_ROUNDS): deadline = time.time() + _MODEL_DISCOVERY_ROUND_SECONDS while time.time() < deadline: - entity_ids = _get_all_model_entity_ids(client, workspace) + entity_ids = _get_all_model_entity_ids(client, workspace, provider_name=provider_name) if entity_ids: break _pause(_MODEL_DISCOVERY_POLL_INTERVAL) @@ -2189,14 +2270,20 @@ def _run_auto_mode( default_model = os.environ.get("NEMO_DEFAULT_MODEL", "").strip() if not default_model and entity_ids: default_model = entity_ids[0] + fast_model = os.environ.get("NEMO_FAST_MODEL", "").strip() or default_model if default_model: - _save_default_model(cli_context, default_model) - console.print(f" {CHECK} Default model: {default_model}") + model_pair = ModelPair(default=default_model, fast=fast_model) + _save_model_pair(cli_context, model_pair) + console.print(f" {CHECK} Default model: {model_pair.default}") + console.print(f" {CHECK} Fast model: {model_pair.fast}") else: + if fast_model: + console.print(f" {WARN} NEMO_FAST_MODEL is ignored until a default model is available") console.print(f" {WARN} No default model set (no models discovered yet)") - console.print(" Run [cyan]nemo setup[/cyan] again after models sync, or set via env var:") + console.print(" Run [cyan]nemo setup[/cyan] again after models sync, or set the models via env vars:") console.print(" [cyan]export NEMO_DEFAULT_MODEL=[/cyan]") + console.print(" [cyan]export NEMO_FAST_MODEL=[/cyan]") _maybe_install_skills( auto=True, @@ -2273,17 +2360,19 @@ def _run_interactive_mode( else: console.print(f" {WARN} No models discovered yet (provider may still be syncing)") - console.print("\n[bold]Step 5: Choose default model[/bold]\n") + console.print("\n[bold]Step 5: Choose agent models[/bold]\n") fallback_model_choices = _get_all_model_choices(client, workspace) if not models else [] if not models and fallback_model_choices: console.print(f" {WARN} Models from existing providers are available, but not from '{provider_name}' yet.") - default_model = _select_default_model(client, workspace) if models else None - if default_model: - _save_default_model(cli_context, default_model) - console.print(f" {CHECK} Default model set to {_display_model_name(default_model)}") + model_pair = _select_model_pair(client, workspace, provider_name=provider_name) if models else None + default_model = model_pair.default if model_pair else None + if model_pair: + _save_model_pair(cli_context, model_pair) + console.print(f" {CHECK} Default model set to {_display_model_name(model_pair.default)}") + console.print(f" {CHECK} Fast model set to {_display_model_name(model_pair.fast)}") else: - console.print(f" {WARN} No default model set for this provider yet.") + console.print(f" {WARN} No agent models set for this provider yet.") console.print(" Run [cyan]nemo setup[/cyan] again after models sync") console.print("\n[bold]Step 6: Install skills[/bold]\n") @@ -2305,7 +2394,13 @@ def _run_interactive_mode( headers=_platform_request_headers(cli_context), ) - _print_onboarding(base_url, provider_name, default_model, demo_deployed=demo_deployed) + _print_onboarding( + base_url, + provider_name, + default_model, + fast_model=model_pair.fast if model_pair else None, + demo_deployed=demo_deployed, + ) except UserCancelled: console.print(f"\n{WARN} Setup cancelled.") @@ -2372,6 +2467,7 @@ def _print_onboarding( provider_name: str, default_model: str | None, *, + fast_model: str | None = None, demo_deployed: bool = False, ) -> None: """Print setup summary, then present goal-oriented onboarding paths.""" @@ -2382,6 +2478,8 @@ def _print_onboarding( console.print(f" Provider: {provider_name}") if default_model: console.print(f" Default model: {_display_model_name(default_model)}") + if fast_model: + console.print(f" Fast model: {_display_model_name(fast_model)}") if demo_deployed: console.print(f" Demo agent: {_DEMO_AGENT_NAME}") diff --git a/sdk/python/nemo-platform/src/nemo_platform/config/config.py b/sdk/python/nemo-platform/src/nemo_platform/config/config.py index 8c5f3e4b20..7eef6c623d 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/config/config.py +++ b/sdk/python/nemo-platform/src/nemo_platform/config/config.py @@ -63,6 +63,13 @@ class _RuntimeAccessTokenSource: label: str +def _resolve_model_pair(default_model: str | None, fast_model: str | None) -> tuple[str | None, str | None]: + """Apply environment overrides and the fast-to-default fallback consistently.""" + effective_default = os.environ.get("NEMO_DEFAULT_MODEL") or default_model + effective_fast = os.environ.get("NEMO_FAST_MODEL") or fast_model or effective_default + return effective_default, effective_fast + + class Config(BaseModel): """ Configuration manager for nemo_platform. @@ -72,10 +79,10 @@ class Config(BaseModel): 2. Runtime overrides (from env vars, CLI, or code) 3. Resolution of effective configuration - Environment variables (prefix NMP_): + Environment variables: NMP_CURRENT_CONTEXT, NMP_WORKSPACE, NMP_OUTPUT_FORMAT, NMP_TIMESTAMP_FORMAT, NMP_PAGE_SIZE, NMP_COLOR_OUTPUT, - NMP_BASE_URL, NMP_ACCESS_TOKEN + NMP_BASE_URL, NMP_ACCESS_TOKEN, NEMO_DEFAULT_MODEL, NEMO_FAST_MODEL """ # Runtime overrides - can be set via env vars @@ -486,8 +493,12 @@ def resolve(self) -> Context: if effective_workspace is None: effective_workspace = DEFAULT_WORKSPACE - # Resolve default model: env var > config file - effective_default_model = os.environ.get("NEMO_DEFAULT_MODEL") or context.default_model + # The default model is the quality-oriented model for agent workloads. + # Existing single-model contexts also supply the fast role. + effective_default_model, effective_fast_model = _resolve_model_pair( + context.default_model, + context.fast_model, + ) # Create the resolved Context (runtime model) from the ContextDefinition (config file model) return Context( @@ -496,6 +507,7 @@ def resolve(self) -> Context: user=user, # Fully resolved User with authentication credentials workspace=effective_workspace, default_model=effective_default_model, + fast_model=effective_fast_model, preferences=prefs, ) @@ -545,12 +557,17 @@ def _create_default_config(self) -> Context: prefs.color_output = self.color_output # Return resolved Context (runtime model) + effective_default_model, effective_fast_model = _resolve_model_pair( + context_def.default_model, + context_def.fast_model, + ) return Context( context_name=context_name, cluster=cluster, user=user, workspace=context_def.workspace or DEFAULT_WORKSPACE, - default_model=os.environ.get("NEMO_DEFAULT_MODEL") or context_def.default_model, + default_model=effective_default_model, + fast_model=effective_fast_model, preferences=prefs, ) diff --git a/sdk/python/nemo-platform/src/nemo_platform/config/models.py b/sdk/python/nemo-platform/src/nemo_platform/config/models.py index f9a6500686..f94ac4982c 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/config/models.py +++ b/sdk/python/nemo-platform/src/nemo_platform/config/models.py @@ -140,6 +140,7 @@ class ContextDefinition(BaseModel): user: str = Field(..., min_length=1, description="Reference to user name (string, not resolved)") workspace: str | None = Field(default=None, description="Default workspace") default_model: str | None = Field(default=None, description="Default model entity ID for inference") + fast_model: str | None = Field(default=None, description="Low-latency model entity ID for agent workloads") preferences: Preferences = Field(default_factory=Preferences, description="Context-specific preferences") metadata: dict[str, Any] = Field(default_factory=dict, description="Additional context metadata") @@ -174,6 +175,7 @@ class ConfigParams(TypedDict, total=False): workspace: str default_model: str + fast_model: str output_format: OutputFormat timestamp_format: TimestampFormat @@ -316,6 +318,8 @@ def ensure_context( context.workspace = params["workspace"] if "default_model" in params: context.default_model = params["default_model"] + if "fast_model" in params: + context.fast_model = params["fast_model"] if "output_format" in params: context.preferences.output_format = params["output_format"] if "timestamp_format" in params: @@ -383,4 +387,5 @@ class Context(BaseModel): user: User | None = Field(default=None, description="Resolved user with authentication credentials") workspace: str = Field(..., description="Active workspace (required)") default_model: str | None = Field(default=None, description="Default model entity ID for inference") + fast_model: str | None = Field(default=None, description="Low-latency model entity ID for agent workloads") preferences: Preferences = Field(..., description="Effective preferences") diff --git a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_setup.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_setup.py index d237db00cf..14980e8b20 100644 --- a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_setup.py +++ b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_setup.py @@ -34,6 +34,7 @@ KNOWN_PROVIDERS, ONBOARDING_PATHS, KeyValidationResult, + ModelPair, _agent_config_path, _agents_plugin_available, _auto_setup, @@ -42,6 +43,7 @@ _check_ollama_running, _check_platform_reachable, _check_platform_reachable_with_retries, + _configure_local_connection, _create_provider, _deploy_demo_agent, _detect_coding_agents, @@ -62,9 +64,10 @@ _render_onboarding_card, _resolve_provider_for_url, _resolve_setup_workspace, + _run_auto_mode, _run_interactive_mode, _save_data_dir, - _select_default_model, + _select_model_pair, _start_services_background, _validate_api_key, _verify_platform_health, @@ -419,8 +422,8 @@ def test_no_auth_header_format_skips_auth_fields(self): workspace="default", ) call_kwargs = client.inference.providers.create.call_args.kwargs - assert "auth_header_format" not in call_kwargs assert "required_extra_headers" not in call_kwargs + assert "auth_header_format" not in call_kwargs def test_provider_without_secret(self): """Providers without secrets (e.g. Ollama) should omit api_key_secret_name.""" @@ -455,16 +458,16 @@ def _bypass_key_validation(): @pytest.mark.usefixtures("_bypass_key_validation") class TestAutoSetup: - def test_no_env_vars_returns_false(self): + def test_no_env_vars_returns_none(self): client = _make_mock_client() with patch.dict("os.environ", {}, clear=True): - assert _auto_setup(client, "default") is False + assert _auto_setup(client, "default") is None def test_openai_key_creates_provider(self): client = _make_mock_client() with patch.dict("os.environ", {"OPENAI_API_KEY": "sk-test123"}, clear=True): result = _auto_setup(client, "default") - assert result is True + assert result == "openai" client.mock_secrets.create_secret.assert_called_once() client.inference.providers.create.assert_called_once() create_kwargs = client.inference.providers.create.call_args @@ -475,7 +478,7 @@ def test_nvidia_key_creates_build_provider(self): client = _make_mock_client() with patch.dict("os.environ", {"NVIDIA_API_KEY": "nvapi-test"}, clear=True): result = _auto_setup(client, "default") - assert result is True + assert result == "nvidia-build" create_kwargs = client.inference.providers.create.call_args assert create_kwargs.kwargs["name"] == "nvidia-build" @@ -487,7 +490,7 @@ def test_nemo_default_key_with_url(self): } with patch.dict("os.environ", env, clear=True): result = _auto_setup(client, "default") - assert result is True + assert result == "openai" create_kwargs = client.inference.providers.create.call_args assert create_kwargs.kwargs["name"] == "openai" @@ -499,7 +502,7 @@ def test_nemo_default_key_with_custom_url(self): } with patch.dict("os.environ", env, clear=True): result = _auto_setup(client, "default") - assert result is True + assert result == "my-custom-llm-example-com" create_kwargs = client.inference.providers.create.call_args assert create_kwargs.kwargs["name"] == "my-custom-llm-example-com" assert create_kwargs.kwargs["host_url"] == "https://my-custom-llm.example.com/v1" @@ -508,7 +511,7 @@ def test_existing_provider_updated_not_recreated(self): client = _make_mock_client(provider_exists=True, secret_exists=True) with patch.dict("os.environ", {"OPENAI_API_KEY": "sk-test123"}, clear=True): result = _auto_setup(client, "default") - assert result is True + assert result == "openai" client.inference.providers.create.assert_not_called() client.inference.providers.update.assert_called_once() @@ -522,7 +525,7 @@ def test_priority_order(self): } with patch.dict("os.environ", env, clear=True): result = _auto_setup(client, "default") - assert result is True + assert result == "anthropic" create_kwargs = client.inference.providers.create.call_args assert create_kwargs.kwargs["name"] == "anthropic" @@ -532,7 +535,7 @@ def test_anthropic_auto_setup_maps_auth_header(self): api_key = "sk-ant-test" with patch.dict("os.environ", {"ANTHROPIC_API_KEY": api_key}, clear=True): result = _auto_setup(client, "default") - assert result is True + assert result == "anthropic" call_kwargs = client.inference.providers.create.call_args.kwargs assert call_kwargs["name"] == "anthropic" assert call_kwargs["api_key_secret_name"] == "anthropic-api-key" @@ -554,9 +557,25 @@ def test_context_definition_has_default_model(self): ctx2 = ContextDefinition(name="test", cluster="c", user="u", default_model="my-model") assert ctx2.default_model == "my-model" + def test_context_definition_has_optional_fast_model(self): + ctx = ContextDefinition(name="test", cluster="c", user="u") + assert ctx.fast_model is None + + configured = ContextDefinition( + name="test", + cluster="c", + user="u", + fast_model="workspace/fast", + ) + assert configured.fast_model == "workspace/fast" + def test_config_params_accepts_default_model(self): - params: ConfigParams = {"default_model": "workspace/openai-gpt-4o"} + params: ConfigParams = { + "default_model": "workspace/openai-gpt-4o", + "fast_model": "workspace/openai-gpt-4o-mini", + } assert params["default_model"] == "workspace/openai-gpt-4o" + assert params["fast_model"] == "workspace/openai-gpt-4o-mini" def test_ensure_context_applies_default_model(self): config_file = ConfigFile( @@ -565,9 +584,13 @@ def test_ensure_context_applies_default_model(self): users=[{"name": "test-user", "type": "no-auth"}], contexts=[{"name": "test", "cluster": "test-cluster", "user": "test-user"}], ) - params: ConfigParams = {"default_model": "workspace/my-model"} + params: ConfigParams = { + "default_model": "workspace/my-model", + "fast_model": "workspace/fast", + } _, _, ctx_def = config_file.ensure_context("test", params) assert ctx_def.default_model == "workspace/my-model" + assert ctx_def.fast_model == "workspace/fast" def test_context_has_default_model(self): ctx = Context( @@ -575,9 +598,79 @@ def test_context_has_default_model(self): cluster=Cluster(name="c", base_url="http://localhost:8080"), workspace="default", default_model="workspace/gpt-4o", + fast_model="workspace/gpt-4o-mini", preferences={}, ) assert ctx.default_model == "workspace/gpt-4o" + assert ctx.fast_model == "workspace/gpt-4o-mini" + + def test_default_model_fills_missing_fast_role(self, tmp_path, monkeypatch): + monkeypatch.delenv("NEMO_DEFAULT_MODEL", raising=False) + monkeypatch.delenv("NEMO_FAST_MODEL", raising=False) + config_file = ConfigFile( + current_context="test", + clusters=[{"name": "test-cluster", "base_url": "http://localhost:8080"}], + users=[{"name": "test-user", "type": "no-auth"}], + contexts=[ + { + "name": "test", + "cluster": "test-cluster", + "user": "test-user", + "default_model": "workspace/legacy-model", + } + ], + ) + + context = Config.create(tmp_path / "config.yaml", config_file).resolve() + + assert context.default_model == "workspace/legacy-model" + assert context.fast_model == "workspace/legacy-model" + + def test_explicit_model_pair_resolves_independently(self, tmp_path, monkeypatch): + monkeypatch.delenv("NEMO_DEFAULT_MODEL", raising=False) + monkeypatch.delenv("NEMO_FAST_MODEL", raising=False) + config_file = ConfigFile( + current_context="test", + clusters=[{"name": "test-cluster", "base_url": "http://localhost:8080"}], + users=[{"name": "test-user", "type": "no-auth"}], + contexts=[ + { + "name": "test", + "cluster": "test-cluster", + "user": "test-user", + "default_model": "workspace/default", + "fast_model": "workspace/fast", + } + ], + ) + + context = Config.create(tmp_path / "config.yaml", config_file).resolve() + + assert context.default_model == "workspace/default" + assert context.fast_model == "workspace/fast" + + def test_model_environment_overrides_resolve_independently(self, tmp_path, monkeypatch): + monkeypatch.setenv("NEMO_DEFAULT_MODEL", "workspace/env-default") + monkeypatch.setenv("NEMO_FAST_MODEL", "workspace/env-fast") + config_file = ConfigFile( + current_context="test", + clusters=[{"name": "test-cluster", "base_url": "http://localhost:8080"}], + users=[{"name": "test-user", "type": "no-auth"}], + contexts=[ + { + "name": "test", + "cluster": "test-cluster", + "user": "test-user", + "default_model": "workspace/config-default", + "fast_model": "workspace/config-fast", + } + ], + ) + + context = Config.create(tmp_path / "config.yaml", config_file).resolve() + + assert context.default_model == "workspace/env-default" + assert context.fast_model == "workspace/env-fast" # --------------------------------------------------------------------------- @@ -918,6 +1011,72 @@ def test_updates_selected_context_and_runtime_url_override(self, tmp_path, monke assert str(dev_cluster.base_url) == "https://new-dev.example.com/" assert cli_context.overrides["base_url"] == "https://new-dev.example.com" + def test_local_connection_preserves_remote_context_and_uses_no_auth(self, tmp_path, monkeypatch): + config_path = tmp_path / "config.yaml" + monkeypatch.setenv("NMP_CONFIG_FILE", str(config_path)) + Config.write( + { + "base_url": "https://remote.example.com", + "access_token": "remote-token", + "refresh_token": "remote-refresh-token", + "current_context": "dev", + }, + context_name="dev", + ) + cli_context = MagicMock() + cli_context.overrides = { + "current_context": "dev", + "base_url": "https://remote.example.com", + } + + _configure_local_connection(cli_context, "local-workspace") + + config = Config.load(config_path=config_path) + config_file = config.get_config_file() + assert config_file.current_context == "local" + remote = next(context for context in config_file.contexts if context.name == "dev") + remote_cluster = next(cluster for cluster in config_file.clusters if cluster.name == remote.cluster) + remote_user = next(user for user in config_file.users if user.name == remote.user) + assert str(remote_cluster.base_url) == "https://remote.example.com/" + assert isinstance(remote_user, OAuthUser) + + local = config.resolve() + assert local.context_name == "local" + assert str(local.cluster.base_url) == "http://localhost:8080/" + assert local.workspace == "local-workspace" + assert isinstance(local.user, NoAuthUser) + assert cli_context.overrides == { + "current_context": "local", + "base_url": "http://localhost:8080", + } + cli_context.reset_sdk_context.assert_called_once_with() + + def test_local_connection_does_not_overwrite_existing_local_context(self, tmp_path, monkeypatch): + config_path = tmp_path / "config.yaml" + monkeypatch.setenv("NMP_CONFIG_FILE", str(config_path)) + Config.write( + { + "base_url": "https://remote.example.com", + "access_token": "remote-token", + "current_context": "local", + }, + context_name="local", + ) + cli_context = MagicMock() + cli_context.overrides = {} + + _configure_local_connection(cli_context, "local-workspace") + + config_file = Config.load(config_path=config_path).get_config_file() + original = next(context for context in config_file.contexts if context.name == "local") + original_cluster = next(cluster for cluster in config_file.clusters if cluster.name == original.cluster) + assert str(original_cluster.base_url) == "https://remote.example.com/" + assert config_file.current_context == "local-2" + assert cli_context.overrides == { + "current_context": "local-2", + "base_url": "http://localhost:8080", + } + def test_preserves_active_workspace_when_flag_not_explicit(self): ctx = MagicMock(spec=typer.Context) ctx.get_parameter_source.return_value = ParameterSource.DEFAULT @@ -1080,7 +1239,7 @@ def test_bootstrap_seeds_cluster_when_only_data_dir_persisted(self, tmp_path, mo ``_bootstrap_config_if_missing`` used to short-circuit on ``config_path.exists()``, leaving the file without any cluster. Later - steps (`_save_default_model` → ``Config.write`` → ``ensure_context``) + steps (`_save_model_pair` → ``Config.write`` → ``ensure_context``) then raised ``Cluster 'default-cluster' does not exist and no base_url provided to create it!`` when the user picked a default model. @@ -1719,7 +1878,7 @@ def test_auto_setup_updates_existing_anthropic_auth(self): api_key = "sk-ant-updated" with patch.dict("os.environ", {"ANTHROPIC_API_KEY": api_key}, clear=True): result = _auto_setup(client, "default") - assert result is True + assert result == "anthropic" call_kwargs = client.inference.providers.update.call_args.kwargs assert call_kwargs["api_key_secret_name"] == "anthropic-api-key" assert call_kwargs["auth_header_format"] == "X-Api-Key: {{ auth_secret }}" @@ -1733,7 +1892,7 @@ def test_auto_setup_updates_existing_provider_secret_binding(self): client = _make_mock_client(provider_exists=True, secret_exists=False) with patch.dict("os.environ", {"OPENAI_API_KEY": "sk-new-key"}, clear=True): result = _auto_setup(client, "default") - assert result is True + assert result == "openai" client.mock_secrets.create_secret.assert_called_once() client.inference.providers.update.assert_called_once() call_kwargs = client.inference.providers.update.call_args.kwargs @@ -1744,7 +1903,7 @@ def test_auto_setup_updates_existing_secret_value(self): client = _make_mock_client(provider_exists=True, secret_exists=True) with patch.dict("os.environ", {"NVIDIA_API_KEY": "nvapi-new"}, clear=True): result = _auto_setup(client, "default") - assert result is True + assert result == "nvidia-build" client.mock_secrets.update_secret.assert_called_once() update_call = client.mock_secrets.update_secret.call_args assert update_call.kwargs["name"] == "nvidia-build-api-key" @@ -1759,9 +1918,47 @@ def test_auto_setup_updates_existing_secret_value(self): # --------------------------------------------------------------------------- -class TestInteractiveDefaultModelSelection: +class TestInteractiveModelPairSelection: _MOD = "nemo_platform.cli.commands.setup" + def test_run_scopes_picker_to_registered_provider(self): + client = MagicMock() + cli_context = MagicMock() + model_pair = ModelPair( + default="default/claude-sonnet-4-6", + fast="default/claude-haiku-4-5-20251001", + ) + + with ( + patch( + f"{self._MOD}._interactive_collect_provider", + return_value=("anthropic", "https://api.anthropic.com", None, None, None), + ), + patch(f"{self._MOD}._register_provider_interactive"), + patch( + f"{self._MOD}._wait_for_models", + return_value=[ + "default/claude-sonnet-4-6", + "default/claude-haiku-4-5-20251001", + ], + ), + patch(f"{self._MOD}._select_model_pair", return_value=model_pair) as select_model_pair, + patch(f"{self._MOD}._save_model_pair"), + patch(f"{self._MOD}._maybe_install_skills"), + patch(f"{self._MOD}._maybe_deploy_agent", return_value=False), + patch(f"{self._MOD}._print_onboarding"), + ): + _run_interactive_mode( + cli_context, + client, + "default", + "http://localhost:8080", + install_skills=False, + deploy_agent=False, + ) + + select_model_pair.assert_called_once_with(client, "default", provider_name="anthropic") + def test_skips_default_model_picker_when_new_provider_has_no_models(self): """When the new provider is still syncing, setup should not show a misleading picker.""" client = MagicMock() @@ -1774,8 +1971,8 @@ def test_skips_default_model_picker_when_new_provider_has_no_models(self): ), patch(f"{self._MOD}._register_provider_interactive"), patch(f"{self._MOD}._wait_for_models", return_value=[]), - patch(f"{self._MOD}._select_default_model") as mock_select_default_model, - patch(f"{self._MOD}._save_default_model"), + patch(f"{self._MOD}._select_model_pair") as mock_select_model_pair, + patch(f"{self._MOD}._save_model_pair"), patch(f"{self._MOD}._maybe_install_skills"), patch(f"{self._MOD}._maybe_deploy_agent"), patch(f"{self._MOD}._print_onboarding"), @@ -1790,11 +1987,11 @@ def test_skips_default_model_picker_when_new_provider_has_no_models(self): deploy_agent=False, ) - mock_select_default_model.assert_not_called() + mock_select_model_pair.assert_not_called() printed_lines = [call.args[0] for call in mock_console.print.call_args_list if call.args] assert any("No models discovered yet (provider may still be syncing)" in line for line in printed_lines) - assert any("Step 5: Choose default model" in line for line in printed_lines) - assert any("No default model set for this provider yet." in line for line in printed_lines) + assert any("Step 5: Choose agent models" in line for line in printed_lines) + assert any("No agent models set for this provider yet." in line for line in printed_lines) assert any("Run [cyan]nemo setup[/cyan] again after models sync" in line for line in printed_lines) def test_warns_when_only_existing_provider_models_are_available(self): @@ -1818,8 +2015,8 @@ def test_warns_when_only_existing_provider_models_are_available(self): ) ], ), - patch(f"{self._MOD}._select_default_model") as mock_select_default_model, - patch(f"{self._MOD}._save_default_model"), + patch(f"{self._MOD}._select_model_pair") as mock_select_model_pair, + patch(f"{self._MOD}._save_model_pair"), patch(f"{self._MOD}._maybe_install_skills"), patch(f"{self._MOD}._maybe_deploy_agent"), patch(f"{self._MOD}._print_onboarding"), @@ -1834,15 +2031,15 @@ def test_warns_when_only_existing_provider_models_are_available(self): deploy_agent=False, ) - mock_select_default_model.assert_not_called() + mock_select_model_pair.assert_not_called() printed_lines = [call.args[0] for call in mock_console.print.call_args_list if call.args] assert any( "Models from existing providers are available, but not from 'my-ollama-custom' yet." in line for line in printed_lines ) - def test_picker_labels_include_provider_names(self): - """Default model choices should show which provider each model comes from.""" + def test_picker_only_includes_models_from_selected_provider(self): + """Models from another provider in the workspace must not enter the picker.""" client = MagicMock() provider_a = MagicMock() provider_a.name = "nvidia-build" @@ -1856,20 +2053,167 @@ def test_picker_labels_include_provider_names(self): ] client.inference.providers.list.return_value = MagicMock(data=[provider_a, provider_b]) - with patch(f"{self._MOD}.prompt_select", return_value="default/qwen2.5:1.5b") as mock_prompt_select: - result = _select_default_model(client, "default") + with patch( + f"{self._MOD}.prompt_select", + side_effect=["default/qwen2.5:1.5b", "default/qwen2.5:1.5b"], + ) as mock_prompt_select: + result = _select_model_pair(client, "default", provider_name="my-ollama-custom") - assert result == "default/qwen2.5:1.5b" - assert mock_prompt_select.call_args.kwargs["choices"] == [ - ( - "default/meta-llama-3-1-8b-instruct", - "meta-llama-3-1-8b-instruct (nvidia-build)", - ), + assert result == ModelPair( + default="default/qwen2.5:1.5b", + fast="default/qwen2.5:1.5b", + ) + expected_choices = [ ( "default/qwen2.5:1.5b", "qwen2.5:1.5b (my-ollama-custom)", ), ] + assert mock_prompt_select.call_count == 2 + assert all(call.kwargs["choices"] == expected_choices for call in mock_prompt_select.call_args_list) + default_call, fast_call = mock_prompt_select.call_args_list + assert default_call.args[0] == "Choose your default model (used for quality-critical agent work):" + assert default_call.kwargs["default"] == "default/qwen2.5:1.5b" + assert default_call.kwargs["hint"] == "Press Enter to accept the default." + assert fast_call.args[0] == "Choose your fast model (used for latency-sensitive agent work):" + assert fast_call.kwargs["default"] == "default/qwen2.5:1.5b" + assert fast_call.kwargs["hint"] == "Press Enter to reuse the default model." + + +class TestAutoModelPairSelection: + def test_entity_discovery_filters_other_workspace_providers(self): + client = MagicMock() + openai = MagicMock() + openai.name = "openai" + openai.served_models = [MagicMock(model_entity_id="default/gpt-4.1")] + anthropic = MagicMock() + anthropic.name = "anthropic" + anthropic.served_models = [MagicMock(model_entity_id="default/claude-sonnet-4-6")] + client.inference.providers.list.return_value = MagicMock(data=[openai, anthropic]) + + result = setup_commands._get_all_model_entity_ids( + client, + "default", + provider_name="anthropic", + ) + + assert result == ["default/claude-sonnet-4-6"] + + def test_explicit_default_and_fast_models_are_saved_independently(self): + cli_context = MagicMock() + with ( + patch.dict( + "os.environ", + { + "NEMO_DEFAULT_MODEL": "default/quality", + "NEMO_FAST_MODEL": "default/fast", + }, + clear=True, + ), + patch(f"{SETUP_MOD}._auto_setup", return_value="anthropic"), + patch(f"{SETUP_MOD}._get_all_model_entity_ids", return_value=["default/discovered"]), + patch(f"{SETUP_MOD}._save_model_pair") as save_pair, + patch(f"{SETUP_MOD}._maybe_install_skills"), + patch(f"{SETUP_MOD}._maybe_deploy_agent"), + patch(f"{SETUP_MOD}._verify_platform_health", return_value=True), + ): + _run_auto_mode( + cli_context, + MagicMock(), + "default", + "http://localhost:8080", + install_skills=False, + deploy_agent=False, + ) + + save_pair.assert_called_once_with(cli_context, ModelPair(default="default/quality", fast="default/fast")) + + def test_discovered_default_is_scoped_to_auto_configured_provider(self): + cli_context = MagicMock() + client = MagicMock() + with ( + patch.dict("os.environ", {}, clear=True), + patch(f"{SETUP_MOD}._auto_setup", return_value="anthropic"), + patch( + f"{SETUP_MOD}._get_all_model_entity_ids", + return_value=["default/claude-sonnet-4-6"], + ) as get_model_ids, + patch(f"{SETUP_MOD}._save_model_pair") as save_pair, + patch(f"{SETUP_MOD}._maybe_install_skills"), + patch(f"{SETUP_MOD}._maybe_deploy_agent"), + patch(f"{SETUP_MOD}._verify_platform_health", return_value=True), + ): + _run_auto_mode( + cli_context, + client, + "default", + "http://localhost:8080", + install_skills=False, + deploy_agent=False, + ) + + get_model_ids.assert_called_once_with(client, "default", provider_name="anthropic") + save_pair.assert_called_once_with( + cli_context, + ModelPair( + default="default/claude-sonnet-4-6", + fast="default/claude-sonnet-4-6", + ), + ) + + def test_first_discovered_model_is_default_and_fast_fallback(self): + cli_context = MagicMock() + with ( + patch.dict("os.environ", {}, clear=True), + patch(f"{SETUP_MOD}._auto_setup", return_value="openai"), + patch( + f"{SETUP_MOD}._get_all_model_entity_ids", + return_value=["default/a-model", "default/z-model"], + ), + patch(f"{SETUP_MOD}._save_model_pair") as save_pair, + patch(f"{SETUP_MOD}._maybe_install_skills"), + patch(f"{SETUP_MOD}._maybe_deploy_agent"), + patch(f"{SETUP_MOD}._verify_platform_health", return_value=True), + ): + _run_auto_mode( + cli_context, + MagicMock(), + "default", + "http://localhost:8080", + install_skills=False, + deploy_agent=False, + ) + + save_pair.assert_called_once_with( + cli_context, + ModelPair(default="default/a-model", fast="default/a-model"), + ) + + def test_fast_override_without_default_warns_and_is_not_saved(self): + cli_context = MagicMock() + with ( + patch.dict("os.environ", {"NEMO_FAST_MODEL": "default/fast"}, clear=True), + patch(f"{SETUP_MOD}._MODEL_DISCOVERY_MAX_ROUNDS", 1), + patch(f"{SETUP_MOD}._MODEL_DISCOVERY_ROUND_SECONDS", 0), + patch(f"{SETUP_MOD}._auto_setup", return_value="openai"), + patch(f"{SETUP_MOD}._get_all_model_entity_ids", return_value=[]), + patch(f"{SETUP_MOD}._save_model_pair") as save_pair, + patch(f"{SETUP_MOD}._maybe_install_skills"), + patch(f"{SETUP_MOD}._maybe_deploy_agent"), + patch(f"{SETUP_MOD}._verify_platform_health", return_value=True), + patch(f"{SETUP_MOD}.console.print") as print_message, + ): + _run_auto_mode( + cli_context, + MagicMock(), + "default", + "http://localhost:8080", + install_skills=False, + deploy_agent=False, + ) + + save_pair.assert_not_called() + assert any("NEMO_FAST_MODEL is ignored" in call.args[0] for call in print_message.call_args_list) # --------------------------------------------------------------------------- @@ -2756,6 +3100,7 @@ def _patch_setup_command( run_interactive=stack.enter_context(patch(f"{SETUP_MOD}._run_interactive_mode")), prompt_remote=None, configure_remote=None, + configure_local=stack.enter_context(patch(f"{SETUP_MOD}._configure_local_connection")), ensure_auth=None, config_write=None, run_auto=None, @@ -2850,7 +3195,6 @@ def test_start_local_persists_default_url_and_starts_services(self): ctx, cli_context = _make_setup_command_ctx(base_url="https://remote.example.com") with _patch_setup_command( maybe_start_services=["start_local", "ready"], - include_config_write=True, ) as mocks: setup_command(ctx) @@ -2861,10 +3205,7 @@ def test_start_local_persists_default_url_and_starts_services(self): start_services=True, timeout=_SERVICE_STARTUP_TIMEOUT_SECONDS, ) - mocks.bootstrap.assert_any_call(DEFAULT_BASE_URL, "default") - mocks.config_write.assert_called_once_with({"base_url": DEFAULT_BASE_URL}, context_name="default") - assert cli_context.overrides["base_url"] == DEFAULT_BASE_URL - cli_context.reset_sdk_context.assert_called() + mocks.configure_local.assert_called_once_with(cli_context, "default") assert mocks.run_interactive.call_args.args[3] == DEFAULT_BASE_URL