nemo-platform-plugin 0.1.0__py3-none-any.whl

nemo_platform_plugin/functions/frames.py

@@ -0,0 +1,69 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Convention frames for NDJSON-streaming :class:`NemoFunction`s.
+
+Plugin authors may use these directly or define their own frame
+shapes — the route adapter doesn't enforce a particular schema on
+user-yielded frames. The one exception is :class:`Heartbeat`: the
+adapter emits a heartbeat frame on idle so proxies don't time the
+connection out and so callers can tell "server is still working" apart
+from "connection died".
+
+The discriminator field name is ``kind``. The canonical values agreed
+in `plan-functions.md` are ``"log" | "progress" | "data" | "error" |
+"done" | "heartbeat"``. This module ships the three structural ones
+(:class:`Heartbeat`, :class:`Done`, :class:`Error`); domain frames
+(``log`` / ``progress`` / ``data``) are typically plugin-specific so
+plugin authors define them themselves.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Literal
+
+from pydantic import BaseModel
+
+
+class Heartbeat(BaseModel):
+    """Idle-tick frame emitted by the route adapter every ~5 s of silence.
+
+    The route adapter inserts these unconditionally on streams that go
+    quiet — they are not produced by user code. Plugins relying on
+    strict frame typing on the receiving side should ignore unknown
+    ``kind`` values or include :class:`Heartbeat` in their own
+    discriminated union.
+    """
+
+    kind: Literal["heartbeat"] = "heartbeat"
+
+
+class Done(BaseModel):
+    """Terminator frame for a successful stream.
+
+    Convention only — yield this as the *last* frame in a stream. The
+    route adapter does not synthesise it; if your stream just ends
+    without a ``Done``, callers see a clean EOF instead of an explicit
+    success marker. The marker is useful when the same stream may end
+    with :class:`Error` instead.
+    """
+
+    kind: Literal["done"] = "done"
+
+
+class Error(BaseModel):
+    """Terminator frame for an unsuccessful stream.
+
+    Convention only — yield this when an error occurs mid-stream and
+    you want to surface a structured failure rather than letting the
+    underlying exception propagate (which would close the connection
+    abruptly). Pair with raising the underlying exception for
+    server-side observability.
+    """
+
+    kind: Literal["error"] = "error"
+    message: str
+    details: dict[str, Any] | None = None
+
+
+__all__ = ["Done", "Error", "Heartbeat"]

nemo_platform_plugin/functions/routes.py

