# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Typed wrappers around the JSON-based NeMo Relay execution APIs. The core runtime operates on JSON-like values. This module adds a codec layer so callers can work with typed Python objects at the boundaries while the middleware pipeline continues to operate on JSON. Example:: from dataclasses import dataclass import nemo_relay.typed as typed @dataclass class SearchArgs: query: str @dataclass class SearchResult: answer: str async def tool_impl(args: SearchArgs) -> SearchResult: return SearchResult(answer=args.query.upper()) result = await typed.tool_execute( "search", SearchArgs(query="hello"), tool_impl, args_codec=typed.DataclassCodec(SearchArgs), result_codec=typed.DataclassCodec(SearchResult), ) """ from __future__ import annotations import base64 import dataclasses import importlib import inspect import io import json import pickle import typing import weakref from typing import AsyncIterator, Awaitable, Callable, Generic, Protocol, TypeVar, cast, overload from nemo_relay import Json, llm, tools from nemo_relay._native import LLMRequest, LlmStream, ScopeHandle from nemo_relay.codecs import LlmCodec, LlmResponseCodec T = TypeVar("T") TArgs = TypeVar("TArgs") TResult = TypeVar("TResult") TResponse = TypeVar("TResponse") TResponseChunk = TypeVar("TResponseChunk") if typing.TYPE_CHECKING: from _typeshed import DataclassInstance else: DataclassInstance = object class _SupportsModelDump(Protocol): def model_dump(self, *, mode: str = "python") -> Json: ... class _SupportsModelValidate(Protocol[T]): @classmethod def model_validate(cls, data: Json) -> T: ... _RUNTIME_TYPE_REGISTRY: weakref.WeakValueDictionary[str, type[object]] = weakref.WeakValueDictionary() _NO_RECONSTRUCTION = object() class _RestrictedUnpickler(pickle.Unpickler): """Unpickle only the builtin containers used by the fallback codec.""" _ALLOWED_GLOBALS = { ("builtins", "dict"), ("builtins", "frozenset"), ("builtins", "list"), ("builtins", "set"), ("builtins", "tuple"), } def find_class(self, module: str, name: str) -> object: if (module, name) not in self._ALLOWED_GLOBALS: raise pickle.UnpicklingError(f"global '{module}.{name}' is not allowed") return super().find_class(module, name) def _serialize_dataclass_instance(value: DataclassInstance) -> Json: return { field_info.name: _serialize_value(field_value) for field_info in dataclasses.fields(value) if (field_value := getattr(value, field_info.name)) is not None } def _serialize_value(value: object) -> Json: if dataclasses.is_dataclass(value) and not isinstance(value, type): return _serialize_dataclass_instance(cast(DataclassInstance, value)) if isinstance(value, list): return [_serialize_value(item) for item in value] if isinstance(value, dict): return {cast(str, key): _serialize_value(item) for key, item in value.items()} if hasattr(value, "model_dump"): try: return cast(_SupportsModelDump, value).model_dump(mode="json") except Exception: pass # Intentional: fall through to the default string encoding return cast(Json, value) def _register_runtime_type(type_obj: type[object]) -> str: token = f"{type_obj.__module__}.{type_obj.__qualname__}:{id(type_obj)}" _RUNTIME_TYPE_REGISTRY[token] = type_obj return token def _resolve_runtime_type(token: object) -> type[object] | None: if not isinstance(token, str): return None resolved = _RUNTIME_TYPE_REGISTRY.get(token) return resolved if isinstance(resolved, type) else None def _resolve_importable_type(path: object) -> type[object] | None: if not isinstance(path, str) or not path: return None parts = path.split(".") for split_at in range(len(parts) - 1, 0, -1): mod_name = ".".join(parts[:split_at]) qualname_parts = parts[split_at:] if "" in qualname_parts: continue try: resolved: object = importlib.import_module(mod_name) for part in qualname_parts: resolved = getattr(resolved, part) except (ImportError, AttributeError): continue if isinstance(resolved, type): return resolved return None # --------------------------------------------------------------------------- # Codec protocol and built-in implementations # --------------------------------------------------------------------------- class Codec(Generic[T]): """Bidirectional conversion protocol between a Python type and JSON. Implementations convert between ergonomic Python objects and the JSON-compatible values used by the underlying NeMo Relay runtime. """ def to_json(self, value: T) -> Json: """Convert a typed value to a JSON-serializable object. Args: value: Typed Python object to serialize. Returns: Json: JSON-compatible representation of ``value``. """ raise NotImplementedError def from_json(self, data: Json) -> T: """Reconstruct a typed value from a JSON-serializable object. Args: data: JSON-compatible value to deserialize. Returns: T: Typed Python object reconstructed from ``data``. """ raise NotImplementedError class JsonPassthrough(Codec[Json]): """Identity codec for callers already working with JSON values.""" def to_json(self, value: Json) -> Json: """Return ``value`` unchanged. Args: value: JSON-compatible value to forward into the runtime. Returns: Json: The original ``value``. """ return value def from_json(self, data: Json) -> Json: """Return ``data`` unchanged. Args: data: JSON-compatible value produced by the runtime. Returns: Json: The original ``data``. """ return data class PydanticCodec(Codec[T]): """Codec for models exposing ``model_dump`` and ``model_validate``. Args: model_cls: Pydantic-compatible model class implementing ``model_validate()``. Example:: codec = PydanticCodec(MyModel) """ def __init__(self, model_cls: type[T]) -> None: self._cls = model_cls def to_json(self, value: T) -> Json: """Serialize a Pydantic model to a JSON-serializable dict. Args: value: Pydantic model instance to serialize. Returns: Json: JSON-compatible dictionary produced by ``model_dump()``. """ return cast(_SupportsModelDump, value).model_dump(mode="json") def from_json(self, data: Json) -> T: """Deserialize JSON data into a Pydantic model. Args: data: JSON-compatible value to validate. Returns: T: Model instance produced by ``model_validate()``. """ return cast(_SupportsModelValidate[T], self._cls).model_validate(data) class DataclassCodec(Codec[T]): """Codec for ``dataclasses.dataclass`` models. Args: dc_cls: Dataclass type to serialize and deserialize. Example:: from dataclasses import dataclass @dataclass class Point: x: int y: int codec = DataclassCodec(Point) """ def __init__(self, dc_cls: type[T]) -> None: self._cls = dc_cls def to_json(self, value: T) -> Json: """Serialize a dataclass instance to a JSON-compatible dictionary. Args: value: Dataclass instance to serialize. Returns: Json: JSON-compatible dictionary representation of ``value``. """ return _serialize_dataclass_instance(cast(DataclassInstance, value)) def from_json(self, data: Json) -> T: """Deserialize JSON data into a dataclass instance. Args: data: JSON-compatible mapping used as keyword arguments for the dataclass constructor. Returns: T: Dataclass instance reconstructed from ``data``. """ return cast(Callable[..., T], self._cls)(**cast(dict[str, object], data)) class BestEffortAnyCodec(Codec[object]): """Best-effort codec for arbitrary Python values. It prefers JSON-native encodings, then dataclass or Pydantic-style model support, and finally falls back to pickle or string encoding when needed. Pickle decoding is restricted to a small allowlist of builtin containers; arbitrary classes and callable globals are never reconstructed. Notes: This codec favors resilience over strictness. If it cannot preserve the original type exactly, it degrades to a string instead of raising. """ def to_json(self, value: object) -> Json: """Serialize an arbitrary Python value to a JSON-serializable form. Args: value: Arbitrary Python object to serialize. Returns: Json: Tagged JSON-compatible value that ``from_json()`` can reconstruct on a best-effort basis. Tries, in order: Pydantic ``model_dump()``, ``dataclasses.asdict()``, native JSON encoding, pickle fallback, and finally ``str()`` as a last resort. Each encoding is tagged with a ``__nv_*__`` key so that ``from_json`` can reconstruct the original type. """ try: if hasattr(value, "model_dump"): dumped = cast(_SupportsModelDump, value).model_dump(mode="json") return { "__nv_pydantic__": f"{value.__class__.__module__}.{value.__class__.__qualname__}", "__nv_runtime_type__": _register_runtime_type(value.__class__), "data": dumped, } except Exception: pass # Intentional: fall through to the next encoding strategy # Dataclass if dataclasses.is_dataclass(value) and not isinstance(value, type): return { "__nv_dataclass__": f"{value.__class__.__module__}.{value.__class__.__qualname__}", "__nv_runtime_type__": _register_runtime_type(value.__class__), "data": _serialize_dataclass_instance(cast(DataclassInstance, value)), } # Try JSON encoding directly try: return cast(Json, json.loads(json.dumps(value))) except Exception: # Fallback: pickle try: pickled = pickle.dumps(value) encoded = base64.b64encode(pickled).decode("ascii") return { "__nv_pickle__": f"{value.__class__.__module__}.{value.__class__.__qualname__}", "data": encoded, } except Exception: # As last resort, do string (may be lossy, but not error) return { "__nv_fallback_str__": f"{value.__class__.__module__}.{value.__class__.__qualname__}", "data": str(value), } @staticmethod def _resolve_tagged_type(data: dict[str, Json], tag: str) -> type[object] | None: return _resolve_runtime_type(data.get("__nv_runtime_type__")) or _resolve_importable_type(data[tag]) @staticmethod def _mapping_payload(data: dict[str, Json]) -> dict[str, object] | None: payload = data.get("data") if isinstance(payload, dict): return cast(dict[str, object], payload) return None @staticmethod def _string_payload(data: dict[str, Json]) -> str | None: payload = data.get("data") return payload if isinstance(payload, str) else None def _reconstruct_pydantic(self, data: dict[str, Json]) -> object: if "__nv_pydantic__" not in data: return _NO_RECONSTRUCTION try: cls = self._resolve_tagged_type(data, "__nv_pydantic__") if cls is not None and hasattr(cls, "model_validate"): return cast(_SupportsModelValidate[object], cls).model_validate(data["data"]) except Exception: pass return _NO_RECONSTRUCTION def _reconstruct_dataclass(self, data: dict[str, Json]) -> object: if "__nv_dataclass__" not in data: return _NO_RECONSTRUCTION try: cls = self._resolve_tagged_type(data, "__nv_dataclass__") payload = self._mapping_payload(data) if cls is not None and dataclasses.is_dataclass(cls) and payload is not None: return cls(**payload) except Exception: pass return _NO_RECONSTRUCTION @staticmethod def _reconstruct_pickle(data: dict[str, Json]) -> object: if "__nv_pickle__" not in data: return _NO_RECONSTRUCTION try: payload = BestEffortAnyCodec._string_payload(data) if payload is None: return _NO_RECONSTRUCTION decoded = base64.b64decode(payload, validate=True) return _RestrictedUnpickler(io.BytesIO(decoded)).load() except Exception: return _NO_RECONSTRUCTION @staticmethod def _reconstruct_fallback_string(data: dict[str, Json]) -> object: if "__nv_fallback_str__" in data: payload = BestEffortAnyCodec._string_payload(data) if payload is not None: return payload return _NO_RECONSTRUCTION def from_json(self, data: Json) -> object: """Reconstruct a Python value from its tagged JSON representation. Args: data: Tagged JSON-compatible value produced by ``to_json()`` or a plain JSON-compatible value. Returns: object: Best-effort reconstruction of the original Python value. Recognises the ``__nv_pydantic__``, ``__nv_dataclass__``, ``__nv_pickle__``, and ``__nv_fallback_str__`` tags produced by ``to_json`` and dispatches to the appropriate reconstruction strategy. Pickle tags are decoded only for allowlisted builtin containers. Falls through to returning the raw data if no tag is recognised. """ if isinstance(data, dict) and "data" in data: for reconstructor in ( self._reconstruct_pydantic, self._reconstruct_dataclass, self._reconstruct_pickle, self._reconstruct_fallback_string, ): result = reconstructor(data) if result is not _NO_RECONSTRUCTION: return result return data # --------------------------------------------------------------------------- # Typed execute wrappers # --------------------------------------------------------------------------- @overload async def tool_execute( name: str, args: TArgs, func: Callable[[TArgs], Awaitable[TResult]], args_codec: Codec[TArgs], result_codec: Codec[TResult], *, handle: ScopeHandle | None = None, attributes: int | None = None, data: Json | None = None, metadata: Json | None = None, ) -> TResult: ... @overload async def tool_execute( name: str, args: TArgs, func: Callable[[TArgs], TResult], args_codec: Codec[TArgs], result_codec: Codec[TResult], *, handle: ScopeHandle | None = None, attributes: int | None = None, data: Json | None = None, metadata: Json | None = None, ) -> TResult: ... async def tool_execute( name: str, args: TArgs, func: Callable[[TArgs], TResult] | Callable[[TArgs], Awaitable[TResult]], args_codec: Codec[TArgs], result_codec: Codec[TResult], *, handle: ScopeHandle | None = None, attributes: int | None = None, data: Json | None = None, metadata: Json | None = None, ) -> TResult: """Run ``nemo_relay.tools.execute`` with typed arguments and results. Args: name: Tool name recorded on emitted lifecycle events. args: Typed arguments to serialize before entering the runtime. func: Tool implementation invoked with deserialized typed arguments. The implementation may be synchronous or asynchronous. args_codec: Codec used to convert ``args`` to and from JSON. result_codec: Codec used to convert the tool result to and from JSON. handle: Optional parent scope handle. When omitted, the current scope becomes the parent. attributes: Optional native tool attributes attached to the start event. data: Optional JSON payload recorded on the emitted start event. metadata: Optional JSON metadata recorded on the emitted start event. Returns: TResult: The decoded typed result produced by ``func``. Example:: from dataclasses import dataclass from nemo_relay.typed import DataclassCodec, tool_execute @dataclass class SearchArgs: query: str @dataclass class SearchResult: answer: str async def tool_impl(args: SearchArgs) -> SearchResult: return SearchResult(answer=args.query.upper()) result = await tool_execute( "search", SearchArgs(query="hello"), tool_impl, args_codec=DataclassCodec(SearchArgs), result_codec=DataclassCodec(SearchResult), handle=None, attributes=None, data={"path": "typed"}, metadata={"request_id": "req-1"}, ) """ json_args = args_codec.to_json(args) async def _json_func(json_args_inner: Json) -> Json: typed_args = args_codec.from_json(json_args_inner) result: TResult | Awaitable[TResult] = func(typed_args) if inspect.isawaitable(result): return result_codec.to_json(await typing.cast(Awaitable[TResult], result)) return result_codec.to_json(typing.cast(TResult, result)) json_result = await tools.execute( name, json_args, _json_func, handle=handle, attributes=attributes, data=data, metadata=metadata, ) return result_codec.from_json(json_result) @overload async def llm_execute( name: str, request: LLMRequest, func: Callable[[LLMRequest], Awaitable[TResponse]], response_json_codec: Codec[TResponse], *, handle: ScopeHandle | None = None, attributes: int | None = None, data: Json | None = None, metadata: Json | None = None, model_name: str | None = None, codec: LlmCodec | None = None, response_codec: LlmResponseCodec | None = None, ) -> TResponse: ... @overload async def llm_execute( name: str, request: LLMRequest, func: Callable[[LLMRequest], TResponse], response_json_codec: Codec[TResponse], *, handle: ScopeHandle | None = None, attributes: int | None = None, data: Json | None = None, metadata: Json | None = None, model_name: str | None = None, codec: LlmCodec | None = None, response_codec: LlmResponseCodec | None = None, ) -> TResponse: ... async def llm_execute( name: str, request: LLMRequest, func: Callable[[LLMRequest], TResponse] | Callable[[LLMRequest], Awaitable[TResponse]], response_json_codec: Codec[TResponse], *, handle: ScopeHandle | None = None, attributes: int | None = None, data: Json | None = None, metadata: Json | None = None, model_name: str | None = None, codec: LlmCodec | None = None, response_codec: LlmResponseCodec | None = None, ) -> TResponse: """Run ``nemo_relay.llm.execute`` and decode the returned response type. Args: name: Provider or logical call name recorded on emitted events. request: Raw ``LLMRequest`` passed through the managed LLM pipeline. func: Provider callback invoked with the possibly intercepted request. The implementation may be synchronous or asynchronous. response_json_codec: Codec used to convert the provider response to and from JSON. handle: Optional parent scope handle. When omitted, the current scope becomes the parent. attributes: Optional native LLM attributes attached to the start event. data: Optional JSON payload recorded on the emitted start event. metadata: Optional JSON metadata recorded on the emitted start event. model_name: Optional normalized model name to record separately from the provider-specific request payload. codec: Optional request codec used to expose ``AnnotatedLLMRequest`` to request intercepts. response_codec: Optional observability codec used to attach an annotated response to the emitted ``LLMEnd`` event. Returns: TResponse: The decoded typed response produced by ``func``. Example:: import nemo_relay from dataclasses import dataclass from nemo_relay.typed import DataclassCodec, llm_execute @dataclass class MyResponse: text: str async def llm_impl(request: nemo_relay.LLMRequest): return {"text": "hello"} typed_response = await llm_execute( "demo-provider", nemo_relay.LLMRequest({}, {"messages": [{"role": "user", "content": "hi"}]}), llm_impl, response_json_codec=DataclassCodec(MyResponse), handle=None, attributes=None, data={"path": "typed"}, metadata={"request_id": "req-2"}, model_name="demo-model", codec=None, response_codec=None, ) """ async def _json_func(request_inner: LLMRequest) -> Json: result: TResponse | Awaitable[TResponse] = func(request_inner) if inspect.isawaitable(result): return response_json_codec.to_json(await typing.cast(Awaitable[TResponse], result)) return response_json_codec.to_json(typing.cast(TResponse, result)) json_result = await llm.execute( name, request, _json_func, handle=handle, attributes=attributes, data=data, metadata=metadata, model_name=model_name, codec=codec, response_codec=response_codec, ) return response_json_codec.from_json(json_result) # SONAR_IGNORE_START async def llm_stream_execute( name: str, request: LLMRequest, func: Callable[[LLMRequest], AsyncIterator[TResponseChunk]], collector: Callable[[TResponseChunk], None], finalizer: Callable[[], TResponse], chunk_json_codec: Codec[TResponseChunk], response_json_codec: Codec[TResponse], *, handle: ScopeHandle | None = None, attributes: int | None = None, data: Json | None = None, metadata: Json | None = None, model_name: str | None = None, codec: LlmCodec | None = None, response_codec: LlmResponseCodec | None = None, ) -> LlmStream: # SONAR_IGNORE_STOP """Run ``nemo_relay.llm.stream_execute`` with typed chunks and final output. Args: name: Provider or logical call name recorded on emitted events. request: Raw ``LLMRequest`` passed through the managed LLM pipeline. func: Async generator invoked with the request and yielding typed chunks. collector: Callback invoked for each decoded typed chunk after it passes through the runtime's streaming intercept chain. finalizer: Callback invoked after the stream completes to build the final typed response recorded on the emitted ``LLMEnd`` event. chunk_json_codec: Codec used to convert streamed chunks to and from JSON. response_json_codec: Codec used to convert the final aggregated response to and from JSON. handle: Optional parent scope handle. When omitted, the current scope becomes the parent. attributes: Optional native LLM attributes attached to the start event. data: Optional JSON payload recorded on the emitted start event. metadata: Optional JSON metadata recorded on the emitted start event. model_name: Optional normalized model name to record separately from the provider-specific request payload. codec: Optional request codec used to expose ``AnnotatedLLMRequest`` to request intercepts. response_codec: Optional observability codec used to attach an annotated response to the emitted ``LLMEnd`` event. Returns: LlmStream: Async iterator that yields the streamed JSON chunks. Notes: ``collector`` receives typed chunks, but the returned stream still yields the raw JSON chunk values produced by the underlying runtime. Example:: import nemo_relay from dataclasses import dataclass from nemo_relay.typed import DataclassCodec, llm_stream_execute @dataclass class MyChunk: token: str @dataclass class MyResponse: text: str collected = [] async def stream_impl(request: nemo_relay.LLMRequest): yield MyChunk(token="hel") yield MyChunk(token="lo") def collect_chunk(chunk: MyChunk) -> None: collected.append(chunk) def finish_response() -> MyResponse: return MyResponse(text="".join(chunk.token for chunk in collected)) stream = await llm_stream_execute( "demo-provider", nemo_relay.LLMRequest({}, {"messages": [{"role": "user", "content": "hi"}]}), stream_impl, collector=collect_chunk, finalizer=finish_response, chunk_json_codec=DataclassCodec(MyChunk), response_json_codec=DataclassCodec(MyResponse), handle=None, attributes=None, data={"path": "typed-stream"}, metadata={"request_id": "req-3"}, model_name="demo-model", codec=None, response_codec=None, ) async for chunk in stream: print(chunk) """ async def _json_func(request_inner: LLMRequest) -> AsyncIterator[Json]: async for typed_chunk in func(request_inner): yield chunk_json_codec.to_json(typed_chunk) def _json_collector(json_chunk: Json) -> None: collector(chunk_json_codec.from_json(json_chunk)) def _json_finalizer() -> Json: return response_json_codec.to_json(finalizer()) return await llm.stream_execute( name, request, _json_func, _json_collector, _json_finalizer, handle=handle, attributes=attributes, data=data, metadata=metadata, model_name=model_name, codec=codec, response_codec=response_codec, ) __all__ = [ "Codec", "BestEffortAnyCodec", "DataclassCodec", "JsonPassthrough", "PydanticCodec", "tool_execute", "llm_execute", "llm_stream_execute", ]