From 9e0e812cd578d3a4522aea16a7840dc962f40cdb Mon Sep 17 00:00:00 2001 From: Manjunath Janardhan Date: Wed, 30 Sep 2026 10:49:24 +0530 Subject: [PATCH 1/5] Python: Add NVIDIA NIM embedding client package Adds agent-framework-nvidia, an alpha package with NvidiaEmbeddingClient and RawNvidiaEmbeddingClient for embedding models hosted on NVIDIA NIM. NIM is OpenAI-compatible, so the clients build on RawOpenAIEmbeddingClient. NvidiaEmbeddingOptions adds input_type ("passage"/"query") and truncate. Asymmetric retrieval models embed stored documents and search text differently, so input_type is set per call through the vector-store embeddings_options. Both options are sent as extra_body, and invalid values are rejected before the request. The hosted endpoint accepts at most 256 inputs and a 2 MiB body per request, so larger calls, such as an upsert of many records, are sent as consecutive requests within both limits that keep vector order and sum usage. RawOpenAIEmbeddingClient gains a protected _prepare_extra_request_options hook for fields outside the OpenAI embedding schema. The default adds nothing, so OpenAI and Azure OpenAI requests are unchanged. Includes a package-local sample, unit tests, and an integration test that is skipped without NVIDIA_API_KEY. Feature-usage index 89. --- docs/specs/feature-usage-bit-registry.md | 3 +- python/PACKAGE_STATUS.md | 1 + python/packages/nvidia/AGENTS.md | 48 ++ python/packages/nvidia/LICENSE | 21 + python/packages/nvidia/README.md | 86 ++++ .../nvidia/agent_framework_nvidia/__init__.py | 23 + .../_embedding_client.py | 302 ++++++++++++ .../agent_framework_nvidia/_feature_usage.py | 9 + .../nvidia/agent_framework_nvidia/py.typed | 0 python/packages/nvidia/pyproject.toml | 70 +++ python/packages/nvidia/samples/README.md | 10 + .../nvidia/samples/nvidia_embeddings.py | 103 ++++ .../nvidia/test_nvidia_embedding_client.py | 442 ++++++++++++++++++ python/packages/openai/AGENTS.md | 8 + .../_embedding_client.py | 17 + .../openai/test_openai_embedding_client.py | 48 ++ python/pyproject.toml | 1 + python/uv.lock | 16 + 18 files changed, 1207 insertions(+), 1 deletion(-) create mode 100644 python/packages/nvidia/AGENTS.md create mode 100644 python/packages/nvidia/LICENSE create mode 100644 python/packages/nvidia/README.md create mode 100644 python/packages/nvidia/agent_framework_nvidia/__init__.py create mode 100644 python/packages/nvidia/agent_framework_nvidia/_embedding_client.py create mode 100644 python/packages/nvidia/agent_framework_nvidia/_feature_usage.py create mode 100644 python/packages/nvidia/agent_framework_nvidia/py.typed create mode 100644 python/packages/nvidia/pyproject.toml create mode 100644 python/packages/nvidia/samples/README.md create mode 100644 python/packages/nvidia/samples/nvidia_embeddings.py create mode 100644 python/packages/nvidia/tests/nvidia/test_nvidia_embedding_client.py diff --git a/docs/specs/feature-usage-bit-registry.md b/docs/specs/feature-usage-bit-registry.md index 939ed2748ed..65ee24303a3 100644 --- a/docs/specs/feature-usage-bit-registry.md +++ b/docs/specs/feature-usage-bit-registry.md @@ -188,7 +188,8 @@ only to approved first-party endpoints. | 86 | `hosting.responses` | OpenAI Responses hosting converters | `agent_framework_hosting_responses.responses_to_run` | | 87 | `hosting.telegram` | Telegram hosting converters | `agent_framework_hosting_telegram.telegram_to_run` | | 88 | `lab` | Experimental Agent Framework Lab features | `agent_framework.lab` feature entry points | -| 89–127 | _reserved_ | future packages | — | +| 89 | `nvidia` | NVIDIA NIM embedding client | `agent_framework_nvidia.NvidiaEmbeddingClient` | +| 90–127 | _reserved_ | future packages | — | ## Index table — .NET (`agent-framework-dotnet`, version 1) diff --git a/python/PACKAGE_STATUS.md b/python/PACKAGE_STATUS.md index b8701ea2d8c..1830f70369f 100644 --- a/python/PACKAGE_STATUS.md +++ b/python/PACKAGE_STATUS.md @@ -47,6 +47,7 @@ Status is grouped into these buckets: | `agent-framework-mistral` | `python/packages/mistral` | `beta` | | `agent-framework-mongodb` | `python/packages/mongodb` | `alpha` | | `agent-framework-monty` | `python/packages/monty` | `beta` | +| `agent-framework-nvidia` | `python/packages/nvidia` | `alpha` | | `agent-framework-ollama` | `python/packages/ollama` | `beta` | | `agent-framework-openai` | `python/packages/openai` | `released` | | `agent-framework-orchestrations` | `python/packages/orchestrations` | `released` | diff --git a/python/packages/nvidia/AGENTS.md b/python/packages/nvidia/AGENTS.md new file mode 100644 index 00000000000..aa28aa4fb55 --- /dev/null +++ b/python/packages/nvidia/AGENTS.md @@ -0,0 +1,48 @@ +# NVIDIA Package (agent-framework-nvidia) + +Integration with NVIDIA NIM for embedding generation. + +## Implementation Notes + +- NVIDIA NIM exposes an OpenAI-compatible endpoint, so this package builds on + `RawOpenAIEmbeddingClient` from `agent-framework-openai` rather than reimplementing the wire + format. This mirrors `agent-framework-foundry-local`, which builds on the OpenAI Chat + Completions client for the same reason. +- The default endpoint is `https://integrate.api.nvidia.com/v1`. `base_url` overrides it, which is + the hook for a self-hosted NIM container. +- Model availability is per-account. An unentitled model returns `404 ... Not found for account`, + and a retired model returns `410 Gone`, so neither is a request-shape error. +- `input_type` and `truncate` are outside the OpenAI embedding schema. The client sends them as + `extra_body` by overriding `_prepare_extra_request_options`, and rejects invalid values before + the request. Callers set `input_type` per call through vector-store `embeddings_options`, the + same route the Gemini client uses for `task_type`. +- The hosted endpoint rejects requests with more than 256 inputs (`MAX_BATCH_SIZE`) or a body over + 2 MiB (`413`, budgeted by `MAX_BATCH_BYTES`), and a vector-store upsert embeds every record in one + call, so `get_embeddings` splits larger inputs into consecutive requests, as the Gemini client does + for Enterprise. + +## Main Classes + +- **`NvidiaEmbeddingClient`** - Embedding client for NVIDIA NIM models, with telemetry +- **`RawNvidiaEmbeddingClient`** - The same client without the telemetry layer +- **`NvidiaEmbeddingOptions`** - Options TypedDict for NVIDIA-specific embedding parameters +- **`NvidiaEmbeddingSettings`** - TypedDict settings for NVIDIA configuration + +## Usage + +```python +from agent_framework_nvidia import NvidiaEmbeddingClient + +# Requires NVIDIA_API_KEY environment variable (or pass api_key= directly) +client = NvidiaEmbeddingClient(model="nvidia/nemotron-3-embed-1b") +result = await client.get_embeddings(["Hello, world!"]) +print(result[0].vector) +``` + +## Import Path + +The package is alpha, so there is no `agent_framework.nvidia` namespace in core yet. + +```python +from agent_framework_nvidia import NvidiaEmbeddingClient +``` diff --git a/python/packages/nvidia/LICENSE b/python/packages/nvidia/LICENSE new file mode 100644 index 00000000000..9e841e7a26e --- /dev/null +++ b/python/packages/nvidia/LICENSE @@ -0,0 +1,21 @@ + MIT License + + Copyright (c) Microsoft Corporation. + + Permission is hereby granted, free of charge, to any person obtaining a copy + of this software and associated documentation files (the "Software"), to deal + in the Software without restriction, including without limitation the rights + to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + copies of the Software, and to permit persons to whom the Software is + furnished to do so, subject to the following conditions: + + The above copyright notice and this permission notice shall be included in all + copies or substantial portions of the Software. + + THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + SOFTWARE diff --git a/python/packages/nvidia/README.md b/python/packages/nvidia/README.md new file mode 100644 index 00000000000..575c56d5708 --- /dev/null +++ b/python/packages/nvidia/README.md @@ -0,0 +1,86 @@ +# Get Started with Microsoft Agent Framework NVIDIA NIM + +An alpha NVIDIA NIM embedding integration for Microsoft Agent Framework. Please install this package: + +```bash +pip install agent-framework-nvidia --pre +``` + +and see the [README](https://github.com/microsoft/agent-framework/tree/main/python/README.md) for more information. + +See the [NVIDIA samples](samples/README.md) for a runnable example. + +## Embedding Client + +The `NvidiaEmbeddingClient` provides embedding generation using models hosted on +[NVIDIA NIM](https://build.nvidia.com/). NIM exposes an OpenAI-compatible surface, so this client +builds on the OpenAI embedding client rather than reimplementing the wire format. + +### Quick Start + +```python +from agent_framework_nvidia import NvidiaEmbeddingClient + +# Using environment variables (NVIDIA_API_KEY, NVIDIA_EMBEDDING_MODEL) +# Parameters can also be passed directly: +# NvidiaEmbeddingClient(model="nvidia/nemotron-3-embed-1b", api_key="your-api-key") +client = NvidiaEmbeddingClient() + +result = await client.get_embeddings(["Hello, world!", "How are you?"]) +for embedding in result: + print(f"Dimensions: {embedding.dimensions}") + print(f"Vector: {embedding.vector[:5]}...") + +# Close the HTTP client the embedding client created. A client passed as async_client is left open. +await client.close() +``` + +### Configuration + +| Environment Variable | Description | +|---|---| +| `NVIDIA_API_KEY` | Your NVIDIA API key, from [build.nvidia.com](https://build.nvidia.com/) | +| `NVIDIA_EMBEDDING_MODEL` | Embedding model name (e.g., `nvidia/nemotron-3-embed-1b`) | +| `NVIDIA_BASE_URL` | Optional endpoint override, for a self-hosted NIM container | + +The service throttles bursts with `429`, which the OpenAI SDK retries twice by default. For heavy +parallel ingestion, pass a pre-configured `async_client`, for example +`AsyncOpenAI(..., max_retries=5)`, to change retries or timeouts. + +### Options + +`NvidiaEmbeddingOptions` adds two NVIDIA fields to the OpenAI embedding options: + +| Option | Values | Purpose | +|---|---|---| +| `input_type` | `"passage"`, `"query"` | Asymmetric retrieval models embed documents and search text differently | +| `truncate` | `"NONE"`, `"START"`, `"END"` | Reject over-long input, or drop tokens from that end | + +For retrieval, embed stored documents as passages and search text as queries. Vector collections +forward `embeddings_options` to each request: + +```python +await collection.upsert(notes, embeddings_options={"input_type": "passage", "truncate": "END"}) +results = await collection.search("How do I call a model?", embeddings_options={"input_type": "query"}) +``` + +`create_vector_search_tool(..., embeddings_options={"input_type": "query"})` does the same for the +tool an agent calls. `nvidia/nemotron-3-embed-1b` accepts only 2048 dimensions, so declare that on +the vector field. + +The hosted endpoint accepts at most 256 inputs and a 2 MiB body per request, so the client sends +larger calls, such as an upsert of many records, as consecutive requests within both limits and +returns the vectors in order. + +Without `truncate`, the service shortens input over the model's token limit itself, but it rejects any +input over 65,536 characters, and one such document fails the whole call, so nothing in that upsert +is written. Set `truncate="END"` (or `"START"`) when ingesting documents of unknown length. +`truncate="NONE"` rejects input over the token limit instead. + +### Choosing a model + +Model availability is per-account: the endpoint lists more embedding models than a given API key +is entitled to call, and an unentitled model returns `404 ... Not found for account` rather than a +permission error. Retired models return `410 Gone`. If a model name from the documentation does +not work, list the models available to your key rather than assuming the name is wrong. An invalid key +returns `403 Forbidden` with `Authorization failed`, not `401`. diff --git a/python/packages/nvidia/agent_framework_nvidia/__init__.py b/python/packages/nvidia/agent_framework_nvidia/__init__.py new file mode 100644 index 00000000000..67cc8aac471 --- /dev/null +++ b/python/packages/nvidia/agent_framework_nvidia/__init__.py @@ -0,0 +1,23 @@ +# Copyright (c) Microsoft. All rights reserved. + +import importlib.metadata + +from ._embedding_client import ( + NvidiaEmbeddingClient, + NvidiaEmbeddingOptions, + NvidiaEmbeddingSettings, + RawNvidiaEmbeddingClient, +) + +try: + __version__ = importlib.metadata.version(__name__) +except importlib.metadata.PackageNotFoundError: + __version__ = "0.0.0" # Fallback for development mode + +__all__ = [ + "NvidiaEmbeddingClient", + "NvidiaEmbeddingOptions", + "NvidiaEmbeddingSettings", + "RawNvidiaEmbeddingClient", + "__version__", +] diff --git a/python/packages/nvidia/agent_framework_nvidia/_embedding_client.py b/python/packages/nvidia/agent_framework_nvidia/_embedding_client.py new file mode 100644 index 00000000000..c78ecb0f483 --- /dev/null +++ b/python/packages/nvidia/agent_framework_nvidia/_embedding_client.py @@ -0,0 +1,302 @@ +# Copyright (c) Microsoft. All rights reserved. + +from __future__ import annotations + +import json +import sys +from collections.abc import Iterator, Mapping, Sequence +from typing import Any, ClassVar, Generic, Literal, TypedDict + +from agent_framework import ( + Embedding, + GeneratedEmbeddings, + SecretString, + UsageDetails, + add_usage_details, + load_settings, +) +from agent_framework.observability import EmbeddingTelemetryLayer +from agent_framework_openai import OpenAIEmbeddingOptions +from agent_framework_openai._embedding_client import RawOpenAIEmbeddingClient +from openai import AsyncOpenAI + +from ._feature_usage import FeatureIndex + +if sys.version_info >= (3, 13): + from typing import TypeVar # pragma: no cover +else: + from typing_extensions import TypeVar # pragma: no cover + +__all__ = [ + "NvidiaEmbeddingClient", + "NvidiaEmbeddingOptions", + "NvidiaEmbeddingSettings", + "RawNvidiaEmbeddingClient", +] + +DEFAULT_BASE_URL = "https://integrate.api.nvidia.com/v1" +"""NVIDIA's hosted NIM endpoint, which exposes an OpenAI-compatible surface.""" + + +MAX_BATCH_SIZE = 256 +"""The most inputs NVIDIA's hosted endpoint accepts in one embeddings request.""" + +MAX_BATCH_BYTES = 2_000_000 +"""Budget for the inputs of one request. The hosted endpoint rejects bodies over 2 MiB with ``413``, +so this leaves room for the other request fields.""" + + +def _batches(values: Sequence[str]) -> Iterator[Sequence[str]]: + """Split ``values`` into consecutive runs within both ``MAX_BATCH_SIZE`` and ``MAX_BATCH_BYTES``. + + Each input is costed at its ASCII-escaped JSON length, which is never less than the bytes it + occupies in the request. An input over the byte budget on its own still gets its own request. + """ + start, size = 0, 0 + for index, value in enumerate(values): + cost = len(json.dumps(value)) + 1 + if index > start and (index - start == MAX_BATCH_SIZE or size + cost > MAX_BATCH_BYTES): + yield values[start:index] + start, size = index, 0 + size += cost + if start < len(values): + yield values[start:] + + +_INPUT_TYPES = ("query", "passage") +_TRUNCATE_MODES = ("NONE", "START", "END") + + +class NvidiaEmbeddingOptions(OpenAIEmbeddingOptions, total=False): + """Request settings for NVIDIA embedding generation. + + NVIDIA accepts the OpenAI embedding options, including ``encoding_format`` and ``user``, + plus the two keys below. + + Keys: + input_type: ``"passage"`` for documents being stored and ``"query"`` for search + text. Asymmetric retrieval models embed the two differently, so set it per + call, for example through ``embeddings_options`` on vector upsert and search. + Omitted, the service uses its default and the key is not sent. + truncate: How the service handles long input. ``"START"`` or ``"END"`` drop tokens + from that end, and ``"NONE"`` rejects input over the model's token limit + (4,096 for ``nvidia/nemotron-3-embed-1b``). Omitted, the service shortens + over-limit input itself but still rejects any input over 65,536 characters, + and one such input fails the whole call, so set ``"START"`` or ``"END"`` when + ingesting documents of unknown length. + """ + + input_type: Literal["query", "passage"] + truncate: Literal["NONE", "START", "END"] + + +NvidiaEmbeddingOptionsT = TypeVar( + "NvidiaEmbeddingOptionsT", + bound=TypedDict, # type: ignore[valid-type] + default="NvidiaEmbeddingOptions", + covariant=True, +) + + +class NvidiaEmbeddingSettings(TypedDict, total=False): + """NVIDIA embedding settings. + + Settings are resolved in this order: explicit keyword arguments, values from an + explicitly provided .env file, then environment variables with the prefix + 'NVIDIA_'. + + Keys: + api_key: The API key for NVIDIA's hosted endpoint. + (Env var NVIDIA_API_KEY) + embedding_model: The embedding model to use, for example + ``nvidia/nemotron-3-embed-1b``. (Env var NVIDIA_EMBEDDING_MODEL) + base_url: Override for the endpoint. Defaults to NVIDIA's hosted NIM surface, + and is the hook for pointing at a self-hosted NIM container. + (Env var NVIDIA_BASE_URL) + """ + + api_key: SecretString | None + embedding_model: str | None + base_url: str | None + + +class RawNvidiaEmbeddingClient( + RawOpenAIEmbeddingClient[NvidiaEmbeddingOptionsT], + Generic[NvidiaEmbeddingOptionsT], +): + """NVIDIA NIM embedding client without telemetry. + + NVIDIA's NIM endpoint is OpenAI-compatible, so this builds on the OpenAI embedding + client rather than reimplementing the wire format, in the same way + ``agent_framework_foundry_local`` builds on the OpenAI Chat Completions client. + """ + + OTEL_PROVIDER_NAME: ClassVar[str] = "nvidia" + _FEATURE_USAGE_INDEX: ClassVar[int | None] = FeatureIndex.NVIDIA + + def __init__( + self, + *, + model: str | None = None, + api_key: str | SecretString | None = None, + base_url: str | None = None, + async_client: AsyncOpenAI | None = None, + additional_properties: dict[str, Any] | None = None, + env_file_path: str | None = None, + env_file_encoding: str | None = None, + ) -> None: + """Initialize a raw NVIDIA embedding client. + + Keyword Args: + model: The embedding model to use. Falls back to NVIDIA_EMBEDDING_MODEL. + api_key: The API key. Falls back to NVIDIA_API_KEY. + base_url: Endpoint override, for a self-hosted NIM container. Falls back to + NVIDIA_BASE_URL, then to NVIDIA's hosted endpoint. + async_client: A pre-configured client to use instead of one built from the + settings above. It stays caller-owned: ``close()`` leaves it open. + additional_properties: Extra properties carried on the client. + env_file_path: Path to a .env file to read settings from. + env_file_encoding: Encoding of that .env file. + + Raises: + ValueError: If no model is supplied, or if no API key is available and no + pre-configured client was passed. + """ + settings = load_settings( + NvidiaEmbeddingSettings, + env_prefix="NVIDIA_", + required_fields=["embedding_model"], + embedding_model=model, + api_key=api_key, + base_url=base_url, + env_file_path=env_file_path, + env_file_encoding=env_file_encoding, + ) + model_setting: str = settings["embedding_model"] # type: ignore[assignment] + + owns_client = async_client is None + if async_client is None: + key = settings.get("api_key") + if key is None: + raise ValueError( + "An NVIDIA API key is required. Set NVIDIA_API_KEY, pass api_key, " + "or supply a pre-configured async_client." + ) + async_client = AsyncOpenAI( + api_key=key.get_secret_value(), + base_url=settings.get("base_url") or DEFAULT_BASE_URL, + ) + + super().__init__( + model=model_setting, + async_client=async_client, + additional_properties=additional_properties, + ) + self._owns_client = owns_client + + async def close(self) -> None: + """Close the HTTP client this instance created. A supplied ``async_client`` stays open for its owner.""" + if self._owns_client: + await self.client.close() + + async def get_embeddings( + self, + values: Sequence[str], + *, + options: NvidiaEmbeddingOptionsT | None = None, + ) -> GeneratedEmbeddings[list[float], NvidiaEmbeddingOptionsT]: + """Generate one embedding per value, splitting large inputs into several requests. + + The hosted endpoint rejects a request with more than ``MAX_BATCH_SIZE`` inputs or a body + over 2 MiB, and a vector-store upsert embeds every record in one call, so larger inputs + are sent as consecutive requests within both limits. Vectors keep the order of + ``values``, usage is summed, and a failed request fails the whole call. + + Args: + values: The texts to embed. + options: Embedding options, applied to every request. + + Returns: + One embedding per value, in order. + """ + batches = list(_batches(values)) + if len(batches) <= 1: + return await super().get_embeddings(values, options=options) + embeddings: list[Embedding[list[float]]] = [] + usage: UsageDetails | None = None + for values_batch in batches: + batch = await super().get_embeddings(values_batch, options=options) + embeddings.extend(batch) + if batch.usage: + usage = add_usage_details(usage, batch.usage) + return GeneratedEmbeddings(embeddings, options=options, usage=usage) + + def _prepare_extra_request_options(self, options: Mapping[str, Any]) -> dict[str, Any]: + """Send ``input_type`` and ``truncate``, which are outside the OpenAI schema, as ``extra_body``. + + Raises: + ValueError: If ``input_type`` or ``truncate`` is not a value NVIDIA accepts. + """ + extra_body: dict[str, Any] = {} + input_type = options.get("input_type") + if input_type is not None: + if input_type not in _INPUT_TYPES: + raise ValueError(f"Unsupported NVIDIA embedding input_type: {input_type!r}.") + extra_body["input_type"] = input_type + truncate = options.get("truncate") + if truncate is not None: + if truncate not in _TRUNCATE_MODES: + raise ValueError(f"Unsupported NVIDIA embedding truncate mode: {truncate!r}.") + extra_body["truncate"] = truncate + return {"extra_body": extra_body} if extra_body else {} + + +class NvidiaEmbeddingClient( + EmbeddingTelemetryLayer[str, list[float], NvidiaEmbeddingOptionsT], + RawNvidiaEmbeddingClient[NvidiaEmbeddingOptionsT], + Generic[NvidiaEmbeddingOptionsT], +): + """NVIDIA NIM embedding client with telemetry support.""" + + OTEL_PROVIDER_NAME: ClassVar[str] = "nvidia" + + def __init__( + self, + *, + model: str | None = None, + api_key: str | SecretString | None = None, + base_url: str | None = None, + async_client: AsyncOpenAI | None = None, + additional_properties: dict[str, Any] | None = None, + otel_provider_name: str | None = None, + env_file_path: str | None = None, + env_file_encoding: str | None = None, + ) -> None: + """Initialize an NVIDIA embedding client with telemetry. + + Keyword Args: + model: The embedding model to use. Falls back to NVIDIA_EMBEDDING_MODEL. + api_key: The API key. Falls back to NVIDIA_API_KEY. + base_url: Endpoint override, for a self-hosted NIM container. Falls back to + NVIDIA_BASE_URL, then to NVIDIA's hosted endpoint. + async_client: A pre-configured client to use instead of one built from the + settings above. It stays caller-owned: ``close()`` leaves it open. + additional_properties: Extra properties carried on the client. + otel_provider_name: Override for the provider name reported in telemetry. + env_file_path: Path to a .env file to read settings from. + env_file_encoding: Encoding of that .env file. + + Raises: + ValueError: If no model is supplied, or if no API key is available and no + pre-configured client was passed. + """ + super().__init__( + model=model, + api_key=api_key, + base_url=base_url, + async_client=async_client, + additional_properties=additional_properties, + otel_provider_name=otel_provider_name, + env_file_path=env_file_path, + env_file_encoding=env_file_encoding, + ) diff --git a/python/packages/nvidia/agent_framework_nvidia/_feature_usage.py b/python/packages/nvidia/agent_framework_nvidia/_feature_usage.py new file mode 100644 index 00000000000..516b33043f4 --- /dev/null +++ b/python/packages/nvidia/agent_framework_nvidia/_feature_usage.py @@ -0,0 +1,9 @@ +# Copyright (c) Microsoft. All rights reserved. + +from enum import IntEnum + + +class FeatureIndex(IntEnum): + """NVIDIA-owned feature-usage indexes.""" + + NVIDIA = 89 diff --git a/python/packages/nvidia/agent_framework_nvidia/py.typed b/python/packages/nvidia/agent_framework_nvidia/py.typed new file mode 100644 index 00000000000..e69de29bb2d diff --git a/python/packages/nvidia/pyproject.toml b/python/packages/nvidia/pyproject.toml new file mode 100644 index 00000000000..2c46cc74d87 --- /dev/null +++ b/python/packages/nvidia/pyproject.toml @@ -0,0 +1,70 @@ +[project] +name = "agent-framework-nvidia" +description = "NVIDIA NIM embeddings for Microsoft Agent Framework." +authors = [{ name = "Microsoft", email = "af-support@microsoft.com" }] +readme = "README.md" +requires-python = ">=3.10" +version = "1.0.0a260929" +license-files = ["LICENSE"] +urls.homepage = "https://aka.ms/agent-framework" +urls.source = "https://github.com/microsoft/agent-framework/tree/main/python" +urls.issues = "https://github.com/microsoft/agent-framework/issues" +classifiers = [ + "License :: OSI Approved :: MIT License", + "Development Status :: 3 - Alpha", + "Intended Audience :: Developers", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Programming Language :: Python :: 3.14", + "Typing :: Typed", +] +dependencies = [ + "agent-framework-core>=1.19.0,<2", + # input_type and truncate need RawOpenAIEmbeddingClient._prepare_extra_request_options. + # 1.14.4 predates it and silently drops them, so raise this floor to the release that adds it. + "agent-framework-openai>=1.14.4,<2", +] + +[tool.ruff] +extend = "../../pyproject.toml" + +[tool.ruff.lint.per-file-ignores] +"samples/**" = ["D", "INP", "TD", "ERA001", "RUF", "S", "T201", "CPY"] +"tests/**" = ["D", "INP", "TD", "ERA001", "RUF", "S"] + +[tool.pyright] +extends = "../../pyproject.toml" +include = ["agent_framework_nvidia"] + +[tool.pytest.ini_options] +testpaths = ["tests"] +addopts = "-ra -q -r fEX" +asyncio_mode = "auto" +asyncio_default_fixture_loop_scope = "function" +markers = [ + "integration: marks tests as integration tests that require external services", + "flaky: marks tests as flaky and eligible for automatic retry", +] +timeout = 60 + +[tool.coverage.run] +omit = ["**/__init__.py"] + +[tool.poe] +executor.type = "uv" +include = "../../shared_tasks.toml" + +[tool.poe.tasks.test] +help = "Run NVIDIA connector unit tests." +cmd = 'pytest -m "not integration" --cov=agent_framework_nvidia --cov-report=term-missing:skip-covered tests' + +[tool.poe.tasks.test-integration] +help = "Run tests against the NVIDIA service using NVIDIA_API_KEY." +cmd = 'pytest -m integration tests' + +[build-system] +requires = ["flit-core>=3.11,<4.0"] +build-backend = "flit_core.buildapi" diff --git a/python/packages/nvidia/samples/README.md b/python/packages/nvidia/samples/README.md new file mode 100644 index 00000000000..ee4872db368 --- /dev/null +++ b/python/packages/nvidia/samples/README.md @@ -0,0 +1,10 @@ +# NVIDIA samples + +`nvidia_embeddings.py` generates single and batched embeddings, then indexes two +notes in an in-memory collection and retrieves the closer one. Stored notes are +embedded with `input_type="passage"` and the search text with `input_type="query"`, +which asymmetric models such as `nvidia/nemotron-3-embed-1b` embed differently. + +Set `NVIDIA_API_KEY` to a key from [build.nvidia.com](https://build.nvidia.com/). +See the [package README](../README.md) for the other settings and for choosing a +model. diff --git a/python/packages/nvidia/samples/nvidia_embeddings.py b/python/packages/nvidia/samples/nvidia_embeddings.py new file mode 100644 index 00000000000..1cab9d1f2bc --- /dev/null +++ b/python/packages/nvidia/samples/nvidia_embeddings.py @@ -0,0 +1,103 @@ +# Copyright (c) Microsoft. All rights reserved. + +import asyncio +from dataclasses import dataclass +from typing import Annotated + +from agent_framework import InMemoryCollection, VectorStoreField, vectorstoremodel + +from agent_framework_nvidia import NvidiaEmbeddingClient + +"""Generate embeddings with models hosted on NVIDIA NIM, then retrieve with them. + +Set ``NVIDIA_API_KEY`` to a key from https://build.nvidia.com/, then run from python/: + uv run --package agent-framework-nvidia python packages/nvidia/samples/nvidia_embeddings.py + +Model availability is per account. If a model name returns ``404 ... Not found for account``, +it exists but your key is not entitled to it; a retired model returns ``410 Gone``. List the +models available to your key rather than assuming the name is wrong. +""" + + +@vectorstoremodel(collection_name="nvidia-notes") +@dataclass +class Note: + id: Annotated[str, VectorStoreField("key")] + text: Annotated[str, VectorStoreField("data")] + vector: Annotated[str | list[float] | None, VectorStoreField("vector", dimensions=2048)] = None + + +async def main() -> None: + """Generate embeddings with NVIDIA NIM.""" + # The API key is read from NVIDIA_API_KEY. + client = NvidiaEmbeddingClient(model="nvidia/nemotron-3-embed-1b") + + try: + # 1. Generate a single embedding. + result = await client.get_embeddings(["Hello, world!"]) + print(f"Single embedding dimensions: {result[0].dimensions}") + print(f"First 5 values: {result[0].vector[:5]}") + print(f"Model: {result[0].model}") + print(f"Usage: {result.usage}") + print() + + # 2. Generate embeddings for multiple inputs. + texts = [ + "The weather is sunny today.", + "It is raining outside.", + "Machine learning is fascinating.", + ] + result = await client.get_embeddings(texts) + print(f"Batch of {len(result)} embeddings, each with {result[0].dimensions} dimensions") + print(f"First embedding vector: {result[0].vector[:5]}") + print() + + # 3. Retrieval: embed stored documents as passages and search text as a query. + # Asymmetric models such as nemotron embed the two differently, and the vector + # store forwards input_type to each embedding request. + collection: InMemoryCollection[str, Note] = InMemoryCollection(Note, embedding_generator=client) + await collection.ensure_collection_exists() + sources = { + "one": "NIM serves models behind an OpenAI-compatible API.", + "two": "The weather is sunny today.", + } + # The vector field holds the text to embed until the collection replaces it with the vector. + # truncate="END" keeps one over-long document from failing the whole upsert. + notes = [Note(note_id, text, text) for note_id, text in sources.items()] + await collection.upsert(notes, embeddings_options={"input_type": "passage", "truncate": "END"}) + results = await collection.search( + "How do I call a model hosted on NIM?", embeddings_options={"input_type": "query"}, top=1 + ) + async for item in results: + print(f"Best match: {item['record'].text}") + + # 4. Point at a self-hosted NIM container instead of the hosted endpoint. + # base_url also reads from NVIDIA_BASE_URL. + # local_client = NvidiaEmbeddingClient( + # model="nvidia/nemotron-3-embed-1b", + # base_url="http://localhost:8000/v1", + # api_key="not-used-by-local-nim", + # ) + + finally: + await client.close() + + +if __name__ == "__main__": + asyncio.run(main()) + + +""" +Sample output: +Single embedding dimensions: 2048 +First 5 values: [0.4655469059944153, -0.027859577909111977, -0.1295825093984604, -0.0473068505525589, +-0.009855593554675579] +Model: nvidia/nemotron-3-embed-1b +Usage: {'input_token_count': 4, 'total_token_count': 4} + +Batch of 3 embeddings, each with 2048 dimensions +First embedding vector: [0.22611355781555176, -0.003007160732522607, -0.0892254188656807, +-0.027936991304159164, 0.01019396260380745] + +Best match: NIM serves models behind an OpenAI-compatible API. +""" diff --git a/python/packages/nvidia/tests/nvidia/test_nvidia_embedding_client.py b/python/packages/nvidia/tests/nvidia/test_nvidia_embedding_client.py new file mode 100644 index 00000000000..9340916a6f7 --- /dev/null +++ b/python/packages/nvidia/tests/nvidia/test_nvidia_embedding_client.py @@ -0,0 +1,442 @@ +# Copyright (c) Microsoft. All rights reserved. + +import json +import math +import os +from dataclasses import dataclass +from typing import Annotated, Any, cast + +import httpx2 +import pytest +from agent_framework import InMemoryCollection, VectorStoreField, create_vector_search_tool, vectorstoremodel +from agent_framework._settings import SecretString +from openai import AsyncOpenAI + +from agent_framework_nvidia import NvidiaEmbeddingClient, RawNvidiaEmbeddingClient +from agent_framework_nvidia._embedding_client import ( # pyright: ignore[reportPrivateUsage] + DEFAULT_BASE_URL, + MAX_BATCH_BYTES, + MAX_BATCH_SIZE, + _batches, # pyright: ignore[reportPrivateUsage] +) +from agent_framework_nvidia._feature_usage import FeatureIndex # pyright: ignore[reportPrivateUsage] + +# region: helpers + + +@pytest.fixture(autouse=True) +def _clear_nvidia_env(request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch) -> None: + """Keep a developer's real NVIDIA settings out of the unit tests. + + Without this a populated environment can make a construction test pass for the wrong + reason, and can make the missing-setting tests fail only on other people's machines. + Integration tests keep the environment, since it holds their credentials. + """ + if request.node.get_closest_marker("integration"): + return + for name in ("NVIDIA_API_KEY", "NVIDIA_EMBEDDING_MODEL", "NVIDIA_BASE_URL"): + monkeypatch.delenv(name, raising=False) + + +def make_embeddings_payload(vectors: list[list[float]], model: str = "test-model") -> dict[str, Any]: + return { + "object": "list", + "model": model, + "data": [{"object": "embedding", "index": i, "embedding": v} for i, v in enumerate(vectors)], + "usage": {"prompt_tokens": 4, "total_tokens": 4}, + } + + +class MockNvidia: + """Captures the requests the client actually puts on the wire.""" + + def __init__(self, payload: dict[str, Any]) -> None: + self._payload = payload + self.requests: list[httpx2.Request] = [] + self.bodies: list[dict[str, Any]] = [] + + def handler(self, request: httpx2.Request) -> httpx2.Response: + self.requests.append(request) + self.bodies.append(json.loads(request.content)) + return httpx2.Response(200, json=self._payload) + + +def make_mock_client(model: str = "test-model") -> tuple[NvidiaEmbeddingClient, MockNvidia]: + server = MockNvidia(make_embeddings_payload([[0.1, 0.2, 0.3]], model=model)) + async_client = AsyncOpenAI( + api_key="test-key", + base_url=DEFAULT_BASE_URL, + http_client=httpx2.AsyncClient(transport=httpx2.MockTransport(server.handler)), + ) + return NvidiaEmbeddingClient(model=model, async_client=async_client), server + + +# region: construction and settings + + +def test_construction_from_env(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("NVIDIA_EMBEDDING_MODEL", "nvidia/nemotron-3-embed-1b") + monkeypatch.setenv("NVIDIA_API_KEY", "test-key") + + client = NvidiaEmbeddingClient() + + assert client.model == "nvidia/nemotron-3-embed-1b" + + +def test_defaults_to_the_nvidia_endpoint(monkeypatch: pytest.MonkeyPatch) -> None: + """The whole point of the package: an NVIDIA key must not be sent to OpenAI. + + The client is built on the OpenAI client, whose own default base URL is OpenAI's, so + this asserts the override rather than trusting it. + """ + monkeypatch.setenv("NVIDIA_EMBEDDING_MODEL", "test-model") + monkeypatch.setenv("NVIDIA_API_KEY", "test-key") + + client = NvidiaEmbeddingClient() + + assert "integrate.api.nvidia.com" in str(client.client.base_url) + assert "api.openai.com" not in str(client.client.base_url) + + +def test_base_url_override_from_env(monkeypatch: pytest.MonkeyPatch) -> None: + """Self-hosted NIM containers are reached by pointing base_url at them.""" + monkeypatch.setenv("NVIDIA_EMBEDDING_MODEL", "test-model") + monkeypatch.setenv("NVIDIA_API_KEY", "test-key") + monkeypatch.setenv("NVIDIA_BASE_URL", "http://localhost:8000/v1") + + client = NvidiaEmbeddingClient() + + assert "localhost:8000" in str(client.client.base_url) + + +def test_explicit_arguments_win_over_the_environment(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("NVIDIA_EMBEDDING_MODEL", "from-env") + monkeypatch.setenv("NVIDIA_API_KEY", "test-key") + monkeypatch.setenv("NVIDIA_BASE_URL", "http://from-env/v1") + + client = NvidiaEmbeddingClient(model="from-arg", base_url="http://from-arg/v1") + + assert client.model == "from-arg" + assert "from-arg" in str(client.client.base_url) + + +@pytest.mark.parametrize("api_key", ["test-key", SecretString("test-key")], ids=["str", "secret"]) +def test_api_key_accepts_str_or_secretstring(api_key: str | SecretString) -> None: + client = NvidiaEmbeddingClient(model="test-model", api_key=api_key) + + assert client.model == "test-model" + + +def test_missing_model_raises() -> None: + with pytest.raises(Exception, match="model"): + NvidiaEmbeddingClient(api_key="test-key") + + +def test_missing_api_key_raises_and_names_the_setting() -> None: + with pytest.raises(ValueError, match="NVIDIA_API_KEY"): + NvidiaEmbeddingClient(model="test-model") + + +def test_injected_client_does_not_require_an_api_key() -> None: + """A caller supplying their own client has already handled auth.""" + async_client = AsyncOpenAI(api_key="test-key", base_url="http://injected/v1") + + client = NvidiaEmbeddingClient(model="test-model", async_client=async_client) + + assert "injected" in str(client.client.base_url) + + +async def test_close_closes_a_client_it_created() -> None: + client = NvidiaEmbeddingClient(model="test-model", api_key="test-key") + + await client.close() + + assert client.client.is_closed() + + +async def test_close_leaves_an_injected_client_open() -> None: + async_client = AsyncOpenAI(api_key="test-key", base_url="http://injected/v1") + client = NvidiaEmbeddingClient(model="test-model", async_client=async_client) + + await client.close() + + assert not async_client.is_closed() + await async_client.close() + + +# region: wiring + + +def test_provider_identity() -> None: + client = NvidiaEmbeddingClient(model="test-model", api_key="test-key") + + assert client.OTEL_PROVIDER_NAME == "nvidia" + assert NvidiaEmbeddingClient._FEATURE_USAGE_INDEX == FeatureIndex.NVIDIA # pyright: ignore[reportPrivateUsage] + + +def test_telemetry_layer_wraps_the_raw_client() -> None: + """The layer order decides whether calls are traced, so pin it.""" + names = [c.__name__ for c in NvidiaEmbeddingClient.__mro__] + + assert names.index("EmbeddingTelemetryLayer") < names.index("RawNvidiaEmbeddingClient") + assert names.index("RawNvidiaEmbeddingClient") < names.index("RawOpenAIEmbeddingClient") + + +async def test_raw_client_sends_nvidia_options_without_telemetry() -> None: + server = MockNvidia(make_embeddings_payload([[0.1, 0.2, 0.3]])) + async_client = AsyncOpenAI( + api_key="test-key", + base_url=DEFAULT_BASE_URL, + http_client=httpx2.AsyncClient(transport=httpx2.MockTransport(server.handler)), + ) + client = RawNvidiaEmbeddingClient(model="test-model", async_client=async_client) + + await client.get_embeddings(["hello"], options={"input_type": "query"}) + + assert "EmbeddingTelemetryLayer" not in [c.__name__ for c in type(client).__mro__] + assert server.bodies[-1]["input_type"] == "query" + + +# region: requests + + +async def test_get_embeddings_returns_vectors() -> None: + client, server = make_mock_client() + + result = await client.get_embeddings(["hello"]) + + assert len(result) == 1 + assert result[0].vector == [0.1, 0.2, 0.3] + assert server.bodies[-1]["input"] == ["hello"] + + +async def test_requests_go_to_the_nvidia_host() -> None: + client, server = make_mock_client() + + await client.get_embeddings(["hello"]) + + assert server.requests[-1].url.host == "integrate.api.nvidia.com" + + +# region: NVIDIA options + + +async def test_input_type_and_truncate_reach_the_wire() -> None: + client, server = make_mock_client() + + await client.get_embeddings(["hello"], options={"input_type": "passage", "truncate": "END"}) + + assert server.bodies[-1]["input_type"] == "passage" + assert server.bodies[-1]["truncate"] == "END" + + +async def test_omitted_options_send_no_nvidia_fields() -> None: + client, server = make_mock_client() + + await client.get_embeddings(["hello"]) + + assert "input_type" not in server.bodies[-1] + assert "truncate" not in server.bodies[-1] + + +@pytest.mark.parametrize( + ("options", "message"), + [ + ({"input_type": "document"}, "input_type"), + ({"truncate": "end"}, "truncate"), + ], + ids=["input_type", "truncate"], +) +async def test_invalid_options_fail_before_the_request(options: dict[str, Any], message: str) -> None: + client, server = make_mock_client() + + with pytest.raises(ValueError, match=message): + await client.get_embeddings(["hello"], options=cast(Any, options)) + + assert server.requests == [] + + +@pytest.mark.filterwarnings("ignore::agent_framework._feature_stage.ExperimentalWarning") +async def test_vector_search_tool_routes_passage_and_query_to_nvidia() -> None: + @vectorstoremodel(collection_name="nvidia-input-type-test") + @dataclass + class Note: + id: Annotated[str, VectorStoreField("key")] + vector: Annotated[str | list[float] | None, VectorStoreField("vector", dimensions=3)] = None + + client, server = make_mock_client() + collection: InMemoryCollection[str, Note] = InMemoryCollection(Note, embedding_generator=client) + await collection.ensure_collection_exists() + + await collection.upsert([Note("one", "NIM hosts models")], embeddings_options={"input_type": "passage"}) + search_tool = create_vector_search_tool(collection, embeddings_options={"input_type": "query"}, top=1) + await search_tool(query="What hosts models?") + + upsert_body, search_body = server.bodies + assert (upsert_body["input_type"], upsert_body["dimensions"]) == ("passage", 3) + assert (search_body["input_type"], search_body["dimensions"]) == ("query", 3) + + +# region: batching + + +class BatchingNvidia: + """Returns one vector per input, numbered across requests, and can fail a chosen request.""" + + def __init__(self, fail_request: int | None = None) -> None: + self.batch_sizes: list[int] = [] + self.bodies: list[dict[str, Any]] = [] + self._fail_request = fail_request + self._next = 0 + + def handler(self, request: httpx2.Request) -> httpx2.Response: + body = json.loads(request.content) + self.bodies.append(body) + self.batch_sizes.append(len(body["input"])) + if len(self.batch_sizes) == self._fail_request: + return httpx2.Response( + 400, json={"object": "error", "message": "rejected", "type": "invalid_request_error"} + ) + vectors = [[float(self._next + i), 0.0, 0.0] for i in range(len(body["input"]))] + self._next += len(vectors) + payload = make_embeddings_payload(vectors) + payload["usage"] = {"prompt_tokens": len(vectors), "total_tokens": len(vectors)} + return httpx2.Response(200, json=payload) + + +def make_batching_client(server: BatchingNvidia) -> NvidiaEmbeddingClient: + async_client = AsyncOpenAI( + api_key="test-key", + base_url=DEFAULT_BASE_URL, + http_client=httpx2.AsyncClient(transport=httpx2.MockTransport(server.handler)), + ) + return NvidiaEmbeddingClient(model="test-model", async_client=async_client) + + +@pytest.mark.parametrize( + ("count", "batch_sizes"), + [(MAX_BATCH_SIZE, [MAX_BATCH_SIZE]), (MAX_BATCH_SIZE + 1, [MAX_BATCH_SIZE, 1])], + ids=["at-limit", "over-limit"], +) +async def test_requests_stay_within_the_batch_limit(count: int, batch_sizes: list[int]) -> None: + server = BatchingNvidia() + client = make_batching_client(server) + + await client.get_embeddings([f"text {i}" for i in range(count)]) + + assert server.batch_sizes == batch_sizes + + +async def test_batched_results_keep_order_usage_and_options() -> None: + server = BatchingNvidia() + client = make_batching_client(server) + + result = await client.get_embeddings([f"text {i}" for i in range(600)], options={"input_type": "passage"}) + + assert server.batch_sizes == [256, 256, 88] + assert [embedding.vector[0] for embedding in result] == [float(i) for i in range(600)] + assert result.usage == {"input_token_count": 600, "total_token_count": 600} + assert all(body["input_type"] == "passage" for body in server.bodies) + + +async def test_empty_input_sends_no_request() -> None: + server = BatchingNvidia() + client = make_batching_client(server) + + result = await client.get_embeddings([]) + + assert list(result) == [] + assert server.batch_sizes == [] + + +async def test_invalid_options_fail_before_any_batch_is_sent() -> None: + server = BatchingNvidia() + client = make_batching_client(server) + + with pytest.raises(ValueError, match="input_type"): + await client.get_embeddings([f"text {i}" for i in range(600)], options=cast(Any, {"input_type": "document"})) + + assert server.batch_sizes == [] + + +async def test_a_failed_batch_fails_the_whole_call() -> None: + server = BatchingNvidia(fail_request=2) + client = make_batching_client(server) + + with pytest.raises(Exception, match="rejected"): + await client.get_embeddings([f"text {i}" for i in range(600)]) + + assert server.batch_sizes == [256, 256] + + +@pytest.mark.parametrize( + ("values", "batch_sizes"), + [ + (["a" * 600_000] * 5, [3, 2]), + (["a" * (MAX_BATCH_BYTES + 1), "b"], [1, 1]), + (["é" * 400_000, "b"], [1, 1]), + (["short"] * (MAX_BATCH_SIZE * 2), [MAX_BATCH_SIZE, MAX_BATCH_SIZE]), + ], + ids=["by-bytes", "oversized-alone", "non-ascii-costed-as-escaped", "by-count"], +) +def test_batches_respect_count_and_byte_limits(values: list[str], batch_sizes: list[int]) -> None: + batches = list(_batches(values)) + + assert [len(batch) for batch in batches] == batch_sizes + assert [value for batch in batches for value in batch] == values + + +async def test_large_documents_are_split_by_request_size() -> None: + server = BatchingNvidia() + client = make_batching_client(server) + + result = await client.get_embeddings(["a" * 600_000] * 5) + + assert server.batch_sizes == [3, 2] + assert len(result) == 5 + + +@pytest.mark.filterwarnings("ignore::agent_framework._feature_stage.ExperimentalWarning") +async def test_upsert_larger_than_the_batch_limit() -> None: + @vectorstoremodel(collection_name="nvidia-batch-test") + @dataclass + class Note: + id: Annotated[str, VectorStoreField("key")] + vector: Annotated[str | list[float] | None, VectorStoreField("vector", dimensions=3)] = None + + server = BatchingNvidia() + client = make_batching_client(server) + collection: InMemoryCollection[str, Note] = InMemoryCollection(Note, embedding_generator=client) + await collection.ensure_collection_exists() + + await collection.upsert( + [Note(str(i), f"note {i}") for i in range(300)], embeddings_options={"input_type": "passage"} + ) + + assert server.batch_sizes == [256, 44] + stored = await collection.get(["0", "299"], include_vectors=True) + assert [note.vector for note in stored] == [[0.0, 0.0, 0.0], [299.0, 0.0, 0.0]] + + +# region: integration + + +@pytest.mark.flaky +@pytest.mark.integration +@pytest.mark.skipif(not os.getenv("NVIDIA_API_KEY"), reason="Set NVIDIA_API_KEY to run.") +async def test_nvidia_embedding_integration() -> None: + """The service embeds a passage and a query differently, which a mock cannot show. + + Repeated identical requests are not bit-for-bit equal, so compare cosine similarity: + the same text measures about 1.0 under one input type and about 0.66 across the two. + """ + client = NvidiaEmbeddingClient(model=os.getenv("NVIDIA_EMBEDDING_MODEL") or "nvidia/nemotron-3-embed-1b") + try: + passage = await client.get_embeddings(["What does NIM serve?"], options={"input_type": "passage"}) + query = await client.get_embeddings(["What does NIM serve?"], options={"input_type": "query"}) + finally: + await client.close() + + a, b = passage[0].vector, query[0].vector + similarity = sum(x * y for x, y in zip(a, b, strict=True)) / (math.hypot(*a) * math.hypot(*b)) + assert similarity < 0.99 diff --git a/python/packages/openai/AGENTS.md b/python/packages/openai/AGENTS.md index 2dc541ca802..e45e2eb1f94 100644 --- a/python/packages/openai/AGENTS.md +++ b/python/packages/openai/AGENTS.md @@ -61,6 +61,14 @@ Both default to `None` (no-op → byte-identical stock OpenAI behavior) and are **both**: the parser surfaces the field for display, the preparer sends it back. Prefer a dedicated client (e.g. `agent-framework-mistral`) when an endpoint diverges substantially. +## Adapting the embedding client to OpenAI-compatible endpoints + +`RawOpenAIEmbeddingClient.get_embeddings` forwards only the OpenAI embedding schema (`dimensions`, +`encoding_format`, `user`) and drops other option keys. A dedicated client for an OpenAI-compatible +service overrides `_prepare_extra_request_options(options)` to return extra `embeddings.create` +keyword arguments, typically `extra_body`, and to raise `ValueError` for invalid values. It runs +before the request is marked or sent. `agent-framework-nvidia` uses it for `input_type` and `truncate`. + Native Responses and Chat Completions refusals remain ordinary text content with the experimental `additional_properties["model_output_kind"] == "refusal"` marker. Assistant history reconstructs the provider's native refusal field without splitting mixed text/refusal turns; non-assistant diff --git a/python/packages/openai/agent_framework_openai/_embedding_client.py b/python/packages/openai/agent_framework_openai/_embedding_client.py index c3ab22bdd13..1b8afa2f9e4 100644 --- a/python/packages/openai/agent_framework_openai/_embedding_client.py +++ b/python/packages/openai/agent_framework_openai/_embedding_client.py @@ -251,6 +251,21 @@ def service_url(self) -> str: """Get the URL of the service.""" return str(self.client.base_url) if self.client else "Unknown" + def _prepare_extra_request_options(self, options: Mapping[str, Any]) -> dict[str, Any]: + """Return extra keyword arguments for the ``embeddings.create`` request. + + Only the OpenAI embedding schema is forwarded by default. Subclasses for + OpenAI-compatible services override this to send provider-specific fields, + typically as ``extra_body``, and to reject invalid values before the request. + + Args: + options: The embedding options for the current call. + + Returns: + Keyword arguments merged into the request. The default adds none. + """ + return {} + async def get_embeddings( self, values: Sequence[str], @@ -276,6 +291,7 @@ async def get_embeddings( model = opts.get("model") or self.model if not model: raise ValueError("model is required") + extra_request_options = self._prepare_extra_request_options(opts) kwargs: dict[str, Any] = {"input": list(values), "model": model} if self._FEATURE_USAGE_INDEX is not None: @@ -286,6 +302,7 @@ async def get_embeddings( kwargs["encoding_format"] = encoding_format if user := opts.get("user"): kwargs["user"] = user + kwargs.update(extra_request_options) response = await self.client.embeddings.create(**kwargs) diff --git a/python/packages/openai/tests/openai/test_openai_embedding_client.py b/python/packages/openai/tests/openai/test_openai_embedding_client.py index ba57b8ee754..d8a5c86d667 100644 --- a/python/packages/openai/tests/openai/test_openai_embedding_client.py +++ b/python/packages/openai/tests/openai/test_openai_embedding_client.py @@ -200,6 +200,54 @@ async def test_openai_options_passthrough_encoding_format(openai_unit_test_env: assert call_kwargs["encoding_format"] == "base64" +async def test_openai_does_not_forward_unknown_options(openai_unit_test_env: dict[str, str]) -> None: + client = OpenAIEmbeddingClient() + client.client = MagicMock() + client.client.embeddings = MagicMock() + client.client.embeddings.create = AsyncMock(return_value=_make_openai_response(embeddings=[[0.1]])) + + await client.get_embeddings(["test"], options=cast(Any, {"input_type": "query"})) + + call_kwargs = client.client.embeddings.create.call_args[1] + assert set(call_kwargs) == {"input", "model"} + + +async def test_subclass_extra_request_options_are_forwarded(openai_unit_test_env: dict[str, str]) -> None: + class CompatibleEmbeddingClient(OpenAIEmbeddingClient): + def _prepare_extra_request_options(self, options: Any) -> dict[str, Any]: + return {"extra_body": {"input_type": options["input_type"]}} + + client = CompatibleEmbeddingClient() + client.client = MagicMock() + client.client.embeddings = MagicMock() + client.client.embeddings.create = AsyncMock(return_value=_make_openai_response(embeddings=[[0.1]])) + + await client.get_embeddings(["test"], options=cast(Any, {"input_type": "query", "dimensions": 8})) + + call_kwargs = client.client.embeddings.create.call_args[1] + assert call_kwargs["extra_body"] == {"input_type": "query"} + assert call_kwargs["dimensions"] == 8 + + +async def test_rejected_extra_request_options_send_nothing(openai_unit_test_env: dict[str, str]) -> None: + class CompatibleEmbeddingClient(OpenAIEmbeddingClient): + def _prepare_extra_request_options(self, options: Any) -> dict[str, Any]: + raise ValueError("unsupported option") + + client = CompatibleEmbeddingClient() + client.client = MagicMock() + client.client.embeddings = MagicMock() + client.client.embeddings.create = AsyncMock(return_value=_make_openai_response(embeddings=[[0.1]])) + with telemetry._feature_mask_lock: + telemetry._feature_mask = 0 + + with pytest.raises(ValueError, match="unsupported option"): + await client.get_embeddings(["test"]) + + client.client.embeddings.create.assert_not_awaited() + assert get_feature_token() is None + + async def test_openai_base64_decoding(openai_unit_test_env: dict[str, str]) -> None: import base64 import struct diff --git a/python/pyproject.toml b/python/pyproject.toml index 9de798cbbba..a85c3b0805e 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -103,6 +103,7 @@ agent-framework-mem0 = { workspace = true } agent-framework-mistral = { workspace = true } agent-framework-mongodb = { workspace = true } agent-framework-monty = { workspace = true } +agent-framework-nvidia = { workspace = true } agent-framework-ollama = { workspace = true } agent-framework-openai = { workspace = true } agent-framework-orchestrations = { workspace = true } diff --git a/python/uv.lock b/python/uv.lock index 580a37e3565..dbb3144dd1a 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -58,6 +58,7 @@ members = [ "agent-framework-mistral", "agent-framework-mongodb", "agent-framework-monty", + "agent-framework-nvidia", "agent-framework-ollama", "agent-framework-openai", "agent-framework-orchestrations", @@ -904,6 +905,21 @@ requires-dist = [ { name = "pydantic-monty", specifier = ">=0.0.23,<0.1" }, ] +[[package]] +name = "agent-framework-nvidia" +version = "1.0.0a260929" +source = { editable = "packages/nvidia" } +dependencies = [ + { name = "agent-framework-core", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "agent-framework-openai", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, +] + +[package.metadata] +requires-dist = [ + { name = "agent-framework-core", editable = "packages/core" }, + { name = "agent-framework-openai", editable = "packages/openai" }, +] + [[package]] name = "agent-framework-ollama" version = "1.0.0b260813" From a8e3c03db1bf8ce816a8b0e1cafbb56c010c04c9 Mon Sep 17 00:00:00 2001 From: Manjunath Janardhan Date: Thu, 1 Oct 2026 18:49:16 +0530 Subject: [PATCH 2/5] Python: Address NVIDIA package review feedback Move the self-hosted NIM example next to the default constructor in the sample and broaden the package description. --- python/packages/nvidia/pyproject.toml | 2 +- .../nvidia/samples/nvidia_embeddings.py | 17 ++++++++--------- 2 files changed, 9 insertions(+), 10 deletions(-) diff --git a/python/packages/nvidia/pyproject.toml b/python/packages/nvidia/pyproject.toml index 2c46cc74d87..78403ffaf8d 100644 --- a/python/packages/nvidia/pyproject.toml +++ b/python/packages/nvidia/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "agent-framework-nvidia" -description = "NVIDIA NIM embeddings for Microsoft Agent Framework." +description = "NVIDIA Integrations for Microsoft Agent Framework." authors = [{ name = "Microsoft", email = "af-support@microsoft.com" }] readme = "README.md" requires-python = ">=3.10" diff --git a/python/packages/nvidia/samples/nvidia_embeddings.py b/python/packages/nvidia/samples/nvidia_embeddings.py index 1cab9d1f2bc..5e14da507b4 100644 --- a/python/packages/nvidia/samples/nvidia_embeddings.py +++ b/python/packages/nvidia/samples/nvidia_embeddings.py @@ -32,6 +32,14 @@ async def main() -> None: # The API key is read from NVIDIA_API_KEY. client = NvidiaEmbeddingClient(model="nvidia/nemotron-3-embed-1b") + # To use a self-hosted NIM container instead of the hosted endpoint, set base_url + # (also read from NVIDIA_BASE_URL): + # client = NvidiaEmbeddingClient( + # model="nvidia/nemotron-3-embed-1b", + # base_url="http://localhost:8000/v1", + # api_key="not-used-by-local-nim", + # ) + try: # 1. Generate a single embedding. result = await client.get_embeddings(["Hello, world!"]) @@ -70,15 +78,6 @@ async def main() -> None: ) async for item in results: print(f"Best match: {item['record'].text}") - - # 4. Point at a self-hosted NIM container instead of the hosted endpoint. - # base_url also reads from NVIDIA_BASE_URL. - # local_client = NvidiaEmbeddingClient( - # model="nvidia/nemotron-3-embed-1b", - # base_url="http://localhost:8000/v1", - # api_key="not-used-by-local-nim", - # ) - finally: await client.close() From ce92674186c0b04566860cdfff2d773be7cf4581 Mon Sep 17 00:00:00 2001 From: Manjunath Janardhan Date: Wed, 7 Oct 2026 18:20:01 +0530 Subject: [PATCH 3/5] Python: Fail closed when agent-framework-openai lacks the NVIDIA options hook Raise ImportError on import when RawOpenAIEmbeddingClient has no _prepare_extra_request_options, so an older agent-framework-openai cannot silently drop input_type and truncate. --- python/packages/nvidia/AGENTS.md | 3 ++- .../agent_framework_nvidia/_embedding_client.py | 9 +++++++++ python/packages/nvidia/pyproject.toml | 3 ++- .../tests/nvidia/test_nvidia_embedding_client.py | 15 +++++++++++++++ 4 files changed, 28 insertions(+), 2 deletions(-) diff --git a/python/packages/nvidia/AGENTS.md b/python/packages/nvidia/AGENTS.md index aa28aa4fb55..5651f2cee6d 100644 --- a/python/packages/nvidia/AGENTS.md +++ b/python/packages/nvidia/AGENTS.md @@ -15,7 +15,8 @@ Integration with NVIDIA NIM for embedding generation. - `input_type` and `truncate` are outside the OpenAI embedding schema. The client sends them as `extra_body` by overriding `_prepare_extra_request_options`, and rejects invalid values before the request. Callers set `input_type` per call through vector-store `embeddings_options`, the - same route the Gemini client uses for `task_type`. + same route the Gemini client uses for `task_type`. An `agent-framework-openai` without that hook + would drop both fields silently, so importing the package raises `ImportError` instead. - The hosted endpoint rejects requests with more than 256 inputs (`MAX_BATCH_SIZE`) or a body over 2 MiB (`413`, budgeted by `MAX_BATCH_BYTES`), and a vector-store upsert embeds every record in one call, so `get_embeddings` splits larger inputs into consecutive requests, as the Gemini client does diff --git a/python/packages/nvidia/agent_framework_nvidia/_embedding_client.py b/python/packages/nvidia/agent_framework_nvidia/_embedding_client.py index c78ecb0f483..bc9f0cafce5 100644 --- a/python/packages/nvidia/agent_framework_nvidia/_embedding_client.py +++ b/python/packages/nvidia/agent_framework_nvidia/_embedding_client.py @@ -34,6 +34,15 @@ "RawNvidiaEmbeddingClient", ] +# An agent-framework-openai without this hook never calls the override below, so input_type and +# truncate would be dropped from every request without an error. Refuse to load instead. Remove this +# once the agent-framework-openai floor in pyproject.toml is a release that has the hook. +if not hasattr(RawOpenAIEmbeddingClient, "_prepare_extra_request_options"): + raise ImportError( + "agent-framework-nvidia requires an agent-framework-openai release whose RawOpenAIEmbeddingClient " + "provides _prepare_extra_request_options. Upgrade agent-framework-openai." + ) + DEFAULT_BASE_URL = "https://integrate.api.nvidia.com/v1" """NVIDIA's hosted NIM endpoint, which exposes an OpenAI-compatible surface.""" diff --git a/python/packages/nvidia/pyproject.toml b/python/packages/nvidia/pyproject.toml index 78403ffaf8d..af5b49cfdc3 100644 --- a/python/packages/nvidia/pyproject.toml +++ b/python/packages/nvidia/pyproject.toml @@ -24,7 +24,8 @@ classifiers = [ dependencies = [ "agent-framework-core>=1.19.0,<2", # input_type and truncate need RawOpenAIEmbeddingClient._prepare_extra_request_options. - # 1.14.4 predates it and silently drops them, so raise this floor to the release that adds it. + # 1.15.0 and earlier lack it, so the package refuses to import with them. Raise this floor to + # the first release that adds it. "agent-framework-openai>=1.14.4,<2", ] diff --git a/python/packages/nvidia/tests/nvidia/test_nvidia_embedding_client.py b/python/packages/nvidia/tests/nvidia/test_nvidia_embedding_client.py index 9340916a6f7..ca23f4c73e3 100644 --- a/python/packages/nvidia/tests/nvidia/test_nvidia_embedding_client.py +++ b/python/packages/nvidia/tests/nvidia/test_nvidia_embedding_client.py @@ -1,5 +1,6 @@ # Copyright (c) Microsoft. All rights reserved. +import importlib.util import json import math import os @@ -10,6 +11,7 @@ import pytest from agent_framework import InMemoryCollection, VectorStoreField, create_vector_search_tool, vectorstoremodel from agent_framework._settings import SecretString +from agent_framework_openai._embedding_client import RawOpenAIEmbeddingClient # pyright: ignore[reportPrivateUsage] from openai import AsyncOpenAI from agent_framework_nvidia import NvidiaEmbeddingClient, RawNvidiaEmbeddingClient @@ -182,6 +184,19 @@ def test_telemetry_layer_wraps_the_raw_client() -> None: assert names.index("RawNvidiaEmbeddingClient") < names.index("RawOpenAIEmbeddingClient") +def test_refuses_to_load_without_the_openai_hook(monkeypatch: pytest.MonkeyPatch) -> None: + """An older agent-framework-openai would silently drop input_type and truncate, so the import must fail.""" + import agent_framework_nvidia._embedding_client as module # pyright: ignore[reportPrivateUsage] + + monkeypatch.delattr(RawOpenAIEmbeddingClient, "_prepare_extra_request_options") + # Load a fresh copy under another name so the module the other tests use is left untouched. + spec = importlib.util.spec_from_file_location("agent_framework_nvidia._embedding_client_probe", module.__file__) + assert spec is not None and spec.loader is not None + + with pytest.raises(ImportError, match="Upgrade agent-framework-openai"): + spec.loader.exec_module(importlib.util.module_from_spec(spec)) + + async def test_raw_client_sends_nvidia_options_without_telemetry() -> None: server = MockNvidia(make_embeddings_payload([[0.1, 0.2, 0.3]])) async_client = AsyncOpenAI( From d4d5f43c69d17b8534a53558b393fc81b4c4d44b Mon Sep 17 00:00:00 2001 From: Manjunath Janardhan Date: Fri, 23 Oct 2026 11:39:15 +0530 Subject: [PATCH 4/5] Python: Add Python 3.15 classifier to the NVIDIA package Also drop the trailing space the main merge left in the PACKAGE_STATUS.md row. --- python/PACKAGE_STATUS.md | 2 +- python/packages/nvidia/pyproject.toml | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/python/PACKAGE_STATUS.md b/python/PACKAGE_STATUS.md index e5d3bf30004..1ee21eb454c 100644 --- a/python/PACKAGE_STATUS.md +++ b/python/PACKAGE_STATUS.md @@ -55,7 +55,7 @@ column records whether `pip install agent-framework` installs the workspace pack | `agent-framework-mistral` | `python/packages/mistral` | `beta` | `yes` | | `agent-framework-mongodb` | `python/packages/mongodb` | `alpha` | `no` | | `agent-framework-monty` | `python/packages/monty` | `beta` | `yes` | -| `agent-framework-nvidia` | `python/packages/nvidia` | `alpha` | `no` | +| `agent-framework-nvidia` | `python/packages/nvidia` | `alpha` | `no` | | `agent-framework-ollama` | `python/packages/ollama` | `beta` | `yes` | | `agent-framework-openai` | `python/packages/openai` | `released` | `yes` | | `agent-framework-oracle` | `python/packages/oracle` | `alpha` | `no` | diff --git a/python/packages/nvidia/pyproject.toml b/python/packages/nvidia/pyproject.toml index af5b49cfdc3..b158657e599 100644 --- a/python/packages/nvidia/pyproject.toml +++ b/python/packages/nvidia/pyproject.toml @@ -19,6 +19,7 @@ classifiers = [ "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", "Programming Language :: Python :: 3.14", + "Programming Language :: Python :: 3.15", "Typing :: Typed", ] dependencies = [ From 4146e58c2df66ff1858736e62bdfdc93305e68dc Mon Sep 17 00:00:00 2001 From: Manjunath Janardhan Date: Sat, 24 Oct 2026 17:42:54 +0530 Subject: [PATCH 5/5] Python: Raise the NVIDIA agent-framework-openai floor to 1.17.0 --- python/packages/nvidia/pyproject.toml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/python/packages/nvidia/pyproject.toml b/python/packages/nvidia/pyproject.toml index b158657e599..8d582dd3615 100644 --- a/python/packages/nvidia/pyproject.toml +++ b/python/packages/nvidia/pyproject.toml @@ -25,9 +25,9 @@ classifiers = [ dependencies = [ "agent-framework-core>=1.19.0,<2", # input_type and truncate need RawOpenAIEmbeddingClient._prepare_extra_request_options. - # 1.15.0 and earlier lack it, so the package refuses to import with them. Raise this floor to - # the first release that adds it. - "agent-framework-openai>=1.14.4,<2", + # Published 1.17.0 and earlier lack it, so the package refuses to import with them. Raise this + # floor to the first release that adds it. + "agent-framework-openai>=1.17.0,<2", ] [tool.ruff]