@@ -0,0 +1,454 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Auto-derived FastAPI route for a :class:`NemoFunction` subclass.
+
+``add_function_routes(MyFunction)`` is the one-liner that mirrors
+:func:`~nemo_platform_plugin.jobs.routes.add_job_routes` for the function tier.
+It returns an :class:`APIRouter` carrying a single ``POST`` route that:
+
+- Validates the request body against
+  :attr:`~nemo_platform_plugin.function.NemoFunction.spec_schema`.
+- Builds a :class:`~nemo_platform_plugin.function_context.FunctionContext` from
+  the workspace path parameter and the optional ``X-Request-ID``
+  header.
+- Resolves keyword-only DI parameters on
+  :meth:`~nemo_platform_plugin.function.NemoFunction.run` by parameter name —
+  ``ctx`` (FunctionContext), ``async_sdk``
+  (``AsyncNeMoPlatform``, resolved from
+  :func:`~nemo_platform_plugin.dependencies.get_sdk_client`), and ``is_local=False``. Functions
+  intentionally do **not** receive a sync ``sdk`` here: the API
+  process is async; sync work belongs inside
+  :func:`anyio.to_thread.run_sync` and bridges back to the loop via
+  :func:`anyio.from_thread.run`.
+- Awaits ``run(spec, **resolved)``.
+- If the result is an async iterator/generator, wraps it in a
+  :class:`StreamingResponse` with media type
+  :data:`NDJSON_MEDIA_TYPE`. Each yielded
+  :class:`pydantic.BaseModel` becomes ``model_dump_json() + "\\n"``;
+  any other value is JSON-encoded the same way.
+- Otherwise returns the value as a normal JSON response (FastAPI
+  serialises ``BaseModel`` instances).
+
+The streaming branch wraps the user's iterator with a heartbeat
+injector that emits a :class:`~nemo_platform_plugin.functions.frames.Heartbeat`
+frame every :data:`HEARTBEAT_INTERVAL_SECONDS` of idle time. Heartbeat
+emission is unconditional on streaming routes — proxies between the
+plugin service and the caller need it to keep the connection open.
+
+Usage::
+
+    from nemo_platform_plugin.functions.routes import add_function_routes
+
+    router = add_function_routes(GreetFunction)
+    app.include_router(
+        router,
+        prefix="/apis/example/v2/workspaces/{workspace}",
+        tags=["Example"],
+    )
+
+The mounted route resolves to
+``POST /apis/example/v2/workspaces/{workspace}/{function-name}``.
+Override the trailing path segment per-class via
+:attr:`NemoFunction.endpoint` when a legacy URL needs to be preserved.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+from collections.abc import AsyncGenerator, AsyncIterator
+from contextlib import suppress
+from typing import TYPE_CHECKING, Any
+
+from fastapi import APIRouter, Depends, Header
+from fastapi.responses import StreamingResponse
+from nemo_platform_plugin.dependencies import get_sdk_client
+from nemo_platform_plugin.function import NemoFunction, returns_async_iterator
+from nemo_platform_plugin.function_context import FunctionContext
+from nemo_platform_plugin.functions.frames import Heartbeat
+from pydantic import BaseModel
+
+if TYPE_CHECKING:
+    pass
+
+logger = logging.getLogger(__name__)
+
+# NDJSON over chunked HTTP. Matches what the SDK and Studio clients
+# expect when ``returns_async_iterator(...)`` is True. DD's existing
+# ``application/jsonl`` is treated as a synonym by clients today;
+# ``application/x-ndjson`` is what `plan-functions.md` standardises on
+# for new functions.
+NDJSON_MEDIA_TYPE = "application/x-ndjson"
+
+# Default idle interval after which the route adapter emits a
+# :class:`Heartbeat` frame. Matches DD `preview`'s existing
+# ``HEARTBEAT_SECONDS = 5`` and `plan-functions.md`'s recommendation.
+# Override per-call via the ``add_function_routes`` keyword argument
+# when wiring a router for a function with very different latency
+# expectations.
+HEARTBEAT_INTERVAL_SECONDS: float = 5.0
+
+# Default mount path. Routers are typically included under a
+# ``/apis/<plugin>/v2/workspaces/{workspace}`` prefix, so the bare
+# ``/{name}`` here resolves to the canonical
+# ``POST /apis/<plugin>/v2/workspaces/{workspace}/{name}``. Functions
+# overriding :attr:`NemoFunction.endpoint` substitute its template
+# instead — see :func:`_resolve_route_path`.
+DEFAULT_FUNCTION_PATH: str = "/{name}"
+
+
+def add_function_routes(
+    function_cls: type[NemoFunction],
+    *,
+    heartbeat_interval_seconds: float = HEARTBEAT_INTERVAL_SECONDS,
+) -> APIRouter:
+    """Mount a single ``POST`` route for *function_cls* on a fresh router.
+
+    Args:
+        function_cls: The :class:`NemoFunction` subclass to expose. Must
+            declare :attr:`~NemoFunction.spec_schema` — that's how the
+            adapter knows what to validate the request body against.
+        heartbeat_interval_seconds: Idle gap before a
+            :class:`Heartbeat` frame is injected on streaming
+            responses. Defaults to :data:`HEARTBEAT_INTERVAL_SECONDS`.
+            Lower values are useful in tests; production callers
+            usually leave the default.
+
+    Returns:
+        An :class:`APIRouter` with one route. The caller mounts it
+        with the URL prefix of their choice (typically
+        ``/apis/<plugin>/v2/workspaces/{workspace}``).
+
+    Raises:
+        TypeError: If ``function_cls`` doesn't declare a
+            ``spec_schema``.
+    """
+    spec_schema = getattr(function_cls, "spec_schema", None)
+    if spec_schema is None:
+        raise TypeError(
+            f"{function_cls.__name__}.spec_schema is None; add_function_routes "
+            f"requires a declared spec_schema. Set it to a Pydantic BaseModel "
+            f"subclass on the NemoFunction class."
+        )
+
+    router = APIRouter()
+    instance = function_cls()
+    run_params = function_cls.run_signature().parameters
+    wants_ctx = "ctx" in run_params
+    wants_async_sdk = "async_sdk" in run_params
+    wants_is_local = "is_local" in run_params
+
+    path = _resolve_route_path(function_cls)
+
+    # The handler signature is built dynamically so we only declare
+    # ``Depends(get_sdk_client)`` for functions that actually opt in.
+    # Otherwise FastAPI would resolve it on every request and the
+    # placeholder raises ``RuntimeError`` when no ``dependency_overrides``
+    # entry has been registered. We also only declare the request-id
+    # header when the function wants ``ctx`` — pure-spec functions stay
+    # purely positional.
+    handler = _build_route_handler(
+        instance=instance,
+        spec_schema=spec_schema,
+        wants_ctx=wants_ctx,
+        wants_async_sdk=wants_async_sdk,
+        wants_is_local=wants_is_local,
+        send_headers_before_first_frame=function_cls.send_headers_before_first_frame,
+        heartbeat_interval_seconds=heartbeat_interval_seconds,
+    )
+    handler.__name__ = f"{function_cls.__name__}__route"
+    handler.__doc__ = function_cls.description or f"Invoke the {function_cls.name} function."
+
+    router.post(
+        path,
+        summary=function_cls.description or f"Invoke {function_cls.name}",
+        # Streaming functions return ``StreamingResponse`` directly, so
+        # leave the response model unconstrained at the route level.
+        # Non-streaming functions get FastAPI's default JSON
+        # serialisation of whatever they return.
+        response_model=None,
+    )(handler)
+
+    return router
+
+
+def _build_route_handler(
+    *,
+    instance: NemoFunction,
+    spec_schema: type[BaseModel],
+    wants_ctx: bool,
+    wants_async_sdk: bool,
+    wants_is_local: bool,
+    send_headers_before_first_frame: bool,
+    heartbeat_interval_seconds: float,
+):
+    """Construct an async route handler whose signature matches *only* the DI the function asks for.
+
+    FastAPI inspects the handler's signature at registration time;
+    every parameter with a ``Depends`` is resolved on every request.
+    Declaring ``Depends(get_sdk_client)`` on a function that doesn't
+    want ``async_sdk`` would force a runtime error on every request
+    until the platform installs ``app.dependency_overrides`` for the
+    SDK. Building the closure body's call-shape from booleans keeps
+    each function's surface as narrow as the function declared.
+    """
+
+    async def _invoke(spec_obj: BaseModel, ctx: FunctionContext | None, async_sdk: Any) -> Any:
+        kwargs: dict[str, Any] = {}
+        if ctx is not None:
+            kwargs["ctx"] = ctx
+        if wants_async_sdk:
+            kwargs["async_sdk"] = async_sdk
+        if wants_is_local:
+            kwargs["is_local"] = False
+
+        result = instance.run(spec_obj, **kwargs)
+        if returns_async_iterator(result):
+            if not send_headers_before_first_frame:
+                result = await _prime_async_iterator(result)
+            stream = _with_heartbeats(result, interval_seconds=heartbeat_interval_seconds)
+            return StreamingResponse(_to_ndjson(stream), media_type=NDJSON_MEDIA_TYPE)
+        return await result
+
+    if wants_ctx and wants_async_sdk:
+
+        async def handler(
+            workspace: str,
+            request_body,
+            x_request_id: str | None = Header(default=None, alias="X-Request-ID"),
+            async_sdk: Any = Depends(get_sdk_client),
+        ) -> Any:
+            ctx = FunctionContext(workspace=workspace, request_id=x_request_id)
+            return await _invoke(request_body, ctx, async_sdk)
+
+    elif wants_ctx:
+
+        async def handler(
+            workspace: str,
+            request_body,
+            x_request_id: str | None = Header(default=None, alias="X-Request-ID"),
+        ) -> Any:
+            ctx = FunctionContext(workspace=workspace, request_id=x_request_id)
+            return await _invoke(request_body, ctx, None)
+
+    elif wants_async_sdk:
+
+        async def handler(
+            workspace: str,
+            request_body,
+            async_sdk: Any = Depends(get_sdk_client),
+        ) -> Any:
+            return await _invoke(request_body, None, async_sdk)
+
+    else:
+
+        async def handler(
+            workspace: str,
+            request_body,
+        ) -> Any:
+            return await _invoke(request_body, None, None)
+
+    # FastAPI inspects ``__annotations__`` (via ``get_type_hints``)
+    # rather than the literal source annotation, and a closure-captured
+    # ``spec_schema`` doesn't resolve from the handler's local scope.
+    # Assigning the Pydantic class directly to the parameter's
+    # annotation ensures FastAPI treats ``request_body`` as a Body
+    # rather than defaulting to a Query parameter.
+    handler.__annotations__["request_body"] = spec_schema
+    return handler
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _resolve_route_path(function_cls: type[NemoFunction]) -> str:
+    """Return the route path the adapter mounts the handler under.
+
+    When :attr:`NemoFunction.endpoint` is set, it's a full template
+    that the operator wants to keep verbatim — the only substitution
+    we apply is ``{name}`` (so the same template can be reused across
+    functions in a plugin without writing the suffix twice). The
+    ``{api}`` and ``{workspace}`` placeholders stay as literal route
+    parameters; FastAPI will refuse to mount the router otherwise.
+
+    The default :data:`DEFAULT_FUNCTION_PATH` template is the trailing
+    ``/{name}`` segment, intended to be mounted under the standard
+    ``/apis/<plugin>/v2/workspaces/{workspace}`` prefix.
+    """
+    template = function_cls.endpoint or DEFAULT_FUNCTION_PATH
+    return template.replace("{name}", function_cls.name)
+
+
+async def _to_ndjson(source: AsyncIterator[Any]) -> AsyncGenerator[str, None]:
+    """Encode each frame from *source* as a single NDJSON line.
+
+    Pydantic models flow through ``model_dump_json()``; every other
+    value goes through ``json.dumps`` so plugin authors can yield
+    plain dicts during prototyping without constructing wrapper
+    models.
+    """
+    async for frame in source:
+        if isinstance(frame, BaseModel):
+            yield frame.model_dump_json() + "\n"
+        else:
+            yield json.dumps(frame, default=str) + "\n"
+
+
+async def _prime_async_iterator(source: AsyncIterator[Any]) -> AsyncIterator[Any]:
+    """Start *source* and wait for its first frame before returning a wrapper.
+
+    This lets route handlers surface validation failures raised before a
+    streaming function's first frame as normal HTTP errors, before
+    ``StreamingResponse`` sends response headers.
+    """
+    queue: asyncio.Queue[Any] = asyncio.Queue()
+    sentinel = object()
+    first_ready = asyncio.Event()
+    producer_error: list[BaseException] = []
+    frames_produced = 0
+
+    async def _drain() -> None:
+        nonlocal frames_produced
+        try:
+            async for frame in source:
+                frames_produced += 1
+                await queue.put(frame)
+                if frames_produced == 1:
+                    first_ready.set()
+        except BaseException as exc:  # noqa: BLE001 — re-raised by the consumer
+            producer_error.append(exc)
+        finally:
+            await queue.put(sentinel)
+            first_ready.set()
+
+    drain_task = asyncio.create_task(_drain())
+    try:
+        await first_ready.wait()
+    except asyncio.CancelledError:
+        if not drain_task.done():
+            drain_task.cancel()
+            with suppress(asyncio.CancelledError):
+                await drain_task
+        raise
+    if frames_produced == 0 and producer_error:
+        if not drain_task.done():
+            drain_task.cancel()
+            with suppress(asyncio.CancelledError):
+                await drain_task
+        else:
+            await drain_task
+        raise producer_error[0]
+
+    return _consume_primed_queue(queue, sentinel, drain_task, producer_error)
+
+
+async def _consume_primed_queue(
+    queue: asyncio.Queue[Any],
+    sentinel: object,
+    drain_task: asyncio.Task[None],
+    producer_error: list[BaseException],
+) -> AsyncIterator[Any]:
+    try:
+        while True:
+            item = await queue.get()
+            if item is sentinel:
+                await drain_task
+                if producer_error:
+                    raise producer_error[0]
+                return
+            yield item
+    finally:
+        if not drain_task.done():
+            drain_task.cancel()
+            with suppress(asyncio.CancelledError):
+                await drain_task
+
+
+async def _with_heartbeats(
+    source: AsyncIterator[Any],
+    *,
+    interval_seconds: float,
+) -> AsyncGenerator[Any, None]:
+    """Wrap *source* with a heartbeat injector.
+
+    Yields whatever *source* yields, with a :class:`Heartbeat` frame
+    inserted whenever the source has been quiet for at least
+    *interval_seconds*. The timer resets on every user-yielded frame
+    so back-to-back work doesn't accumulate spurious heartbeats.
+
+    The implementation drains *source* into an unbounded
+    :class:`asyncio.Queue` from a background task so the heartbeat
+    timer can run on the caller side. When *source* is exhausted (or
+    raises) the background task pushes a sentinel to terminate the
+    output stream cleanly; producer exceptions ride alongside the
+    sentinel and are re-raised on the consumer side so the caller
+    sees the failure rather than a clean EOF on a truncated stream.
+    Cancellation propagates: cancelling the consumer cancels the
+    drain task too.
+
+    *interval_seconds* of ``0`` or less disables heartbeats — useful
+    for tests that want pure pass-through and for plugins that yield
+    fast enough to never hit the idle threshold.
+    """
+    if interval_seconds <= 0:
+        async for frame in source:
+            yield frame
+        return
+
+    queue: asyncio.Queue[Any] = asyncio.Queue()
+    sentinel = object()
+    # Captured outside the queue so the consumer can re-raise it once the
+    # sentinel arrives. Putting the exception on the queue alongside frames
+    # would be ambiguous — a plugin yielding a bare ``Exception`` instance
+    # as a frame would otherwise look indistinguishable from a producer crash.
+    producer_error: list[BaseException] = []
+
+    async def _drain() -> None:
+        try:
+            async for frame in source:
+                await queue.put(frame)
+        except BaseException as exc:  # noqa: BLE001 — re-raised on consumer side
+            producer_error.append(exc)
+        finally:
+            await queue.put(sentinel)
+
+    drain_task = asyncio.create_task(_drain())
+    try:
+        while True:
+            try:
+                item = await asyncio.wait_for(queue.get(), timeout=interval_seconds)
+            except asyncio.TimeoutError:
+                yield Heartbeat()
+                continue
+            if item is sentinel:
+                # Drain finished — surface a producer crash if there was one.
+                # Awaiting the task lets ``CancelledError`` propagate naturally
+                # while making sure any synchronous bookkeeping in ``_drain``
+                # has run to completion before we re-raise.
+                await drain_task
+                if producer_error:
+                    raise producer_error[0]
+                return
+            yield item
+    finally:
+        if not drain_task.done():
+            drain_task.cancel()
+            # Suppress the cancellation exception so we don't mask the
+            # original error if the consumer is also being torn down
+            # because of one.
+            try:
+                await drain_task
+            except (asyncio.CancelledError, Exception):  # noqa: BLE001 — see comment above
+                pass
+
+
+__all__ = [
+    "DEFAULT_FUNCTION_PATH",
+    "HEARTBEAT_INTERVAL_SECONDS",
+    "NDJSON_MEDIA_TYPE",
+    "add_function_routes",
+]