[khazic]

What does this PR do ?

Adds an OpenAI-compatible POST /v1/embeddings endpoint to the Model Server, so memory- and retrieval-based agents can route embedding calls through Gym instead of calling a provider directly.

• Base model server (nemo_gym/base_responses_api_model.py): registers the /v1/embeddings route and adds a non-abstract embeddings() default that returns 501. This mirrors the existing messages() precedent (uniform route surface: every model server advertises the same endpoints; unsupported ones fail at call time). The route is declared non-generating, so an embeddings call under a rollout's token-capture prefix does not mark that rollout incomplete.
• openai_model and inference_provider: override embeddings() as a passthrough to the provider's /v1/embeddings. A single OpenAI-compatible adapter covers OpenAI, NVIDIA's hosted endpoint (integrate.api.nvidia.com/v1), and any other OpenAI-compatible provider (including vLLM served as an embedding model). openai_model sends the call through _call_upstream, so its upstream retry policy and concurrency slots apply as they do for chat and responses.
• NeMoGymAsyncOpenAI gains create_embeddings(), alongside create_chat_completion/create_response/create_tokenize, all going through the global aiohttp client.

The passthrough merges the server's extra_body under the request (a field the request sets wins), replaces model with the configured one, and returns the provider's JSON unchanged. So encoding_format: base64 (the OpenAI SDK's default on the wire) reaches the provider and comes back as base64, and a provider that omits optional fields such as usage still returns 200.

Embeddings calls are not recorded by model-call capture, which covers only /v1/responses, /v1/chat/completions and /v1/messages; the docs say so.

An embedding model is usually distinct from the policy model, so it wires as a separate Model Server instance referenced via ModelServerRef, matching the existing one-server-one-model pattern (no schema change).

Issues

Closes #1943

Usage

gym env start --resources-server <your_resources_server> --model-type openai_model

curl -s http://127.0.0.1:<port>/v1/embeddings \
  -H 'Content-Type: application/json' \
  -d '{"input": "hello world", "model": "text-embedding-3-small"}'

Tests

• tests/unit_tests/test_base_responses_api_model.py: base default returns 501.
• responses_api_models/openai_model/tests/test_app.py and responses_api_models/inference_provider/tests/test_app.py: the provider JSON comes back unchanged (base64 vectors, no usage), the configured model replaces the request's, encoding_format is forwarded as sent, request fields win over extra_body and the config is not mutated, and openai_model retries through upstream_retry_policy.
• tests/unit_tests/test_openai_utils.py: create_embeddings posts to {base_url}/embeddings with the kwargs as the JSON body.
• tests/unit_tests/test_token_id_capture.py: POST /ng-rollout/<id>/training-token-capture/v1/embeddings leaves the rollout complete (fails without the non-generating declaration).

The affected test files pass (731 tests); ruff check/ruff format (v0.9.9) clean.

Real-provider smoke

Both model servers ran in their own process and pointed at a vLLM 0.28 server hosting Qwen/Qwen3-Embedding-0.6B (vllm serve ... --runner pooling --convert embed). Each was queried with curl (float and base64) and with the OpenAI Python SDK, which sends encoding_format=base64 and decodes the reply itself.

Smoke output

=== openai_model ===
$ curl -s http://127.0.0.1:<port>/v1/embeddings -H 'Content-Type: application/json' -d '{"input": "hello world", "model": "text-embedding-3-small"}'
HTTP 200  model=qwen3-embedding-0.6b items=1 encoding=float vector_length=1024 usage={'prompt_tokens': 3, 'total_tokens': 3, 'completion_tokens': 0, 'prompt_tokens_details': None, 'completion_tokens_details': None}
$ curl -s http://127.0.0.1:<port>/v1/embeddings -H 'Content-Type: application/json' -d '{"input": "hello world", "model": "text-embedding-3-small", "encoding_format": "base64"}'
HTTP 200  model=qwen3-embedding-0.6b items=1 encoding=base64 vector_length=1024 usage={'prompt_tokens': 3, 'total_tokens': 3, 'completion_tokens': 0, 'prompt_tokens_details': None, 'completion_tokens_details': None}
base64 decodes to the float vector: True
OpenAI SDK (default wire encoding base64): items=2 vector_length=1024 matches_float_request=True

=== inference_provider ===
$ curl -s http://127.0.0.1:<port>/v1/embeddings -H 'Content-Type: application/json' -d '{"input": "hello world", "model": "text-embedding-3-small"}'
HTTP 200  model=qwen3-embedding-0.6b items=1 encoding=float vector_length=1024 usage={'prompt_tokens': 3, 'total_tokens': 3, 'completion_tokens': 0, 'prompt_tokens_details': None, 'completion_tokens_details': None}
$ curl -s http://127.0.0.1:<port>/v1/embeddings -H 'Content-Type: application/json' -d '{"input": "hello world", "model": "text-embedding-3-small", "encoding_format": "base64"}'
HTTP 200  model=qwen3-embedding-0.6b items=1 encoding=base64 vector_length=1024 usage={'prompt_tokens': 3, 'total_tokens': 3, 'completion_tokens': 0, 'prompt_tokens_details': None, 'completion_tokens_details': None}
base64 decodes to the float vector: True
OpenAI SDK (default wire encoding base64): items=2 vector_length=1024 matches_float_request=True

SMOKE_OK

Rollout evidence

N/A: this adds a model-server route and does not change any environment, agent or verifier behavior.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[khazic]

hello?

[github-actions]

🌿 Preview your docs: https://nvidia-preview-pr-1964-5944df350bc4.docs.buildwithfern.com/nemo/gym

[khazic]

All CI checks are green on the latest commit: DCO, Lint, copyright-check, secrets-detector, the full Server suite (shards 0 through 7), and the Test job all pass. The lone Server suite (shard ${{ matrix.shard }}) entry that shows as failed is just the un-expanded matrix-template placeholder job, not a real shard.

The branch is a little behind main right now; happy to rebase/update it whenever it's convenient. Is there anything else you'd like me to add or adjust here, or is it good to merge?

[yaoyu-33]

Thanks for the PR. The route, the 501 default, and the passthrough adapters follow the existing messages() / chat_completions() patterns, and the new tests pass at this head (test_base_responses_api_model.py, plus the new openai_model and inference_provider embeddings tests).

Process

• GitHub reports this branch as conflicting with main; could you rebase or merge main and re-run CI on the new head?
• Could you add a short real-provider smoke result to the PR body (for example the curl from the Usage section against OpenAI/NVIDIA or a vLLM --task embed server, showing status and vector length)? Every embeddings test here mocks _client.create_embeddings, so no run covers a real provider response.

Observation (not blocking)

• /v1/embeddings is registered only on the plain route, so embeddings calls are not picked up by model-call capture or the rollout-path prefixes that /v1/chat/completions has. A line in the docs saying so would save a surprise later.

Inline comments below: one design question on the forced float / strict response validation, one test gap, one docstring note.

[yaoyu-33]

openai_model/app.py:109

Forcing encoding_format="float" and then validating with the strict CreateEmbeddingResponse has two side effects: a client that explicitly asks for base64 gets float arrays back, and a provider that omits usage, model or data[].object (some OpenAI-compatible servers do) fails validation and comes back as a 500 even though the upstream call succeeded. Would you consider returning the provider's JSON as-is (no forced format, no re-validation)? That keeps the passthrough honest for both cases. If you prefer to keep the typed contract, please document that base64 is not supported. The same applies to inference_provider/app.py:125.

[khazic]

Went with the passthrough in 1100e05 on both servers: no forced encoding_format and no re-validation, so base64 and providers that omit usage both work. openai_model now also goes through _call_upstream.

[yaoyu-33]

openai_model/tests/test_app.py:292

Two cases are untested: (1) openai_model has no extra_body merge test like the one added for inference_provider (request value wins over config for the same key, and self.config.extra_body is not mutated), and (2) NeMoGymAsyncOpenAI.create_embeddings (openai_utils.py:953) is never exercised because every test mocks it, so the {base_url}/embeddings URL and JSON body are not checked. A test that patches _request/_raise_for_status/get_response_json and asserts the URL and body would cover it.

[khazic]

Thanks for adding both in 3d8bbb6. I only updated the expected kwargs of the extra_body test in 1100e05, since encoding_format is no longer forced.

[yaoyu-33]

openai_model/app.py:103

AGENTS.md asks for useful docstrings on public APIs. The base embeddings() has one, but this override, the one in inference_provider and create_embeddings do not. Something like:

Suggested change

	async def embeddings(self, body: dict = Body()) -> CreateEmbeddingResponse:
	async def embeddings(self, body: dict = Body()) -> CreateEmbeddingResponse:
	"""Pass an OpenAI-compatible embeddings request through to the provider; the configured model replaces the request's model."""

[khazic]

Thanks, your docstrings are kept as is.

[khazic]

Thanks for the review and for 3d8bbb6. 1100e05 goes on top of it:

• Embeddings are now a plain passthrough (inline reply on app.py:109).
• The real-provider smoke is in the PR body: curl with float and base64, plus the OpenAI SDK, against a vLLM embeddings server, for both openai_model and inference_provider.
• On the capture observation: besides the docs line, /v1/embeddings is now declared non-generating. Before, a call under the token-capture prefix marked the rollout incomplete; the new test in test_token_id_capture.py covers it.

The branch merges cleanly with main. CI needs a new /ok to test for 1100e05bd89e53c4682b71edab6f9e8cb7ed9dbc.

[yaoyu-33]

/ok to test 1100e05

[yaoyu-33]

LGTM at 1100e05.

[marcromeyn]

/ok to test 5944df3

[github-actions]

🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA-NeMo/Gym/actions/runs/40334275023

[github-actions]

🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA-NeMo/Gym/actions/runs/40344070863

[github-actions]

🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA-NeMo/Gym/actions/runs/40346226624