diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/adapter.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/adapter.py index e44acdf0ec..a08091e8f4 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/adapter.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/adapter.py @@ -11,11 +11,8 @@ from nemo_platform_plugin.client.adapter import client_from_platform - class ExampleClient(_ExampleEndpoints, NemoClient): - pass - - def make_example_client(platform: NeMoPlatform) -> ExampleClient: - return client_from_platform(platform, ExampleClient) + def make_sync_resource(platform: NeMoPlatform) -> NemoClient: + return client_from_platform(platform, NemoClient) """ from __future__ import annotations diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/bound.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/bound.py deleted file mode 100644 index 30de4d270e..0000000000 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/bound.py +++ /dev/null @@ -1,157 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Bound callables returned by endpoint descriptors. - -When an endpoint is accessed as an attribute on a :class:`NemoClient` or -:class:`AsyncNemoClient` instance, its ``__get__`` returns one of these -bound callables. The self-type overloads on ``__call__`` dispatch the -correct argument signature and return type based on ``RequestT`` and -``ResponseT``. -""" - -from __future__ import annotations - -from collections.abc import AsyncIterable, Callable, Iterable -from typing import Generic, Unpack, overload - -from nemo_platform_plugin.client.client import AsyncNemoClient, NemoClient -from nemo_platform_plugin.client.response import ( - AsyncNemoBinaryResponse, - AsyncNemoStreamResponse, - NemoBinaryResponse, - NemoResponse, - NemoStreamResponse, -) -from nemo_platform_plugin.client.types import ( - BinaryContent, - BodyRequestT, - ModelT, - PathT, - PreparedRequest, - RequestT, - ResponseT, - Stream, -) - - -class SyncBoundCall(Generic[PathT, RequestT, ResponseT]): - """Sync callable returned when an :class:`Endpoint` is accessed on a :class:`NemoClient`.""" - - def __init__(self, client: NemoClient, request_fn: Callable[..., PreparedRequest[ResponseT]]) -> None: - self._client = client - self._request_fn = request_fn - - # -- Body (RequestT is BaseModel) × response variants -- - - @overload - def __call__( - self: SyncBoundCall[PathT, BodyRequestT, BinaryContent], payload: BodyRequestT, **kw: Unpack[PathT] - ) -> NemoBinaryResponse: ... - @overload - def __call__( - self: SyncBoundCall[PathT, BodyRequestT, Stream[ModelT]], payload: BodyRequestT, **kw: Unpack[PathT] - ) -> NemoStreamResponse[ModelT]: ... - @overload - def __call__( - self: SyncBoundCall[PathT, BodyRequestT, ResponseT], payload: BodyRequestT, **kw: Unpack[PathT] - ) -> NemoResponse[ResponseT]: ... - - # -- Binary (RequestT is BinaryContent) × response variants -- - - @overload - def __call__( - self: SyncBoundCall[PathT, BinaryContent, BinaryContent], - content: bytes | Iterable[bytes] | AsyncIterable[bytes], - **kw: Unpack[PathT], - ) -> NemoBinaryResponse: ... - @overload - def __call__( - self: SyncBoundCall[PathT, BinaryContent, Stream[ModelT]], - content: bytes | Iterable[bytes] | AsyncIterable[bytes], - **kw: Unpack[PathT], - ) -> NemoStreamResponse[ModelT]: ... - @overload - def __call__( - self: SyncBoundCall[PathT, BinaryContent, ResponseT], - content: bytes | Iterable[bytes] | AsyncIterable[bytes], - **kw: Unpack[PathT], - ) -> NemoResponse[ResponseT]: ... - - # -- No body (RequestT is None) × response variants -- - - @overload - def __call__(self: SyncBoundCall[PathT, None, BinaryContent], **kw: Unpack[PathT]) -> NemoBinaryResponse: ... - @overload - def __call__( - self: SyncBoundCall[PathT, None, Stream[ModelT]], **kw: Unpack[PathT] - ) -> NemoStreamResponse[ModelT]: ... - @overload - def __call__(self: SyncBoundCall[PathT, None, ResponseT], **kw: Unpack[PathT]) -> NemoResponse[ResponseT]: ... - - def __call__(self, *args: object, **kw: object) -> NemoResponse | NemoBinaryResponse | NemoStreamResponse: - return self._client.send(self._request_fn(*args, **kw)) - - -class AsyncBoundCall(Generic[PathT, RequestT, ResponseT]): - """Async callable returned when an :class:`Endpoint` is accessed on an :class:`AsyncNemoClient`.""" - - def __init__(self, client: AsyncNemoClient, request_fn: Callable[..., PreparedRequest[ResponseT]]) -> None: - self._client = client - self._request_fn = request_fn - - # -- Body (RequestT is BaseModel) × response variants -- - - @overload - async def __call__( - self: AsyncBoundCall[PathT, BodyRequestT, BinaryContent], payload: BodyRequestT, **kw: Unpack[PathT] - ) -> AsyncNemoBinaryResponse: ... - @overload - async def __call__( - self: AsyncBoundCall[PathT, BodyRequestT, Stream[ModelT]], payload: BodyRequestT, **kw: Unpack[PathT] - ) -> AsyncNemoStreamResponse[ModelT]: ... - @overload - async def __call__( - self: AsyncBoundCall[PathT, BodyRequestT, ResponseT], payload: BodyRequestT, **kw: Unpack[PathT] - ) -> NemoResponse[ResponseT]: ... - - # -- Binary (RequestT is BinaryContent) × response variants -- - - @overload - async def __call__( - self: AsyncBoundCall[PathT, BinaryContent, BinaryContent], - content: bytes | Iterable[bytes] | AsyncIterable[bytes], - **kw: Unpack[PathT], - ) -> AsyncNemoBinaryResponse: ... - @overload - async def __call__( - self: AsyncBoundCall[PathT, BinaryContent, Stream[ModelT]], - content: bytes | Iterable[bytes] | AsyncIterable[bytes], - **kw: Unpack[PathT], - ) -> AsyncNemoStreamResponse[ModelT]: ... - @overload - async def __call__( - self: AsyncBoundCall[PathT, BinaryContent, ResponseT], - content: bytes | Iterable[bytes] | AsyncIterable[bytes], - **kw: Unpack[PathT], - ) -> NemoResponse[ResponseT]: ... - - # -- No body (RequestT is None) × response variants -- - - @overload - async def __call__( - self: AsyncBoundCall[PathT, None, BinaryContent], **kw: Unpack[PathT] - ) -> AsyncNemoBinaryResponse: ... - @overload - async def __call__( - self: AsyncBoundCall[PathT, None, Stream[ModelT]], **kw: Unpack[PathT] - ) -> AsyncNemoStreamResponse[ModelT]: ... - @overload - async def __call__( - self: AsyncBoundCall[PathT, None, ResponseT], **kw: Unpack[PathT] - ) -> NemoResponse[ResponseT]: ... - - async def __call__( - self, *args: object, **kw: object - ) -> NemoResponse | AsyncNemoBinaryResponse | AsyncNemoStreamResponse: - return await self._client.send(self._request_fn(*args, **kw)) diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/client.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/client.py index ce8a7e5e60..8ea210eed0 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/client.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/client.py @@ -26,10 +26,9 @@ NemoResponse, NemoStreamResponse, ) -from nemo_platform_plugin.client.types import BinaryContent, PreparedRequest, Stream +from nemo_platform_plugin.client.types import BinaryContent, PreparedRequest, ResponseT, Stream from pydantic import BaseModel -ResponseT = TypeVar("ResponseT", bound=BaseModel | None) ModelT = TypeVar("ModelT", bound=BaseModel) DEFAULT_TIMEOUT = 60.0 @@ -80,9 +79,12 @@ def _resolve_path(self, request: PreparedRequest) -> str: return self._base_url + path def _request_headers(self, request: PreparedRequest) -> dict[str, str] | None: + headers: dict[str, str] = {} if request.content_type is not None: - return {"Content-Type": request.content_type} - return None + headers["Content-Type"] = request.content_type + if request.extra_headers: + headers.update(request.extra_headers) + return headers or None def _is_binary(self, request: PreparedRequest) -> bool: return request.response_type is BinaryContent @@ -90,6 +92,13 @@ def _is_binary(self, request: PreparedRequest) -> bool: def _is_stream(self, request: PreparedRequest) -> bool: return get_origin(request.response_type) is Stream + def _resolve_query_params(self, request: PreparedRequest) -> dict[str, str | int | bool] | None: + """Filter out None values from query params for httpx.""" + if request.query_params is None: + return None + filtered = {k: v for k, v in request.query_params.items() if v is not None} + return filtered or None + class NemoClient(BaseNemoClient): """Sync HTTP client for NeMo Platform APIs.""" @@ -110,44 +119,62 @@ def __init__( ) @overload - def send(self, request: PreparedRequest[BinaryContent]) -> NemoBinaryResponse: ... + def send( + self, request: PreparedRequest[BinaryContent], *, headers: dict[str, str] | None = None + ) -> NemoBinaryResponse: ... @overload - def send(self, request: PreparedRequest[Stream[ModelT]]) -> NemoStreamResponse[ModelT]: ... + def send( + self, request: PreparedRequest[Stream[ModelT]], *, headers: dict[str, str] | None = None + ) -> NemoStreamResponse[ModelT]: ... @overload - def send(self, request: PreparedRequest[None]) -> NemoResponse[None]: ... + def send(self, request: PreparedRequest[None], *, headers: dict[str, str] | None = None) -> NemoResponse[None]: ... @overload - def send(self, request: PreparedRequest[ResponseT]) -> NemoResponse[ResponseT]: ... + def send( + self, request: PreparedRequest[ResponseT], *, headers: dict[str, str] | None = None + ) -> NemoResponse[ResponseT]: ... - def send(self, request: PreparedRequest) -> NemoResponse | NemoBinaryResponse | NemoStreamResponse: + def send( + self, request: PreparedRequest, *, headers: dict[str, str] | None = None + ) -> NemoResponse | NemoBinaryResponse | NemoStreamResponse: """Send a prepared request and return a typed response. - The return type is determined by the endpoint's ``ResponseT``. + Args: + request: The prepared request to send. + headers: Optional per-request headers merged on top of client + defaults and content-type headers. For binary and streaming endpoints, the caller should use the response as a context manager to ensure the connection is closed:: - with client.send(DownloadEndpoint.request(...)) as resp: + with client.send(endpoints.download(name="file.csv")) as resp: for chunk in resp: f.write(chunk) """ + if headers: + request = request.with_headers(headers) url = self._resolve_path(request) - headers = self._request_headers(request) + req_headers = self._request_headers(request) + params = self._resolve_query_params(request) if self._is_binary(request): - stream_ctx = self._http.stream(request.method, url, content=request.content, headers=headers) - return NemoBinaryResponse(stream_ctx) + stream_ctx = self._http.stream( + request.method, url, content=request.content, headers=req_headers, params=params + ) + return NemoBinaryResponse(stream_ctx, request) if self._is_stream(request): assert request.response_type is not None - stream_ctx = self._http.stream(request.method, url, content=request.content, headers=headers) + stream_ctx = self._http.stream( + request.method, url, content=request.content, headers=req_headers, params=params + ) model_type = _get_stream_model_type(request.response_type) - return NemoStreamResponse(stream_ctx, model_type) + return NemoStreamResponse(stream_ctx, model_type, request) - raw = self._http.request(request.method, url, content=request.content, headers=headers) + raw = self._http.request(request.method, url, content=request.content, headers=req_headers, params=params) body = None if raw.is_success and request.response_type is not None: body = request.response_type.model_validate(raw.json()) - return NemoResponse(http_response=raw, body=body) + return NemoResponse(http_response=raw, body=body, request=request) class AsyncNemoClient(BaseNemoClient): @@ -172,31 +199,48 @@ def __init__( ) @overload - async def send(self, request: PreparedRequest[BinaryContent]) -> AsyncNemoBinaryResponse: ... + async def send( + self, request: PreparedRequest[BinaryContent], *, headers: dict[str, str] | None = None + ) -> AsyncNemoBinaryResponse: ... @overload - async def send(self, request: PreparedRequest[Stream[ModelT]]) -> AsyncNemoStreamResponse[ModelT]: ... + async def send( + self, request: PreparedRequest[Stream[ModelT]], *, headers: dict[str, str] | None = None + ) -> AsyncNemoStreamResponse[ModelT]: ... @overload - async def send(self, request: PreparedRequest[None]) -> NemoResponse[None]: ... + async def send( + self, request: PreparedRequest[None], *, headers: dict[str, str] | None = None + ) -> NemoResponse[None]: ... @overload - async def send(self, request: PreparedRequest[ResponseT]) -> NemoResponse[ResponseT]: ... + async def send( + self, request: PreparedRequest[ResponseT], *, headers: dict[str, str] | None = None + ) -> NemoResponse[ResponseT]: ... - async def send(self, request: PreparedRequest) -> NemoResponse | AsyncNemoBinaryResponse | AsyncNemoStreamResponse: + async def send( + self, request: PreparedRequest, *, headers: dict[str, str] | None = None + ) -> NemoResponse | AsyncNemoBinaryResponse | AsyncNemoStreamResponse: """Send a prepared request and return a typed response.""" + if headers: + request = request.with_headers(headers) url = self._resolve_path(request) - headers = self._request_headers(request) + req_headers = self._request_headers(request) + params = self._resolve_query_params(request) if self._is_binary(request): - stream_ctx = self._http.stream(request.method, url, content=request.content, headers=headers) - return AsyncNemoBinaryResponse(stream_ctx) + stream_ctx = self._http.stream( + request.method, url, content=request.content, headers=req_headers, params=params + ) + return AsyncNemoBinaryResponse(stream_ctx, request) if self._is_stream(request): assert request.response_type is not None - stream_ctx = self._http.stream(request.method, url, content=request.content, headers=headers) + stream_ctx = self._http.stream( + request.method, url, content=request.content, headers=req_headers, params=params + ) model_type = _get_stream_model_type(request.response_type) - return AsyncNemoStreamResponse(stream_ctx, model_type) + return AsyncNemoStreamResponse(stream_ctx, model_type, request) - raw = await self._http.request(request.method, url, content=request.content, headers=headers) + raw = await self._http.request(request.method, url, content=request.content, headers=req_headers, params=params) body = None if raw.is_success and request.response_type is not None: body = request.response_type.model_validate(raw.json()) - return NemoResponse(http_response=raw, body=body) + return NemoResponse(http_response=raw, body=body, request=request) diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/endpoint.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/endpoint.py index 9835035551..7999a2594c 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/endpoint.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/endpoint.py @@ -1,158 +1,159 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Typed endpoint definitions and factory functions. - -Endpoints are descriptors: when assigned as class attributes on a -:class:`NemoClient` or :class:`AsyncNemoClient` subclass, accessing them -returns a bound callable that sends the request and returns the typed response. - -Define endpoints once in a mixin, then create sync and async client classes:: - - class _ItemEndpoints: - create = post("/items", path_type=WorkspacePath, request_type=CreateItemRequest, response_type=ItemResponse) - get_item = get("/items/{name}", path_type=WorkspaceItemPath, response_type=ItemResponse) - - class ItemsClient(_ItemEndpoints, NemoClient): - pass - - class AsyncItemsClient(_ItemEndpoints, AsyncNemoClient): - pass +"""Typed endpoint definitions using ParamSpec-based decorators. + +Endpoints are declared as decorated methods on a class. The decorator +replaces each method with a callable that builds a :class:`PreparedRequest`, +preserving the original call signature for autocomplete and type checking:: + + class ExampleEndpoints: + @post("/apis/example/v2/workspaces/{workspace}/items") + def create_item(self, body: CreateItemRequest, *, workspace: str) -> Item: + raise NotImplementedError + + @get("/apis/example/hello/{name}") + def hello(self, *, name: str) -> HelloResponse: + raise NotImplementedError + + endpoints = ExampleEndpoints() + req = endpoints.create_item(workspace="default", body=CreateItemRequest(name="x")) + resp = client.send(req) # NemoResponse[Item] + +Parameter conventions: +- ``body`` — JSON request body (Pydantic model, serialized automatically) +- ``content`` — binary request body (raw bytes) +- ``query_params`` — query parameters (dict or TypedDict) +- All other keyword parameters — path parameters (matched to ``{placeholders}`` in the path template) """ from __future__ import annotations -from collections.abc import AsyncIterable, Iterable -from typing import Generic, Unpack, overload +import functools +import inspect +import string +from collections.abc import AsyncIterable, Callable, Iterable +from typing import get_type_hints -from nemo_platform_plugin.client.bound import AsyncBoundCall, SyncBoundCall -from nemo_platform_plugin.client.client import AsyncNemoClient, NemoClient from nemo_platform_plugin.client.types import ( - BinaryContent, - BodyRequestT, - PathT, + P, PreparedRequest, - RequestT, ResponseT, ) from pydantic import BaseModel -class Endpoint(Generic[PathT, RequestT, ResponseT]): - """A typed HTTP endpoint definition. - - Links a path type ``PathT``, request type ``RequestT``, and response type - ``ResponseT`` together with the HTTP method and path template. +def _build_prepared_request( + method: str, + path: str, + sig: inspect.Signature, + path_param_names: set[str], + response_type: type | None, + args: tuple, + kwargs: dict, +) -> PreparedRequest: + """Build a PreparedRequest by binding call arguments to the endpoint signature. - Also a descriptor: when assigned as a class attribute on a - :class:`NemoClient` or :class:`AsyncNemoClient` subclass, accessing it - returns a :class:`SyncBoundCall` or :class:`AsyncBoundCall`. + Uses ``bind_partial`` so that path parameters with client-level defaults + (e.g. ``workspace``) can be omitted by the caller. """ - - def __init__( - self, path: str, method: str, request_type: type[RequestT] | None, response_type: type[ResponseT] | None - ) -> None: - self.path = path - self.method = method - self.request_type = request_type - self.response_type = response_type - - # -- request() overloads: body / binary / no-body ---------------------- - - @overload - def request( - self: Endpoint[PathT, BodyRequestT, ResponseT], payload: BodyRequestT, **path_params: Unpack[PathT] - ) -> PreparedRequest[ResponseT]: ... - @overload - def request( - self: Endpoint[PathT, BinaryContent, ResponseT], - content: bytes | Iterable[bytes] | AsyncIterable[bytes], - **path_params: Unpack[PathT], - ) -> PreparedRequest[ResponseT]: ... - @overload - def request(self: Endpoint[PathT, None, ResponseT], **path_params: Unpack[PathT]) -> PreparedRequest[ResponseT]: ... - - def request(self, *args: object, **path_params: object) -> PreparedRequest: - """Build a :class:`PreparedRequest` from payload/content and path parameters.""" - params = {k: str(v) for k, v in path_params.items()} - content: bytes | Iterable[bytes] | AsyncIterable[bytes] | None - content_type: str | None - - if self.request_type is None: - content = None - content_type = None - elif self.request_type is BinaryContent: - content = args[0] # type: ignore[assignment] - content_type = "application/octet-stream" - else: - payload = args[0] - assert isinstance(payload, BaseModel) - content = payload.model_dump_json().encode() + bound = sig.bind_partial(*args, **kwargs) + bound.apply_defaults() + + path_params: dict[str, str] = {} + query_params: dict[str, str | int | bool | None] | None = None + content: bytes | Iterable[bytes] | AsyncIterable[bytes] | None = None + content_type: str | None = None + + for name, value in bound.arguments.items(): + if name == "self": + continue + if name in path_param_names: + if value is not None: + path_params[name] = str(value) + elif name == "body": + if not isinstance(value, BaseModel): + raise TypeError(f"body must be a BaseModel instance, got {type(value).__name__}") + content = value.model_dump_json().encode() content_type = "application/json" + elif name == "content": + content = value + content_type = "application/octet-stream" + elif name == "query_params": + if value is not None: + query_params = dict(value) + + return PreparedRequest( + path_template=path, + path_params=path_params, + method=method, + content=content, + content_type=content_type, + response_type=response_type, + query_params=query_params, + ) + + +def _make_endpoint(http_method: str, path: str, fn: Callable[P, ResponseT]) -> Callable[P, PreparedRequest[ResponseT]]: + """Create a callable that builds PreparedRequests from the function's signature.""" + sig = inspect.signature(fn) + path_param_names = {field_name for _, field_name, _, _ in string.Formatter().parse(path) if field_name} + hints = get_type_hints(fn) + ret = hints.get("return") + response_type = ret if ret is not None and ret is not type(None) else None - return PreparedRequest( - path_template=self.path, - path_params=params, - method=self.method, - content=content, - content_type=content_type, - response_type=self.response_type, - ) - - # -- Descriptor: sync vs async ----------------------------------------- - - @overload - def __get__(self, obj: NemoClient, objtype: type | None = None) -> SyncBoundCall[PathT, RequestT, ResponseT]: ... - @overload - def __get__( - self, obj: AsyncNemoClient, objtype: type | None = None - ) -> AsyncBoundCall[PathT, RequestT, ResponseT]: ... - - def __get__( - self, obj: NemoClient | AsyncNemoClient | None, objtype: type | None = None - ) -> SyncBoundCall[PathT, RequestT, ResponseT] | AsyncBoundCall[PathT, RequestT, ResponseT]: - assert obj is not None - if isinstance(obj, AsyncNemoClient): - return AsyncBoundCall(obj, self.request) - return SyncBoundCall(obj, self.request) - - def __repr__(self) -> str: - req = self.request_type.__name__ if self.request_type else "None" - resp = self.response_type.__name__ if self.response_type else "None" - return f"Endpoint({self.method} {self.path}, {req} -> {resp})" + @functools.wraps(fn) + def prepare(*args: P.args, **kwargs: P.kwargs) -> PreparedRequest[ResponseT]: + return _build_prepared_request(http_method, path, sig, path_param_names, response_type, args, kwargs) + + return prepare # --------------------------------------------------------------------------- -# Factory functions +# Decorator factories # --------------------------------------------------------------------------- -def get(path: str, path_type: type[PathT], response_type: type[ResponseT]) -> Endpoint[PathT, None, ResponseT]: +def get(path: str) -> Callable[[Callable[P, ResponseT]], Callable[P, PreparedRequest[ResponseT]]]: """Define a GET endpoint (no request body).""" - return Endpoint(path, "GET", None, response_type) + + def decorator(fn: Callable[P, ResponseT]) -> Callable[P, PreparedRequest[ResponseT]]: + return _make_endpoint("GET", path, fn) + + return decorator -def post( - path: str, path_type: type[PathT], request_type: type[RequestT], response_type: type[ResponseT] -) -> Endpoint[PathT, RequestT, ResponseT]: +def post(path: str) -> Callable[[Callable[P, ResponseT]], Callable[P, PreparedRequest[ResponseT]]]: """Define a POST endpoint.""" - return Endpoint(path, "POST", request_type, response_type) + def decorator(fn: Callable[P, ResponseT]) -> Callable[P, PreparedRequest[ResponseT]]: + return _make_endpoint("POST", path, fn) -def put( - path: str, path_type: type[PathT], request_type: type[RequestT], response_type: type[ResponseT] -) -> Endpoint[PathT, RequestT, ResponseT]: + return decorator + + +def put(path: str) -> Callable[[Callable[P, ResponseT]], Callable[P, PreparedRequest[ResponseT]]]: """Define a PUT endpoint.""" - return Endpoint(path, "PUT", request_type, response_type) + def decorator(fn: Callable[P, ResponseT]) -> Callable[P, PreparedRequest[ResponseT]]: + return _make_endpoint("PUT", path, fn) + + return decorator -def patch( - path: str, path_type: type[PathT], request_type: type[RequestT], response_type: type[ResponseT] -) -> Endpoint[PathT, RequestT, ResponseT]: + +def patch(path: str) -> Callable[[Callable[P, ResponseT]], Callable[P, PreparedRequest[ResponseT]]]: """Define a PATCH endpoint.""" - return Endpoint(path, "PATCH", request_type, response_type) + def decorator(fn: Callable[P, ResponseT]) -> Callable[P, PreparedRequest[ResponseT]]: + return _make_endpoint("PATCH", path, fn) + + return decorator + + +def delete(path: str) -> Callable[[Callable[P, ResponseT]], Callable[P, PreparedRequest[ResponseT]]]: + """Define a DELETE endpoint (no request body, optional response body).""" + + def decorator(fn: Callable[P, ResponseT]) -> Callable[P, PreparedRequest[ResponseT]]: + return _make_endpoint("DELETE", path, fn) -def delete(path: str, path_type: type[PathT]) -> Endpoint[PathT, None, None]: - """Define a DELETE endpoint (no request body, no response body).""" - return Endpoint(path, "DELETE", None, None) + return decorator diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/method.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/method.py new file mode 100644 index 0000000000..3c53d074bb --- /dev/null +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/method.py @@ -0,0 +1,80 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Optional convenience layer: turn endpoint methods into client methods. + +Plugin authors define endpoints once in a collection class, then use +``method()`` to bridge them onto a client class:: + + class _ExampleMethods: + hello = method(ExampleEndpoints.hello) + create_item = method(ExampleEndpoints.create_item) + + class ExampleClient(_ExampleMethods, NemoClient): pass + class AsyncExampleClient(_ExampleMethods, AsyncNemoClient): pass + + client = ExampleClient(base_url="...", workspace="default") + resp = client.hello(name="alice") # NemoResponse[HelloResponse] + +The descriptor dispatches sync vs async based on the client type. + +Note: ``ty`` shows ``Unknown |`` on the method types due to unannotated +class attributes (astral-sh/ty#3254). The types themselves are correct +and ``pyright`` resolves them cleanly. +""" + +from __future__ import annotations + +import functools +from collections.abc import Callable +from typing import Any, Coroutine, Generic, overload + +from nemo_platform_plugin.client.client import AsyncNemoClient, NemoClient +from nemo_platform_plugin.client.response import NemoResponse +from nemo_platform_plugin.client.types import P, PreparedRequest, ResponseT + + +class EndpointMethod(Generic[P, ResponseT]): + """Descriptor that binds an endpoint to a client instance. + + When accessed on a :class:`NemoClient`, returns a sync callable. + When accessed on an :class:`AsyncNemoClient`, returns an async callable. + Both preserve the endpoint's full ``ParamSpec`` signature. + """ + + def __init__(self, endpoint_fn: Callable[P, PreparedRequest[ResponseT]]) -> None: + self._endpoint_fn = endpoint_fn + + @overload + def __get__(self, obj: NemoClient, objtype: type | None = None) -> Callable[P, NemoResponse[ResponseT]]: ... + @overload + def __get__( + self, obj: AsyncNemoClient, objtype: type | None = None + ) -> Callable[P, Coroutine[Any, Any, NemoResponse[ResponseT]]]: ... + + def __get__(self, obj: NemoClient | AsyncNemoClient | None, objtype: type | None = None) -> object: + assert obj is not None + if isinstance(obj, AsyncNemoClient): + + @functools.wraps(self._endpoint_fn) + async def async_bound(*args: P.args, **kwargs: P.kwargs) -> NemoResponse[ResponseT]: + return await obj.send(self._endpoint_fn(*args, **kwargs)) + + return async_bound + + @functools.wraps(self._endpoint_fn) + def sync_bound(*args: P.args, **kwargs: P.kwargs) -> NemoResponse[ResponseT]: + return obj.send(self._endpoint_fn(*args, **kwargs)) # type: ignore[return-value] + + return sync_bound + + +def method(endpoint_fn: Callable[P, PreparedRequest[ResponseT]]) -> EndpointMethod[P, ResponseT]: + """Create an :class:`EndpointMethod` descriptor from an endpoint method. + + Usage:: + + class _MyMethods: + create_item = method(MyEndpoints.create_item) + """ + return EndpointMethod(endpoint_fn) diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/response.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/response.py index 25d67ec8c1..3decc8f389 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/response.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/response.py @@ -12,10 +12,10 @@ from typing import Generic, TypeVar import httpx -from nemo_platform_plugin.client.types import BinaryContent, Stream +from nemo_platform_plugin.client.types import PreparedRequest from pydantic import BaseModel -ResponseT = TypeVar("ResponseT", bound=BaseModel | BinaryContent | Stream | None) +ResponseT = TypeVar("ResponseT") ModelT = TypeVar("ModelT", bound=BaseModel) @@ -25,7 +25,7 @@ class NemoResponse(Generic[ResponseT]): Example:: - resp = client.send(GetUserEndpoint.request(workspace="default")) + resp = client.send(endpoints.get_user(workspace="default")) resp.body # UserResponse resp.http_response # full httpx.Response @@ -34,11 +34,12 @@ class NemoResponse(Generic[ResponseT]): http_response: httpx.Response body: ResponseT + request: PreparedRequest def data(self) -> ResponseT: """Return the body if the status is 2xx, otherwise raise.""" if not (200 <= self.http_response.status_code < 300): - raise NemoHTTPError(self.http_response, self.body) + raise NemoHTTPError(self.http_response) return self.body @@ -52,14 +53,15 @@ class NemoBinaryResponse: Use as a context manager:: - with client.send(DownloadEndpoint.request(...)) as resp: + with client.send(endpoints.download(...)) as resp: data = resp.read() # all bytes at once # or: for chunk in resp # iterate chunks """ - def __init__(self, stream_ctx: AbstractContextManager[httpx.Response]) -> None: + def __init__(self, stream_ctx: AbstractContextManager[httpx.Response], request: PreparedRequest) -> None: self._stream_ctx = stream_ctx self._response: httpx.Response | None = None + self.request = request @property def http_response(self) -> httpx.Response: @@ -89,15 +91,21 @@ class NemoStreamResponse(Generic[ModelT]): Use as a context manager:: - with client.send(ChatEndpoint.request(...)) as resp: + with client.send(ChatEndpoint(...)) as resp: for chunk in resp: print(chunk.text) """ - def __init__(self, stream_ctx: AbstractContextManager[httpx.Response], model_type: type[ModelT]) -> None: + def __init__( + self, + stream_ctx: AbstractContextManager[httpx.Response], + model_type: type[ModelT], + request: PreparedRequest, + ) -> None: self._stream_ctx = stream_ctx self._model_type = model_type self._response: httpx.Response | None = None + self.request = request @property def http_response(self) -> httpx.Response: @@ -131,14 +139,15 @@ class AsyncNemoBinaryResponse: Use as an async context manager:: - async with client.send(DownloadEndpoint.request(...)) as resp: + async with client.send(endpoints.download(...)) as resp: data = await resp.read() # all bytes at once # or: async for chunk in resp # iterate chunks """ - def __init__(self, stream_ctx: AbstractAsyncContextManager[httpx.Response]) -> None: + def __init__(self, stream_ctx: AbstractAsyncContextManager[httpx.Response], request: PreparedRequest) -> None: self._stream_ctx = stream_ctx self._response: httpx.Response | None = None + self.request = request @property def http_response(self) -> httpx.Response: @@ -169,15 +178,21 @@ class AsyncNemoStreamResponse(Generic[ModelT]): Use as an async context manager:: - async with client.send(ChatEndpoint.request(...)) as resp: + async with client.send(ChatEndpoint(...)) as resp: async for chunk in resp: print(chunk.text) """ - def __init__(self, stream_ctx: AbstractAsyncContextManager[httpx.Response], model_type: type[ModelT]) -> None: + def __init__( + self, + stream_ctx: AbstractAsyncContextManager[httpx.Response], + model_type: type[ModelT], + request: PreparedRequest, + ) -> None: self._stream_ctx = stream_ctx self._model_type = model_type self._response: httpx.Response | None = None + self.request = request @property def http_response(self) -> httpx.Response: @@ -207,9 +222,28 @@ async def __aexit__( class NemoHTTPError(Exception): - """Raised by :meth:`NemoResponse.data` on non-2xx responses.""" + """Raised by :meth:`NemoResponse.data` on non-2xx responses. + + Attributes: + http_response: The raw httpx response. + status_code: The HTTP status code. + detail: A human-readable error message extracted from the response + body (``{"detail": "..."}`` convention used by FastAPI / NeMo + Platform), or the raw response text as a fallback. + """ - def __init__(self, http_response: httpx.Response, body: object) -> None: + def __init__(self, http_response: httpx.Response) -> None: self.http_response = http_response - self.body = body - super().__init__(f"HTTP {http_response.status_code}") + self.status_code = http_response.status_code + self.detail = self._extract_detail(http_response) + super().__init__(f"HTTP {self.status_code}: {self.detail}") + + @staticmethod + def _extract_detail(resp: httpx.Response) -> str: + try: + body = resp.json() + if isinstance(body, dict) and isinstance(body.get("detail"), str): + return body["detail"] + except Exception: + pass + return resp.text diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/types.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/types.py index 843be75c95..2a0ac85b33 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/types.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/types.py @@ -10,54 +10,36 @@ from __future__ import annotations from collections.abc import AsyncIterable, Iterable -from dataclasses import dataclass -from typing import Generic, NotRequired, TypedDict, TypeVar +from dataclasses import dataclass, replace +from typing import Generic, ParamSpec, TypeVar from pydantic import BaseModel +P = ParamSpec("P") ModelT = TypeVar("ModelT", bound=BaseModel) -BodyRequestT = TypeVar("BodyRequestT", bound=BaseModel) +ResponseT = TypeVar("ResponseT") class BinaryContent: """Marker type: endpoint sends or receives raw bytes. - Use as ``request_type`` for binary uploads or ``response_type`` for - binary downloads:: + Use ``content`` parameter for binary uploads:: - UploadEndpoint = put("/files/{path}", path_type=FilePath, request_type=BinaryContent, response_type=FileResponse) - DownloadEndpoint = get("/files/{path}", path_type=FilePath, response_type=BinaryContent) + @put("/files/{path}") + def UploadEndpoint(content: bytes, *, path: str) -> FileResponse: ... """ class Stream(Generic[ModelT]): """Marker type: endpoint returns a stream of ``ModelT`` objects (SSE/NDJSON). - Used as ``response_type`` in endpoint definitions:: + Used as return type in endpoint definitions:: - ChatEndpoint = post("/chat/{workspace}", path_type=WorkspacePath, request_type=ChatRequest, response_type=Stream[ChatChunk]) + @post("/chat/{workspace}") + def ChatEndpoint(body: ChatRequest, *, workspace: str) -> Stream[ChatChunk]: ... """ -class PathParams(TypedDict): - """Base class for all path parameter types. - - All path TypedDicts must inherit from this so that ``PathT`` is - properly constrained. - """ - - -class WorkspaceParams(PathParams): - """Path params with an optional workspace (filled by client default).""" - - workspace: NotRequired[str] - - -PathT = TypeVar("PathT", bound=PathParams) -RequestT = TypeVar("RequestT", bound=BaseModel | BinaryContent | None) -ResponseT = TypeVar("ResponseT", bound=BaseModel | BinaryContent | Stream | None) - - @dataclass(frozen=True, slots=True) class PreparedRequest(Generic[ResponseT]): """A request ready to be sent — carries the endpoint metadata and payload. @@ -73,3 +55,10 @@ class PreparedRequest(Generic[ResponseT]): content: bytes | Iterable[bytes] | AsyncIterable[bytes] | None content_type: str | None response_type: type[ResponseT] | None + query_params: dict[str, str | int | bool | None] | None = None + extra_headers: dict[str, str] | None = None + + def with_headers(self, headers: dict[str, str]) -> PreparedRequest[ResponseT]: + """Return a new PreparedRequest with additional headers merged in.""" + merged = {**(self.extra_headers or {}), **headers} + return replace(self, extra_headers=merged) diff --git a/packages/nemo_platform_plugin/tests/client/test_client.py b/packages/nemo_platform_plugin/tests/client/test_client.py index 704badfde2..2ea7569ce6 100644 --- a/packages/nemo_platform_plugin/tests/client/test_client.py +++ b/packages/nemo_platform_plugin/tests/client/test_client.py @@ -3,15 +3,13 @@ from __future__ import annotations -from typing import NotRequired from unittest.mock import AsyncMock, MagicMock import httpx import pytest from nemo_platform_plugin.client.client import AsyncNemoClient, NemoClient from nemo_platform_plugin.client.endpoint import delete, get, post -from nemo_platform_plugin.client.response import NemoResponse -from nemo_platform_plugin.client.types import PathParams +from nemo_platform_plugin.client.response import NemoHTTPError, NemoResponse from pydantic import BaseModel BASE = "http://test:8000" @@ -26,30 +24,29 @@ class ItemResponse(BaseModel): name: str -class EmptyPath(PathParams): - pass +@post("/apis/test/v2/items") +def CREATE_ITEM(body: ItemRequest) -> ItemResponse: + raise NotImplementedError -class NamePath(PathParams): - name: str - - -class WorkspacePath(PathParams): - workspace: NotRequired[str] +@get("/apis/test/v2/items/{name}") +def GET_ITEM(*, name: str) -> ItemResponse: + raise NotImplementedError -CREATE_ITEM = post("/apis/test/v2/items", path_type=EmptyPath, request_type=ItemRequest, response_type=ItemResponse) -GET_ITEM = get("/apis/test/v2/items/{name}", path_type=NamePath, response_type=ItemResponse) -DELETE_ITEM = delete("/apis/test/v2/items/{name}", path_type=NamePath) -GET_WS_ITEM = get("/apis/test/v2/workspaces/{workspace}/items", path_type=WorkspacePath, response_type=ItemResponse) +@delete("/apis/test/v2/items/{name}") +def DELETE_ITEM(*, name: str) -> None: + raise NotImplementedError -class StubClient(NemoClient): - pass +@get("/apis/test/v2/workspaces/{workspace}/items") +def GET_WS_ITEM(*, workspace: str | None = None) -> ItemResponse: + raise NotImplementedError -class AsyncStubClient(AsyncNemoClient): - pass +@get("/apis/test/v2/items") +def GET_ITEMS_WITH_PARAMS(*, query_params: dict | None = None) -> ItemResponse: + raise NotImplementedError # --------------------------------------------------------------------------- @@ -65,8 +62,8 @@ def test_send_post() -> None: json={"id": 1, "name": "alice"}, ) - client = StubClient(base_url=BASE, http_client=mock_http) - resp = client.send(CREATE_ITEM.request(ItemRequest(name="alice"))) + client = NemoClient(base_url=BASE, http_client=mock_http) + resp = client.send(CREATE_ITEM(ItemRequest(name="alice"))) assert isinstance(resp, NemoResponse) assert resp.http_response.status_code == 201 @@ -78,6 +75,7 @@ def test_send_post() -> None: f"{BASE}/apis/test/v2/items", content=ItemRequest(name="alice").model_dump_json().encode(), headers={"Content-Type": "application/json"}, + params=None, ) @@ -89,8 +87,8 @@ def test_send_get_with_path_params() -> None: json={"id": 1, "name": "alice"}, ) - client = StubClient(base_url=BASE, http_client=mock_http) - resp = client.send(GET_ITEM.request(name="alice")) + client = NemoClient(base_url=BASE, http_client=mock_http) + resp = client.send(GET_ITEM(name="alice")) assert resp.body.name == "alice" mock_http.request.assert_called_once_with( @@ -98,6 +96,7 @@ def test_send_get_with_path_params() -> None: f"{BASE}/apis/test/v2/items/alice", content=None, headers=None, + params=None, ) @@ -109,8 +108,8 @@ def test_send_delete() -> None: content=b"", ) - client = StubClient(base_url=BASE, http_client=mock_http) - resp = client.send(DELETE_ITEM.request(name="alice")) + client = NemoClient(base_url=BASE, http_client=mock_http) + resp = client.send(DELETE_ITEM(name="alice")) assert resp.http_response.status_code == 204 assert resp.body is None @@ -124,8 +123,8 @@ def test_data_success() -> None: json={"id": 1, "name": "alice"}, ) - client = StubClient(base_url=BASE, http_client=mock_http) - item = client.send(GET_ITEM.request(name="alice")).data() + client = NemoClient(base_url=BASE, http_client=mock_http) + item = client.send(GET_ITEM(name="alice")).data() assert item.name == "alice" @@ -138,8 +137,8 @@ def test_base_url_trailing_slash_stripped() -> None: json={"id": 1, "name": "x"}, ) - client = StubClient(base_url=BASE + "/", http_client=mock_http) - client.send(GET_ITEM.request(name="x")) + client = NemoClient(base_url=BASE + "/", http_client=mock_http) + client.send(GET_ITEM(name="x")) url_called = mock_http.request.call_args[0][1] assert not url_called.startswith(BASE + "//") @@ -159,8 +158,8 @@ async def test_async_send_post() -> None: json={"id": 1, "name": "alice"}, ) - client = AsyncStubClient(base_url=BASE, http_client=mock_http) - resp = await client.send(CREATE_ITEM.request(ItemRequest(name="alice"))) + client = AsyncNemoClient(base_url=BASE, http_client=mock_http) + resp = await client.send(CREATE_ITEM(ItemRequest(name="alice"))) assert resp.http_response.status_code == 201 assert resp.body.name == "alice" @@ -175,8 +174,8 @@ async def test_async_send_get() -> None: json={"id": 1, "name": "alice"}, ) - client = AsyncStubClient(base_url=BASE, http_client=mock_http) - resp = await client.send(GET_ITEM.request(name="alice")) + client = AsyncNemoClient(base_url=BASE, http_client=mock_http) + resp = await client.send(GET_ITEM(name="alice")) assert resp.body.name == "alice" @@ -186,7 +185,7 @@ async def test_async_send_get() -> None: # --------------------------------------------------------------------------- -def test_workspace_default_fills_path() -> None: +def test_workspace_explicit_in_request() -> None: mock_http = MagicMock(spec=httpx.Client) mock_http.request.return_value = httpx.Response( 200, @@ -194,8 +193,24 @@ def test_workspace_default_fills_path() -> None: json={"id": 1, "name": "alice"}, ) - client = StubClient(base_url=BASE, workspace="default", http_client=mock_http) - client.send(GET_WS_ITEM.request()) + client = NemoClient(base_url=BASE, workspace="default", http_client=mock_http) + client.send(GET_WS_ITEM(workspace="default")) + + url_called = mock_http.request.call_args[0][1] + assert "/workspaces/default/" in url_called + + +def test_workspace_default_fills_omitted_path_param() -> None: + mock_http = MagicMock(spec=httpx.Client) + mock_http.request.return_value = httpx.Response( + 200, + request=httpx.Request("GET", f"{BASE}/apis/test/v2/workspaces/default/items"), + json={"id": 1, "name": "alice"}, + ) + + client = NemoClient(base_url=BASE, workspace="default", http_client=mock_http) + # workspace omitted — client default fills it + client.send(GET_WS_ITEM()) url_called = mock_http.request.call_args[0][1] assert "/workspaces/default/" in url_called @@ -209,8 +224,189 @@ def test_workspace_explicit_overrides_default() -> None: json={"id": 1, "name": "alice"}, ) - client = StubClient(base_url=BASE, workspace="default", http_client=mock_http) - client.send(GET_WS_ITEM.request(workspace="other")) + client = NemoClient(base_url=BASE, workspace="default", http_client=mock_http) + client.send(GET_WS_ITEM(workspace="other")) url_called = mock_http.request.call_args[0][1] assert "/workspaces/other/" in url_called + + +# --------------------------------------------------------------------------- +# Query params +# --------------------------------------------------------------------------- + + +def test_query_params_passed_to_httpx() -> None: + mock_http = MagicMock(spec=httpx.Client) + mock_http.request.return_value = httpx.Response( + 200, + request=httpx.Request("GET", f"{BASE}/apis/test/v2/items"), + json={"id": 1, "name": "alice"}, + ) + + client = NemoClient(base_url=BASE, http_client=mock_http) + client.send(GET_ITEMS_WITH_PARAMS(query_params={"page": 2, "page_size": 10})) + + mock_http.request.assert_called_once_with( + "GET", + f"{BASE}/apis/test/v2/items", + content=None, + headers=None, + params={"page": 2, "page_size": 10}, + ) + + +def test_query_params_none_values_filtered() -> None: + mock_http = MagicMock(spec=httpx.Client) + mock_http.request.return_value = httpx.Response( + 200, + request=httpx.Request("GET", f"{BASE}/apis/test/v2/items"), + json={"id": 1, "name": "alice"}, + ) + + client = NemoClient(base_url=BASE, http_client=mock_http) + client.send(GET_ITEMS_WITH_PARAMS(query_params={"page_cursor": None, "page_size": 10})) + + mock_http.request.assert_called_once_with( + "GET", + f"{BASE}/apis/test/v2/items", + content=None, + headers=None, + params={"page_size": 10}, + ) + + +def test_query_params_all_none_becomes_none() -> None: + mock_http = MagicMock(spec=httpx.Client) + mock_http.request.return_value = httpx.Response( + 200, + request=httpx.Request("GET", f"{BASE}/apis/test/v2/items"), + json={"id": 1, "name": "alice"}, + ) + + client = NemoClient(base_url=BASE, http_client=mock_http) + client.send(GET_ITEMS_WITH_PARAMS(query_params={"page_cursor": None})) + + mock_http.request.assert_called_once_with( + "GET", + f"{BASE}/apis/test/v2/items", + content=None, + headers=None, + params=None, + ) + + +# --------------------------------------------------------------------------- +# Error response body parsing +# --------------------------------------------------------------------------- + + +def test_error_response_extracts_detail() -> None: + mock_http = MagicMock(spec=httpx.Client) + mock_http.request.return_value = httpx.Response( + 422, + request=httpx.Request("POST", f"{BASE}/apis/test/v2/items"), + json={"detail": "Validation failed: name is required"}, + ) + + client = NemoClient(base_url=BASE, http_client=mock_http) + resp = client.send(CREATE_ITEM(ItemRequest(name=""))) + + with pytest.raises(NemoHTTPError) as exc_info: + resp.data() + + assert exc_info.value.status_code == 422 + assert exc_info.value.detail == "Validation failed: name is required" + assert "422" in str(exc_info.value) + assert "Validation failed" in str(exc_info.value) + + +def test_error_response_fallback_to_text() -> None: + mock_http = MagicMock(spec=httpx.Client) + mock_http.request.return_value = httpx.Response( + 500, + request=httpx.Request("GET", f"{BASE}/apis/test/v2/items/x"), + text="Internal Server Error", + ) + + client = NemoClient(base_url=BASE, http_client=mock_http) + resp = client.send(GET_ITEM(name="x")) + + with pytest.raises(NemoHTTPError) as exc_info: + resp.data() + + assert exc_info.value.status_code == 500 + assert exc_info.value.detail == "Internal Server Error" + + +def test_error_response_body_is_none() -> None: + """On error, body should be None (not deserialized as the response type).""" + mock_http = MagicMock(spec=httpx.Client) + mock_http.request.return_value = httpx.Response( + 404, + request=httpx.Request("GET", f"{BASE}/apis/test/v2/items/missing"), + json={"detail": "Not found"}, + ) + + client = NemoClient(base_url=BASE, http_client=mock_http) + resp = client.send(GET_ITEM(name="missing")) + + assert resp.body is None + assert resp.http_response.status_code == 404 + + +# --------------------------------------------------------------------------- +# Per-request headers +# --------------------------------------------------------------------------- + + +def test_extra_headers_merged_into_request() -> None: + mock_http = MagicMock(spec=httpx.Client) + mock_http.request.return_value = httpx.Response( + 200, + request=httpx.Request("GET", f"{BASE}/apis/test/v2/items/alice"), + json={"id": 1, "name": "alice"}, + ) + + client = NemoClient(base_url=BASE, http_client=mock_http) + client.send(GET_ITEM(name="alice"), headers={"Accept": "application/octet-stream"}) + + _, kwargs = mock_http.request.call_args + assert kwargs["headers"]["Accept"] == "application/octet-stream" + + +def test_extra_headers_dont_override_content_type() -> None: + mock_http = MagicMock(spec=httpx.Client) + mock_http.request.return_value = httpx.Response( + 201, + request=httpx.Request("POST", f"{BASE}/apis/test/v2/items"), + json={"id": 1, "name": "alice"}, + ) + + client = NemoClient(base_url=BASE, http_client=mock_http) + client.send(CREATE_ITEM(ItemRequest(name="alice")), headers={"X-Custom": "value"}) + + _, kwargs = mock_http.request.call_args + assert kwargs["headers"]["Content-Type"] == "application/json" + assert kwargs["headers"]["X-Custom"] == "value" + + +# --------------------------------------------------------------------------- +# Response carries request +# --------------------------------------------------------------------------- + + +def test_response_carries_prepared_request() -> None: + mock_http = MagicMock(spec=httpx.Client) + mock_http.request.return_value = httpx.Response( + 200, + request=httpx.Request("GET", f"{BASE}/apis/test/v2/items/alice"), + json={"id": 1, "name": "alice"}, + ) + + client = NemoClient(base_url=BASE, http_client=mock_http) + resp = client.send(GET_ITEM(name="alice")) + + assert resp.request is not None + assert resp.request.method == "GET" + assert resp.request.path_params == {"name": "alice"} diff --git a/packages/nemo_platform_plugin/tests/client/test_endpoint.py b/packages/nemo_platform_plugin/tests/client/test_endpoint.py index 72022a849a..42de7be496 100644 --- a/packages/nemo_platform_plugin/tests/client/test_endpoint.py +++ b/packages/nemo_platform_plugin/tests/client/test_endpoint.py @@ -3,8 +3,10 @@ from __future__ import annotations +from typing import NotRequired, TypedDict + from nemo_platform_plugin.client.endpoint import delete, get, patch, post -from nemo_platform_plugin.client.types import PathParams, PreparedRequest +from nemo_platform_plugin.client.types import PreparedRequest from pydantic import BaseModel @@ -17,41 +19,51 @@ class FakeResponse(BaseModel): name: str -class WorkspacePath(PathParams): - workspace: str +@post("/v2/workspaces/{workspace}/items") +def PostEndpoint(body: FakeRequest, *, workspace: str) -> FakeResponse: + raise NotImplementedError -class WorkspaceItemPath(PathParams): - workspace: str - name: str +@get("/v2/workspaces/{workspace}/items/{name}") +def GetEndpoint(*, workspace: str, name: str) -> FakeResponse: + raise NotImplementedError + + +@delete("/v2/workspaces/{workspace}/items/{name}") +def DeleteEndpoint(*, workspace: str, name: str) -> None: + raise NotImplementedError + + +@patch("/items/{id}") +def PatchEndpoint(body: FakeRequest, *, id: str) -> FakeResponse: + raise NotImplementedError + +class ListQueryParams(TypedDict, total=False): + page: NotRequired[int] + page_size: NotRequired[int] -class IdPath(PathParams): - id: str + +@get("/v2/workspaces/{workspace}/items") +def ListEndpoint(*, workspace: str, query_params: ListQueryParams | None = None) -> FakeResponse: + raise NotImplementedError def test_post_endpoint_produces_prepared_request() -> None: - ep = post( - "/v2/workspaces/{workspace}/items", - path_type=WorkspacePath, - request_type=FakeRequest, - response_type=FakeResponse, - ) - payload = FakeRequest(name="alice") - prepared = ep.request(payload, workspace="default") + body = FakeRequest(name="alice") + prepared = PostEndpoint(body, workspace="default") assert isinstance(prepared, PreparedRequest) assert prepared.path_template == "/v2/workspaces/{workspace}/items" assert prepared.path_params == {"workspace": "default"} assert prepared.method == "POST" - assert prepared.content == payload.model_dump_json().encode() + assert prepared.content == body.model_dump_json().encode() assert prepared.content_type == "application/json" assert prepared.response_type is FakeResponse def test_get_endpoint_no_body() -> None: - ep = get("/v2/workspaces/{workspace}/items/{name}", path_type=WorkspaceItemPath, response_type=FakeResponse) - prepared = ep.request(workspace="default", name="item-1") + prepared = GetEndpoint(workspace="default", name="item-1") assert prepared.path_template == "/v2/workspaces/{workspace}/items/{name}" assert prepared.path_params == {"workspace": "default", "name": "item-1"} @@ -62,8 +74,7 @@ def test_get_endpoint_no_body() -> None: def test_delete_endpoint() -> None: - ep = delete("/v2/workspaces/{workspace}/items/{name}", path_type=WorkspaceItemPath) - prepared = ep.request(workspace="default", name="item-1") + prepared = DeleteEndpoint(workspace="default", name="item-1") assert prepared.path_params == {"workspace": "default", "name": "item-1"} assert prepared.method == "DELETE" @@ -71,18 +82,47 @@ def test_delete_endpoint() -> None: def test_patch_endpoint() -> None: - ep = patch("/items/{id}", path_type=IdPath, request_type=FakeRequest, response_type=FakeResponse) - payload = FakeRequest(name="updated") - prepared = ep.request(payload, id="42") + body = FakeRequest(name="updated") + prepared = PatchEndpoint(body, id="42") assert prepared.path_params == {"id": "42"} assert prepared.method == "PATCH" - assert prepared.content == payload.model_dump_json().encode() + assert prepared.content == body.model_dump_json().encode() + + +def test_get_with_query_params() -> None: + prepared = ListEndpoint(workspace="default", query_params={"page": 1, "page_size": 10}) + + assert prepared.path_params == {"workspace": "default"} + assert prepared.query_params == {"page": 1, "page_size": 10} + assert prepared.method == "GET" -def test_endpoint_repr() -> None: - ep = post("/items/{id}", path_type=IdPath, request_type=FakeRequest, response_type=FakeResponse) - r = repr(ep) - assert "/items" in r - assert "FakeRequest" in r - assert "FakeResponse" in r +def test_query_params_default_none() -> None: + prepared = ListEndpoint(workspace="default") + + assert prepared.query_params is None + + +def test_post_with_query_params() -> None: + @post("/v2/items/{workspace}") + def PostWithQuery(body: FakeRequest, *, workspace: str, query_params: dict | None = None) -> FakeResponse: + raise NotImplementedError + + body = FakeRequest(name="alice") + prepared = PostWithQuery(body, workspace="default", query_params={"dry_run": True}) + + assert prepared.query_params == {"dry_run": True} + assert prepared.content == body.model_dump_json().encode() + + +def test_delete_with_response_type() -> None: + @delete("/v2/items/{id}") + def DeleteWithResp(*, id: str) -> FakeResponse: + raise NotImplementedError + + prepared = DeleteWithResp(id="42") + + assert prepared.method == "DELETE" + assert prepared.response_type is FakeResponse + assert prepared.content is None diff --git a/plugins/example-plugin/src/nemo_example_plugin/sdk.py b/plugins/example-plugin/src/nemo_example_plugin/sdk.py index ad18dc1a77..0206928958 100644 --- a/plugins/example-plugin/src/nemo_example_plugin/sdk.py +++ b/plugins/example-plugin/src/nemo_example_plugin/sdk.py @@ -3,60 +3,40 @@ """SDK resources for the example plugin. -Endpoints are defined once in a mixin, then sync/async client classes -inherit the mixin + the appropriate client base. -The descriptor protocol on each endpoint returns the right bound callable. +Endpoints are defined in ``types.endpoints`` as decorated functions. +The client classes expose them as direct methods via ``method()`` wrappers. """ from __future__ import annotations -from nemo_example_plugin.types.endpoints import ( - CountEndpoint, - CreateItemEndpoint, - DeleteItemEndpoint, - DownloadBlobEndpoint, - GetItemEndpoint, - HelloEndpoint, - ListItemsEndpoint, - UpdateItemEndpoint, - UploadBlobEndpoint, -) +from nemo_example_plugin.types import endpoints from nemo_platform import AsyncNeMoPlatform, NeMoPlatform from nemo_platform_plugin.client.adapter import client_from_platform from nemo_platform_plugin.client.client import AsyncNemoClient, NemoClient +from nemo_platform_plugin.client.method import method from nemo_platform_plugin.sdk import NemoPluginSDKResources -# -- Endpoint assignments (defined once) ----------------------------------- - - -class _ExampleEndpoints: - hello = HelloEndpoint - create_item = CreateItemEndpoint - list_items = ListItemsEndpoint - get_item = GetItemEndpoint - update_item = UpdateItemEndpoint - delete_item = DeleteItemEndpoint - count = CountEndpoint - upload_blob = UploadBlobEndpoint - download_blob = DownloadBlobEndpoint +class _ExampleMethods: + hello = method(endpoints.hello) + create_item = method(endpoints.create_item) + list_items = method(endpoints.list_items) + get_item = method(endpoints.get_item) + update_item = method(endpoints.update_item) + delete_item = method(endpoints.delete_item) + count = method(endpoints.count) + upload_blob = method(endpoints.upload_blob) + download_blob = method(endpoints.download_blob) -# -- Client classes: mixin + client base ----------------------------------- - -class ExampleClient(_ExampleEndpoints, NemoClient): +class ExampleClient(_ExampleMethods, NemoClient): """Sync client for the example plugin API.""" -class AsyncExampleClient(_ExampleEndpoints, AsyncNemoClient): +class AsyncExampleClient(_ExampleMethods, AsyncNemoClient): """Async client for the example plugin API.""" -# --------------------------------------------------------------------------- -# Plugin SDK registration — bridges NeMoPlatform to the new typed client -# --------------------------------------------------------------------------- - - def _make_sync_resource(platform: NeMoPlatform) -> ExampleClient: return client_from_platform(platform, ExampleClient) diff --git a/plugins/example-plugin/src/nemo_example_plugin/types/endpoints.py b/plugins/example-plugin/src/nemo_example_plugin/types/endpoints.py index e51afe0789..d9a16e9d3a 100644 --- a/plugins/example-plugin/src/nemo_example_plugin/types/endpoints.py +++ b/plugins/example-plugin/src/nemo_example_plugin/types/endpoints.py @@ -3,15 +3,15 @@ """Typed endpoint definitions for the example plugin. -These are the single source of truth for the HTTP contract. Both the SDK -client and (eventually) server route registration can be derived from them. - -Request and response models are plain Pydantic — they have no knowledge of -the HTTP layer. +These are the single source of truth for the HTTP contract. Each endpoint +is a decorated function that declares its call signature and response type. """ from __future__ import annotations +from abc import abstractmethod +from typing import NotRequired, TypedDict + from nemo_example_plugin.entities import ExampleItem from nemo_example_plugin.types.payloads import ( BlobUploadResponse, @@ -23,69 +23,56 @@ UpdateExampleItemRequest, ) from nemo_platform_plugin.client.endpoint import delete, get, patch, post, put -from nemo_platform_plugin.client.types import BinaryContent, PathParams, Stream, WorkspaceParams +from nemo_platform_plugin.client.types import BinaryContent, Stream -# -- Path parameter types -------------------------------------------------- +class ListItemsQueryParams(TypedDict, total=False): + page: NotRequired[int] + page_size: NotRequired[int] -class NamePath(PathParams): - name: str +@get("/apis/example/hello/{name}") +@abstractmethod +def hello(*, name: str) -> HelloResponse: ... -class WorkspaceItemPath(WorkspaceParams): - name: str +@post("/apis/example/v2/workspaces/{workspace}/items") +@abstractmethod +def create_item(*, workspace: str | None = None, body: CreateExampleItemRequest) -> ExampleItem: ... -# -- Hello ----------------------------------------------------------------- -HelloEndpoint = get("/apis/example/hello/{name}", path_type=NamePath, response_type=HelloResponse) +@get("/apis/example/v2/workspaces/{workspace}/items") +@abstractmethod +def list_items( + *, workspace: str | None = None, query_params: ListItemsQueryParams | None = None +) -> ExampleItemPage: ... -# -- Items CRUD ------------------------------------------------------------ -CreateItemEndpoint = post( - "/apis/example/v2/workspaces/{workspace}/items", - path_type=WorkspaceParams, - request_type=CreateExampleItemRequest, - response_type=ExampleItem, -) +@get("/apis/example/v2/workspaces/{workspace}/items/{name}") +@abstractmethod +def get_item(*, workspace: str | None = None, name: str) -> ExampleItem: ... -ListItemsEndpoint = get( - "/apis/example/v2/workspaces/{workspace}/items", path_type=WorkspaceParams, response_type=ExampleItemPage -) -GetItemEndpoint = get( - "/apis/example/v2/workspaces/{workspace}/items/{name}", path_type=WorkspaceItemPath, response_type=ExampleItem -) +@patch("/apis/example/v2/workspaces/{workspace}/items/{name}") +@abstractmethod +def update_item(*, workspace: str | None = None, name: str, body: UpdateExampleItemRequest) -> ExampleItem: ... -UpdateItemEndpoint = patch( - "/apis/example/v2/workspaces/{workspace}/items/{name}", - path_type=WorkspaceItemPath, - request_type=UpdateExampleItemRequest, - response_type=ExampleItem, -) -DeleteItemEndpoint = delete("/apis/example/v2/workspaces/{workspace}/items/{name}", path_type=WorkspaceItemPath) +@delete("/apis/example/v2/workspaces/{workspace}/items/{name}") +@abstractmethod +def delete_item(*, workspace: str | None = None, name: str) -> None: ... -# -- Functions ------------------------------------------------------------- -CountEndpoint = post( - "/apis/example/v2/workspaces/{workspace}/count", - path_type=WorkspaceParams, - request_type=CountRequest, - response_type=Stream[Tick], -) +@post("/apis/example/v2/workspaces/{workspace}/count") +@abstractmethod +def count(*, workspace: str | None = None, body: CountRequest) -> Stream[Tick]: ... -# -- Binary ---------------------------------------------------------------- -UploadBlobEndpoint = put( - "/apis/example/blob/{name}", - path_type=NamePath, - request_type=BinaryContent, - response_type=BlobUploadResponse, -) +@put("/apis/example/blob/{name}") +@abstractmethod +def upload_blob(*, name: str, content: bytes) -> BlobUploadResponse: ... -DownloadBlobEndpoint = get( - "/apis/example/blob/{name}", - path_type=NamePath, - response_type=BinaryContent, -) + +@get("/apis/example/blob/{name}") +@abstractmethod +def download_blob(*, name: str) -> BinaryContent: ... diff --git a/plugins/example-plugin/tests/test_sdk.py b/plugins/example-plugin/tests/test_sdk.py index 8176d3a4ba..7c326b4a80 100644 --- a/plugins/example-plugin/tests/test_sdk.py +++ b/plugins/example-plugin/tests/test_sdk.py @@ -10,10 +10,12 @@ import httpx import pytest from nemo_example_plugin.sdk import AsyncExampleClient, ExampleClient +from nemo_example_plugin.types import endpoints from nemo_example_plugin.types.payloads import ( CreateExampleItemRequest, UpdateExampleItemRequest, ) +from nemo_platform_plugin.client.client import NemoClient BASE = "http://test:8000" WS = "default" @@ -41,18 +43,18 @@ def _resp(status: int, payload=None) -> httpx.Response: def _sync_client() -> tuple[ExampleClient, MagicMock]: mock_http = MagicMock(spec=httpx.Client) - client = ExampleClient(base_url=BASE, http_client=mock_http) + client = ExampleClient(base_url=BASE, workspace=WS, http_client=mock_http) return client, mock_http def _async_client() -> tuple[AsyncExampleClient, AsyncMock]: mock_http = AsyncMock(spec=httpx.AsyncClient) - client = AsyncExampleClient(base_url=BASE, http_client=mock_http) + client = AsyncExampleClient(base_url=BASE, workspace=WS, http_client=mock_http) return client, mock_http # --------------------------------------------------------------------------- -# hello +# hello — client.method() style # --------------------------------------------------------------------------- @@ -72,7 +74,7 @@ async def test_async_hello() -> None: # --------------------------------------------------------------------------- -# Items CRUD — sync +# Items CRUD — client.method() style (sync) # --------------------------------------------------------------------------- @@ -80,7 +82,7 @@ def test_sync_create_item() -> None: client, mock_http = _sync_client() mock_http.request.return_value = _resp(201, ITEM_PAYLOAD) - resp = client.create_item(CreateExampleItemRequest(name="my-item", title="My Item"), workspace=WS) + resp = client.create_item(body=CreateExampleItemRequest(name="my-item", title="My Item")) item = resp.data() assert item.name == "my-item" @@ -88,11 +90,22 @@ def test_sync_create_item() -> None: mock_http.request.assert_called_once() +def test_sync_create_item_explicit_workspace() -> None: + client, mock_http = _sync_client() + mock_http.request.return_value = _resp(201, ITEM_PAYLOAD) + + resp = client.create_item(workspace="other", body=CreateExampleItemRequest(name="my-item", title="My Item")) + + assert resp.data().name == "my-item" + url_called = mock_http.request.call_args[0][1] + assert "/workspaces/other/" in url_called + + def test_sync_get_item() -> None: client, mock_http = _sync_client() mock_http.request.return_value = _resp(200, ITEM_PAYLOAD) - resp = client.get_item(workspace=WS, name="my-item") + resp = client.get_item(name="my-item") assert resp.data().name == "my-item" @@ -103,19 +116,33 @@ def test_sync_list_items() -> None: 200, {"data": [ITEM_PAYLOAD], "pagination": None, "sort": None, "filter": None} ) - resp = client.list_items(workspace=WS) + resp = client.list_items() page = resp.data() assert len(page.data) == 1 assert page.data[0].name == "my-item" +def test_sync_list_items_with_query_params() -> None: + client, mock_http = _sync_client() + mock_http.request.return_value = _resp( + 200, {"data": [ITEM_PAYLOAD], "pagination": None, "sort": None, "filter": None} + ) + + resp = client.list_items(query_params={"page": 2, "page_size": 5}) + page = resp.data() + + assert len(page.data) == 1 + _, kwargs = mock_http.request.call_args + assert kwargs["params"] == {"page": 2, "page_size": 5} + + def test_sync_update_item() -> None: client, mock_http = _sync_client() updated = {**ITEM_PAYLOAD, "title": "Updated"} mock_http.request.return_value = _resp(200, updated) - resp = client.update_item(UpdateExampleItemRequest(title="Updated"), workspace=WS, name="my-item") + resp = client.update_item(name="my-item", body=UpdateExampleItemRequest(title="Updated")) assert resp.data().title == "Updated" @@ -124,13 +151,28 @@ def test_sync_delete_item() -> None: client, mock_http = _sync_client() mock_http.request.return_value = _resp(204) - client.delete_item(workspace=WS, name="my-item") + client.delete_item(name="my-item") mock_http.request.assert_called_once() # --------------------------------------------------------------------------- -# Items CRUD — async +# Low-level: endpoints + client.send() still works +# --------------------------------------------------------------------------- + + +def test_send_with_endpoint_function() -> None: + mock_http = MagicMock(spec=httpx.Client) + mock_http.request.return_value = _resp(200, {"message": "Hello, alice!"}) + client = NemoClient(base_url=BASE, workspace=WS, http_client=mock_http) + + resp = client.send(endpoints.hello(name="alice")) + + assert resp.data().message == "Hello, alice!" + + +# --------------------------------------------------------------------------- +# Items CRUD — client.method() style (async) # --------------------------------------------------------------------------- @@ -139,7 +181,7 @@ async def test_async_create_item() -> None: client, mock_http = _async_client() mock_http.request.return_value = _resp(201, ITEM_PAYLOAD) - resp = await client.create_item(CreateExampleItemRequest(name="my-item", title="My Item"), workspace=WS) + resp = await client.create_item(body=CreateExampleItemRequest(name="my-item", title="My Item")) assert resp.data().name == "my-item" @@ -149,7 +191,7 @@ async def test_async_get_item() -> None: client, mock_http = _async_client() mock_http.request.return_value = _resp(200, ITEM_PAYLOAD) - resp = await client.get_item(workspace=WS, name="my-item") + resp = await client.get_item(name="my-item") assert resp.data().name == "my-item" @@ -161,7 +203,7 @@ async def test_async_list_items() -> None: 200, {"data": [ITEM_PAYLOAD], "pagination": None, "sort": None, "filter": None} ) - resp = await client.list_items(workspace=WS) + resp = await client.list_items() page = resp.data() assert len(page.data) == 1 @@ -174,7 +216,7 @@ async def test_async_update_item() -> None: updated = {**ITEM_PAYLOAD, "title": "Updated"} mock_http.request.return_value = _resp(200, updated) - resp = await client.update_item(UpdateExampleItemRequest(title="Updated"), workspace=WS, name="my-item") + resp = await client.update_item(name="my-item", body=UpdateExampleItemRequest(title="Updated")) assert resp.data().title == "Updated" @@ -184,6 +226,6 @@ async def test_async_delete_item() -> None: client, mock_http = _async_client() mock_http.request.return_value = _resp(204) - await client.delete_item(workspace=WS, name="my-item") + await client.delete_item(name="my-item") mock_http.request.assert_awaited_once()