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 8c5b083d746..1ee21eb454c 100644 --- a/python/PACKAGE_STATUS.md +++ b/python/PACKAGE_STATUS.md @@ -55,6 +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-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/AGENTS.md b/python/packages/nvidia/AGENTS.md new file mode 100644 index 00000000000..5651f2cee6d --- /dev/null +++ b/python/packages/nvidia/AGENTS.md @@ -0,0 +1,49 @@ +# 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`. 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 + 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..bc9f0cafce5 --- /dev/null +++ b/python/packages/nvidia/agent_framework_nvidia/_embedding_client.py @@ -0,0 +1,311 @@ +# 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", +] + +# 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.""" + + +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..8d582dd3615 --- /dev/null +++ b/python/packages/nvidia/pyproject.toml @@ -0,0 +1,72 @@ +[project] +name = "agent-framework-nvidia" +description = "NVIDIA Integrations 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", + "Programming Language :: Python :: 3.15", + "Typing :: Typed", +] +dependencies = [ + "agent-framework-core>=1.19.0,<2", + # input_type and truncate need RawOpenAIEmbeddingClient._prepare_extra_request_options. + # 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] +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..5e14da507b4 --- /dev/null +++ b/python/packages/nvidia/samples/nvidia_embeddings.py @@ -0,0 +1,102 @@ +# 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") + + # 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!"]) + 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}") + 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..ca23f4c73e3 --- /dev/null +++ b/python/packages/nvidia/tests/nvidia/test_nvidia_embedding_client.py @@ -0,0 +1,457 @@ +# Copyright (c) Microsoft. All rights reserved. + +import importlib.util +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 agent_framework_openai._embedding_client import RawOpenAIEmbeddingClient # pyright: ignore[reportPrivateUsage] +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") + + +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( + 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 f2478efb98d..472c0d780ce 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 df7d14d338d..2555ad9ea37 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -108,6 +108,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-oracle = { workspace = true } diff --git a/python/uv.lock b/python/uv.lock index 84130377c9f..621143e6864 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -55,6 +55,7 @@ members = [ "agent-framework-mistral", "agent-framework-mongodb", "agent-framework-monty", + "agent-framework-nvidia", "agent-framework-ollama", "agent-framework-openai", "agent-framework-oracle", @@ -913,6 +914,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.0b261022"