From 3ba1c021b034af49f2fb8ee5f58601d0a94fc463 Mon Sep 17 00:00:00 2001 From: Kyle Zheng <126034466+KyleZheng1284@users.noreply.github.com> Date: Tue, 30 Jun 2026 05:27:33 +0000 Subject: [PATCH 1/3] Add abstraction for LLM-based tasks Signed-off-by: Kyle Zheng <126034466+KyleZheng1284@users.noreply.github.com> --- .../nemo_retriever/common/params/__init__.py | 4 + .../nemo_retriever/common/params/models.py | 162 ++- .../graph/graph_pipeline_registry.py | 617 ++++++++++-- .../src/nemo_retriever/models/llm/__init__.py | 36 +- .../models/llm/clients/__init__.py | 4 +- .../models/llm/clients/judge.py | 6 +- .../models/llm/clients/litellm.py | 127 +-- .../models/llm/tasks/__init__.py | 17 + .../nemo_retriever/models/llm/tasks/base.py | 80 ++ .../models/llm/tasks/generic.py | 106 ++ .../models/llm/tasks/rag_answer.py | 149 +++ .../models/llm/tasks/summarize.py | 90 ++ .../nemo_retriever/models/llm/text_utils.py | 2 +- .../src/nemo_retriever/models/llm/types.py | 40 +- .../src/nemo_retriever/operators/__init__.py | 11 +- .../operators/generation/__init__.py | 15 + .../operators/generation/base.py | 296 ++++++ .../operators/generation/generic.py | 70 ++ .../operators/generation/summarization.py | 65 ++ .../tools/evaluation/generation.py | 111 ++- nemo_retriever/tests/test_caption.py | 6 + nemo_retriever/tests/test_generation_tasks.py | 925 ++++++++++++++++++ .../tests/test_graph_pipeline_registry.py | 126 ++- 23 files changed, 2795 insertions(+), 270 deletions(-) create mode 100644 nemo_retriever/src/nemo_retriever/models/llm/tasks/__init__.py create mode 100644 nemo_retriever/src/nemo_retriever/models/llm/tasks/base.py create mode 100644 nemo_retriever/src/nemo_retriever/models/llm/tasks/generic.py create mode 100644 nemo_retriever/src/nemo_retriever/models/llm/tasks/rag_answer.py create mode 100644 nemo_retriever/src/nemo_retriever/models/llm/tasks/summarize.py create mode 100644 nemo_retriever/src/nemo_retriever/operators/generation/__init__.py create mode 100644 nemo_retriever/src/nemo_retriever/operators/generation/base.py create mode 100644 nemo_retriever/src/nemo_retriever/operators/generation/generic.py create mode 100644 nemo_retriever/src/nemo_retriever/operators/generation/summarization.py create mode 100644 nemo_retriever/tests/test_generation_tasks.py diff --git a/nemo_retriever/src/nemo_retriever/common/params/__init__.py b/nemo_retriever/src/nemo_retriever/common/params/__init__.py index 57c8918f2..daec10ee1 100644 --- a/nemo_retriever/src/nemo_retriever/common/params/__init__.py +++ b/nemo_retriever/src/nemo_retriever/common/params/__init__.py @@ -20,6 +20,7 @@ from nemo_retriever.common.params.models import LanceDbParams from nemo_retriever.common.params.models import LLMInferenceParams from nemo_retriever.common.params.models import LLMRemoteClientParams +from nemo_retriever.common.params.models import LLMSamplingOverrides from nemo_retriever.common.params.models import ModelRuntimeParams from nemo_retriever.common.params.models import OcrParams from nemo_retriever.common.params.models import PageElementsParams @@ -30,6 +31,7 @@ from nemo_retriever.common.params.models import TabularExtractParams from nemo_retriever.common.params.models import TableParams from nemo_retriever.common.params.models import TextChunkParams +from nemo_retriever.common.params.models import TextGenerationParams from nemo_retriever.common.params.models import MetaJoinKey from nemo_retriever.common.params.models import VdbUploadParams from nemo_retriever.common.params.models import VideoFrameParams @@ -58,6 +60,7 @@ "LanceDbParams", "LLMInferenceParams", "LLMRemoteClientParams", + "LLMSamplingOverrides", "ModelRuntimeParams", "OcrParams", "PageElementsParams", @@ -69,6 +72,7 @@ "TabularExtractParams", "TableParams", "TextChunkParams", + "TextGenerationParams", "MetaJoinKey", "VdbUploadParams", "VideoFrameParams", diff --git a/nemo_retriever/src/nemo_retriever/common/params/models.py b/nemo_retriever/src/nemo_retriever/common/params/models.py index de3d30e84..72e44139c 100644 --- a/nemo_retriever/src/nemo_retriever/common/params/models.py +++ b/nemo_retriever/src/nemo_retriever/common/params/models.py @@ -13,7 +13,7 @@ from upath import UPath from nemo_retriever.tabular_data.sql_database import SQLDatabase -from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator +from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, field_validator, model_serializer, model_validator from nemo_retriever.common.remote_auth import resolve_remote_api_key @@ -50,17 +50,35 @@ class _ParamsModel(BaseModel): model_config = ConfigDict(extra="forbid") + # Keep the explicit no-auth intent after NO_API_KEY is normalized to + # None for existing runtime consumers. Graph persistence uses this + # private provenance to distinguish no-auth from worker-side env lookup. + _no_api_key_fields: set[str] = PrivateAttr(default_factory=set) + @model_validator(mode="after") def _resolve_api_keys(self) -> "_ParamsModel": for field_name in type(self).model_fields: if _is_api_key_field(field_name): value = getattr(self, field_name, None) if value is None: + if field_name in self._no_api_key_fields: + continue setattr(self, field_name, resolve_remote_api_key()) elif value == NO_API_KEY: + self._no_api_key_fields.add(field_name) setattr(self, field_name, None) + else: + self._no_api_key_fields.discard(field_name) return self + def _uses_no_api_key(self, field_name: str) -> bool: + """Return whether an API-key field was explicitly disabled. + + This is an internal persistence hook. Runtime callers continue to + observe None for NO_API_KEY exactly as before. + """ + return field_name in self._no_api_key_fields and getattr(self, field_name, None) is None + def __repr__(self) -> str: parts: list[str] = [] for field_name in type(self).model_fields: @@ -547,14 +565,14 @@ class LLMInferenceParams(_ParamsModel): to any task that invokes an LLM (captioning, summarization, etc.). """ - temperature: float = 1.0 + temperature: Optional[float] = 1.0 top_p: Optional[float] = None max_tokens: int = 1024 @field_validator("temperature") @classmethod - def _check_temperature(cls, v: float) -> float: - if not (0.0 <= v <= 2.0): + def _check_temperature(cls, v: Optional[float]) -> Optional[float]: + if v is not None and not (0.0 <= v <= 2.0): raise ValueError("temperature must be between 0.0 and 2.0") return v @@ -579,7 +597,9 @@ def to_sampling_kwargs(self) -> dict[str, Any]: many backends (vLLM, OpenAI, NIM) change behaviour when the key is present vs. absent. """ - kw: dict[str, Any] = {"temperature": self.temperature, "max_tokens": self.max_tokens} + kw: dict[str, Any] = {"max_tokens": self.max_tokens} + if self.temperature is not None: + kw["temperature"] = self.temperature if self.top_p is not None: kw["top_p"] = self.top_p return kw @@ -618,6 +638,131 @@ def _check_timeout(cls, v: float) -> float: return v +class LLMSamplingOverrides(_ParamsModel): + """Partial sampling overrides resolved on top of task-specific defaults.""" + + temperature: Optional[float] = None + top_p: Optional[float] = None + max_tokens: Optional[int] = None + + @field_validator("temperature") + @classmethod + def _check_temperature(cls, v: Optional[float]) -> Optional[float]: + if v is not None and not (0.0 <= v <= 2.0): + raise ValueError("temperature must be between 0.0 and 2.0") + return v + + @field_validator("top_p") + @classmethod + def _check_top_p(cls, v: Optional[float]) -> Optional[float]: + if v is not None and not (0.0 <= v <= 1.0): + raise ValueError("top_p must be between 0.0 and 1.0") + return v + + @field_validator("max_tokens") + @classmethod + def _check_max_tokens(cls, v: Optional[int]) -> Optional[int]: + if v is not None and v <= 0: + raise ValueError("max_tokens must be > 0") + return v + + @model_validator(mode="after") + def _reject_explicit_null_max_tokens(self) -> "LLMSamplingOverrides": + if "max_tokens" in self.model_fields_set and self.max_tokens is None: + raise ValueError("max_tokens cannot be None; omit it to inherit the task default") + return self + + @model_serializer(mode="plain") + def _serialize_only_explicit_overrides(self) -> dict[str, Any]: + """Preserve omitted-vs-null state across model and JSON round trips.""" + return { + name: getattr(self, name) + for name in ("temperature", "top_p", "max_tokens") + if name in self.model_fields_set + } + + def __eq__(self, other: object) -> bool: + if isinstance(other, LLMSamplingOverrides): + return self.model_fields_set == other.model_fields_set and super().__eq__(other) + return super().__eq__(other) + + def resolve(self, defaults: LLMInferenceParams) -> LLMInferenceParams: + """Apply explicitly supplied fields to defaults.""" + values = defaults.model_dump() + for name in self.model_fields_set: + value = getattr(self, name) + values[name] = value + return LLMInferenceParams(**values) + + +_SAMPLING_UNSET = object() + + +class TextGenerationParams(_ParamsModel): + """Transport, task controls, and partial sampling for text generation.""" + + transport: LLMRemoteClientParams + sampling: LLMSamplingOverrides = Field(default_factory=LLMSamplingOverrides) + prompt: Optional[str] = None + system_prompt: Optional[str] = None + reasoning_enabled: Optional[bool] = None + max_workers: int = Field(default=8, ge=1) + + def resolve_sampling(self, defaults: LLMInferenceParams) -> LLMInferenceParams: + """Resolve explicit sampling fields over a task's defaults.""" + return self.sampling.resolve(defaults) + + @classmethod + def from_kwargs( + cls, + *, + model: str, + api_base: Optional[str] = None, + api_key: Optional[str] = None, + temperature: Any = _SAMPLING_UNSET, + top_p: Any = _SAMPLING_UNSET, + max_tokens: Any = _SAMPLING_UNSET, + extra_params: Optional[dict[str, Any]] = None, + num_retries: int = 3, + timeout: float = 120.0, + rag_system_prompt: Optional[str] = None, + rag_system_prompt_prefix: Optional[str] = None, + reasoning_enabled: Optional[bool] = None, + prompt: Optional[str] = None, + system_prompt: Optional[str] = None, + max_workers: int = 8, + ) -> "TextGenerationParams": + """Construct structured text-generation params from flat kwargs.""" + sampling_values: dict[str, Any] = {} + for name, value in ( + ("temperature", temperature), + ("top_p", top_p), + ("max_tokens", max_tokens), + ): + if value is not _SAMPLING_UNSET: + sampling_values[name] = value + + transport_reasoning = True if reasoning_enabled is None else reasoning_enabled + return cls( + transport=LLMRemoteClientParams( + model=model, + api_base=api_base, + api_key=api_key, + num_retries=num_retries, + timeout=timeout, + extra_params=extra_params or {}, + rag_system_prompt=rag_system_prompt, + rag_system_prompt_prefix=rag_system_prompt_prefix, + reasoning_enabled=transport_reasoning, + ), + sampling=LLMSamplingOverrides(**sampling_values), + prompt=prompt, + system_prompt=system_prompt, + reasoning_enabled=reasoning_enabled, + max_workers=max_workers, + ) + + class CaptionParams(LLMInferenceParams): endpoint_url: Optional[str] = None model_name: str = "nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16" @@ -633,6 +778,13 @@ class CaptionParams(LLMInferenceParams): caption_infographics: bool = False extra_body: dict[str, Any] = Field(default_factory=dict) + @field_validator("temperature") + @classmethod + def _require_temperature(cls, value: Optional[float]) -> float: + if value is None: + raise ValueError("temperature cannot be None for captioning") + return value + class WebhookParams(_ParamsModel): """Configuration for the webhook notification stage. diff --git a/nemo_retriever/src/nemo_retriever/graph/graph_pipeline_registry.py b/nemo_retriever/src/nemo_retriever/graph/graph_pipeline_registry.py index 97c04ae1b..5cba747d5 100644 --- a/nemo_retriever/src/nemo_retriever/graph/graph_pipeline_registry.py +++ b/nemo_retriever/src/nemo_retriever/graph/graph_pipeline_registry.py @@ -48,9 +48,11 @@ def _build(): Union, ) -from nemo_retriever.operators.abstract_operator import AbstractOperator -from nemo_retriever.graph.pipeline_graph import Graph, Node +from pydantic import BaseModel +from nemo_retriever.common.remote_auth import resolve_remote_api_key +from nemo_retriever.graph.pipeline_graph import Graph, Node +from nemo_retriever.operators.abstract_operator import AbstractOperator # --------------------------------------------------------------------------- # Helpers @@ -75,44 +77,379 @@ def _import_class(qualified: str) -> type: return cls +_GRAPH_FORMAT_VERSION = 2 +_PYDANTIC_MODEL_MARKER = "__pydantic_model__" +_PYDANTIC_FIELDS = "fields" +_PYDANTIC_FIELDS_SET = "fields_set" +_SECRET_ENV_MARKER = "__secret_env__" +_SECRET_NO_AUTH_MARKER = "__secret_no_auth__" +_TUPLE_MARKER = "__tuple__" +_FROZENSET_MARKER = "__frozenset__" +_MAPPING_MARKER = "__mapping__" +_OMIT_FIELD = object() + + +class GraphSerializationError(ValueError): + """Raised when graph state cannot be serialized safely and losslessly.""" + + +def _is_api_key_field(field_name: Optional[str]) -> bool: + return bool(field_name) and (field_name == "api_key" or field_name.endswith("_api_key")) + + +def _is_obvious_secret_field(field_name: Optional[str]) -> bool: + if not field_name: + return False + normalized = field_name.lower().replace("-", "_") + compact = normalized.replace("_", "") + if _is_api_key_field(normalized): + return True + if normalized in { + "authorization", + "credential", + "credentials", + "password", + "passwd", + "private_key", + "secret", + "secret_key", + "storage_options", + }: + return True + if set(normalized.split("_")) & {"password", "passwd", "secret"}: + return True + if normalized == "token" or normalized.endswith("_token"): + return True + if compact in { + "authorization", + "credential", + "credentials", + "storageoptions", + }: + return True + return compact.endswith(("apikey", "password", "passwd", "secret", "secretkey", "privatekey", "token")) + + +def _is_empty_secret(value: Any) -> bool: + if value is None or value == "": + return True + return isinstance(value, (dict, list, tuple, set, frozenset)) and not value + + +def _import_qualified_object(qualified: str) -> Any: + """Import a qualified module attribute, including nested attributes.""" + parts = qualified.split(".") + for index in range(len(parts) - 1, 0, -1): + module_name = ".".join(parts[:index]) + try: + obj: Any = importlib.import_module(module_name) + except ModuleNotFoundError: + continue + try: + for part in parts[index:]: + obj = getattr(obj, part) + except AttributeError as exc: + raise ImportError(f"Cannot import qualified object {qualified!r}") from exc + return obj + raise ImportError(f"Cannot import qualified object {qualified!r}") + + +def _model_uses_no_api_key(model: BaseModel, field_name: str) -> bool: + checker = getattr(model, "_uses_no_api_key", None) + if callable(checker): + return bool(checker(field_name)) + return field_name in getattr(model, "_no_api_key_fields", set()) + + +def _contains_pydantic_model(value: Any, seen: Optional[Set[int]] = None) -> bool: + if isinstance(value, BaseModel): + return True + if value is None or isinstance(value, (bool, int, float, str, bytes)): + return False + if seen is None: + seen = set() + value_id = id(value) + if value_id in seen: + return False + seen.add(value_id) + if isinstance(value, dict): + return any(_contains_pydantic_model(item, seen) for pair in value.items() for item in pair) + if isinstance(value, (list, tuple, set, frozenset)): + return any(_contains_pydantic_model(item, seen) for item in value) + return False + + +def _encode_secret( + value: Any, + *, + field_name: str, + path: str, + owner: Optional[BaseModel], + allow_api_key_env: bool, +) -> Any: + if _is_api_key_field(field_name): + if value == "" or (value is None and owner is not None and _model_uses_no_api_key(owner, field_name)): + return {_SECRET_NO_AUTH_MARKER: ""} + if owner is None and allow_api_key_env and isinstance(value, str) and value.strip().startswith("os.environ/"): + return value.strip() + if owner is None and not allow_api_key_env: + if value is None: + return None + raise GraphSerializationError( + f"{path}: refusing to serialize an API key inside an opaque mapping; " + "move it to a typed params field or top-level operator kwarg" + ) + if value is not None and not isinstance(value, str): + raise GraphSerializationError(f"{path}: API-key fields must be strings, null, or the no-auth marker") + scope = "model" if owner is not None else "operator" + return {_SECRET_ENV_MARKER: scope} + if not _is_empty_secret(value): + raise GraphSerializationError(f"{path}: refusing to serialize non-rehydratable secret field {field_name!r}") + return value + + +def _encode_value( + value: Any, + *, + path: str, + field_name: Optional[str] = None, + owner: Optional[BaseModel] = None, + allow_api_key_env: bool = False, +) -> Any: + """Recursively encode ``value`` into lossless, JSON-native graph state.""" + if _is_obvious_secret_field(field_name): + return _encode_secret( + value, + field_name=field_name or "", + path=path, + owner=owner, + allow_api_key_env=allow_api_key_env, + ) + if value is None or isinstance(value, (bool, int, float, str)): + return value + if isinstance(value, BaseModel): + model_type = type(value) + qualified = _qualified_name(model_type) + try: + restored_type = _import_qualified_object(qualified) + except ImportError as exc: + raise GraphSerializationError(f"{path}: Pydantic model {qualified!r} is not rehydratable") from exc + if restored_type is not model_type: + raise GraphSerializationError(f"{path}: Pydantic model {qualified!r} does not round-trip by identity") + dumped = value.model_dump(mode="python") + if not isinstance(dumped, dict): + raise GraphSerializationError( + f"{path}: Pydantic model serializer must return a mapping, got {type(dumped).__name__}" + ) + fields: Dict[str, Any] = {} + for name, dumped_item in dumped.items(): + actual = getattr(value, name, dumped_item) + item = actual if _contains_pydantic_model(actual) else dumped_item + fields[name] = _encode_value( + item, + path=f"{path}.{name}", + field_name=name, + owner=value, + ) + return { + _PYDANTIC_MODEL_MARKER: qualified, + _PYDANTIC_FIELDS: fields, + _PYDANTIC_FIELDS_SET: sorted(value.model_fields_set), + } + if isinstance(value, type): + qualified = _qualified_name(value) + try: + restored = _import_qualified_object(qualified) + except ImportError as exc: + raise GraphSerializationError(f"{path}: type {qualified!r} is not rehydratable") from exc + if restored is not value: + raise GraphSerializationError(f"{path}: type {qualified!r} does not round-trip by identity") + return {"__type_ref__": qualified} + if callable(value) and hasattr(value, "__qualname__"): + module = getattr(value, "__module__", None) or "" + qualified = f"{module}.{value.__qualname__}" + try: + restored = _import_qualified_object(qualified) + except ImportError as exc: + raise GraphSerializationError(f"{path}: callable {qualified!r} is not rehydratable") from exc + if restored is not value: + raise GraphSerializationError(f"{path}: callable {qualified!r} does not round-trip by identity") + return {"__callable_ref__": qualified} + if isinstance(value, Path): + return {"__path__": str(value)} + if isinstance(value, tuple): + return {_TUPLE_MARKER: [_encode_value(item, path=f"{path}[{index}]") for index, item in enumerate(value)]} + if isinstance(value, list): + return [_encode_value(item, path=f"{path}[{index}]") for index, item in enumerate(value)] + if isinstance(value, (set, frozenset)): + encoded = [_encode_value(item, path=f"{path}[{index}]") for index, item in enumerate(value)] + encoded.sort(key=lambda item: json.dumps(item, sort_keys=True)) + marker = _FROZENSET_MARKER if isinstance(value, frozenset) else "__set__" + return {marker: encoded} + if isinstance(value, dict): + encoded_dict: Dict[str, Any] = {} + for key, item in value.items(): + if not isinstance(key, str): + raise GraphSerializationError( + f"{path}: mapping key {key!r} is not a string and cannot round-trip through JSON" + ) + encoded_dict[key] = _encode_value( + item, + path=f"{path}.{key}", + field_name=key, + ) + return {_MAPPING_MARKER: encoded_dict} + raise GraphSerializationError( + f"{path}: unsupported value of type {type(value).__module__}.{type(value).__qualname__}" + ) + + +def _decode_value( + value: Any, + *, + path: str, + format_version: int, + field_name: Optional[str] = None, +) -> Any: + """Recursively restore graph state encoded by v2 or accepted v1 markers.""" + if value is None or isinstance(value, (bool, int, float, str)): + return value + if isinstance(value, list): + return [ + _decode_value(item, path=f"{path}[{index}]", format_version=format_version) + for index, item in enumerate(value) + ] + if not isinstance(value, dict): + if format_version == 1: + return value + raise GraphSerializationError(f"{path}: expected JSON-native graph state") + if format_version >= 2 and _SECRET_ENV_MARKER in value: + if not _is_api_key_field(field_name): + raise GraphSerializationError(f"{path}: secret environment marker is outside an API-key field") + scope = value[_SECRET_ENV_MARKER] + if scope == "model": + return _OMIT_FIELD + if scope == "operator": + return resolve_remote_api_key() + raise GraphSerializationError(f"{path}: invalid API-key environment marker") + if format_version >= 2 and _SECRET_NO_AUTH_MARKER in value: + if not _is_api_key_field(field_name) or value[_SECRET_NO_AUTH_MARKER] != "": + raise GraphSerializationError(f"{path}: invalid no-auth API-key marker") + return "" + if format_version >= 2 and _MAPPING_MARKER in value: + items = value[_MAPPING_MARKER] + if not isinstance(items, dict): + raise GraphSerializationError(f"{path}: malformed mapping envelope") + decoded_mapping: Dict[str, Any] = {} + for key, item in items.items(): + decoded = _decode_value( + item, + path=f"{path}.{key}", + format_version=format_version, + field_name=key, + ) + if decoded is not _OMIT_FIELD: + decoded_mapping[key] = decoded + return decoded_mapping + if format_version >= 2 and _PYDANTIC_MODEL_MARKER in value: + qualified = value[_PYDANTIC_MODEL_MARKER] + fields = value.get(_PYDANTIC_FIELDS) + fields_set = value.get(_PYDANTIC_FIELDS_SET, []) + if not isinstance(qualified, str) or not isinstance(fields, dict): + raise GraphSerializationError(f"{path}: malformed Pydantic model envelope") + try: + model_cls = _import_qualified_object(qualified) + except ImportError as exc: + raise GraphSerializationError(f"{path}: Pydantic model {qualified!r} is not importable") from exc + if not isinstance(model_cls, type) or not issubclass(model_cls, BaseModel): + raise GraphSerializationError(f"{path}: {qualified!r} is not a Pydantic model type") + decoded_fields: Dict[str, Any] = {} + for name, item in fields.items(): + decoded = _decode_value( + item, + path=f"{path}.{name}", + format_version=format_version, + field_name=name, + ) + if decoded is not _OMIT_FIELD: + decoded_fields[name] = decoded + try: + model = model_cls.model_validate(decoded_fields) + except Exception as exc: + raise GraphSerializationError( + f"{path}: failed to validate restored Pydantic model {qualified!r}: {exc}" + ) from exc + if not isinstance(fields_set, list) or not all(isinstance(name, str) for name in fields_set): + raise GraphSerializationError(f"{path}: malformed Pydantic fields_set") + unknown = set(fields_set) - set(type(model).model_fields) + if unknown: + raise GraphSerializationError(f"{path}: Pydantic fields_set contains unknown fields: {sorted(unknown)}") + model.__pydantic_fields_set__ = set(fields_set) + return model + if "__type_ref__" in value: + qualified = value["__type_ref__"] + try: + restored = _import_qualified_object(qualified) + except (ImportError, TypeError) as exc: + if format_version == 1: + return value + raise GraphSerializationError(f"{path}: type reference {qualified!r} is not importable") from exc + if not isinstance(restored, type): + if format_version == 1: + return value + raise GraphSerializationError(f"{path}: type reference {qualified!r} is not a type") + return restored + if "__callable_ref__" in value: + qualified = value["__callable_ref__"] + try: + restored = _import_qualified_object(qualified) + except (ImportError, TypeError) as exc: + if format_version == 1: + return value + raise GraphSerializationError(f"{path}: callable reference {qualified!r} is not importable") from exc + if not callable(restored): + if format_version == 1: + return value + raise GraphSerializationError(f"{path}: callable reference {qualified!r} is not callable") + return restored + if "__path__" in value: + return Path(value["__path__"]) + if "__set__" in value: + return { + _decode_value(item, path=f"{path}[{index}]", format_version=format_version) + for index, item in enumerate(value["__set__"]) + } + if format_version >= 2 and _FROZENSET_MARKER in value: + return frozenset( + _decode_value(item, path=f"{path}[{index}]", format_version=format_version) + for index, item in enumerate(value[_FROZENSET_MARKER]) + ) + if format_version >= 2 and _TUPLE_MARKER in value: + return tuple( + _decode_value(item, path=f"{path}[{index}]", format_version=format_version) + for index, item in enumerate(value[_TUPLE_MARKER]) + ) + return { + key: _decode_value( + item, + path=f"{path}.{key}", + format_version=format_version, + field_name=key, + ) + for key, item in value.items() + } + + class _RegistryJSONEncoder(json.JSONEncoder): - """JSON encoder that handles common non-serializable types found in operator kwargs.""" + """Compatibility encoder delegating non-native values to the v2 codec.""" def default(self, obj: Any) -> Any: - if isinstance(obj, type): - return {"__type_ref__": _qualified_name(obj)} - if callable(obj) and hasattr(obj, "__qualname__"): - module = getattr(obj, "__module__", None) or "" - return {"__callable_ref__": f"{module}.{obj.__qualname__}"} - if isinstance(obj, Path): - return {"__path__": str(obj)} - if isinstance(obj, (set, frozenset)): - return {"__set__": sorted(obj, key=str)} - if isinstance(obj, bytes): - return {"__bytes_len__": len(obj), "__repr__": repr(obj[:64])} - if hasattr(obj, "__dict__"): - safe_attrs = {} - for k, v in obj.__dict__.items(): - if not k.startswith("_"): - try: - json.dumps(v, cls=_RegistryJSONEncoder) - safe_attrs[k] = v - except (TypeError, ValueError): - safe_attrs[k] = repr(v) - return { - "__object__": _qualified_name(type(obj)), - "__attrs__": safe_attrs, - } - return super().default(obj) + return _encode_value(obj, path="$json") def _safe_serialize_value(value: Any) -> Any: - """Best-effort conversion of *value* into something JSON-safe.""" - try: - json.dumps(value, cls=_RegistryJSONEncoder) - return value - except (TypeError, ValueError, OverflowError): - return repr(value) + """Encode a value without lossy repr fallback (compatibility helper).""" + return _encode_value(value, path="$value") # --------------------------------------------------------------------------- @@ -191,6 +528,48 @@ def list_all_kwargs(graph: Graph) -> Dict[str, Dict[str, Any]]: # --------------------------------------------------------------------------- +def _redact_display_value( + value: Any, + *, + field_name: Optional[str] = None, + seen: Optional[Set[int]] = None, +) -> Any: + """Return a recursively redacted copy suitable only for diagnostics.""" + if _is_obvious_secret_field(field_name) and not _is_empty_secret(value): + return "***" + if value is None or isinstance(value, (bool, int, float, str, bytes)): + return value + + if seen is None: + seen = set() + value_id = id(value) + if value_id in seen: + return "" + seen.add(value_id) + + if isinstance(value, BaseModel): + return { + name: _redact_display_value( + getattr(value, name), + field_name=name, + seen=seen, + ) + for name in type(value).model_fields + } + if isinstance(value, dict): + return {key: _redact_display_value(item, field_name=str(key), seen=seen) for key, item in value.items()} + if isinstance(value, (list, tuple)): + redacted = [_redact_display_value(item, seen=seen) for item in value] + return tuple(redacted) if isinstance(value, tuple) else redacted + if isinstance(value, (set, frozenset)): + return [_redact_display_value(item, seen=seen) for item in value] + return value + + +def _display_repr(field_name: str, value: Any) -> str: + return repr(_redact_display_value(value, field_name=field_name)) + + def format_graph_tree( graph: Graph, *, @@ -243,7 +622,7 @@ def _render(node: Node, prefix: str, is_last: bool, is_root: bool) -> None: if show_kwargs and node.operator_kwargs: kw_prefix = prefix + ("" if is_root else (" " if is_last else "│ ")) for key, val in sorted(node.operator_kwargs.items()): - val_repr = repr(val) if not isinstance(val, str) else f"'{val}'" + val_repr = _display_repr(key, val) if len(val_repr) > max_value_width: val_repr = val_repr[: max_value_width - 3] + "..." lines.append(f"{kw_prefix} ╰ {key} = {val_repr}") @@ -269,7 +648,7 @@ def format_node_details(node: Node) -> str: f" Kwargs ({len(node.operator_kwargs)}):", ] for key, val in sorted(node.operator_kwargs.items()): - val_repr = repr(val) + val_repr = _display_repr(key, val) if len(val_repr) > 200: val_repr = val_repr[:197] + "..." lines.append(f" {key:30s} = {val_repr}") @@ -475,15 +854,15 @@ def format(self) -> str: if nd.kwargs_added: lines.append(" + Added kwargs:") for k, v in sorted(nd.kwargs_added.items()): - lines.append(f" {k} = {repr(v)}") + lines.append(f" {k} = {_display_repr(k, v)}") if nd.kwargs_removed: lines.append(" - Removed kwargs:") for k, v in sorted(nd.kwargs_removed.items()): - lines.append(f" {k} = {repr(v)}") + lines.append(f" {k} = {_display_repr(k, v)}") if nd.kwargs_changed: lines.append(" ~ Changed kwargs:") for k, (old, new) in sorted(nd.kwargs_changed.items()): - lines.append(f" {k}: {repr(old)} -> {repr(new)}") + lines.append(f" {k}: {_display_repr(k, old)} -> {_display_repr(k, new)}") if nd.children_a_only: lines.append(f" Children only in A: {nd.children_a_only}") if nd.children_b_only: @@ -610,28 +989,33 @@ def print_diff(graph_a: Graph, graph_b: Graph) -> None: # --------------------------------------------------------------------------- -def _serialize_node(node: Node) -> dict: +def _serialize_node(node: Node, *, path: str) -> dict: """Serialize a single node to a JSON-compatible dict.""" - safe_kwargs = {} - for k, v in node.operator_kwargs.items(): - safe_kwargs[k] = _safe_serialize_value(v) + safe_kwargs: Dict[str, Any] = {} + for key, value in node.operator_kwargs.items(): + if not isinstance(key, str): + raise GraphSerializationError(f"{path}.operator_kwargs: kwarg names must be strings") + safe_kwargs[key] = _encode_value( + value, + path=f"{path}.operator_kwargs.{key}", + field_name=key, + allow_api_key_env=True, + ) return { "name": node.name, "operator_class": _qualified_name(node.operator_class), "operator_kwargs": safe_kwargs, - "children": [_serialize_node(child) for child in node.children], + "children": [ + _serialize_node(child, path=f"{path}.children[{index}]") for index, child in enumerate(node.children) + ], } def serialize_graph(graph: Graph) -> dict: - """Serialize a graph to a JSON-compatible dictionary. - - The result can be passed to :func:`json.dumps` (with the - :class:`_RegistryJSONEncoder`) and later restored via - :func:`deserialize_graph`. - """ + """Serialize a graph to a versioned, recursively JSON-native dictionary.""" return { - "roots": [_serialize_node(root) for root in graph.roots], + "format_version": _GRAPH_FORMAT_VERSION, + "roots": [_serialize_node(root, path=f"roots[{index}]({root.name})") for index, root in enumerate(graph.roots)], "metadata": { "node_count": node_count(graph), "max_depth": max_depth(graph), @@ -661,56 +1045,79 @@ def postprocess(self, data: Any, **kwargs: Any) -> Any: return data -def _restore_special_values(kwargs: dict) -> dict: - """Walk a kwargs dict and restore ``__type_ref__``, ``__path__``, etc.""" +def _restore_special_values( + kwargs: dict, + *, + format_version: int = 1, + path: str = "operator_kwargs", +) -> dict: + """Recursively restore encoded operator kwargs, including v1 markers.""" cleaned: Dict[str, Any] = {} - for k, v in kwargs.items(): - if isinstance(v, dict): - if "__type_ref__" in v: - try: - cleaned[k] = _import_class(v["__type_ref__"]) - except ImportError: - cleaned[k] = v - continue - if "__callable_ref__" in v: - try: - cleaned[k] = _import_class(v["__callable_ref__"]) - except ImportError: - cleaned[k] = v - continue - if "__path__" in v: - cleaned[k] = Path(v["__path__"]) - continue - if "__set__" in v: - cleaned[k] = set(v["__set__"]) - continue - cleaned[k] = v + for key, value in kwargs.items(): + decoded = _decode_value( + value, + path=f"{path}.{key}", + format_version=format_version, + field_name=key, + ) + if decoded is not _OMIT_FIELD: + cleaned[key] = decoded return cleaned -def _deserialize_node(data: dict) -> Node: +def _deserialize_node(data: dict, *, format_version: int, path: str) -> Node: """Reconstruct a :class:`Node` from its serialized dict.""" cls = _import_class(data["operator_class"]) raw_kwargs = data.get("operator_kwargs", {}) - cleaned = _restore_special_values(raw_kwargs) + if not isinstance(raw_kwargs, dict): + raise GraphSerializationError(f"{path}.operator_kwargs: expected a mapping") + cleaned = _restore_special_values( + raw_kwargs, + format_version=format_version, + path=f"{path}.operator_kwargs", + ) try: op = cls(**cleaned) - except Exception: + except Exception as exc: + if format_version >= 2: + raise GraphSerializationError( + f"{path}: failed to construct operator " f"{data['operator_class']!r}: {exc}" + ) from exc op = _PlaceholderOperator(original_class=data["operator_class"], original_kwargs=cleaned) node = Node(op, name=data.get("name"), operator_class=cls, operator_kwargs=cleaned) - for child_data in data.get("children", []): - child_node = _deserialize_node(child_data) + for index, child_data in enumerate(data.get("children", [])): + child_node = _deserialize_node( + child_data, + format_version=format_version, + path=f"{path}.children[{index}]", + ) node.children.append(child_node) return node +def _read_format_version(data: dict) -> int: + version = data.get("format_version", 1) + if isinstance(version, bool) or not isinstance(version, int): + raise GraphSerializationError("format_version must be an integer") + if version not in (1, _GRAPH_FORMAT_VERSION): + raise GraphSerializationError(f"unsupported graph format_version: {version}") + return version + + def deserialize_graph(data: dict) -> Graph: - """Reconstruct a :class:`Graph` from a dict produced by :func:`serialize_graph`.""" + """Reconstruct a graph from v2 data or a versionless v1 payload.""" + if not isinstance(data, dict): + raise GraphSerializationError("serialized graph must be a mapping") + format_version = _read_format_version(data) graph = Graph() - for root_data in data.get("roots", []): - root_node = _deserialize_node(root_data) + for index, root_data in enumerate(data.get("roots", [])): + root_node = _deserialize_node( + root_data, + format_version=format_version, + path=f"roots[{index}]", + ) graph.roots.append(root_node) return graph @@ -722,7 +1129,7 @@ def save_graph(graph: Graph, path: Union[str, Path], *, indent: int = 2) -> Path """ path = Path(path) payload = serialize_graph(graph) - path.write_text(json.dumps(payload, cls=_RegistryJSONEncoder, indent=indent, default=repr)) + path.write_text(json.dumps(payload, indent=indent)) return path @@ -977,23 +1384,25 @@ def build_with_overrides(self, name: str, overrides: Dict[str, Dict[str, Any]]) def save_all(self, path: Union[str, Path], *, indent: int = 2) -> Path: """Serialize every registered graph to a single JSON file. - The file contains ``{name: {roots, metadata, blueprint}}`` for each - registered graph. Returns the resolved path. + Version 2 stores graphs under a versioned ``graphs`` mapping. Returns + the resolved path. """ path = Path(path) - payload: Dict[str, Any] = {} + graphs_payload: Dict[str, Any] = {} for name, bp in self._blueprints.items(): graph = bp.build() entry = serialize_graph(graph) entry["blueprint"] = { + "name": bp.name, "description": bp.description, "version": bp.version, "tags": bp.tags, "created_at": bp.created_at, "updated_at": bp.updated_at, } - payload[name] = entry - path.write_text(json.dumps(payload, cls=_RegistryJSONEncoder, indent=indent, default=repr)) + graphs_payload[name] = entry + payload = {"format_version": _GRAPH_FORMAT_VERSION, "graphs": graphs_payload} + path.write_text(json.dumps(payload, indent=indent)) return path def load_all(self, path: Union[str, Path], *, overwrite: bool = False) -> List[str]: @@ -1004,8 +1413,20 @@ def load_all(self, path: Union[str, Path], *, overwrite: bool = False) -> List[s """ path = Path(path) payload = json.loads(path.read_text()) + if not isinstance(payload, dict): + raise GraphSerializationError("serialized graph registry must be a mapping") + version_marker = payload.get("format_version") + if isinstance(version_marker, int) and not isinstance(version_marker, bool): + _read_format_version(payload) + entries = payload.get("graphs") + if not isinstance(entries, dict): + raise GraphSerializationError("version 2 graph registry requires a graphs mapping") + else: + # Versionless v1 registries were a direct name -> graph mapping, + # and graph names were unrestricted (including "format_version"). + entries = payload loaded: List[str] = [] - for name, entry in payload.items(): + for name, entry in entries.items(): bp_meta = entry.get("blueprint", {}) graph_data = {k: v for k, v in entry.items() if k != "blueprint"} @@ -1020,6 +1441,11 @@ def _factory(_gd: dict = graph_data) -> Graph: tags=bp_meta.get("tags", []), overwrite=overwrite, ) + restored_bp = self.get_blueprint(name) + if isinstance(bp_meta.get("created_at"), str): + restored_bp.created_at = bp_meta["created_at"] + if isinstance(bp_meta.get("updated_at"), str): + restored_bp.updated_at = bp_meta["updated_at"] loaded.append(name) return loaded @@ -1029,6 +1455,7 @@ def save_graph(self, name: str, path: Union[str, Path], *, indent: int = 2) -> P bp = self.get_blueprint(name) payload = serialize_graph(graph) payload["blueprint"] = { + "name": bp.name, "description": bp.description, "version": bp.version, "tags": bp.tags, @@ -1036,7 +1463,7 @@ def save_graph(self, name: str, path: Union[str, Path], *, indent: int = 2) -> P "updated_at": bp.updated_at, } path = Path(path) - path.write_text(json.dumps(payload, cls=_RegistryJSONEncoder, indent=indent, default=repr)) + path.write_text(json.dumps(payload, indent=indent)) return path def load_graph(self, path: Union[str, Path], *, name: Optional[str] = None, overwrite: bool = False) -> str: @@ -1047,7 +1474,12 @@ def load_graph(self, path: Union[str, Path], *, name: Optional[str] = None, over """ path = Path(path) payload = json.loads(path.read_text()) + if not isinstance(payload, dict): + raise GraphSerializationError("serialized graph must be a mapping") + _read_format_version(payload) bp_meta = payload.get("blueprint", {}) + if not isinstance(bp_meta, dict): + raise GraphSerializationError("blueprint metadata must be a mapping") graph_data = {k: v for k, v in payload.items() if k != "blueprint"} resolved_name = name or bp_meta.get("name") or path.stem @@ -1062,6 +1494,11 @@ def _factory(_gd: dict = graph_data) -> Graph: tags=bp_meta.get("tags", []), overwrite=overwrite, ) + restored_bp = self.get_blueprint(resolved_name) + if isinstance(bp_meta.get("created_at"), str): + restored_bp.created_at = bp_meta["created_at"] + if isinstance(bp_meta.get("updated_at"), str): + restored_bp.updated_at = bp_meta["updated_at"] return resolved_name diff --git a/nemo_retriever/src/nemo_retriever/models/llm/__init__.py b/nemo_retriever/src/nemo_retriever/models/llm/__init__.py index 837771e83..7ad463aae 100644 --- a/nemo_retriever/src/nemo_retriever/models/llm/__init__.py +++ b/nemo_retriever/src/nemo_retriever/models/llm/__init__.py @@ -27,9 +27,9 @@ Public surface contract ----------------------- The names in ``__all__`` below are the frozen public API of this -module. External callers should import from ``nemo_retriever.llm`` -rather than reaching into submodules (``llm.clients.litellm``, -``llm.text_utils``) directly -- those submodule paths are implementation +module. External callers should import from ``nemo_retriever.models.llm`` +rather than reaching into submodules (``models.llm.clients.litellm``, +``models.llm.text_utils``) directly -- those submodule paths are implementation details and may be reorganised in future releases without notice. The Protocols + result dataclasses + concrete clients + re-exported params models listed here are the supported integration points. @@ -39,17 +39,31 @@ AnswerJudge, AnswerRequest, AnswerResult, + CompletionClient, + GeneratedTextResult, + GenerationRequest, GenerationResult, JudgeResult, LLMClient, RetrievalResult, RetrieverStrategy, ) -from nemo_retriever.common.params.models import LLMInferenceParams, LLMRemoteClientParams +from nemo_retriever.common.params.models import ( + LLMInferenceParams, + LLMRemoteClientParams, + LLMSamplingOverrides, + TextGenerationParams, +) +from nemo_retriever.models.llm.tasks import ( + GenerationTask, + GenericPromptTask, + RagAnswerTask, + SummarizeTask, +) _LAZY_IMPORTS = { - "LiteLLMClient": "nemo_retriever.llm.clients.litellm", - "LLMJudge": "nemo_retriever.llm.clients.judge", + "LiteLLMClient": "nemo_retriever.models.llm.clients.litellm", + "LLMJudge": "nemo_retriever.models.llm.clients.judge", } @@ -65,18 +79,28 @@ def __getattr__(name: str): __all__ = [ # Protocols "AnswerJudge", + "CompletionClient", + "GenerationTask", "LLMClient", "RetrieverStrategy", # Request/result models "AnswerRequest", "AnswerResult", + "GeneratedTextResult", + "GenerationRequest", "GenerationResult", "JudgeResult", "RetrievalResult", + # Tasks + "GenericPromptTask", + "RagAnswerTask", + "SummarizeTask", # Concrete clients (lazy-loaded) "LLMJudge", "LiteLLMClient", # Transport / sampling params (re-exported for ergonomics) "LLMInferenceParams", "LLMRemoteClientParams", + "LLMSamplingOverrides", + "TextGenerationParams", ] diff --git a/nemo_retriever/src/nemo_retriever/models/llm/clients/__init__.py b/nemo_retriever/src/nemo_retriever/models/llm/clients/__init__.py index 572a239bc..c6905145b 100644 --- a/nemo_retriever/src/nemo_retriever/models/llm/clients/__init__.py +++ b/nemo_retriever/src/nemo_retriever/models/llm/clients/__init__.py @@ -5,14 +5,14 @@ """ Concrete LLM client implementations and a lightweight client registry. -The ``llm.clients`` package hosts every concrete :class:`LLMClient` +The ``models.llm.clients`` package hosts every concrete :class:`LLMClient` implementation in its own submodule (``litellm.py``, ``judge.py``, ...) so that adding a new client means adding a new file rather than extending a monolithic module. To keep the public import path stable, the registered client classes plus the internal prompt-helper (``_build_rag_prompt``) are re-exported from this package's namespace. Any caller that imports -``from nemo_retriever.llm.clients import LiteLLMClient`` will therefore +``from nemo_retriever.models.llm.clients import LiteLLMClient`` will therefore continue to work unchanged after the module-to-package refactor. """ diff --git a/nemo_retriever/src/nemo_retriever/models/llm/clients/judge.py b/nemo_retriever/src/nemo_retriever/models/llm/clients/judge.py index a8f63f976..de7e76b34 100644 --- a/nemo_retriever/src/nemo_retriever/models/llm/clients/judge.py +++ b/nemo_retriever/src/nemo_retriever/models/llm/clients/judge.py @@ -7,7 +7,7 @@ ``LLMJudge`` reproduces the dual-judge logic of ragas' :class:`ragas.metrics.collections.AnswerAccuracy` directly on -:class:`~nemo_retriever.llm.clients.litellm.LiteLLMClient`, with no ragas / +:class:`~nemo_retriever.models.llm.clients.litellm.LiteLLMClient`, with no ragas / instructor / openai dependency: * Two paraphrased judge prompts (verbatim from ragas), each rating the answer @@ -224,10 +224,10 @@ class LLMJudge: Configuration is split into two Pydantic objects: - * ``transport``: :class:`~nemo_retriever.params.LLMRemoteClientParams` owns + * ``transport``: :class:`~nemo_retriever.common.params.LLMRemoteClientParams` owns the endpoint, api_key, retries, and timeout. ``num_retries`` is reused as the per-judge attempt budget for obtaining a valid ``0/2/4`` rating. - * ``sampling``: :class:`~nemo_retriever.params.LLMInferenceParams` owns + * ``sampling``: :class:`~nemo_retriever.common.params.LLMInferenceParams` owns ``temperature`` / ``top_p`` / ``max_tokens``. Defaults to ``temperature=0.1, max_tokens=4096`` for judge consistency. diff --git a/nemo_retriever/src/nemo_retriever/models/llm/clients/litellm.py b/nemo_retriever/src/nemo_retriever/models/llm/clients/litellm.py index d5eb9ac87..42a10b18d 100644 --- a/nemo_retriever/src/nemo_retriever/models/llm/clients/litellm.py +++ b/nemo_retriever/src/nemo_retriever/models/llm/clients/litellm.py @@ -12,88 +12,22 @@ from __future__ import annotations -from copy import deepcopy import logging import time from typing import Any, Optional -from nemo_retriever.models.llm.text_utils import strip_think_tags +from nemo_retriever.models.llm.tasks.rag_answer import ( + RagAnswerTask, + _build_rag_prompt as _task_build_rag_prompt, + _deep_merge_dicts, + _format_rag_system_prompt, +) from nemo_retriever.models.llm.types import GenerationResult from nemo_retriever.common.params.models import LLMInferenceParams, LLMRemoteClientParams logger = logging.getLogger(__name__) - -_RAG_SYSTEM_PROMPT = ( - "You are a precise question-answering assistant. " - "Answer the question using ONLY the information provided in the context below. " - "If the context does not contain enough information to answer, say so clearly. " - "Be concise and factual." -) - -_RAG_USER_TEMPLATE = """\ -Context: -{context} - -Question: {query} - -Answer:""" - -_NO_REASONING_SYSTEM_DIRECTIVE = "/no_think" -_NO_REASONING_EXTRA_PARAMS = {"chat_template_kwargs": {"enable_thinking": False}} - - -def _format_rag_system_prompt( - *, - rag_system_prompt: Optional[str] = None, - rag_system_prompt_prefix: Optional[str] = None, -) -> str: - """Resolve the system prompt used for RAG answer generation.""" - prompt = (rag_system_prompt if rag_system_prompt is not None else _RAG_SYSTEM_PROMPT).strip() - prefix = (rag_system_prompt_prefix or "").strip() - if not prefix: - return prompt - if not prompt: - return prefix - return f"{prefix}\n{prompt}" - - -def _build_rag_prompt( - query: str, - chunks: list[str], - *, - formatted_rag_system_prompt: str, -) -> list[dict]: - """Build the OpenAI-style messages list for a RAG prompt.""" - context = "\n\n---\n\n".join(chunks) if chunks else "(no context retrieved)" - user_content = _RAG_USER_TEMPLATE.format(context=context, query=query) - return [ - {"role": "system", "content": formatted_rag_system_prompt}, - {"role": "user", "content": user_content}, - ] - - -def _deep_merge_dicts(left: dict[str, Any], right: dict[str, Any]) -> dict[str, Any]: - """Return a recursive merge where ``right`` wins without mutating inputs.""" - merged = deepcopy(left) - for key, value in right.items(): - if isinstance(merged.get(key), dict) and isinstance(value, dict): - merged[key] = _deep_merge_dicts(merged[key], value) - else: - merged[key] = deepcopy(value) - return merged - - -def _with_no_reasoning_controls(messages: list[dict]) -> list[dict]: - """Add no-reasoning prompt metadata understood by current Nemotron LLM NIMs.""" - updated = [dict(message) for message in messages] - if updated and updated[0].get("role") == "system": - content = str(updated[0].get("content") or "").strip() - if _NO_REASONING_SYSTEM_DIRECTIVE not in content: - content = f"{_NO_REASONING_SYSTEM_DIRECTIVE}\n{content}" if content else _NO_REASONING_SYSTEM_DIRECTIVE - updated[0]["content"] = content - return updated - updated.insert(0, {"role": "system", "content": _NO_REASONING_SYSTEM_DIRECTIVE}) - return updated +# Backwards-compatible helper export retained for existing callers. +_build_rag_prompt = _task_build_rag_prompt class LiteLLMClient: @@ -110,9 +44,9 @@ class LiteLLMClient: Configuration is split into two orthogonal Pydantic objects: - * ``transport``: :class:`~nemo_retriever.params.LLMRemoteClientParams` + * ``transport``: :class:`~nemo_retriever.common.params.LLMRemoteClientParams` owns provider endpoint, authentication, retry, and timeout. - * ``sampling``: :class:`~nemo_retriever.params.LLMInferenceParams` + * ``sampling``: :class:`~nemo_retriever.common.params.LLMInferenceParams` owns ``temperature``, ``top_p``, and ``max_tokens``. Use :meth:`from_kwargs` for a flat, backwards-compatible constructor. @@ -149,7 +83,7 @@ def from_kwargs( model: str = _DEFAULT_MODEL, api_base: Optional[str] = None, api_key: Optional[str] = None, - temperature: float = 0.0, + temperature: Optional[float] = 0.0, top_p: Optional[float] = None, max_tokens: int = 4096, extra_params: Optional[dict[str, Any]] = None, @@ -238,34 +172,17 @@ def generate( reasoning_enabled: Optional[bool] = None, ) -> GenerationResult: """Generate an answer for the given query using retrieved chunks as context.""" - messages = _build_rag_prompt( - query, - chunks, - formatted_rag_system_prompt=self._formatted_rag_system_prompt, - ) - request_extra_params: dict[str, Any] | None = None effective_reasoning_enabled = ( self.transport.reasoning_enabled if reasoning_enabled is None else reasoning_enabled ) - if not effective_reasoning_enabled: - messages = _with_no_reasoning_controls(messages) - request_extra_params = _NO_REASONING_EXTRA_PARAMS - try: - raw_answer, latency = self.complete(messages, extra_params=request_extra_params) - answer = strip_think_tags(raw_answer) - if not answer: - return GenerationResult( - answer="", - latency_s=latency, - model=self.transport.model, - error="thinking_truncated", - ) - return GenerationResult(answer=answer, latency_s=latency, model=self.transport.model) - except Exception as exc: - logger.debug("Generation failed for model=%s: %s", self.transport.model, exc) - return GenerationResult( - answer="", - latency_s=0.0, - model=self.transport.model, - error=str(exc), - ) + task = RagAnswerTask( + system_prompt=self._formatted_rag_system_prompt, + reasoning_enabled=effective_reasoning_enabled, + ) + result = task.execute(self, query=query, chunks=chunks) + return GenerationResult( + answer=result.text, + latency_s=result.latency_s, + model=result.model, + error=result.error, + ) diff --git a/nemo_retriever/src/nemo_retriever/models/llm/tasks/__init__.py b/nemo_retriever/src/nemo_retriever/models/llm/tasks/__init__.py new file mode 100644 index 000000000..775ec4ef9 --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/models/llm/tasks/__init__.py @@ -0,0 +1,17 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-26, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Reusable text-generation task strategies.""" + +from nemo_retriever.models.llm.tasks.base import GenerationTask +from nemo_retriever.models.llm.tasks.generic import GenericPromptTask +from nemo_retriever.models.llm.tasks.rag_answer import RagAnswerTask +from nemo_retriever.models.llm.tasks.summarize import SummarizeTask + +__all__ = [ + "GenerationTask", + "GenericPromptTask", + "RagAnswerTask", + "SummarizeTask", +] diff --git a/nemo_retriever/src/nemo_retriever/models/llm/tasks/base.py b/nemo_retriever/src/nemo_retriever/models/llm/tasks/base.py new file mode 100644 index 000000000..4d7782d83 --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/models/llm/tasks/base.py @@ -0,0 +1,80 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-26, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Reusable request/response lifecycle for text-generation tasks.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Any, ClassVar, Optional + +from nemo_retriever.common.params.models import LLMInferenceParams +from nemo_retriever.models.llm.types import ( + CompletionClient, + GeneratedTextResult, + GenerationRequest, +) + + +class GenerationTask(ABC): + """Stateless strategy that turns logical inputs into one completion call.""" + + required_inputs: tuple[str, ...] = () + _default_sampling: ClassVar[dict[str, Any]] = { + "temperature": 1.0, + "top_p": None, + "max_tokens": 1024, + } + empty_output_error: ClassVar[str] = "empty_output" + + @property + def default_sampling(self) -> LLMInferenceParams: + """Return a fresh copy of this task's sampling defaults.""" + return LLMInferenceParams(**self._default_sampling) + + @abstractmethod + def build_request(self, **inputs: object) -> GenerationRequest: + """Build one provider-neutral request from logical task inputs.""" + + def parse(self, raw_text: str) -> str: + """Parse completion text into the task's text result.""" + return raw_text.strip() + + def _preflight_error(self, **inputs: object) -> Optional[str]: + """Return an error code when no provider request should be made.""" + return None + + def execute(self, client: CompletionClient, **inputs: object) -> GeneratedTextResult: + """Build, execute, and parse one request without leaking row failures.""" + latency_s = 0.0 + try: + preflight_error = self._preflight_error(**inputs) + if preflight_error is not None: + return GeneratedTextResult( + text="", + latency_s=0.0, + model=client.model, + error=preflight_error, + ) + + request = self.build_request(**inputs) + raw_text, latency_s = client.complete( + request.messages, + max_tokens=request.max_tokens, + extra_params=request.extra_params, + ) + text = self.parse(raw_text) + return GeneratedTextResult( + text=text, + latency_s=latency_s, + model=client.model, + error=None if text else self.empty_output_error, + ) + except Exception as exc: + return GeneratedTextResult( + text="", + latency_s=latency_s, + model=client.model, + error=str(exc), + ) diff --git a/nemo_retriever/src/nemo_retriever/models/llm/tasks/generic.py b/nemo_retriever/src/nemo_retriever/models/llm/tasks/generic.py new file mode 100644 index 000000000..5bc27234b --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/models/llm/tasks/generic.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-26, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Validated prompt-template task for general text generation.""" + +from __future__ import annotations + +from string import Formatter +from typing import Any, ClassVar, Optional, Sequence + +from nemo_retriever.models.llm.tasks.base import GenerationTask +from nemo_retriever.models.llm.tasks.rag_answer import _apply_reasoning_control +from nemo_retriever.models.llm.text_utils import strip_think_tags +from nemo_retriever.models.llm.types import GenerationRequest + + +def _validate_prompt_template(prompt: str, input_names: tuple[str, ...]) -> None: + """Require a one-to-one match between simple fields and declared inputs.""" + if not isinstance(prompt, str): + raise TypeError("prompt must be a string") + if not input_names: + raise ValueError("input_names must declare at least one input") + if len(set(input_names)) != len(input_names): + raise ValueError("input_names must not contain duplicates") + if any(not name.isidentifier() for name in input_names): + raise ValueError("input_names must contain valid Python identifiers") + + fields: list[str] = [] + try: + for _, field_name, format_spec, conversion in Formatter().parse(prompt): + if field_name is None: + continue + if ( + not field_name + or not field_name.isidentifier() + or "." in field_name + or "[" in field_name + or format_spec + or conversion + ): + raise ValueError("prompt placeholders must be simple names") + fields.append(field_name) + except ValueError as exc: + if str(exc) == "prompt placeholders must be simple names": + raise + raise ValueError(f"invalid prompt template: {exc}") from exc + + if not fields: + raise ValueError("prompt must contain at least one declared placeholder") + declared = set(input_names) + referenced = set(fields) + missing = declared - referenced + undeclared = referenced - declared + if missing: + raise ValueError(f"prompt is missing declared placeholders: {sorted(missing)}") + if undeclared: + raise ValueError(f"prompt contains undeclared placeholders: {sorted(undeclared)}") + + +class GenericPromptTask(GenerationTask): + """Render declared row inputs into a validated prompt template.""" + + _default_sampling: ClassVar[dict[str, Any]] = { + "temperature": 1.0, + "top_p": None, + "max_tokens": 1024, + } + + def __init__( + self, + *, + prompt: str, + input_names: Sequence[str], + system_prompt: Optional[str] = None, + reasoning_enabled: Optional[bool] = None, + ) -> None: + if isinstance(input_names, str): + raise TypeError("input_names must be a sequence of names, not a string") + names = tuple(input_names) + _validate_prompt_template(prompt, names) + self.prompt = prompt + self.required_inputs = names + self.system_prompt = system_prompt + self.reasoning_enabled = reasoning_enabled + + def build_request(self, **inputs: object) -> GenerationRequest: + """Render declared inputs and build one completion request.""" + missing = [name for name in self.required_inputs if name not in inputs] + if missing: + raise KeyError(f"missing required inputs: {missing}") + values = {name: inputs[name] for name in self.required_inputs} + user_content = self.prompt.format(**values) + messages: list[dict[str, Any]] = [] + if self.system_prompt is not None: + messages.append({"role": "system", "content": self.system_prompt}) + messages.append({"role": "user", "content": user_content}) + messages, extra_params = _apply_reasoning_control(messages, self.reasoning_enabled) + return GenerationRequest(messages=messages, extra_params=extra_params) + + def parse(self, raw_text: str) -> str: + """Remove visible model reasoning from generated text.""" + return strip_think_tags(raw_text) + + +__all__ = ["GenericPromptTask"] diff --git a/nemo_retriever/src/nemo_retriever/models/llm/tasks/rag_answer.py b/nemo_retriever/src/nemo_retriever/models/llm/tasks/rag_answer.py new file mode 100644 index 000000000..ebf7f9355 --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/models/llm/tasks/rag_answer.py @@ -0,0 +1,149 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-26, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Retrieval-augmented answer-generation task and prompt helpers.""" + +from __future__ import annotations + +from copy import deepcopy +from dataclasses import dataclass +from typing import Any, ClassVar, Optional + +from nemo_retriever.models.llm.tasks.base import GenerationTask +from nemo_retriever.models.llm.text_utils import strip_think_tags +from nemo_retriever.models.llm.types import GenerationRequest + +_RAG_SYSTEM_PROMPT = ( + "You are a precise question-answering assistant. " + "Answer the question using ONLY the information provided in the context below. " + "If the context does not contain enough information to answer, say so clearly. " + "Be concise and factual." +) + +_RAG_USER_TEMPLATE = """\ +Context: +{context} + +Question: {query} + +Answer:""" + +_NO_REASONING_SYSTEM_DIRECTIVE = "/no_think" +_NO_REASONING_EXTRA_PARAMS = {"chat_template_kwargs": {"enable_thinking": False}} + + +def _format_rag_system_prompt( + *, + rag_system_prompt: Optional[str] = None, + rag_system_prompt_prefix: Optional[str] = None, +) -> str: + """Resolve the system prompt used for RAG answer generation.""" + prompt = (rag_system_prompt if rag_system_prompt is not None else _RAG_SYSTEM_PROMPT).strip() + prefix = (rag_system_prompt_prefix or "").strip() + if not prefix: + return prompt + if not prompt: + return prefix + return f"{prefix}\n{prompt}" + + +def _build_rag_prompt( + query: str, + chunks: list[str], + *, + formatted_rag_system_prompt: str, +) -> list[dict[str, Any]]: + """Build the OpenAI-style messages list for a RAG prompt.""" + context = "\n\n---\n\n".join(chunks) if chunks else "(no context retrieved)" + user_content = _RAG_USER_TEMPLATE.format(context=context, query=query) + return [ + {"role": "system", "content": formatted_rag_system_prompt}, + {"role": "user", "content": user_content}, + ] + + +def _deep_merge_dicts(left: dict[str, Any], right: dict[str, Any]) -> dict[str, Any]: + """Return a recursive merge where ``right`` wins without mutating inputs.""" + merged = deepcopy(left) + for key, value in right.items(): + if isinstance(merged.get(key), dict) and isinstance(value, dict): + merged[key] = _deep_merge_dicts(merged[key], value) + else: + merged[key] = deepcopy(value) + return merged + + +def _with_no_reasoning_controls(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: + """Add no-reasoning prompt metadata understood by current Nemotron NIMs.""" + updated = [dict(message) for message in messages] + if updated and updated[0].get("role") == "system": + content = str(updated[0].get("content") or "").strip() + if _NO_REASONING_SYSTEM_DIRECTIVE not in content: + content = f"{_NO_REASONING_SYSTEM_DIRECTIVE}\n{content}" if content else _NO_REASONING_SYSTEM_DIRECTIVE + updated[0]["content"] = content + return updated + updated.insert(0, {"role": "system", "content": _NO_REASONING_SYSTEM_DIRECTIVE}) + return updated + + +def _apply_reasoning_control( + messages: list[dict[str, Any]], + reasoning_enabled: Optional[bool], +) -> tuple[list[dict[str, Any]], Optional[dict[str, Any]]]: + """Apply task-level reasoning controls to messages and request extras.""" + if reasoning_enabled is not False: + return messages, None + return _with_no_reasoning_controls(messages), deepcopy(_NO_REASONING_EXTRA_PARAMS) + + +@dataclass(frozen=True) +class RagAnswerTask(GenerationTask): + """Generate a grounded answer from a query and retrieved text chunks.""" + + prompt: Optional[str] = None + system_prompt: Optional[str] = None + system_prompt_prefix: Optional[str] = None + reasoning_enabled: Optional[bool] = None + + required_inputs: ClassVar[tuple[str, ...]] = ("query", "chunks") + _default_sampling: ClassVar[dict[str, Any]] = { + "temperature": 0.0, + "top_p": None, + "max_tokens": 4096, + } + empty_output_error: ClassVar[str] = "thinking_truncated" + + def build_request(self, **inputs: object) -> GenerationRequest: + """Build a grounded answer request, including optional reasoning controls.""" + query = inputs["query"] + chunks = inputs["chunks"] + if not isinstance(query, str): + raise TypeError("query must be a string") + if not isinstance(chunks, list) or not all(isinstance(chunk, str) for chunk in chunks): + raise TypeError("chunks must be a list of strings") + + formatted_system_prompt = _format_rag_system_prompt( + rag_system_prompt=self.system_prompt, + rag_system_prompt_prefix=self.system_prompt_prefix, + ) + messages = _build_rag_prompt( + query, + chunks, + formatted_rag_system_prompt=formatted_system_prompt, + ) + if self.prompt is not None: + context = "\n\n---\n\n".join(chunks) if chunks else "(no context retrieved)" + messages[-1]["content"] = self.prompt.format(context=context, query=query) + + per_request_reasoning = inputs.get("reasoning_enabled") + effective_reasoning = self.reasoning_enabled if per_request_reasoning is None else bool(per_request_reasoning) + messages, extra_params = _apply_reasoning_control(messages, effective_reasoning) + return GenerationRequest(messages=messages, extra_params=extra_params) + + def parse(self, raw_text: str) -> str: + """Remove visible model reasoning from the answer.""" + return strip_think_tags(raw_text) + + +__all__ = ["RagAnswerTask"] diff --git a/nemo_retriever/src/nemo_retriever/models/llm/tasks/summarize.py b/nemo_retriever/src/nemo_retriever/models/llm/tasks/summarize.py new file mode 100644 index 000000000..beb54cd0d --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/models/llm/tasks/summarize.py @@ -0,0 +1,90 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-26, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Single-request text summarization task.""" + +from __future__ import annotations + +from dataclasses import dataclass +from string import Formatter +from typing import Any, ClassVar, Optional + +from nemo_retriever.models.llm.tasks.base import GenerationTask +from nemo_retriever.models.llm.tasks.rag_answer import _apply_reasoning_control +from nemo_retriever.models.llm.text_utils import strip_think_tags +from nemo_retriever.models.llm.types import GenerationRequest + +_SUMMARIZE_SYSTEM_PROMPT = ( + "You are a precise summarization assistant. Produce a faithful, concise summary " + "without adding information that is not present in the source text." +) +_SUMMARIZE_USER_TEMPLATE = "Summarize the following text:\n\n{text}" + + +def _summary_prompt_fields(prompt: str) -> set[str]: + """Validate and return fields used by a custom summarization prompt.""" + fields: set[str] = set() + try: + parsed = Formatter().parse(prompt) + for _, field_name, format_spec, conversion in parsed: + if field_name is None: + continue + if field_name != "text" or format_spec or conversion: + raise ValueError("summarization prompt may only use the simple {text} placeholder") + fields.add(field_name) + except ValueError as exc: + if str(exc) == "summarization prompt may only use the simple {text} placeholder": + raise + raise ValueError(f"invalid summarization prompt: {exc}") from exc + return fields + + +@dataclass(frozen=True) +class SummarizeTask(GenerationTask): + """Summarize one text value without truncation or hidden map-reduce.""" + + prompt: Optional[str] = None + system_prompt: Optional[str] = None + reasoning_enabled: Optional[bool] = None + + required_inputs: ClassVar[tuple[str, ...]] = ("text",) + _default_sampling: ClassVar[dict[str, Any]] = { + "temperature": 0.0, + "top_p": None, + "max_tokens": 1024, + } + + def __post_init__(self) -> None: + if self.prompt is not None: + _summary_prompt_fields(self.prompt) + + def _preflight_error(self, **inputs: object) -> Optional[str]: + text = inputs.get("text") + if isinstance(text, str) and not text.strip(): + return "empty_input" + return None + + def build_request(self, **inputs: object) -> GenerationRequest: + """Build one faithful-summary request for the supplied text.""" + text = inputs["text"] + if not isinstance(text, str): + raise TypeError("text must be a string") + + prompt = self.prompt if self.prompt is not None else _SUMMARIZE_USER_TEMPLATE + fields = _summary_prompt_fields(prompt) + user_content = prompt.format(text=text) if fields else f"{prompt}\n\n{text}" + system_content = self.system_prompt if self.system_prompt is not None else _SUMMARIZE_SYSTEM_PROMPT + messages = [ + {"role": "system", "content": system_content}, + {"role": "user", "content": user_content}, + ] + messages, extra_params = _apply_reasoning_control(messages, self.reasoning_enabled) + return GenerationRequest(messages=messages, extra_params=extra_params) + + def parse(self, raw_text: str) -> str: + """Remove visible model reasoning from the summary.""" + return strip_think_tags(raw_text) + + +__all__ = ["SummarizeTask"] diff --git a/nemo_retriever/src/nemo_retriever/models/llm/text_utils.py b/nemo_retriever/src/nemo_retriever/models/llm/text_utils.py index 34b8634a2..e022a4267 100644 --- a/nemo_retriever/src/nemo_retriever/models/llm/text_utils.py +++ b/nemo_retriever/src/nemo_retriever/models/llm/text_utils.py @@ -4,7 +4,7 @@ """Shared text-processing utilities for LLM output hygiene. -Pure-stdlib module. Lives under ``nemo_retriever.llm`` so that the +Pure-stdlib module. Lives under ``nemo_retriever.models.llm`` so that the lightweight SDK surface (``LiteLLMClient``, ``Retriever.answer``) does not pull in ``pandas`` or any evaluation dependencies just to clean ```` tags out of a model response. diff --git a/nemo_retriever/src/nemo_retriever/models/llm/types.py b/nemo_retriever/src/nemo_retriever/models/llm/types.py index 4fd56da1e..9093eeacf 100644 --- a/nemo_retriever/src/nemo_retriever/models/llm/types.py +++ b/nemo_retriever/src/nemo_retriever/models/llm/types.py @@ -39,6 +39,25 @@ def generate( ) -> "GenerationResult": ... +@runtime_checkable +class CompletionClient(Protocol): + """Minimal client contract consumed by reusable generation tasks.""" + + @property + def model(self) -> str: + """Return the model identifier used for generated results.""" + ... + + def complete( + self, + messages: list[dict[str, Any]], + max_tokens: Optional[int] = None, + extra_params: Optional[dict[str, Any]] = None, + ) -> tuple[str, float]: + """Return generated text and wall-clock latency in seconds.""" + ... + + @runtime_checkable class AnswerJudge(Protocol): """Pluggable answer scoring interface.""" @@ -64,6 +83,25 @@ class GenerationResult: error: Optional[str] = None +@dataclass +class GenerationRequest: + """Provider-neutral request produced by a generation task.""" + + messages: list[dict[str, Any]] + max_tokens: Optional[int] = None + extra_params: Optional[dict[str, Any]] = None + + +@dataclass +class GeneratedTextResult: + """Task-neutral result from a single text-generation request.""" + + text: str + latency_s: float + model: str + error: Optional[str] = None + + @dataclass class JudgeResult: """Result from a single judge evaluation. @@ -104,7 +142,7 @@ class AnswerResult(BaseModel): to produce it and -- when a ``reference`` answer and/or ``judge`` are supplied -- the Tier-1 / Tier-2 / Tier-3 scoring artefacts produced by :mod:`nemo_retriever.evaluation.scoring` and - :class:`~nemo_retriever.llm.clients.judge.LLMJudge`. + :class:`~nemo_retriever.models.llm.clients.judge.LLMJudge`. Attributes: query: The question that was answered. diff --git a/nemo_retriever/src/nemo_retriever/operators/__init__.py b/nemo_retriever/src/nemo_retriever/operators/__init__.py index aaf3ba1ec..fcfc2b73e 100644 --- a/nemo_retriever/src/nemo_retriever/operators/__init__.py +++ b/nemo_retriever/src/nemo_retriever/operators/__init__.py @@ -8,14 +8,23 @@ not eagerly pull the compatibility aliases (which keeps package initialization free of import cycles during the bucket reorganization). """ + from __future__ import annotations -__all__ = ["AbstractOperator", "CPUOperator", "GPUOperator", "ExplodeContentActor", "_BatchEmbedActor"] +__all__ = [ + "AbstractOperator", + "CPUOperator", + "GPUOperator", + "TextGenerationOperator", + "ExplodeContentActor", + "_BatchEmbedActor", +] _LAZY = { "AbstractOperator": "nemo_retriever.operators.abstract_operator", "CPUOperator": "nemo_retriever.operators.cpu_operator", "GPUOperator": "nemo_retriever.operators.gpu_operator", + "TextGenerationOperator": "nemo_retriever.operators.generation", "ExplodeContentActor": "nemo_retriever.operators.graph_ops.content_operators", "_BatchEmbedActor": "nemo_retriever.operators.embed.operators", } diff --git a/nemo_retriever/src/nemo_retriever/operators/generation/__init__.py b/nemo_retriever/src/nemo_retriever/operators/generation/__init__.py new file mode 100644 index 000000000..b45bebaba --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/operators/generation/__init__.py @@ -0,0 +1,15 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Reusable operators for row-oriented text generation.""" + +from nemo_retriever.operators.generation.base import TextGenerationOperator +from nemo_retriever.operators.generation.generic import GenericGenerationOperator +from nemo_retriever.operators.generation.summarization import SummarizationOperator + +__all__ = [ + "GenericGenerationOperator", + "SummarizationOperator", + "TextGenerationOperator", +] diff --git a/nemo_retriever/src/nemo_retriever/operators/generation/base.py b/nemo_retriever/src/nemo_retriever/operators/generation/base.py new file mode 100644 index 000000000..d63454bd2 --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/operators/generation/base.py @@ -0,0 +1,296 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Reusable DataFrame operator for text-generation tasks.""" + +from __future__ import annotations + +import inspect +import logging +from abc import abstractmethod +from collections.abc import Mapping +from concurrent.futures import Future, ThreadPoolExecutor, as_completed +from typing import Any, ClassVar + +import pandas as pd + +from pydantic import BaseModel +from nemo_retriever.common.params import TextGenerationParams +from nemo_retriever.models.llm.clients import LiteLLMClient +from nemo_retriever.models.llm.tasks import GenerationTask +from nemo_retriever.models.llm.types import CompletionClient, GeneratedTextResult +from nemo_retriever.operators.abstract_operator import AbstractOperator +from nemo_retriever.operators.cpu_operator import CPUOperator + +logger = logging.getLogger(__name__) + + +class TextGenerationOperator(AbstractOperator, CPUOperator): + """Base operator for one text-generation request per DataFrame row. + + Subclasses bind a concrete :class:`GenerationTask` through + :meth:`_create_task`. The base owns the common DataFrame concerns: column + validation, bounded threaded execution, positional result ordering, and + standard output metadata. + + ``input_columns`` maps each task-level input name to a physical DataFrame + column. Results are tracked by row position rather than index label so + duplicate DataFrame indices remain safe. + """ + + required_columns: ClassVar[tuple[str, ...]] = () + output_columns: ClassVar[tuple[str, ...]] = () + + def __init__( + self, + params: TextGenerationParams, + *, + input_columns: Mapping[str, str], + output_column: str, + latency_column: str | None = None, + model_column: str | None = None, + error_column: str | None = None, + overwrite: bool = False, + client: CompletionClient | None = None, + ) -> None: + logical_columns = dict(input_columns) + self._validate_input_mapping(logical_columns) + + resolved_latency_column = latency_column if latency_column is not None else f"{output_column}_latency_s" + resolved_model_column = model_column if model_column is not None else f"{output_column}_model" + resolved_error_column = error_column if error_column is not None else f"{output_column}_error" + output_columns = ( + output_column, + resolved_latency_column, + resolved_model_column, + resolved_error_column, + ) + self._validate_output_columns(output_columns) + + super().__init__() + + self._params = params + self._input_columns = logical_columns + self._output_column = output_column + self._latency_column = resolved_latency_column + self._model_column = resolved_model_column + self._error_column = resolved_error_column + self._latency_column_arg = latency_column + self._model_column_arg = model_column + self._error_column_arg = error_column + self._overwrite = overwrite + self._max_workers = params.max_workers + self._configured_model = params.transport.model + + self.required_columns = tuple(dict.fromkeys(logical_columns.values())) + self.output_columns = output_columns + + self._task = self._create_task(params, tuple(logical_columns)) + missing_inputs = [name for name in self._task.required_inputs if name not in logical_columns] + if missing_inputs: + raise ValueError(f"{type(self).__name__} is missing task input mappings: {missing_inputs}") + + if client is None: + sampling = params.resolve_sampling(self._task.default_sampling) + self._client: CompletionClient = LiteLLMClient(transport=params.transport, sampling=sampling) + else: + self._client = client + + @abstractmethod + def _create_task( + self, + params: TextGenerationParams, + logical_inputs: tuple[str, ...], + ) -> GenerationTask: + """Create the stateless task bound to this concrete operator.""" + ... + + @abstractmethod + def _get_generation_constructor_kwargs(self) -> dict[str, Any]: + """Return reconstructible public constructor state for this operator.""" + ... + + @classmethod + def _contains_runtime_object( + cls, + value: Any, + targets: tuple[object, ...], + seen: set[int], + ) -> bool: + if any(value is target for target in targets): + return True + if isinstance(value, (str, bytes, bytearray, memoryview)): + return False + value_id = id(value) + if value_id in seen: + return False + seen.add(value_id) + if isinstance(value, BaseModel): + return any( + cls._contains_runtime_object(getattr(value, name), targets, seen) for name in type(value).model_fields + ) + if isinstance(value, Mapping): + return any(cls._contains_runtime_object(item, targets, seen) for pair in value.items() for item in pair) + if isinstance(value, (list, tuple, set, frozenset)): + return any(cls._contains_runtime_object(item, targets, seen) for item in value) + return False + + def get_constructor_kwargs(self) -> dict[str, Any]: + """Return validated graph state without capturing a live task or client.""" + kwargs = dict(self._get_generation_constructor_kwargs()) + forbidden_keys = {"client", "task", "_client", "_task"}.intersection(kwargs) + if forbidden_keys: + raise TypeError( + f"{type(self).__name__} graph constructor hook returned runtime-only keys: " f"{sorted(forbidden_keys)}" + ) + if self._contains_runtime_object( + kwargs, + (self._client, self._task), + set(), + ): + raise TypeError(f"{type(self).__name__} graph constructor hook captured a live client or task") + + signature = inspect.signature(type(self).__init__) + try: + signature.bind(None, **kwargs) + except TypeError as exc: + raise TypeError(f"{type(self).__name__} returned invalid graph constructor kwargs: {exc}") from exc + return kwargs + + @staticmethod + def _validate_input_mapping(input_columns: Mapping[str, str]) -> None: + if not input_columns: + raise ValueError("input_columns must contain at least one task input mapping") + for logical_name, column_name in input_columns.items(): + if not isinstance(logical_name, str) or not logical_name: + raise ValueError("input_columns task input names must be non-empty strings") + if not isinstance(column_name, str) or not column_name: + raise ValueError("input_columns DataFrame column names must be non-empty strings") + + @staticmethod + def _validate_output_columns(output_columns: tuple[str, ...]) -> None: + if any(not isinstance(column, str) or not column for column in output_columns): + raise ValueError("output column names must be non-empty strings") + if len(set(output_columns)) != len(output_columns): + raise ValueError(f"output column names must be distinct: {list(output_columns)}") + + @staticmethod + def _label_positions(data: pd.DataFrame, label: str) -> list[int]: + return [int(position) for position in data.columns.get_indexer_for([label]) if position >= 0] + + def _validate_and_resolve_dataframe( + self, + data: Any, + ) -> tuple[pd.DataFrame, dict[str, int]]: + if not isinstance(data, pd.DataFrame): + raise TypeError(f"{type(self).__name__} requires a pandas DataFrame") + + input_positions: dict[str, int] = {} + missing: list[str] = [] + ambiguous_inputs: list[str] = [] + for logical_name, column_name in self._input_columns.items(): + positions = self._label_positions(data, column_name) + if not positions: + missing.append(column_name) + elif len(positions) > 1: + ambiguous_inputs.append(column_name) + else: + input_positions[logical_name] = positions[0] + if missing: + missing = list(dict.fromkeys(missing)) + raise ValueError(f"{type(self).__name__} requires missing columns: {missing}") + if ambiguous_inputs: + ambiguous_inputs = list(dict.fromkeys(ambiguous_inputs)) + raise ValueError( + f"{type(self).__name__} mapped input columns are ambiguous because their labels " + f"are duplicated: {ambiguous_inputs}" + ) + + if not self._overwrite: + collisions = [column for column in self.output_columns if self._label_positions(data, column)] + if collisions: + raise ValueError( + f"{type(self).__name__} output columns already exist: {collisions}; " + "set overwrite=True to replace them" + ) + else: + ambiguous_outputs = [ + column for column in self.output_columns if len(self._label_positions(data, column)) > 1 + ] + if ambiguous_outputs: + raise ValueError( + f"{type(self).__name__} cannot overwrite ambiguous duplicate output " f"labels: {ambiguous_outputs}" + ) + return data, input_positions + + def preprocess(self, data: Any, **kwargs: Any) -> pd.DataFrame: + df, _ = self._validate_and_resolve_dataframe(data) + return df + + def _execute_task(self, inputs: dict[str, Any]) -> GeneratedTextResult: + """Execute the configured task; subclasses may adapt legacy clients.""" + return self._task.execute(self._client, **inputs) + + def _execute_row(self, position: int, inputs: dict[str, Any]) -> tuple[int, GeneratedTextResult]: + return position, self._execute_task(inputs) + + def _failure_model(self) -> str: + try: + model = self._client.model + except Exception: + return self._configured_model + return model if isinstance(model, str) and model else self._configured_model + + def _failure_result(self, exc: Exception) -> GeneratedTextResult: + return GeneratedTextResult( + text="", + latency_s=0.0, + model=self._failure_model(), + error=str(exc), + ) + + def process(self, data: Any, **kwargs: Any) -> pd.DataFrame: + df, input_positions = self._validate_and_resolve_dataframe(data) + results: list[GeneratedTextResult | None] = [None] * len(df) + + if len(df): + futures: dict[Future[tuple[int, GeneratedTextResult]], int] = {} + with ThreadPoolExecutor(max_workers=min(self._max_workers, len(df))) as pool: + for position in range(len(df)): + inputs = { + name: df.iat[position, column_position] for name, column_position in input_positions.items() + } + future = pool.submit(self._execute_row, position, inputs) + futures[future] = position + + for future in as_completed(futures): + position = futures[future] + try: + result_position, result = future.result() + if result_position != position: + raise RuntimeError( + f"generation result position {result_position} does not match " + f"submitted position {position}" + ) + results[position] = result + except Exception as exc: + logger.warning("Row %d generation failed: %s", position, exc) + results[position] = self._failure_result(exc) + + # Every non-empty row is assigned either a task result or a failure + # result above. The cast-free local assertion catches future changes to + # that invariant before partially writing output columns. + if any(result is None for result in results): + raise RuntimeError("generation completed without a result for every row") + completed_results = [result for result in results if result is not None] + + out = df.copy() + out[self._output_column] = [result.text for result in completed_results] + out[self._latency_column] = [result.latency_s for result in completed_results] + out[self._model_column] = [result.model for result in completed_results] + out[self._error_column] = [result.error for result in completed_results] + return out + + def postprocess(self, data: Any, **kwargs: Any) -> Any: + return data diff --git a/nemo_retriever/src/nemo_retriever/operators/generation/generic.py b/nemo_retriever/src/nemo_retriever/operators/generation/generic.py new file mode 100644 index 000000000..69ff2ff96 --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/operators/generation/generic.py @@ -0,0 +1,70 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Generic prompt-template operator built on the generation layer.""" + +from __future__ import annotations + +from collections.abc import Mapping + +from nemo_retriever.common.params import TextGenerationParams +from nemo_retriever.models.llm.tasks import GenerationTask, GenericPromptTask +from nemo_retriever.models.llm.types import CompletionClient +from nemo_retriever.operators.generation.base import TextGenerationOperator + + +class GenericGenerationOperator(TextGenerationOperator): + """Generate text from a validated prompt template and mapped row inputs.""" + + def __init__( + self, + params: TextGenerationParams, + input_columns: Mapping[str, str], + output_column: str = "generated_text", + *, + latency_column: str | None = None, + model_column: str | None = None, + error_column: str | None = None, + overwrite: bool = False, + client: CompletionClient | None = None, + ) -> None: + normalized_input_columns = dict(input_columns) + super().__init__( + params, + input_columns=normalized_input_columns, + output_column=output_column, + latency_column=latency_column, + model_column=model_column, + error_column=error_column, + overwrite=overwrite, + client=client, + ) + + def _get_generation_constructor_kwargs(self) -> dict[str, object]: + return { + "params": self._params, + "input_columns": dict(self._input_columns), + "output_column": self._output_column, + "latency_column": self._latency_column_arg, + "model_column": self._model_column_arg, + "error_column": self._error_column_arg, + "overwrite": self._overwrite, + } + + def _create_task( + self, + params: TextGenerationParams, + logical_inputs: tuple[str, ...], + ) -> GenerationTask: + if params.prompt is None: + raise ValueError("GenericGenerationOperator requires params.prompt") + reasoning_enabled = ( + params.reasoning_enabled if params.reasoning_enabled is not None else params.transport.reasoning_enabled + ) + return GenericPromptTask( + prompt=params.prompt, + input_names=logical_inputs, + system_prompt=params.system_prompt, + reasoning_enabled=reasoning_enabled, + ) diff --git a/nemo_retriever/src/nemo_retriever/operators/generation/summarization.py b/nemo_retriever/src/nemo_retriever/operators/generation/summarization.py new file mode 100644 index 000000000..a18d14f0c --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/operators/generation/summarization.py @@ -0,0 +1,65 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Summarization operator built on the reusable generation layer.""" + +from __future__ import annotations + +from nemo_retriever.common.params import TextGenerationParams +from nemo_retriever.models.llm.tasks import GenerationTask, SummarizeTask +from nemo_retriever.models.llm.types import CompletionClient +from nemo_retriever.operators.generation.base import TextGenerationOperator + + +class SummarizationOperator(TextGenerationOperator): + """Summarize the text in one DataFrame column per row.""" + + def __init__( + self, + params: TextGenerationParams, + input_column: str = "text", + output_column: str = "summary", + *, + latency_column: str | None = None, + model_column: str | None = None, + error_column: str | None = None, + overwrite: bool = False, + client: CompletionClient | None = None, + ) -> None: + super().__init__( + params, + input_columns={"text": input_column}, + output_column=output_column, + latency_column=latency_column, + model_column=model_column, + error_column=error_column, + overwrite=overwrite, + client=client, + ) + + def _get_generation_constructor_kwargs(self) -> dict[str, object]: + return { + "params": self._params, + "input_column": self._input_columns["text"], + "output_column": self._output_column, + "latency_column": self._latency_column_arg, + "model_column": self._model_column_arg, + "error_column": self._error_column_arg, + "overwrite": self._overwrite, + } + + def _create_task( + self, + params: TextGenerationParams, + logical_inputs: tuple[str, ...], + ) -> GenerationTask: + del logical_inputs + reasoning_enabled = ( + params.reasoning_enabled if params.reasoning_enabled is not None else params.transport.reasoning_enabled + ) + return SummarizeTask( + prompt=params.prompt, + system_prompt=params.system_prompt, + reasoning_enabled=reasoning_enabled, + ) diff --git a/nemo_retriever/src/nemo_retriever/tools/evaluation/generation.py b/nemo_retriever/src/nemo_retriever/tools/evaluation/generation.py index bda4e5f3f..d85fdc9bd 100644 --- a/nemo_retriever/src/nemo_retriever/tools/evaluation/generation.py +++ b/nemo_retriever/src/nemo_retriever/tools/evaluation/generation.py @@ -6,18 +6,15 @@ from __future__ import annotations -import logging -from concurrent.futures import ThreadPoolExecutor, as_completed from typing import Any, ClassVar, Optional -from nemo_retriever.operators.graph_ops.eval_operator import EvalOperator -from nemo_retriever.models.llm.clients import LiteLLMClient -from nemo_retriever.models.llm.types import GenerationResult +from nemo_retriever.common.params import TextGenerationParams +from nemo_retriever.models.llm.tasks import GenerationTask, RagAnswerTask +from nemo_retriever.models.llm.types import CompletionClient, GeneratedTextResult, LLMClient +from nemo_retriever.operators.generation import TextGenerationOperator -logger = logging.getLogger(__name__) - -class QAGenerationOperator(EvalOperator): +class QAGenerationOperator(TextGenerationOperator): """Generate answers for each row using a single LLM. Input DataFrame must have ``query`` and ``context`` columns. @@ -35,7 +32,7 @@ def __init__( *, api_base: Optional[str] = None, api_key: Optional[str] = None, - temperature: float = 0.0, + temperature: Optional[float] = 0.0, top_p: Optional[float] = None, max_tokens: int = 4096, extra_params: Optional[dict[str, Any]] = None, @@ -46,7 +43,7 @@ def __init__( rag_system_prompt_prefix: Optional[str] = None, reasoning_enabled: bool = True, ) -> None: - super().__init__( + params = TextGenerationParams.from_kwargs( model=model, api_base=api_base, api_key=api_key, @@ -61,46 +58,62 @@ def __init__( rag_system_prompt_prefix=rag_system_prompt_prefix, reasoning_enabled=reasoning_enabled, ) - self._client = LiteLLMClient.from_kwargs( - model=model, - api_base=api_base, - api_key=api_key, - temperature=temperature, - top_p=top_p, - max_tokens=max_tokens, - extra_params=extra_params, - num_retries=num_retries, - timeout=timeout, - rag_system_prompt=rag_system_prompt, - rag_system_prompt_prefix=rag_system_prompt_prefix, - reasoning_enabled=reasoning_enabled, + super().__init__( + params, + input_columns={"query": "query", "chunks": "context"}, + output_column="answer", + latency_column="latency_s", + model_column="model", + error_column="gen_error", + overwrite=True, ) - self._max_workers = max_workers + self._qa_constructor_kwargs = { + "model": model, + "api_base": api_base, + "api_key": api_key, + "temperature": temperature, + "top_p": top_p, + "max_tokens": max_tokens, + "extra_params": extra_params, + "num_retries": num_retries, + "timeout": timeout, + "max_workers": max_workers, + "rag_system_prompt": rag_system_prompt, + "rag_system_prompt_prefix": rag_system_prompt_prefix, + "reasoning_enabled": reasoning_enabled, + } - def process(self, data: Any, **kwargs: Any) -> Any: - df = data - results: list = [None] * len(df) + def _get_generation_constructor_kwargs(self) -> dict[str, Any]: + """Preserve the legacy flat QA constructor contract for graph workers.""" + return dict(self._qa_constructor_kwargs) - def _generate(idx: int, query: str, context: list[str]): - gen = self._client.generate(query, context) - return idx, gen + def _execute_task(self, inputs: dict[str, Any]) -> GeneratedTextResult: + """Prefer completion tasks while adapting legacy generate-only clients.""" + client = self._client + if isinstance(client, CompletionClient): + return super()._execute_task(inputs) + if isinstance(client, LLMClient): + result = client.generate(inputs["query"], inputs["chunks"]) + return GeneratedTextResult( + text=result.answer, + latency_s=result.latency_s, + model=result.model, + error=result.error, + ) + raise TypeError("QAGenerationOperator client must implement CompletionClient or LLMClient") - with ThreadPoolExecutor(max_workers=self._max_workers) as pool: - futures = { - pool.submit(_generate, i, row["query"], row["context"]): i for i, (_, row) in enumerate(df.iterrows()) - } - for future in as_completed(futures): - try: - idx, gen = future.result() - results[idx] = gen - except Exception as exc: - idx = futures[future] - logger.warning("Row %d generation failed: %s", idx, exc) - results[idx] = GenerationResult(answer="", latency_s=0.0, model=self._client.model, error=str(exc)) - - out = df.copy() - out["answer"] = [r.answer for r in results] - out["latency_s"] = [r.latency_s for r in results] - out["model"] = [r.model for r in results] - out["gen_error"] = [r.error for r in results] - return out + def _create_task( + self, + params: TextGenerationParams, + logical_inputs: tuple[str, ...], + ) -> GenerationTask: + del logical_inputs + reasoning_enabled = ( + params.reasoning_enabled if params.reasoning_enabled is not None else params.transport.reasoning_enabled + ) + return RagAnswerTask( + prompt=params.prompt, + system_prompt=params.transport.rag_system_prompt, + system_prompt_prefix=params.transport.rag_system_prompt_prefix, + reasoning_enabled=reasoning_enabled, + ) diff --git a/nemo_retriever/tests/test_caption.py b/nemo_retriever/tests/test_caption.py index 090621721..7b2fdfdab 100644 --- a/nemo_retriever/tests/test_caption.py +++ b/nemo_retriever/tests/test_caption.py @@ -218,6 +218,12 @@ def test_validation_inherited(self): with pytest.raises(ValueError, match="temperature must be between"): CaptionParams(temperature=-1.0) + def test_null_temperature_rejected(self): + from nemo_retriever.common.params import CaptionParams + + with pytest.raises(ValueError, match="cannot be None for captioning"): + CaptionParams(temperature=None) + class TestCaptionImageParamThreading: """Verify top_p and max_tokens flow through to the model / client.""" diff --git a/nemo_retriever/tests/test_generation_tasks.py b/nemo_retriever/tests/test_generation_tasks.py new file mode 100644 index 000000000..eff2d91da --- /dev/null +++ b/nemo_retriever/tests/test_generation_tasks.py @@ -0,0 +1,925 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Focused offline tests for the reusable text-generation task layer.""" + +from __future__ import annotations + +from collections.abc import Callable +import json +import threading +import time +from typing import Any + +import pandas as pd +import pytest + + +class FakeCompletionClient: + """Small thread-safe completion client used by task and operator tests.""" + + def __init__( + self, + handler: Callable[[list[dict], int | None, dict[str, Any] | None], tuple[str, float]] | None = None, + *, + model: str = "fake/model", + ) -> None: + self._model = model + self._handler = handler or (lambda messages, max_tokens, extra_params: ("generated", 0.25)) + self._lock = threading.Lock() + self.calls: list[tuple[list[dict], int | None, dict[str, Any] | None]] = [] + + @property + def model(self) -> str: + return self._model + + def complete( + self, + messages: list[dict], + max_tokens: int | None = None, + extra_params: dict[str, Any] | None = None, + ) -> tuple[str, float]: + with self._lock: + self.calls.append((messages, max_tokens, extra_params)) + return self._handler(messages, max_tokens, extra_params) + + +def _params(**kwargs: Any): + from nemo_retriever.common.params import TextGenerationParams + + kwargs.setdefault("api_key", "") + return TextGenerationParams.from_kwargs(model="fake/model", **kwargs) + + +class TestTextGenerationParams: + def test_flat_constructor_composes_transport_and_partial_sampling(self): + params = _params( + api_base="http://llm.test/v1", + api_key="secret", + temperature=0.4, + top_p=0.8, + max_tokens=77, + num_retries=5, + timeout=12.0, + extra_params={"seed": 3}, + prompt="Prompt {text}", + system_prompt="System", + reasoning_enabled=False, + max_workers=3, + ) + + assert params.transport.model == "fake/model" + assert params.transport.api_base == "http://llm.test/v1" + assert params.transport.api_key == "secret" + assert params.transport.num_retries == 5 + assert params.transport.timeout == 12.0 + assert params.transport.extra_params == {"seed": 3} + assert params.sampling.temperature == 0.4 + assert params.sampling.top_p == 0.8 + assert params.sampling.max_tokens == 77 + assert params.prompt == "Prompt {text}" + assert params.system_prompt == "System" + assert params.reasoning_enabled is False + assert params.max_workers == 3 + + def test_sampling_overrides_only_explicit_fields(self): + from nemo_retriever.common.params import LLMInferenceParams + + task_defaults = LLMInferenceParams(temperature=0.0, top_p=0.35, max_tokens=4096) + + inherited = _params().resolve_sampling(task_defaults) + overridden = _params(temperature=0.7, max_tokens=123).resolve_sampling(task_defaults) + + assert inherited == task_defaults + assert overridden.temperature == 0.7 + assert overridden.top_p == 0.35 + assert overridden.max_tokens == 123 + + def test_api_key_is_resolved_but_redacted_from_display(self, monkeypatch): + from nemo_retriever.common.params import models as params_models + + monkeypatch.setattr(params_models, "resolve_remote_api_key", lambda: "resolved-secret") + params = params_models.TextGenerationParams.from_kwargs(model="m") + + assert params.transport.api_key == "resolved-secret" + assert "resolved-secret" not in repr(params) + assert "resolved-secret" not in str(params) + + def test_no_api_key_survives_nested_validation(self, monkeypatch): + from nemo_retriever.common.params import models as params_models + + monkeypatch.setattr( + params_models, + "resolve_remote_api_key", + lambda: "environment-secret", + ) + params = params_models.TextGenerationParams.from_kwargs(model="m", api_key="") + + assert params.transport.api_key is None + + def test_max_workers_must_be_positive(self): + with pytest.raises(ValueError, match="max_workers"): + _params(max_workers=0) + + +class TestGenerationTasks: + def test_summary_builds_request_and_executes(self): + from nemo_retriever.models.llm.tasks import SummarizeTask + + client = FakeCompletionClient(lambda messages, max_tokens, extra_params: (" concise summary ", 0.4)) + task = SummarizeTask() + + request = task.build_request(text="Source text") + result = task.execute(client, text="Source text") + + assert [message["role"] for message in request.messages] == ["system", "user"] + assert request.messages[-1]["content"] == "Summarize the following text:\n\nSource text" + assert result.text == "concise summary" + assert result.latency_s == 0.4 + assert result.model == "fake/model" + assert result.error is None + assert client.calls == [(request.messages, request.max_tokens, request.extra_params)] + + def test_summary_empty_input_short_circuits(self): + from nemo_retriever.models.llm.tasks import SummarizeTask + + client = FakeCompletionClient() + result = SummarizeTask().execute(client, text=" \n") + + assert result.text == "" + assert result.error == "empty_input" + assert result.model == "fake/model" + assert client.calls == [] + + def test_task_converts_transport_exception_to_result(self): + from nemo_retriever.models.llm.tasks import SummarizeTask + + def fail(messages, max_tokens, extra_params): + raise RuntimeError("service unavailable") + + result = SummarizeTask().execute(FakeCompletionClient(fail), text="source") + + assert result.text == "" + assert result.latency_s == 0.0 + assert result.model == "fake/model" + assert result.error == "service unavailable" + + def test_rag_request_applies_no_reasoning_controls_and_think_cleanup(self): + from nemo_retriever.models.llm.tasks import RagAnswerTask + + client = FakeCompletionClient( + lambda messages, max_tokens, extra_params: ("private final answer ", 0.2) + ) + task = RagAnswerTask(reasoning_enabled=False) + + request = task.build_request(query="What?", chunks=["First", "Second"]) + result = task.execute(client, query="What?", chunks=["First", "Second"]) + + assert request.messages[0]["content"].startswith("/no_think\n") + assert "First\n\n---\n\nSecond" in request.messages[-1]["content"] + assert request.extra_params == {"chat_template_kwargs": {"enable_thinking": False}} + assert result.text == "final answer" + assert result.error is None + + def test_rag_think_only_output_uses_compatibility_error(self): + from nemo_retriever.models.llm.tasks import RagAnswerTask + + client = FakeCompletionClient(lambda messages, max_tokens, extra_params: ("unfinished", 0.3)) + result = RagAnswerTask().execute(client, query="q", chunks=[]) + + assert result.text == "" + assert result.latency_s == 0.3 + assert result.error == "thinking_truncated" + + def test_generic_template_supports_declared_fields_and_escaped_braces(self): + from nemo_retriever.models.llm.tasks import GenericPromptTask + + task = GenericPromptTask( + prompt="Keep {{literal}} and greet {name}", + input_names=("name",), + system_prompt="Be brief.", + ) + request = task.build_request(name="Ada") + + assert request.messages == [ + {"role": "system", "content": "Be brief."}, + {"role": "user", "content": "Keep {literal} and greet Ada"}, + ] + + @pytest.mark.parametrize( + ("prompt", "input_names"), + [ + ("constant text", ("name",)), + ("hello {name}", ("name", "unused")), + ("hello {other}", ("name",)), + ("hello {user.name}", ("user",)), + ("hello {items[0]}", ("items",)), + ("hello {name", ("name",)), + ], + ) + def test_generic_template_rejects_invalid_field_contracts(self, prompt, input_names): + from nemo_retriever.models.llm.tasks import GenericPromptTask + + with pytest.raises(ValueError): + GenericPromptTask(prompt=prompt, input_names=input_names) + + def test_generic_missing_runtime_input_is_an_error_result(self): + from nemo_retriever.models.llm.tasks import GenericPromptTask + + client = FakeCompletionClient() + result = GenericPromptTask(prompt="hello {name}", input_names=("name",)).execute(client) + + assert result.text == "" + assert result.error + assert client.calls == [] + + +class TestTextGenerationOperators: + def test_base_is_exported_from_canonical_operator_package(self): + from nemo_retriever.operators import TextGenerationOperator + from nemo_retriever.operators.generation import TextGenerationOperator as DirectExport + + assert TextGenerationOperator is DirectExport + + def test_summary_happy_path_uses_namespaced_metadata(self): + from nemo_retriever.operators.generation import SummarizationOperator + + client = FakeCompletionClient(lambda messages, max_tokens, extra_params: ("short", 0.1)) + source = pd.DataFrame({"text": ["long source"]}) + + out = SummarizationOperator(_params(), client=client).run(source) + + assert out.to_dict(orient="records") == [ + { + "text": "long source", + "summary": "short", + "summary_latency_s": 0.1, + "summary_model": "fake/model", + "summary_error": None, + } + ] + assert list(source.columns) == ["text"] + + def test_missing_and_colliding_columns_are_batch_errors(self): + from nemo_retriever.operators.generation import SummarizationOperator + + op = SummarizationOperator(_params(), client=FakeCompletionClient()) + + with pytest.raises(ValueError, match="missing columns"): + op.run(pd.DataFrame({"body": ["source"]})) + with pytest.raises(ValueError, match="already exist"): + op.run(pd.DataFrame({"text": ["source"], "summary": ["old"]})) + + def test_overwrite_allows_existing_output_column(self): + from nemo_retriever.operators.generation import SummarizationOperator + + client = FakeCompletionClient(lambda messages, max_tokens, extra_params: ("new", 0.1)) + op = SummarizationOperator(_params(), overwrite=True, client=client) + + out = op.run(pd.DataFrame({"text": ["source"], "summary": ["old"]})) + + assert out["summary"].tolist() == ["new"] + + def test_duplicate_indices_order_and_mixed_failures(self): + from nemo_retriever.operators.generation import SummarizationOperator + + def respond(messages, max_tokens, extra_params): + content = messages[-1]["content"] + if "boom" in content: + raise RuntimeError("offline for boom") + if "slow" in content: + time.sleep(0.03) + return "SLOW", 0.3 + return "FAST", 0.1 + + source = pd.DataFrame({"text": ["slow", "boom", "fast"]}, index=[7, 7, 2]) + out = SummarizationOperator(_params(max_workers=3), client=FakeCompletionClient(respond)).run(source) + + assert out.index.tolist() == [7, 7, 2] + assert out["text"].tolist() == ["slow", "boom", "fast"] + assert out["summary"].tolist() == ["SLOW", "", "FAST"] + assert out["summary_error"].tolist()[0] is None + assert out["summary_error"].tolist()[1] == "offline for boom" + assert out["summary_error"].tolist()[2] is None + assert out["summary_model"].tolist() == ["fake/model"] * 3 + + def test_empty_batch_adds_schema_without_calling_client(self): + from nemo_retriever.operators.generation import SummarizationOperator + + client = FakeCompletionClient() + out = SummarizationOperator(_params(), client=client).run(pd.DataFrame({"text": pd.Series(dtype=str)})) + + assert list(out.columns) == [ + "text", + "summary", + "summary_latency_s", + "summary_model", + "summary_error", + ] + assert out.empty + assert client.calls == [] + + def test_generic_operator_maps_multiple_logical_inputs(self): + from nemo_retriever.operators.generation import GenericGenerationOperator + + client = FakeCompletionClient(lambda messages, max_tokens, extra_params: ("bonjour", 0.2)) + params = _params(prompt="Translate {text} to {language}.", system_prompt="Translator") + op = GenericGenerationOperator( + params, + input_columns={"text": "body", "language": "target"}, + client=client, + ) + + out = op.run(pd.DataFrame({"body": ["hello"], "target": ["French"]})) + + assert out["generated_text"].tolist() == ["bonjour"] + assert client.calls[0][0] == [ + {"role": "system", "content": "Translator"}, + {"role": "user", "content": "Translate hello to French."}, + ] + + def test_injected_and_built_clients_are_not_graph_constructor_state(self): + from nemo_retriever.graph.pipeline_graph import Node + from nemo_retriever.operators.generation import SummarizationOperator + + client = FakeCompletionClient() + injected = SummarizationOperator(_params(), client=client) + built = SummarizationOperator(_params()) + + assert injected._client is client + for operator in (injected, built): + kwargs = Node(operator).operator_kwargs + assert "client" not in kwargs + assert "task" not in kwargs + assert client not in kwargs.values() + + +class TestQACompatibility: + def test_qa_run_preserves_schema_order_and_overwrites_legacy_outputs(self): + from nemo_retriever.tools.evaluation.generation import QAGenerationOperator + + client = FakeCompletionClient(lambda messages, max_tokens, extra_params: ("grounded answer", 0.6)) + operator = QAGenerationOperator(model="fake/model", api_key="") + operator._client = client + + source = pd.DataFrame({"query": ["question"], "context": [["context"]]}) + out = operator.run(source) + + assert list(out.columns) == [ + "query", + "context", + "answer", + "latency_s", + "model", + "gen_error", + ] + assert out.loc[0, ["answer", "latency_s", "model", "gen_error"]].tolist() == [ + "grounded answer", + 0.6, + "fake/model", + None, + ] + + existing = out.assign( + answer="stale", + latency_s=-1.0, + model="stale/model", + gen_error="stale error", + ) + overwritten = operator.run(existing) + + assert list(overwritten.columns) == list(existing.columns) + assert overwritten.loc[0, ["answer", "latency_s", "model", "gen_error"]].tolist() == [ + "grounded answer", + 0.6, + "fake/model", + None, + ] + + def test_qa_graph_kwargs_remain_flat_and_reconstructible(self): + from nemo_retriever.graph.pipeline_graph import Node + from nemo_retriever.tools.evaluation.generation import QAGenerationOperator + + operator = QAGenerationOperator( + "fake/model", + api_base="http://llm.test/v1", + api_key="", + temperature=0.2, + top_p=0.7, + max_tokens=321, + extra_params={"seed": 4}, + num_retries=5, + timeout=11.0, + max_workers=2, + rag_system_prompt="Use only context.", + rag_system_prompt_prefix="Prefix", + reasoning_enabled=False, + ) + kwargs = Node(operator).operator_kwargs + + assert set(kwargs) == { + "model", + "api_base", + "api_key", + "temperature", + "top_p", + "max_tokens", + "extra_params", + "num_retries", + "timeout", + "max_workers", + "rag_system_prompt", + "rag_system_prompt_prefix", + "reasoning_enabled", + } + assert not ({"params", "input_columns", "output_column", "task", "client"} & kwargs.keys()) + + reconstructed = QAGenerationOperator(**kwargs) + assert reconstructed._client.transport.model == "fake/model" + assert reconstructed._client.sampling.temperature == 0.2 + assert reconstructed._client.sampling.top_p == 0.7 + assert reconstructed._client.sampling.max_tokens == 321 + assert reconstructed.required_columns == ("query", "context") + assert reconstructed.output_columns == ("answer", "latency_s", "model", "gen_error") + + def test_litellm_generate_preserves_legacy_result_sentinels(self): + from nemo_retriever.models.llm.clients import LiteLLMClient + from nemo_retriever.models.llm.types import GenerationResult + + client = LiteLLMClient.from_kwargs(model="fake/model", api_key="") + client.complete = lambda messages, max_tokens=None, extra_params=None: ( + "reasoning only", + 0.45, + ) + truncated = client.generate("question", ["context"]) + + assert isinstance(truncated, GenerationResult) + assert truncated == GenerationResult( + answer="", + latency_s=0.45, + model="fake/model", + error="thinking_truncated", + ) + + def fail(messages, max_tokens=None, extra_params=None): + raise RuntimeError("transport unavailable") + + client.complete = fail + failed = client.generate("question", ["context"]) + assert failed == GenerationResult( + answer="", + latency_s=0.0, + model="fake/model", + error="transport unavailable", + ) + + def test_rag_prompt_helper_reexport_preserves_exact_messages(self): + from nemo_retriever.models.llm.clients import _build_rag_prompt + + empty = _build_rag_prompt("Where?", [], formatted_rag_system_prompt="System") + populated = _build_rag_prompt( + "Where?", + ["first", "second"], + formatted_rag_system_prompt="System", + ) + + assert empty == [ + {"role": "system", "content": "System"}, + { + "role": "user", + "content": "Context:\n(no context retrieved)\n\nQuestion: Where?\n\nAnswer:", + }, + ] + assert populated[-1]["content"] == ("Context:\nfirst\n\n---\n\nsecond\n\nQuestion: Where?\n\nAnswer:") + + +class TestSamplingRemediation: + def test_omitted_and_explicit_null_overrides_survive_round_trip(self): + from nemo_retriever.common.params import LLMInferenceParams, LLMSamplingOverrides + + defaults = LLMInferenceParams(temperature=0.4, top_p=0.7, max_tokens=99) + omitted = LLMSamplingOverrides() + cleared = LLMSamplingOverrides(temperature=None, top_p=None) + + omitted_restored = LLMSamplingOverrides.model_validate_json(omitted.model_dump_json()) + cleared_restored = LLMSamplingOverrides.model_validate_json(cleared.model_dump_json()) + + assert omitted.model_dump() == {} + assert omitted_restored.resolve(defaults) == defaults + assert omitted != cleared + assert _params() != _params(temperature=None) + assert cleared.model_dump() == {"temperature": None, "top_p": None} + resolved = cleared_restored.resolve(defaults) + assert resolved.temperature is None + assert resolved.top_p is None + assert resolved.to_sampling_kwargs() == {"max_tokens": 99} + + def test_nested_text_params_preserve_partial_sampling_state(self): + from nemo_retriever.common.params import ( + LLMInferenceParams, + TextGenerationParams, + ) + + params = _params(top_p=None) + restored = TextGenerationParams.model_validate(params.model_dump()) + defaults = LLMInferenceParams(temperature=0.3, top_p=0.7, max_tokens=10) + + assert restored.sampling.model_fields_set == {"top_p"} + assert restored.resolve_sampling(defaults).top_p is None + assert restored.resolve_sampling(defaults).temperature == 0.3 + + def test_explicit_null_max_tokens_is_invalid(self): + from nemo_retriever.common.params import LLMSamplingOverrides + + with pytest.raises(ValueError, match="omit it"): + LLMSamplingOverrides(max_tokens=None) + + +class TestOperatorRemediation: + def test_mixed_numeric_columns_preserve_exact_integer_prompt_value(self): + from nemo_retriever.operators.generation import GenericGenerationOperator + + large = 2**60 + 1 + seen: list[str] = [] + + def respond(messages, max_tokens, extra_params): + seen.append(messages[-1]["content"]) + return "ok", 0.1 + + source = pd.DataFrame( + { + "large": pd.Series([large], dtype="int64"), + "ratio": pd.Series([0.5], dtype="float64"), + } + ) + operator = GenericGenerationOperator( + _params(prompt="{large}|{ratio}"), + input_columns={"large": "large", "ratio": "ratio"}, + client=FakeCompletionClient(respond), + ) + + out = operator.run(source) + + assert seen == [f"{large}|0.5"] + assert out["generated_text"].tolist() == ["ok"] + + def test_ambiguous_mapped_input_label_is_rejected(self): + from nemo_retriever.operators.generation import SummarizationOperator + + source = pd.DataFrame([["first", "second"]], columns=["text", "text"]) + operator = SummarizationOperator(_params(), client=FakeCompletionClient()) + + with pytest.raises(ValueError, match="mapped input columns are ambiguous"): + operator.run(source) + + def test_ambiguous_output_label_is_rejected_when_overwriting(self): + from nemo_retriever.operators.generation import SummarizationOperator + + source = pd.DataFrame( + [["source", "old one", "old two"]], + columns=["text", "summary", "summary"], + ) + operator = SummarizationOperator(_params(), overwrite=True, client=FakeCompletionClient()) + + with pytest.raises(ValueError, match="ambiguous duplicate output"): + operator.run(source) + + def test_reconstruction_hook_rejects_nested_runtime_client(self): + from nemo_retriever.graph.pipeline_graph import Node + from nemo_retriever.operators.generation import SummarizationOperator + + class InvalidOperator(SummarizationOperator): + def _get_generation_constructor_kwargs(self): + kwargs = super()._get_generation_constructor_kwargs() + kwargs["input_column"] = { + "nested": [self._client], + } + return kwargs + + operator = InvalidOperator( + _params(), + client=FakeCompletionClient(), + ) + + with pytest.raises(TypeError, match="captured a live client or task"): + Node(operator) + + +class TestQALegacyClientRemediation: + def test_generate_only_client_without_model_is_called(self): + from nemo_retriever.models.llm.types import GenerationResult + from nemo_retriever.tools.evaluation.generation import QAGenerationOperator + + class LegacyClient: + def __init__(self): + self.calls: list[tuple[str, list[str]]] = [] + + def generate(self, query, chunks, *, reasoning_enabled=None): + self.calls.append((query, chunks)) + return GenerationResult("legacy answer", 0.4, "legacy/model") + + client = LegacyClient() + operator = QAGenerationOperator(model="configured/model", api_key="") + operator._client = client + + out = operator.run(pd.DataFrame({"query": ["q"], "context": [["c"]]})) + + assert client.calls == [("q", ["c"])] + assert out.loc[0, "answer"] == "legacy answer" + assert out.loc[0, "model"] == "legacy/model" + assert out.loc[0, "gen_error"] is None + + def test_generate_only_failure_uses_configured_model(self): + from nemo_retriever.tools.evaluation.generation import QAGenerationOperator + + class FailingLegacyClient: + def generate(self, query, chunks, *, reasoning_enabled=None): + raise RuntimeError("legacy unavailable") + + operator = QAGenerationOperator(model="configured/model", api_key="") + operator._client = FailingLegacyClient() + + out = operator.run(pd.DataFrame({"query": ["q"], "context": [["c"]]})) + + assert out.loc[0, "answer"] == "" + assert out.loc[0, "model"] == "configured/model" + assert out.loc[0, "gen_error"] == "legacy unavailable" + + def test_dual_protocol_client_prefers_completion_contract(self): + from nemo_retriever.models.llm.types import GenerationResult + from nemo_retriever.tools.evaluation.generation import QAGenerationOperator + + class DualClient(FakeCompletionClient): + def __init__(self): + super().__init__(lambda messages, max_tokens, extra_params: ("completion answer", 0.2)) + self.generate_calls = 0 + + def generate(self, query, chunks, *, reasoning_enabled=None): + self.generate_calls += 1 + return GenerationResult("legacy answer", 0.3, self.model) + + client = DualClient() + operator = QAGenerationOperator(model="configured/model", api_key="") + operator._client = client + + out = operator.run(pd.DataFrame({"query": ["q"], "context": [["c"]]})) + + assert out.loc[0, "answer"] == "completion answer" + assert len(client.calls) == 1 + assert client.generate_calls == 0 + + +class TestPublicLLMExports: + def test_direct_and_star_imports_resolve_canonical_clients(self): + from nemo_retriever.models.llm import LLMJudge, LiteLLMClient + + namespace: dict[str, Any] = {} + exec("from nemo_retriever.models.llm import *", namespace) + + assert LiteLLMClient.__module__ == "nemo_retriever.models.llm.clients.litellm" + assert LLMJudge.__module__ == "nemo_retriever.models.llm.clients.judge" + assert namespace["LiteLLMClient"] is LiteLLMClient + assert namespace["LLMJudge"] is LLMJudge + + def test_type_contract_import_does_not_eagerly_load_clients(self): + import subprocess + import sys + + code = ( + "import sys; " + "from nemo_retriever.models.llm import CompletionClient; " + "assert 'nemo_retriever.models.llm.clients.litellm' not in sys.modules; " + "assert 'litellm' not in sys.modules; " + "from nemo_retriever.models.llm import LiteLLMClient; " + "assert 'nemo_retriever.models.llm.clients.litellm' in sys.modules; " + "assert 'litellm' not in sys.modules" + ) + + subprocess.run( + [sys.executable, "-c", code], + check=True, + ) + + +class TestGenerationGraphPersistence: + @staticmethod + def _single_root_graph(operator): + from nemo_retriever.graph.pipeline_graph import Graph + + graph = Graph() + graph.add_root(operator) + return graph + + def test_typed_params_are_secret_free_and_rehydrate_from_environment( + self, + monkeypatch, + tmp_path, + ): + from nemo_retriever.common.params import TextGenerationParams + from nemo_retriever.graph.graph_pipeline_registry import ( + deserialize_graph, + load_graph, + save_graph, + serialize_graph, + ) + from nemo_retriever.operators.generation import SummarizationOperator + + params = TextGenerationParams.from_kwargs( + model="fake/model", + api_key="constructor-secret", + temperature=None, + top_p=0.7, + ) + graph = self._single_root_graph(SummarizationOperator(params)) + + payload = serialize_graph(graph) + encoded = json.dumps(payload) + + assert payload["format_version"] == 2 + assert "constructor-secret" not in encoded + assert "__pydantic_model__" in encoded + assert "__secret_env__" in encoded + + monkeypatch.setenv("NVIDIA_API_KEY", "worker-secret") + restored = deserialize_graph(payload) + restored_operator = restored.roots[0].operator + + assert isinstance(restored_operator, SummarizationOperator) + assert restored_operator._params.transport.api_key == "worker-secret" + assert restored_operator._params.sampling.model_fields_set == { + "temperature", + "top_p", + } + assert restored_operator._client.sampling.temperature is None + assert restored_operator._client.sampling.top_p == 0.7 + + path = tmp_path / "generation.json" + save_graph(graph, path) + assert "constructor-secret" not in path.read_text() + assert isinstance(load_graph(path).roots[0].operator, SummarizationOperator) + + def test_no_auth_survives_clone_with_environment_key_present(self, monkeypatch): + from nemo_retriever.graph.graph_pipeline_registry import clone_graph + from nemo_retriever.operators.generation import GenericGenerationOperator + + monkeypatch.setenv("NVIDIA_API_KEY", "must-not-be-used") + operator = GenericGenerationOperator( + _params(prompt="Hello {name}", api_key=""), + input_columns={"name": "name"}, + ) + + cloned = clone_graph(self._single_root_graph(operator)) + cloned_operator = cloned.roots[0].operator + + assert isinstance(cloned_operator, GenericGenerationOperator) + assert cloned_operator._params.transport.api_key is None + assert cloned_operator._params.transport._uses_no_api_key("api_key") + + def test_qa_and_registry_round_trip_to_concrete_operators( + self, + monkeypatch, + tmp_path, + ): + from nemo_retriever.graph.graph_pipeline_registry import ( + GraphPipelineRegistry, + deserialize_graph, + serialize_graph, + ) + from nemo_retriever.tools.evaluation.generation import QAGenerationOperator + + operator = QAGenerationOperator( + model="fake/model", + api_key="qa-constructor-secret", + ) + graph = self._single_root_graph(operator) + payload = serialize_graph(graph) + + monkeypatch.setenv("NVIDIA_API_KEY", "qa-worker-secret") + restored = deserialize_graph(payload) + assert isinstance(restored.roots[0].operator, QAGenerationOperator) + assert restored.roots[0].operator._client.transport.api_key == "qa-worker-secret" + + registry = GraphPipelineRegistry() + registry.register_graph( + "qa", + lambda: graph, + description="QA graph", + tags=["generation"], + ) + path = tmp_path / "registry.json" + registry.save_all(path) + raw = json.loads(path.read_text()) + + assert raw["format_version"] == 2 + assert set(raw["graphs"]) == {"qa"} + assert "qa-constructor-secret" not in path.read_text() + + loaded_registry = GraphPipelineRegistry() + assert loaded_registry.load_all(path) == ["qa"] + rebuilt = loaded_registry.build("qa") + assert isinstance(rebuilt.roots[0].operator, QAGenerationOperator) + + def test_opaque_nested_api_key_is_rejected_with_context(self): + from nemo_retriever.graph.graph_pipeline_registry import ( + GraphSerializationError, + serialize_graph, + ) + from nemo_retriever.operators.generation import SummarizationOperator + + graph = self._single_root_graph(SummarizationOperator(_params())) + graph.roots[0].operator_kwargs["opaque"] = { + "api_key": "must-not-leak", + } + + with pytest.raises( + GraphSerializationError, + match="opaque mapping", + ): + serialize_graph(graph) + + def test_graph_diagnostics_redact_flat_and_nested_credentials(self): + from nemo_retriever.graph.graph_pipeline_registry import ( + diff_graphs, + format_full_report, + format_graph_tree, + format_node_details, + ) + from nemo_retriever.tools.evaluation.generation import QAGenerationOperator + + first = QAGenerationOperator( + model="fake/model", + api_key="FIRST-SECRET", + extra_params={ + "headers": { + "Authorization": "Bearer INNER-FIRST", + }, + }, + ) + second = QAGenerationOperator( + model="fake/model", + api_key="SECOND-SECRET", + extra_params={ + "headers": { + "Authorization": "Bearer INNER-SECOND", + }, + }, + ) + first_graph = self._single_root_graph(first) + second_graph = self._single_root_graph(second) + + rendered = [ + format_graph_tree(first_graph, show_kwargs=True), + format_node_details(first_graph.roots[0]), + format_full_report(first_graph), + diff_graphs(first_graph, second_graph).format(), + ] + + for report in rendered: + assert "FIRST-SECRET" not in report + assert "SECOND-SECRET" not in report + assert "Bearer INNER-FIRST" not in report + assert "Bearer INNER-SECOND" not in report + assert "***" in report + + def test_non_rehydratable_token_fields_are_rejected(self): + from nemo_retriever.common.params.models import ASRParams + from nemo_retriever.graph.graph_pipeline_registry import ( + GraphSerializationError, + _safe_serialize_value, + serialize_graph, + ) + from nemo_retriever.operators.generation import SummarizationOperator + + with pytest.raises( + GraphSerializationError, + match="non-rehydratable secret field", + ): + _safe_serialize_value(ASRParams(auth_token="typed-secret")) + + graph = self._single_root_graph(SummarizationOperator(_params())) + graph.roots[0].operator_kwargs["auth_token"] = "flat-secret" + + with pytest.raises( + GraphSerializationError, + match="non-rehydratable secret field", + ): + serialize_graph(graph) + + def test_v2_constructor_failures_raise_instead_of_using_placeholder(self): + from nemo_retriever.graph.graph_pipeline_registry import ( + GraphSerializationError, + deserialize_graph, + serialize_graph, + ) + from nemo_retriever.operators.generation import SummarizationOperator + + graph = self._single_root_graph(SummarizationOperator(_params())) + payload = serialize_graph(graph) + payload["roots"][0]["operator_kwargs"]["unexpected"] = True + + with pytest.raises( + GraphSerializationError, + match="failed to construct operator", + ): + deserialize_graph(payload) diff --git a/nemo_retriever/tests/test_graph_pipeline_registry.py b/nemo_retriever/tests/test_graph_pipeline_registry.py index 740be9720..e73761ab3 100644 --- a/nemo_retriever/tests/test_graph_pipeline_registry.py +++ b/nemo_retriever/tests/test_graph_pipeline_registry.py @@ -11,12 +11,12 @@ from typing import Any import pytest - -from nemo_retriever.operators.abstract_operator import AbstractOperator -from nemo_retriever.graph.pipeline_graph import Graph, Node +from pydantic import BaseModel +from nemo_retriever.common.params import LLMRemoteClientParams from nemo_retriever.graph.graph_pipeline_registry import ( GraphBlueprint, GraphDiff, + GraphSerializationError, GraphPipelineRegistry, _PlaceholderOperator, _import_class, @@ -48,6 +48,8 @@ walk_nodes, ) +from nemo_retriever.graph.pipeline_graph import Graph, Node +from nemo_retriever.operators.abstract_operator import AbstractOperator # --------------------------------------------------------------------------- # Operator stubs (mirrors test_pipeline_graph.py conventions) @@ -99,6 +101,25 @@ def postprocess(self, data: Any, **kw: Any) -> Any: return data +class ParamsContainer(BaseModel): + children: dict[str, LLMRemoteClientParams] + + +class ParamsContainerOp(AbstractOperator): + def __init__(self, params: ParamsContainer) -> None: + super().__init__(params=params) + self.params = params + + def preprocess(self, data: Any, **kw: Any) -> Any: + return data + + def process(self, data: Any, **kw: Any) -> Any: + return data + + def postprocess(self, data: Any, **kw: Any) -> Any: + return data + + # --------------------------------------------------------------------------- # Helpers for building test graphs # --------------------------------------------------------------------------- @@ -167,10 +188,9 @@ def test_primitive_passes_through(self): assert _safe_serialize_value(42) == 42 assert _safe_serialize_value("hello") == "hello" - def test_non_serializable_becomes_repr(self): - result = _safe_serialize_value(object()) - assert isinstance(result, str) - assert "object" in result + def test_non_serializable_raises_instead_of_becoming_repr(self): + with pytest.raises(GraphSerializationError, match="unsupported value"): + _safe_serialize_value(object()) # ===================================================================== @@ -1200,3 +1220,95 @@ def my_factory(): assert callable(my_factory) assert my_factory() is not None assert isinstance(my_factory(), Graph) + + +class TestLegacyRegistryCompatibility: + def test_v1_graph_may_be_named_format_version(self, tmp_path): + payload = { + "format_version": { + "roots": [], + "metadata": {}, + "blueprint": { + "name": "format_version", + "description": "legacy graph name", + }, + }, + } + path = tmp_path / "legacy-registry.json" + path.write_text(json.dumps(payload)) + + registry = GraphPipelineRegistry() + loaded = registry.load_all(path) + + assert loaded == ["format_version"] + assert registry.build("format_version").roots == [] + + +class TestTypedCodecRegressions: + def test_nested_model_container_preserves_no_auth(self, monkeypatch): + monkeypatch.setenv("NVIDIA_API_KEY", "ENV-SECRET") + child = LLMRemoteClientParams( + model="model", + api_key="", + ) + container = ParamsContainer(children={"child": child}) + graph = Graph() + graph.add_root(ParamsContainerOp(container)) + + payload = serialize_graph(graph) + encoded = json.dumps(payload) + restored = deserialize_graph(payload) + restored_child = restored.roots[0].operator.params.children["child"] + + assert "ENV-SECRET" not in encoded + assert "__secret_no_auth__" in encoded + assert restored_child.api_key is None + assert restored_child._uses_no_api_key("api_key") + + def test_function_local_pydantic_model_is_rejected(self): + class LocalModel(BaseModel): + value: int + + with pytest.raises( + GraphSerializationError, + match="Pydantic model .* is not rehydratable", + ): + _safe_serialize_value(LocalModel(value=1)) + + @pytest.mark.parametrize( + "field_name", + ["refresh_token", "owner_token", "apiKey", "accessToken"], + ) + def test_common_token_and_camel_case_secrets_are_rejected(self, field_name): + with pytest.raises( + GraphSerializationError, + match="non-rehydratable secret field", + ): + _safe_serialize_value( + {field_name: "MUST-NOT-LEAK"}, + ) + + def test_flat_operator_api_key_rehydrates_on_worker(self, monkeypatch): + from nemo_retriever.operators.graph_ops.subquery_operator import ( + SubQueryGeneratorOperator, + ) + + original = SubQueryGeneratorOperator(llm_model="model", api_key="ORIGINAL") + graph = Graph() + graph.add_root(original) + payload = serialize_graph(graph) + + assert "ORIGINAL" not in json.dumps(payload) + monkeypatch.setenv("NVIDIA_API_KEY", "WORKER-KEY") + restored = deserialize_graph(payload).roots[0].operator + assert restored._api_key == "WORKER-KEY" + assert restored._resolve_api_key() == "WORKER-KEY" + + monkeypatch.setenv("CUSTOM_LLM_KEY", "CUSTOM-WORKER-KEY") + reference = "os.environ/CUSTOM_LLM_KEY" + referenced = SubQueryGeneratorOperator(llm_model="model", api_key=reference) + referenced_graph = Graph() + referenced_graph.add_root(referenced) + restored_reference = deserialize_graph(serialize_graph(referenced_graph)).roots[0].operator + assert restored_reference._api_key == reference + assert restored_reference._resolve_api_key() == "CUSTOM-WORKER-KEY" From adf009ef086ea46b4f66e186aa39c5695191aabe Mon Sep 17 00:00:00 2001 From: Kyle Zheng <126034466+KyleZheng1284@users.noreply.github.com> Date: Thu, 2 Jul 2026 22:49:33 +0000 Subject: [PATCH 2/3] Harden text generation abstraction and graph persistence Signed-off-by: Kyle Zheng <126034466+KyleZheng1284@users.noreply.github.com> --- docs/docs/extraction/api-keys.md | 12 + .../nemo-retriever-api-reference.md | 26 + .../nemo_retriever/common/params/models.py | 241 +++++- .../graph/graph_pipeline_registry.py | 774 +++++++++++++++--- .../src/nemo_retriever/models/llm/__init__.py | 26 +- .../models/llm/clients/litellm.py | 60 +- .../models/llm/tasks/__init__.py | 3 +- .../nemo_retriever/models/llm/tasks/base.py | 166 +++- .../models/llm/tasks/generic.py | 15 +- .../src/nemo_retriever/models/llm/types.py | 44 +- .../operators/generation/base.py | 100 ++- .../operators/generation/generic.py | 39 +- .../operators/generation/summarization.py | 32 +- .../tools/evaluation/generation.py | 100 ++- .../tests/test_generation_hardening.py | 325 ++++++++ nemo_retriever/tests/test_generation_tasks.py | 244 +++++- .../tests/test_graph_pipeline_registry.py | 744 ++++++++++++++++- nemo_retriever/tests/test_llm_params.py | 274 ++++++- 18 files changed, 2851 insertions(+), 374 deletions(-) create mode 100644 nemo_retriever/tests/test_generation_hardening.py diff --git a/docs/docs/extraction/api-keys.md b/docs/docs/extraction/api-keys.md index 050b4dc01..e6bcd3f8f 100644 --- a/docs/docs/extraction/api-keys.md +++ b/docs/docs/extraction/api-keys.md @@ -27,6 +27,18 @@ For a full list of related variables, see [Environment configuration variables]( The `NVIDIA_API_KEY` from build.nvidia.com is not the same string as your NGC personal key used for Helm and `nvcr.io` access. Do not substitute one for the other unless your tooling explicitly documents that mapping. +## Credential references in persisted graphs + +Persisted pipeline graphs never contain literal API keys. Configure a graph with an explicit worker-side environment reference such as: + +```python +api_key="os.environ/NVIDIA_API_KEY" +``` + +Use the provider's own variable name, for example `os.environ/OPENAI_API_KEY` for an OpenAI model. The reference is stored in graph JSON and resolved only when the operator is constructed or invoked on the worker. + +Literal keys remain available for non-persisted local execution, but attempting to serialize one raises an error. This prevents graph persistence from silently substituting an NVIDIA credential for another provider's key. + ## NGC personal key (Helm and `nvcr.io`) Many public assets on NGC can be used without authentication. For a Kubernetes deployment, the cluster must still pull NIM and microservice images from `nvcr.io` and may need NGC API access; the Helm chart expects credentials derived from an NGC personal key. diff --git a/docs/docs/extraction/nemo-retriever-api-reference.md b/docs/docs/extraction/nemo-retriever-api-reference.md index 246d99dd6..5b9281bba 100644 --- a/docs/docs/extraction/nemo-retriever-api-reference.md +++ b/docs/docs/extraction/nemo-retriever-api-reference.md @@ -8,6 +8,32 @@ To tune splitter throughput from the CLI, use `--pdf-split-batch-size` (Ray acto **Python client (`pdf_split_config`):** Only `create_ingestor(run_mode="service")` implements `.pdf_split_config(pages_per_chunk=...)`, which records page-chunking settings in the request pipeline spec for the remote gateway. Local graph ingest (`run_mode="inprocess"` or `"batch"`) raises `NotImplementedError` if you call this method; PDFs are split automatically on the default ingest path without client-side configuration. +## One-shot text generation + +`TextGenerationOperator` is the reusable base for synchronous, one-request-per-row text generation. It is a provisional text-only API: it does not support tool calls, agent loops, streaming, multiple choices, or structured domain results. + +Concrete operators construct an immutable `GenerationTask` and provide reconstructible constructor state. Runtime task and client objects must not be included in graph constructor arguments. A custom completion client must be safe for concurrent calls or report that it does not support concurrent calls so the operator serializes access. + +Embedding and captioning remain separate operator families because they use modality grouping, native batching, and specialized CPU/GPU lifecycles. + +## Persisted graphs are trusted configuration + +Graph loading imports operator classes and invokes their constructors. Load graph JSON only from trusted sources; do not expose graph payloads, callable references, or class names as model- or user-controlled agent tools. + +Version 2 graph files preserve shared-node DAG identity and reject cycles. Constructor state must consist of supported JSON-native values, typed Pydantic models, paths, sets and tuples, or importable type/callable references. Runtime data such as DataFrames and opaque client objects is not persistable. + +API keys are never written into graph JSON. Use an explicit environment reference in persisted configuration: + +```python +QAGenerationOperator( + model="openai/gpt-4o-mini", + api_key="os.environ/OPENAI_API_KEY", +) +``` + +Serializing a graph containing a literal API key fails with a contextual error instead of guessing which provider credential should be used on a worker. + + ::: nemo_retriever.ingestor options: filters: diff --git a/nemo_retriever/src/nemo_retriever/common/params/models.py b/nemo_retriever/src/nemo_retriever/common/params/models.py index 72e44139c..82b085b71 100644 --- a/nemo_retriever/src/nemo_retriever/common/params/models.py +++ b/nemo_retriever/src/nemo_retriever/common/params/models.py @@ -4,16 +4,25 @@ from __future__ import annotations -from typing import Any, Literal, Optional, Sequence, Tuple -from urllib.parse import urlparse - +import os +import re import warnings +from typing import Any, ClassVar, Literal, Optional, Sequence, Tuple +from urllib.parse import urlparse from upath import UPath from nemo_retriever.tabular_data.sql_database import SQLDatabase -from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, field_validator, model_serializer, model_validator +from pydantic import ( + BaseModel, + ConfigDict, + Field, + PrivateAttr, + field_validator, + model_serializer, + model_validator, +) from nemo_retriever.common.remote_auth import resolve_remote_api_key @@ -26,12 +35,144 @@ _REDACTED = "***" +ENVIRONMENT_REFERENCE_PREFIX = "os.environ/" +_ENVIRONMENT_VARIABLE_PATTERN = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") +PROTECTED_LLM_REQUEST_KEYS = frozenset( + { + "model", + "messages", + "api_key", + "api_base", + "timeout", + "num_retries", + "temperature", + "top_p", + "max_tokens", + "tools", + "tool_choice", + "parallel_tool_calls", + "functions", + "function_call", + "stream", + "n", + } +) + + +def environment_reference_name(value: object) -> Optional[str]: + """Return and validate the variable name in an ``os.environ/NAME`` reference.""" + if not isinstance(value, str): + return None + stripped = value.strip() + if not stripped.startswith(ENVIRONMENT_REFERENCE_PREFIX): + return None + if stripped != value: + raise ValueError("environment references must not contain surrounding whitespace") + name = stripped.removeprefix(ENVIRONMENT_REFERENCE_PREFIX) + if not _ENVIRONMENT_VARIABLE_PATTERN.fullmatch(name): + raise ValueError("environment references must use os.environ/VARIABLE_NAME with a valid variable name") + return name + + +def resolve_environment_reference(value: Optional[str]) -> Optional[str]: + """Resolve an explicit environment reference, leaving literal values unchanged.""" + name = environment_reference_name(value) + if name is None: + return value + resolved = (os.environ.get(name) or "").strip() + if not resolved: + raise ValueError(f"required credential environment variable {name!r} is not set") + return resolved + + +def validate_llm_extra_params(extra_params: Optional[dict[str, Any]], *, source: str) -> None: + """Reject extension parameters that would replace protected request state.""" + if extra_params is None: + return + if not isinstance(extra_params, dict): + raise TypeError(f"{source} must be a dictionary") + protected = sorted(PROTECTED_LLM_REQUEST_KEYS.intersection(extra_params)) + if protected: + raise ValueError(f"{source} may not override protected request fields: {', '.join(protected)}") + def _is_api_key_field(field_name: str) -> bool: """Return True when ``field_name`` should be masked in ``repr`` / logs.""" return field_name == "api_key" or field_name.endswith("_api_key") +def _is_secret_display_field(field_name: str) -> bool: + """Return whether a nested diagnostic field may contain credentials.""" + normalized = field_name.lower().replace("-", "_") + compact = normalized.replace("_", "") + parts = set(normalized.split("_")) + if _is_api_key_field(normalized): + return True + if parts & { + "authorization", + "bearer", + "cookie", + "credentials", + "password", + "passwd", + "secret", + "token", + }: + return True + return ( + "accesskey" in compact + or "accountkey" in compact + or compact.startswith(("authorization", "bearer", "cookie", "credential", "password", "secret")) + or compact.endswith(("token", "privatekey", "secretkey")) + or compact.endswith("apikey") + ) + + +def _redact_param_display( + value: Any, + *, + field_name: Optional[str] = None, + seen: Optional[set[int]] = None, +) -> Any: + """Build a recursively redacted, repr-safe diagnostic value.""" + if field_name is not None and _is_secret_display_field(field_name): + return _REDACTED + if value is None or isinstance(value, (bool, int, float, str)): + return value + if seen is None: + seen = set() + value_id = id(value) + if value_id in seen: + return "" + seen.add(value_id) + if isinstance(value, BaseModel): + return { + name: _redact_param_display( + getattr(value, name), + field_name=name, + seen=seen, + ) + for name in type(value).model_fields + } + if isinstance(value, dict): + redacted: dict[Any, Any] = {} + for key, item in value.items(): + display_key = ( + key + if isinstance(key, (str, int, float, bool)) or key is None + else f"<{type(key).__module__}.{type(key).__qualname__}>" + ) + redacted[display_key] = _redact_param_display( + item, + field_name=key if isinstance(key, str) else None, + seen=seen, + ) + return redacted + if isinstance(value, (list, tuple, set, frozenset)): + return [_redact_param_display(item, seen=seen) for item in value] + return f"<{type(value).__module__}.{type(value).__qualname__}>" + + class _ParamsModel(BaseModel): """Shared base for all remote-transport Pydantic params models. @@ -48,12 +189,15 @@ class _ParamsModel(BaseModel): so no downstream consumer needs changes. """ - model_config = ConfigDict(extra="forbid") + model_config = ConfigDict(extra="forbid", hide_input_in_errors=True) # Keep the explicit no-auth intent after NO_API_KEY is normalized to # None for existing runtime consumers. Graph persistence uses this # private provenance to distinguish no-auth from worker-side env lookup. _no_api_key_fields: set[str] = PrivateAttr(default_factory=set) + _api_key_env_references: dict[str, str] = PrivateAttr(default_factory=dict) + _api_key_env_values: dict[str, str] = PrivateAttr(default_factory=dict) + _auto_resolve_unset_api_keys: ClassVar[bool] = True @model_validator(mode="after") def _resolve_api_keys(self) -> "_ParamsModel": @@ -63,14 +207,55 @@ def _resolve_api_keys(self) -> "_ParamsModel": if value is None: if field_name in self._no_api_key_fields: continue - setattr(self, field_name, resolve_remote_api_key()) + if not type(self)._auto_resolve_unset_api_keys: + self._api_key_env_references.pop(field_name, None) + self._api_key_env_values.pop(field_name, None) + continue + resolved = resolve_remote_api_key() + setattr(self, field_name, resolved) + source_name = next( + (name for name in ("NVIDIA_API_KEY", "NGC_API_KEY") if (os.environ.get(name) or "").strip()), + None, + ) + if resolved and source_name: + self._api_key_env_references[field_name] = f"{ENVIRONMENT_REFERENCE_PREFIX}{source_name}" + self._api_key_env_values[field_name] = resolved + else: + self._api_key_env_references.pop(field_name, None) + self._api_key_env_values.pop(field_name, None) elif value == NO_API_KEY: self._no_api_key_fields.add(field_name) + self._api_key_env_references.pop(field_name, None) + self._api_key_env_values.pop(field_name, None) setattr(self, field_name, None) else: self._no_api_key_fields.discard(field_name) + explicit_reference = environment_reference_name(value) + if explicit_reference is not None: + self._api_key_env_references[field_name] = value + if type(self)._auto_resolve_unset_api_keys: + value = resolve_environment_reference(value) + setattr(self, field_name, value) + self._api_key_env_values[field_name] = value + else: + prior_value = self._api_key_env_values.get(field_name) + if prior_value != value: + self._api_key_env_references.pop(field_name, None) + self._api_key_env_values.pop(field_name, None) return self + def _api_key_env_reference(self, field_name: str) -> Optional[str]: + """Return the exact environment reference that supplied an API key.""" + value = getattr(self, field_name, None) + if isinstance(value, str) and value.startswith(ENVIRONMENT_REFERENCE_PREFIX): + return value + reference = self._api_key_env_references.get(field_name) + if reference is None: + return None + if self._api_key_env_values.get(field_name) != value: + return None + return reference + def _uses_no_api_key(self, field_name: str) -> bool: """Return whether an API-key field was explicitly disabled. @@ -83,12 +268,16 @@ def __repr__(self) -> str: parts: list[str] = [] for field_name in type(self).model_fields: value = getattr(self, field_name, None) - if _is_api_key_field(field_name) and value: - parts.append(f"{field_name}={_REDACTED}") + if _is_api_key_field(field_name): + if value: + parts.append(f"{field_name}={_REDACTED}") + else: + parts.append(f"{field_name}={value!r}") elif field_name == "storage_options" and value: parts.append(f"{field_name}={_REDACTED}") else: - parts.append(f"{field_name}={value!r}") + display_value = _redact_param_display(value, field_name=field_name) + parts.append(f"{field_name}={display_value!r}") return f"{type(self).__name__}({', '.join(parts)})" __str__ = __repr__ @@ -383,8 +572,14 @@ def _auto_enable_features(self) -> "ExtractParams": raise ValueError("ocr_lang is only supported when ocr_version='v2'.") if not self.use_page_elements: consumers = [ - ("use_table_structure", self.use_table_structure and self.extract_tables), - ("use_graphic_elements", self.use_graphic_elements and self.extract_charts), + ( + "use_table_structure", + self.use_table_structure and self.extract_tables, + ), + ( + "use_graphic_elements", + self.use_graphic_elements and self.extract_charts, + ), ] enabled = [name for name, on in consumers if on] if enabled: @@ -431,7 +626,10 @@ class EmbedParams(_ParamsModel): @field_validator("local_ingest_embed_backend", mode="before") @classmethod def _validate_local_ingest_embed_backend(cls, v: str) -> str: - from nemo_retriever.models import _LOCAL_INGEST_EMBED_BACKENDS, normalize_backend + from nemo_retriever.models import ( + _LOCAL_INGEST_EMBED_BACKENDS, + normalize_backend, + ) return normalize_backend( str(v) if v is not None else None, @@ -440,7 +638,12 @@ def _validate_local_ingest_embed_backend(cls, v: str) -> str: default="vllm", ) - @field_validator("embed_modality", "text_elements_modality", "structured_elements_modality", mode="before") + @field_validator( + "embed_modality", + "text_elements_modality", + "structured_elements_modality", + mode="before", + ) @classmethod def _validate_modality(cls, v: str | None) -> str | None: if v is None: @@ -609,10 +812,12 @@ class LLMRemoteClientParams(_ParamsModel): """Transport / connection parameters for any remote LLM client. Pairs with :class:`LLMInferenceParams` (sampling) to fully specify a - call. ``api_key`` is auto-resolved from the environment by - :class:`_ParamsModel` when left as ``None``. + call. ``api_key=None`` is left unset so LiteLLM can perform provider-native + environment lookup on the worker. """ + _auto_resolve_unset_api_keys: ClassVar[bool] = False + model: str api_base: Optional[str] = None api_key: Optional[str] = None @@ -623,6 +828,12 @@ class LLMRemoteClientParams(_ParamsModel): rag_system_prompt_prefix: Optional[str] = None reasoning_enabled: bool = True + @field_validator("extra_params") + @classmethod + def _check_extra_params(cls, value: dict[str, Any]) -> dict[str, Any]: + validate_llm_extra_params(value, source="LLMRemoteClientParams.extra_params") + return value + @field_validator("num_retries") @classmethod def _check_retries(cls, v: int) -> int: diff --git a/nemo_retriever/src/nemo_retriever/graph/graph_pipeline_registry.py b/nemo_retriever/src/nemo_retriever/graph/graph_pipeline_registry.py index 5cba747d5..342afd47a 100644 --- a/nemo_retriever/src/nemo_retriever/graph/graph_pipeline_registry.py +++ b/nemo_retriever/src/nemo_retriever/graph/graph_pipeline_registry.py @@ -31,6 +31,7 @@ def _build(): import importlib import json +import re from collections import OrderedDict from dataclasses import dataclass, field from datetime import datetime, timezone @@ -50,7 +51,6 @@ def _build(): from pydantic import BaseModel -from nemo_retriever.common.remote_auth import resolve_remote_api_key from nemo_retriever.graph.pipeline_graph import Graph, Node from nemo_retriever.operators.abstract_operator import AbstractOperator @@ -81,31 +81,46 @@ def _import_class(qualified: str) -> type: _PYDANTIC_MODEL_MARKER = "__pydantic_model__" _PYDANTIC_FIELDS = "fields" _PYDANTIC_FIELDS_SET = "fields_set" -_SECRET_ENV_MARKER = "__secret_env__" _SECRET_NO_AUTH_MARKER = "__secret_no_auth__" _TUPLE_MARKER = "__tuple__" _FROZENSET_MARKER = "__frozenset__" _MAPPING_MARKER = "__mapping__" _OMIT_FIELD = object() +_ENV_REFERENCE_PREFIX = "os.environ/" +_ENV_VARIABLE_RE = re.compile(r"[A-Za-z_][A-Za-z0-9_]*\Z") class GraphSerializationError(ValueError): """Raised when graph state cannot be serialized safely and losslessly.""" +def _normalize_field_name(field_name: str) -> str: + snake_case = re.sub(r"(?<=[a-z0-9])(?=[A-Z])", "_", field_name) + return snake_case.replace("-", "_").lower() + + def _is_api_key_field(field_name: Optional[str]) -> bool: - return bool(field_name) and (field_name == "api_key" or field_name.endswith("_api_key")) + if not field_name: + return False + normalized = _normalize_field_name(field_name) + compact = normalized.replace("_", "") + return normalized == "api_key" or normalized.endswith("_api_key") or compact.endswith("apikey") def _is_obvious_secret_field(field_name: Optional[str]) -> bool: if not field_name: return False - normalized = field_name.lower().replace("-", "_") + normalized = _normalize_field_name(field_name) compact = normalized.replace("_", "") if _is_api_key_field(normalized): return True if normalized in { + "access_key", + "account_key", "authorization", + "bearer", + "cookie", + "cookies", "credential", "credentials", "password", @@ -113,27 +128,70 @@ def _is_obvious_secret_field(field_name: Optional[str]) -> bool: "private_key", "secret", "secret_key", - "storage_options", }: return True - if set(normalized.split("_")) & {"password", "passwd", "secret"}: + parts = set(normalized.split("_")) + if parts & { + "authorization", + "bearer", + "cookie", + "cookies", + "credential", + "credentials", + "password", + "passwd", + "secret", + "token", + }: + return True + if {"access", "key"} <= parts or {"account", "key"} <= parts: return True if normalized == "token" or normalized.endswith("_token"): return True if compact in { + "accesskey", + "accountkey", "authorization", + "bearer", + "cookie", + "cookies", "credential", "credentials", - "storageoptions", }: return True - return compact.endswith(("apikey", "password", "passwd", "secret", "secretkey", "privatekey", "token")) + if "accesskey" in compact or "accountkey" in compact: + return True + return compact.endswith( + ( + "apikey", + "password", + "passwd", + "secret", + "secretkey", + "privatekey", + "token", + "cookie", + "bearer", + "credential", + ) + ) + + +def _is_storage_option_secret_field(field_name: Optional[str]) -> bool: + if not field_name: + return False + normalized = _normalize_field_name(field_name) + return normalized == "key" or _is_obvious_secret_field(field_name) def _is_empty_secret(value: Any) -> bool: - if value is None or value == "": + if value is None: return True - return isinstance(value, (dict, list, tuple, set, frozenset)) and not value + if isinstance(value, str): + return value == "" + if isinstance(value, (dict, list, tuple, set, frozenset)): + return len(value) == 0 + return False def _import_qualified_object(qualified: str) -> Any: @@ -154,13 +212,57 @@ def _import_qualified_object(qualified: str) -> Any: raise ImportError(f"Cannot import qualified object {qualified!r}") -def _model_uses_no_api_key(model: BaseModel, field_name: str) -> bool: +def _model_uses_no_api_key(model: BaseModel, field_name: str, *, path: str) -> bool: checker = getattr(model, "_uses_no_api_key", None) if callable(checker): - return bool(checker(field_name)) + provenance_error: Optional[GraphSerializationError] = None + try: + uses_no_auth = bool(checker(field_name)) + except Exception as exc: + exc_type = f"{type(exc).__module__}.{type(exc).__qualname__}" + provenance_error = GraphSerializationError(f"{path}: no-auth credential provenance failed ({exc_type})") + if provenance_error is not None: + raise provenance_error + return uses_no_auth return field_name in getattr(model, "_no_api_key_fields", set()) +def _normalize_env_reference(value: str, *, path: str) -> str: + """Return a canonical os.environ/NAME reference or raise safely.""" + candidate = value.strip() + if not candidate.startswith(_ENV_REFERENCE_PREFIX): + raise GraphSerializationError( + f"{path}: API keys cannot be persisted literally; use " f"{_ENV_REFERENCE_PREFIX}VARIABLE_NAME" + ) + variable = candidate.removeprefix(_ENV_REFERENCE_PREFIX) + if not _ENV_VARIABLE_RE.fullmatch(variable): + raise GraphSerializationError( + f"{path}: invalid environment reference; expected " f"{_ENV_REFERENCE_PREFIX}VARIABLE_NAME" + ) + return f"{_ENV_REFERENCE_PREFIX}{variable}" + + +def _model_api_key_env_reference(model: BaseModel, field_name: str, *, path: str) -> Optional[str]: + """Read optional credential provenance without depending on a model base.""" + getter = getattr(model, "_api_key_env_reference", None) + if not callable(getter): + return None + provenance_error: Optional[GraphSerializationError] = None + try: + reference = getter(field_name) + except Exception as exc: + exc_type = f"{type(exc).__module__}.{type(exc).__qualname__}" + provenance_error = GraphSerializationError(f"{path}: API-key environment provenance failed ({exc_type})") + reference = None + if provenance_error is not None: + raise provenance_error + if reference is None: + return None + if not isinstance(reference, str): + raise GraphSerializationError(f"{path}: API-key environment provenance must be a string or null") + return _normalize_env_reference(reference, path=path) + + def _contains_pydantic_model(value: Any, seen: Optional[Set[int]] = None) -> bool: if isinstance(value, BaseModel): return True @@ -188,10 +290,17 @@ def _encode_secret( allow_api_key_env: bool, ) -> Any: if _is_api_key_field(field_name): - if value == "" or (value is None and owner is not None and _model_uses_no_api_key(owner, field_name)): + if (isinstance(value, str) and value == "") or ( + value is None and owner is not None and _model_uses_no_api_key(owner, field_name, path=path) + ): return {_SECRET_NO_AUTH_MARKER: ""} - if owner is None and allow_api_key_env and isinstance(value, str) and value.strip().startswith("os.environ/"): - return value.strip() + if owner is not None: + reference = _model_api_key_env_reference(owner, field_name, path=path) + if reference is not None: + return reference + if isinstance(value, str) and value.strip().startswith(_ENV_REFERENCE_PREFIX): + if owner is not None or allow_api_key_env: + return _normalize_env_reference(value, path=path) if owner is None and not allow_api_key_env: if value is None: return None @@ -201,11 +310,14 @@ def _encode_secret( ) if value is not None and not isinstance(value, str): raise GraphSerializationError(f"{path}: API-key fields must be strings, null, or the no-auth marker") - scope = "model" if owner is not None else "operator" - return {_SECRET_ENV_MARKER: scope} - if not _is_empty_secret(value): - raise GraphSerializationError(f"{path}: refusing to serialize non-rehydratable secret field {field_name!r}") - return value + if value is None: + return None + raise GraphSerializationError( + f"{path}: API keys cannot be persisted literally; use " f"{_ENV_REFERENCE_PREFIX}VARIABLE_NAME" + ) + if value is None or (isinstance(value, str) and value == ""): + return value + raise GraphSerializationError(f"{path}: refusing to serialize non-rehydratable secret field {field_name!r}") def _encode_value( @@ -215,9 +327,12 @@ def _encode_value( field_name: Optional[str] = None, owner: Optional[BaseModel] = None, allow_api_key_env: bool = False, + inside_storage_options: bool = False, ) -> Any: """Recursively encode ``value`` into lossless, JSON-native graph state.""" - if _is_obvious_secret_field(field_name): + normalized_field = _normalize_field_name(field_name) if field_name else None + inside_storage_options = inside_storage_options or normalized_field == "storage_options" + if _is_obvious_secret_field(field_name) or (inside_storage_options and _is_storage_option_secret_field(field_name)): return _encode_secret( value, field_name=field_name or "", @@ -236,7 +351,15 @@ def _encode_value( raise GraphSerializationError(f"{path}: Pydantic model {qualified!r} is not rehydratable") from exc if restored_type is not model_type: raise GraphSerializationError(f"{path}: Pydantic model {qualified!r} does not round-trip by identity") - dumped = value.model_dump(mode="python") + serialization_error: Optional[GraphSerializationError] = None + try: + dumped = value.model_dump(mode="python") + except Exception as exc: + exc_type = f"{type(exc).__module__}.{type(exc).__qualname__}" + serialization_error = GraphSerializationError(f"{path}: Pydantic model serialization failed ({exc_type})") + dumped = None + if serialization_error is not None: + raise serialization_error if not isinstance(dumped, dict): raise GraphSerializationError( f"{path}: Pydantic model serializer must return a mapping, got {type(dumped).__name__}" @@ -250,6 +373,7 @@ def _encode_value( path=f"{path}.{name}", field_name=name, owner=value, + inside_storage_options=inside_storage_options, ) return { _PYDANTIC_MODEL_MARKER: qualified, @@ -278,11 +402,34 @@ def _encode_value( if isinstance(value, Path): return {"__path__": str(value)} if isinstance(value, tuple): - return {_TUPLE_MARKER: [_encode_value(item, path=f"{path}[{index}]") for index, item in enumerate(value)]} + return { + _TUPLE_MARKER: [ + _encode_value( + item, + path=f"{path}[{index}]", + inside_storage_options=inside_storage_options, + ) + for index, item in enumerate(value) + ] + } if isinstance(value, list): - return [_encode_value(item, path=f"{path}[{index}]") for index, item in enumerate(value)] + return [ + _encode_value( + item, + path=f"{path}[{index}]", + inside_storage_options=inside_storage_options, + ) + for index, item in enumerate(value) + ] if isinstance(value, (set, frozenset)): - encoded = [_encode_value(item, path=f"{path}[{index}]") for index, item in enumerate(value)] + encoded = [ + _encode_value( + item, + path=f"{path}[{index}]", + inside_storage_options=inside_storage_options, + ) + for index, item in enumerate(value) + ] encoded.sort(key=lambda item: json.dumps(item, sort_keys=True)) marker = _FROZENSET_MARKER if isinstance(value, frozenset) else "__set__" return {marker: encoded} @@ -290,13 +437,15 @@ def _encode_value( encoded_dict: Dict[str, Any] = {} for key, item in value.items(): if not isinstance(key, str): + key_type = f"{type(key).__module__}.{type(key).__qualname__}" raise GraphSerializationError( - f"{path}: mapping key {key!r} is not a string and cannot round-trip through JSON" + f"{path}: mapping key of type {key_type} is not a string and cannot round-trip through JSON" ) encoded_dict[key] = _encode_value( item, path=f"{path}.{key}", field_name=key, + inside_storage_options=inside_storage_options, ) return {_MAPPING_MARKER: encoded_dict} raise GraphSerializationError( @@ -304,41 +453,80 @@ def _encode_value( ) +def _validate_value_envelope( + value: dict, + *, + required: Set[str], + optional: Set[str] = frozenset(), + path: str, + label: str, +) -> None: + keys = set(value) + if not required <= keys or not keys <= required | optional: + raise GraphSerializationError(f"{path}: malformed {label} envelope") + + def _decode_value( value: Any, *, path: str, format_version: int, field_name: Optional[str] = None, + inside_storage_options: bool = False, ) -> Any: """Recursively restore graph state encoded by v2 or accepted v1 markers.""" + normalized_field = _normalize_field_name(field_name) if field_name else None + inside_storage_options = inside_storage_options or normalized_field == "storage_options" + is_api_key = _is_api_key_field(field_name) + is_secret = _is_obvious_secret_field(field_name) or ( + inside_storage_options and _is_storage_option_secret_field(field_name) + ) + if format_version >= 2 and is_secret and not is_api_key: + if value is None or (isinstance(value, str) and value == ""): + return value + raise GraphSerializationError(f"{path}: serialized secret fields must be empty") + if format_version >= 2 and is_api_key: + if value is None: + return None + if isinstance(value, str): + return _normalize_env_reference(value, path=path) + if not (isinstance(value, dict) and _SECRET_NO_AUTH_MARKER in value): + raise GraphSerializationError(f"{path}: invalid encoded API-key value") if value is None or isinstance(value, (bool, int, float, str)): return value if isinstance(value, list): return [ - _decode_value(item, path=f"{path}[{index}]", format_version=format_version) + _decode_value( + item, + path=f"{path}[{index}]", + format_version=format_version, + inside_storage_options=inside_storage_options, + ) for index, item in enumerate(value) ] if not isinstance(value, dict): if format_version == 1: return value raise GraphSerializationError(f"{path}: expected JSON-native graph state") - if format_version >= 2 and _SECRET_ENV_MARKER in value: - if not _is_api_key_field(field_name): - raise GraphSerializationError(f"{path}: secret environment marker is outside an API-key field") - scope = value[_SECRET_ENV_MARKER] - if scope == "model": - return _OMIT_FIELD - if scope == "operator": - return resolve_remote_api_key() - raise GraphSerializationError(f"{path}: invalid API-key environment marker") if format_version >= 2 and _SECRET_NO_AUTH_MARKER in value: + _validate_value_envelope( + value, + required={_SECRET_NO_AUTH_MARKER}, + path=path, + label="no-auth API-key", + ) if not _is_api_key_field(field_name) or value[_SECRET_NO_AUTH_MARKER] != "": raise GraphSerializationError(f"{path}: invalid no-auth API-key marker") return "" if format_version >= 2 and _MAPPING_MARKER in value: + _validate_value_envelope( + value, + required={_MAPPING_MARKER}, + path=path, + label="mapping", + ) items = value[_MAPPING_MARKER] - if not isinstance(items, dict): + if not isinstance(items, dict) or not all(isinstance(key, str) for key in items): raise GraphSerializationError(f"{path}: malformed mapping envelope") decoded_mapping: Dict[str, Any] = {} for key, item in items.items(): @@ -347,15 +535,27 @@ def _decode_value( path=f"{path}.{key}", format_version=format_version, field_name=key, + inside_storage_options=inside_storage_options, ) if decoded is not _OMIT_FIELD: decoded_mapping[key] = decoded return decoded_mapping if format_version >= 2 and _PYDANTIC_MODEL_MARKER in value: + _validate_value_envelope( + value, + required={_PYDANTIC_MODEL_MARKER, _PYDANTIC_FIELDS}, + optional={_PYDANTIC_FIELDS_SET}, + path=path, + label="Pydantic model", + ) qualified = value[_PYDANTIC_MODEL_MARKER] - fields = value.get(_PYDANTIC_FIELDS) + fields = value[_PYDANTIC_FIELDS] fields_set = value.get(_PYDANTIC_FIELDS_SET, []) - if not isinstance(qualified, str) or not isinstance(fields, dict): + if ( + not isinstance(qualified, str) + or not isinstance(fields, dict) + or not all(isinstance(name, str) for name in fields) + ): raise GraphSerializationError(f"{path}: malformed Pydantic model envelope") try: model_cls = _import_qualified_object(qualified) @@ -370,15 +570,21 @@ def _decode_value( path=f"{path}.{name}", format_version=format_version, field_name=name, + inside_storage_options=inside_storage_options, ) if decoded is not _OMIT_FIELD: decoded_fields[name] = decoded + validation_error: Optional[GraphSerializationError] = None try: model = model_cls.model_validate(decoded_fields) except Exception as exc: - raise GraphSerializationError( - f"{path}: failed to validate restored Pydantic model {qualified!r}: {exc}" - ) from exc + exc_type = f"{type(exc).__module__}.{type(exc).__qualname__}" + validation_error = GraphSerializationError( + f"{path}: failed to validate restored Pydantic model {qualified!r} ({exc_type})" + ) + model = None + if validation_error is not None: + raise validation_error if not isinstance(fields_set, list) or not all(isinstance(name, str) for name in fields_set): raise GraphSerializationError(f"{path}: malformed Pydantic fields_set") unknown = set(fields_set) - set(type(model).model_fields) @@ -387,7 +593,16 @@ def _decode_value( model.__pydantic_fields_set__ = set(fields_set) return model if "__type_ref__" in value: + if format_version >= 2: + _validate_value_envelope( + value, + required={"__type_ref__"}, + path=path, + label="type reference", + ) qualified = value["__type_ref__"] + if format_version >= 2 and not isinstance(qualified, str): + raise GraphSerializationError(f"{path}: malformed type-reference envelope") try: restored = _import_qualified_object(qualified) except (ImportError, TypeError) as exc: @@ -400,7 +615,16 @@ def _decode_value( raise GraphSerializationError(f"{path}: type reference {qualified!r} is not a type") return restored if "__callable_ref__" in value: + if format_version >= 2: + _validate_value_envelope( + value, + required={"__callable_ref__"}, + path=path, + label="callable reference", + ) qualified = value["__callable_ref__"] + if format_version >= 2 and not isinstance(qualified, str): + raise GraphSerializationError(f"{path}: malformed callable-reference envelope") try: restored = _import_qualified_object(qualified) except (ImportError, TypeError) as exc: @@ -413,28 +637,80 @@ def _decode_value( raise GraphSerializationError(f"{path}: callable reference {qualified!r} is not callable") return restored if "__path__" in value: + if format_version >= 2: + _validate_value_envelope( + value, + required={"__path__"}, + path=path, + label="path", + ) + if not isinstance(value["__path__"], str): + raise GraphSerializationError(f"{path}: malformed path envelope") return Path(value["__path__"]) if "__set__" in value: + if format_version >= 2: + _validate_value_envelope( + value, + required={"__set__"}, + path=path, + label="set", + ) + if not isinstance(value["__set__"], list): + raise GraphSerializationError(f"{path}: malformed set envelope") return { - _decode_value(item, path=f"{path}[{index}]", format_version=format_version) + _decode_value( + item, + path=f"{path}[{index}]", + format_version=format_version, + inside_storage_options=inside_storage_options, + ) for index, item in enumerate(value["__set__"]) } if format_version >= 2 and _FROZENSET_MARKER in value: + _validate_value_envelope( + value, + required={_FROZENSET_MARKER}, + path=path, + label="frozenset", + ) + if not isinstance(value[_FROZENSET_MARKER], list): + raise GraphSerializationError(f"{path}: malformed frozenset envelope") return frozenset( - _decode_value(item, path=f"{path}[{index}]", format_version=format_version) + _decode_value( + item, + path=f"{path}[{index}]", + format_version=format_version, + inside_storage_options=inside_storage_options, + ) for index, item in enumerate(value[_FROZENSET_MARKER]) ) if format_version >= 2 and _TUPLE_MARKER in value: + _validate_value_envelope( + value, + required={_TUPLE_MARKER}, + path=path, + label="tuple", + ) + if not isinstance(value[_TUPLE_MARKER], list): + raise GraphSerializationError(f"{path}: malformed tuple envelope") return tuple( - _decode_value(item, path=f"{path}[{index}]", format_version=format_version) + _decode_value( + item, + path=f"{path}[{index}]", + format_version=format_version, + inside_storage_options=inside_storage_options, + ) for index, item in enumerate(value[_TUPLE_MARKER]) ) + if format_version >= 2: + raise GraphSerializationError(f"{path}: unrecognized encoded mapping") return { key: _decode_value( item, path=f"{path}.{key}", format_version=format_version, field_name=key, + inside_storage_options=inside_storage_options, ) for key, item in value.items() } @@ -533,11 +809,17 @@ def _redact_display_value( *, field_name: Optional[str] = None, seen: Optional[Set[int]] = None, + inside_storage_options: bool = False, ) -> Any: """Return a recursively redacted copy suitable only for diagnostics.""" - if _is_obvious_secret_field(field_name) and not _is_empty_secret(value): + normalized_field = _normalize_field_name(field_name) if field_name else None + inside_storage_options = inside_storage_options or normalized_field == "storage_options" + is_secret = _is_obvious_secret_field(field_name) or ( + inside_storage_options and _is_storage_option_secret_field(field_name) + ) + if is_secret and not _is_empty_secret(value): return "***" - if value is None or isinstance(value, (bool, int, float, str, bytes)): + if value is None or isinstance(value, (bool, int, float, str)): return value if seen is None: @@ -553,17 +835,48 @@ def _redact_display_value( getattr(value, name), field_name=name, seen=seen, + inside_storage_options=inside_storage_options, ) for name in type(value).model_fields } if isinstance(value, dict): - return {key: _redact_display_value(item, field_name=str(key), seen=seen) for key, item in value.items()} + redacted_mapping: Dict[Any, Any] = {} + for key, item in value.items(): + if isinstance(key, (str, int, float, bool)) or key is None: + display_key = key + nested_field = key if isinstance(key, str) else None + else: + key_type = type(key) + display_key = f"<{key_type.__module__}.{key_type.__qualname__}>" + nested_field = None + redacted_mapping[display_key] = _redact_display_value( + item, + field_name=nested_field, + seen=seen, + inside_storage_options=inside_storage_options, + ) + return redacted_mapping if isinstance(value, (list, tuple)): - redacted = [_redact_display_value(item, seen=seen) for item in value] + redacted = [ + _redact_display_value( + item, + seen=seen, + inside_storage_options=inside_storage_options, + ) + for item in value + ] return tuple(redacted) if isinstance(value, tuple) else redacted if isinstance(value, (set, frozenset)): - return [_redact_display_value(item, seen=seen) for item in value] - return value + return [ + _redact_display_value( + item, + seen=seen, + inside_storage_options=inside_storage_options, + ) + for item in value + ] + value_type = type(value) + return f"<{value_type.__module__}.{value_type.__qualname__}>" def _display_repr(field_name: str, value: Any) -> str: @@ -629,7 +942,12 @@ def _render(node: Node, prefix: str, is_last: bool, is_root: bool) -> None: child_prefix = prefix + ("" if is_root else (" " if is_last else "│ ")) for i, child in enumerate(node.children): - _render(child, child_prefix, is_last=(i == len(node.children) - 1), is_root=False) + _render( + child, + child_prefix, + is_last=(i == len(node.children) - 1), + is_root=False, + ) for i, root in enumerate(graph.roots): if i > 0: @@ -888,10 +1206,14 @@ def _diff_kwargs(kwargs_a: dict, kwargs_b: dict) -> Tuple[dict, dict, dict]: elif in_b and not in_a: added[key] = kwargs_b[key] else: - try: - equal = kwargs_a[key] == kwargs_b[key] - except Exception: - equal = repr(kwargs_a[key]) == repr(kwargs_b[key]) + if kwargs_a[key] is kwargs_b[key]: + equal = True + else: + try: + comparison = kwargs_a[key] == kwargs_b[key] + equal = bool(comparison) + except Exception: + equal = False if not equal: changed[key] = (kwargs_a[key], kwargs_b[key]) return added, removed, changed @@ -989,8 +1311,22 @@ def print_diff(graph_a: Graph, graph_b: Graph) -> None: # --------------------------------------------------------------------------- -def _serialize_node(node: Node, *, path: str) -> dict: - """Serialize a single node to a JSON-compatible dict.""" +def _serialize_operator_state(node: Node, *, path: str) -> dict: + """Serialize one node's operator state without traversing its edges.""" + if not isinstance(node.name, str): + raise GraphSerializationError(f"{path}.name: expected a string") + if not isinstance(node.operator_class, type) or not issubclass(node.operator_class, AbstractOperator): + raise GraphSerializationError(f"{path}.operator_class: expected an AbstractOperator class") + qualified = _qualified_name(node.operator_class) + try: + restored_class = _import_class(qualified) + except (ImportError, TypeError) as exc: + raise GraphSerializationError(f"{path}.operator_class: {qualified!r} is not importable") from exc + if restored_class is not node.operator_class: + raise GraphSerializationError(f"{path}.operator_class: {qualified!r} does not round-trip by identity") + if not isinstance(node.operator_kwargs, dict): + raise GraphSerializationError(f"{path}.operator_kwargs: expected a mapping") + safe_kwargs: Dict[str, Any] = {} for key, value in node.operator_kwargs.items(): if not isinstance(key, str): @@ -1003,31 +1339,89 @@ def _serialize_node(node: Node, *, path: str) -> dict: ) return { "name": node.name, - "operator_class": _qualified_name(node.operator_class), + "operator_class": qualified, "operator_kwargs": safe_kwargs, - "children": [ - _serialize_node(child, path=f"{path}.children[{index}]") for index, child in enumerate(node.children) - ], + "children": [], } +def _normalized_max_depth(root_ids: Sequence[str], nodes: Dict[str, dict]) -> int: + if not root_ids: + return 0 + memo: Dict[str, int] = {} + + def _depth(node_id: str) -> int: + if node_id not in memo: + children = nodes[node_id]["children"] + memo[node_id] = 0 if not children else 1 + max(_depth(child_id) for child_id in children) + return memo[node_id] + + return max(_depth(root_id) for root_id in root_ids) + + def serialize_graph(graph: Graph) -> dict: - """Serialize a graph to a versioned, recursively JSON-native dictionary.""" + """Serialize a DAG to deterministic, normalized version-2 graph state.""" + if not isinstance(graph, Graph): + raise GraphSerializationError(f"graph must be a Graph, got {type(graph).__name__}") + + object_ids: Dict[int, str] = {} + nodes: Dict[str, dict] = {} + active: List[int] = [] + complete: Set[int] = set() + + def _visit(node: Node, *, path: str) -> str: + if not isinstance(node, Node): + raise GraphSerializationError(f"{path}: expected a Node, got {type(node).__name__}") + object_id = id(node) + node_id = object_ids.get(object_id) + if node_id is None: + node_id = f"node_{len(object_ids)}" + object_ids[object_id] = node_id + nodes[node_id] = {} + + if object_id in active: + cycle_start = active.index(object_id) + cycle = [object_ids[item] for item in active[cycle_start:]] + [node_id] + raise GraphSerializationError(f"{path}: graph cycle detected: {' -> '.join(cycle)}") + if object_id in complete: + return node_id + + active.append(object_id) + state = _serialize_operator_state(node, path=f"nodes.{node_id}") + child_ids = [ + _visit(child, path=f"nodes.{node_id}.children[{index}]") for index, child in enumerate(node.children) + ] + if len(set(child_ids)) != len(child_ids): + raise GraphSerializationError(f"nodes.{node_id}.children: duplicate child node IDs are not allowed") + state["children"] = child_ids + nodes[node_id] = state + active.pop() + complete.add(object_id) + return node_id + + root_ids = [_visit(root, path=f"roots[{index}]") for index, root in enumerate(graph.roots)] + if len(set(root_ids)) != len(root_ids): + raise GraphSerializationError("roots: duplicate root node IDs are not allowed") return { "format_version": _GRAPH_FORMAT_VERSION, - "roots": [_serialize_node(root, path=f"roots[{index}]({root.name})") for index, root in enumerate(graph.roots)], + "roots": root_ids, + "nodes": nodes, "metadata": { - "node_count": node_count(graph), - "max_depth": max_depth(graph), - "serialized_at": datetime.now(timezone.utc).isoformat(), + "node_count": len(nodes), + "max_depth": _normalized_max_depth(root_ids, nodes), }, } class _PlaceholderOperator(AbstractOperator): - """Stand-in used when the real operator class cannot be instantiated during deserialization.""" + """Stand-in used when a legacy operator cannot be instantiated.""" - def __init__(self, original_class: str = "", original_kwargs: Optional[dict] = None, **kwargs: Any) -> None: + def __init__( + self, + original_class: str = "", + original_kwargs: Optional[dict] = None, + **kwargs: Any, + ) -> None: super().__init__(**kwargs) self._original_class = original_class self._original_kwargs = original_kwargs or {} @@ -1065,38 +1459,158 @@ def _restore_special_values( return cleaned -def _deserialize_node(data: dict, *, format_version: int, path: str) -> Node: - """Reconstruct a :class:`Node` from its serialized dict.""" - cls = _import_class(data["operator_class"]) +def _deserialize_v1_node(data: Any, *, path: str) -> Node: + """Reconstruct one nested node from the versionless legacy format.""" + if not isinstance(data, dict): + raise GraphSerializationError(f"{path}: expected a node mapping") + qualified = data.get("operator_class") + if not isinstance(qualified, str): + raise GraphSerializationError(f"{path}.operator_class: expected a string") + try: + cls = _import_class(qualified) + except (ImportError, TypeError) as exc: + raise GraphSerializationError(f"{path}.operator_class: {qualified!r} is not importable") from exc raw_kwargs = data.get("operator_kwargs", {}) if not isinstance(raw_kwargs, dict): raise GraphSerializationError(f"{path}.operator_kwargs: expected a mapping") - cleaned = _restore_special_values( - raw_kwargs, - format_version=format_version, - path=f"{path}.operator_kwargs", - ) + cleaned = _restore_special_values(raw_kwargs, format_version=1, path=f"{path}.operator_kwargs") try: op = cls(**cleaned) - except Exception as exc: - if format_version >= 2: - raise GraphSerializationError( - f"{path}: failed to construct operator " f"{data['operator_class']!r}: {exc}" - ) from exc - op = _PlaceholderOperator(original_class=data["operator_class"], original_kwargs=cleaned) + except Exception: + op = _PlaceholderOperator(original_class=qualified, original_kwargs=cleaned) node = Node(op, name=data.get("name"), operator_class=cls, operator_kwargs=cleaned) - for index, child_data in enumerate(data.get("children", [])): - child_node = _deserialize_node( - child_data, - format_version=format_version, - path=f"{path}.children[{index}]", - ) - node.children.append(child_node) + children = data.get("children", []) + if not isinstance(children, list): + raise GraphSerializationError(f"{path}.children: expected a list") + for index, child_data in enumerate(children): + node.children.append(_deserialize_v1_node(child_data, path=f"{path}.children[{index}]")) return node +def _validate_v2_topology(data: dict) -> Tuple[List[str], Dict[str, dict]]: + required_top_level = {"format_version", "roots", "nodes"} + allowed_top_level = required_top_level | {"metadata"} + top_level_keys = set(data) + if not required_top_level <= top_level_keys: + raise GraphSerializationError("version 2 graph is missing required top-level fields") + if not top_level_keys <= allowed_top_level: + raise GraphSerializationError("version 2 graph contains unknown top-level fields") + + metadata = data.get("metadata") + if metadata is not None: + if not isinstance(metadata, dict) or set(metadata) != {"node_count", "max_depth"}: + raise GraphSerializationError("metadata: expected node_count and max_depth") + if any(isinstance(value, bool) or not isinstance(value, int) or value < 0 for value in metadata.values()): + raise GraphSerializationError("metadata: node_count and max_depth must be non-negative integers") + + roots = data.get("roots") + raw_nodes = data.get("nodes") + if not isinstance(roots, list): + raise GraphSerializationError("roots: version 2 requires a list of node IDs") + if not all(isinstance(node_id, str) and node_id for node_id in roots): + raise GraphSerializationError("roots: every node ID must be a non-empty string") + if len(set(roots)) != len(roots): + raise GraphSerializationError("roots: duplicate root node IDs are not allowed") + if not isinstance(raw_nodes, dict): + raise GraphSerializationError("nodes: version 2 requires a node-ID mapping") + + required_fields = {"name", "operator_class", "operator_kwargs", "children"} + nodes: Dict[str, dict] = {} + for node_id, record in raw_nodes.items(): + path = f"nodes.{node_id}" + if not isinstance(node_id, str) or not node_id: + raise GraphSerializationError("nodes: every node ID must be a non-empty string") + if not isinstance(record, dict): + raise GraphSerializationError(f"{path}: expected a node mapping") + missing = required_fields - set(record) + unknown = set(record) - required_fields + if missing: + raise GraphSerializationError(f"{path}: missing fields: {sorted(missing)}") + if unknown: + raise GraphSerializationError(f"{path}: unknown fields: {sorted(unknown)}") + if not isinstance(record["name"], str): + raise GraphSerializationError(f"{path}.name: expected a string") + if not isinstance(record["operator_class"], str): + raise GraphSerializationError(f"{path}.operator_class: expected a string") + if not isinstance(record["operator_kwargs"], dict): + raise GraphSerializationError(f"{path}.operator_kwargs: expected a mapping") + children = record["children"] + if not isinstance(children, list) or not all(isinstance(child_id, str) and child_id for child_id in children): + raise GraphSerializationError(f"{path}.children: expected a list of non-empty node IDs") + if len(set(children)) != len(children): + raise GraphSerializationError(f"{path}.children: duplicate child node IDs are not allowed") + nodes[node_id] = record + + for index, root_id in enumerate(roots): + if root_id not in nodes: + raise GraphSerializationError(f"roots[{index}]: unknown node ID {root_id!r}") + for node_id, record in nodes.items(): + for index, child_id in enumerate(record["children"]): + if child_id not in nodes: + raise GraphSerializationError(f"nodes.{node_id}.children[{index}]: unknown node ID {child_id!r}") + + state: Dict[str, int] = {} + reachable: Set[str] = set() + stack: List[str] = [] + + def _visit(node_id: str) -> None: + marker = state.get(node_id, 0) + if marker == 1: + cycle_start = stack.index(node_id) + cycle = stack[cycle_start:] + [node_id] + raise GraphSerializationError(f"nodes: graph cycle detected: {' -> '.join(cycle)}") + if marker == 2: + return + state[node_id] = 1 + reachable.add(node_id) + stack.append(node_id) + for child_id in nodes[node_id]["children"]: + _visit(child_id) + stack.pop() + state[node_id] = 2 + + for root_id in roots: + _visit(root_id) + unreachable = sorted(set(nodes) - reachable) + if unreachable: + raise GraphSerializationError(f"nodes: unreachable node IDs: {unreachable}") + if metadata is not None: + if metadata["node_count"] != len(nodes): + raise GraphSerializationError("metadata.node_count does not match the normalized graph") + if metadata["max_depth"] != _normalized_max_depth(roots, nodes): + raise GraphSerializationError("metadata.max_depth does not match the normalized graph") + return roots, nodes + + +def _instantiate_v2_node(node_id: str, record: dict) -> Node: + path = f"nodes.{node_id}" + qualified = record["operator_class"] + try: + cls = _import_class(qualified) + except (ImportError, TypeError) as exc: + raise GraphSerializationError(f"{path}.operator_class: {qualified!r} is not importable") from exc + if not isinstance(cls, type) or not issubclass(cls, AbstractOperator): + raise GraphSerializationError(f"{path}.operator_class: {qualified!r} is not an AbstractOperator class") + + cleaned = _restore_special_values( + record["operator_kwargs"], + format_version=_GRAPH_FORMAT_VERSION, + path=f"{path}.operator_kwargs", + ) + construction_error: Optional[GraphSerializationError] = None + try: + op = cls(**cleaned) + except Exception as exc: + exc_type = f"{type(exc).__module__}.{type(exc).__qualname__}" + construction_error = GraphSerializationError(f"{path}: failed to construct operator {qualified!r} ({exc_type})") + op = None + if construction_error is not None: + raise construction_error + return Node(op, name=record["name"], operator_class=cls, operator_kwargs=cleaned) + + def _read_format_version(data: dict) -> int: version = data.get("format_version", 1) if isinstance(version, bool) or not isinstance(version, int): @@ -1107,18 +1621,27 @@ def _read_format_version(data: dict) -> int: def deserialize_graph(data: dict) -> Graph: - """Reconstruct a graph from v2 data or a versionless v1 payload.""" + """Load trusted v2 graph state, or a nested versionless-v1 payload. + + Loading imports the recorded operator classes and runs their constructors. + Only deserialize graph artifacts from a trusted source. + """ if not isinstance(data, dict): raise GraphSerializationError("serialized graph must be a mapping") format_version = _read_format_version(data) graph = Graph() - for index, root_data in enumerate(data.get("roots", [])): - root_node = _deserialize_node( - root_data, - format_version=format_version, - path=f"roots[{index}]", - ) - graph.roots.append(root_node) + if format_version == 1: + roots = data.get("roots", []) + if not isinstance(roots, list): + raise GraphSerializationError("roots: expected a list") + graph.roots = [_deserialize_v1_node(root_data, path=f"roots[{index}]") for index, root_data in enumerate(roots)] + return graph + + root_ids, records = _validate_v2_topology(data) + nodes = {node_id: _instantiate_v2_node(node_id, record) for node_id, record in records.items()} + for node_id, record in records.items(): + nodes[node_id].children = [nodes[child_id] for child_id in record["children"]] + graph.roots = [nodes[root_id] for root_id in root_ids] return graph @@ -1133,10 +1656,31 @@ def save_graph(graph: Graph, path: Union[str, Path], *, indent: int = 2) -> Path return path +def _reject_duplicate_json_keys(pairs: List[Tuple[str, Any]]) -> dict: + result: Dict[str, Any] = {} + for key, value in pairs: + if key in result: + raise GraphSerializationError("serialized graph contains a duplicate JSON object key") + result[key] = value + return result + + +def _load_json_payload(path: Path) -> Any: + try: + return json.loads(path.read_text(), object_pairs_hook=_reject_duplicate_json_keys) + except json.JSONDecodeError as exc: + raise GraphSerializationError(f"invalid graph JSON at line {exc.lineno}, column {exc.colno}") from exc + + def load_graph(path: Union[str, Path]) -> Graph: - """Load a graph from a JSON file produced by :func:`save_graph`.""" - path = Path(path) - payload = json.loads(path.read_text()) + """Load trusted graph state from a module or registry graph JSON file.""" + payload = _load_json_payload(Path(path)) + if not isinstance(payload, dict): + raise GraphSerializationError("serialized graph must be a mapping") + if "blueprint" in payload: + if not isinstance(payload["blueprint"], dict): + raise GraphSerializationError("blueprint metadata must be a mapping") + payload = {key: value for key, value in payload.items() if key != "blueprint"} return deserialize_graph(payload) @@ -1412,12 +1956,14 @@ def load_all(self, path: Union[str, Path], *, overwrite: bool = False) -> List[s stored structure. Returns the list of graph names loaded. """ path = Path(path) - payload = json.loads(path.read_text()) + payload = _load_json_payload(path) if not isinstance(payload, dict): raise GraphSerializationError("serialized graph registry must be a mapping") version_marker = payload.get("format_version") if isinstance(version_marker, int) and not isinstance(version_marker, bool): _read_format_version(payload) + if set(payload) != {"format_version", "graphs"}: + raise GraphSerializationError("version 2 graph registry contains unknown top-level fields") entries = payload.get("graphs") if not isinstance(entries, dict): raise GraphSerializationError("version 2 graph registry requires a graphs mapping") @@ -1425,10 +1971,19 @@ def load_all(self, path: Union[str, Path], *, overwrite: bool = False) -> List[s # Versionless v1 registries were a direct name -> graph mapping, # and graph names were unrestricted (including "format_version"). entries = payload - loaded: List[str] = [] + prepared: List[Tuple[str, dict, dict]] = [] for name, entry in entries.items(): + if not isinstance(name, str) or not isinstance(entry, dict): + raise GraphSerializationError("serialized registry entries must map graph names to graph mappings") bp_meta = entry.get("blueprint", {}) - graph_data = {k: v for k, v in entry.items() if k != "blueprint"} + if not isinstance(bp_meta, dict): + raise GraphSerializationError("blueprint metadata must be a mapping") + graph_data = {key: value for key, value in entry.items() if key != "blueprint"} + deserialize_graph(graph_data) + prepared.append((name, bp_meta, graph_data)) + + loaded: List[str] = [] + for name, bp_meta, graph_data in prepared: def _factory(_gd: dict = graph_data) -> Graph: return deserialize_graph(_gd) @@ -1466,14 +2021,20 @@ def save_graph(self, name: str, path: Union[str, Path], *, indent: int = 2) -> P path.write_text(json.dumps(payload, indent=indent)) return path - def load_graph(self, path: Union[str, Path], *, name: Optional[str] = None, overwrite: bool = False) -> str: + def load_graph( + self, + path: Union[str, Path], + *, + name: Optional[str] = None, + overwrite: bool = False, + ) -> str: """Load a single graph from a JSON file and register it. If *name* is not provided, the blueprint name stored in the file is used (falls back to the file stem). Returns the registered name. """ path = Path(path) - payload = json.loads(path.read_text()) + payload = _load_json_payload(path) if not isinstance(payload, dict): raise GraphSerializationError("serialized graph must be a mapping") _read_format_version(payload) @@ -1481,6 +2042,7 @@ def load_graph(self, path: Union[str, Path], *, name: Optional[str] = None, over if not isinstance(bp_meta, dict): raise GraphSerializationError("blueprint metadata must be a mapping") graph_data = {k: v for k, v in payload.items() if k != "blueprint"} + deserialize_graph(graph_data) resolved_name = name or bp_meta.get("name") or path.stem def _factory(_gd: dict = graph_data) -> Graph: diff --git a/nemo_retriever/src/nemo_retriever/models/llm/__init__.py b/nemo_retriever/src/nemo_retriever/models/llm/__init__.py index 7ad463aae..25758df72 100644 --- a/nemo_retriever/src/nemo_retriever/models/llm/__init__.py +++ b/nemo_retriever/src/nemo_retriever/models/llm/__init__.py @@ -18,28 +18,27 @@ Per-component API keys (``api_key``) and base URLs (``api_base``) are passed directly on ``LiteLLMClient.from_kwargs`` / ``LLMJudge.from_kwargs`` or via ``Retriever(embed_kwargs={"api_key": ..., "embedding_endpoint": ...})``. When -``api_key`` is left ``None`` the shared ``_ParamsModel`` validator -resolves ``NVIDIA_API_KEY`` / ``NGC_API_KEY`` from the environment. This -keeps the common single-provider path a one-liner while still allowing -multiple independent endpoints to coexist -- each component takes its -own ``(api_base, api_key, model)`` triple. +``api_key`` is left ``None``, LiteLLM performs provider-native environment +lookup. An explicit ``os.environ/VARIABLE_NAME`` value resolves that variable +immediately before the provider call. Literal keys are supported for local, +non-persisted execution but graph persistence rejects them. Public surface contract ----------------------- -The names in ``__all__`` below are the frozen public API of this -module. External callers should import from ``nemo_retriever.models.llm`` +The one-shot text generation names in ``__all__`` below are a provisional v1 +API. They support synchronous, single-turn text only; tools, streaming, +multiple choices, and structured domain responses are not supported. Other +established names retain their existing compatibility commitments. External +callers should import from ``nemo_retriever.models.llm`` rather than reaching into submodules (``models.llm.clients.litellm``, ``models.llm.text_utils``) directly -- those submodule paths are implementation -details and may be reorganised in future releases without notice. The -Protocols + result dataclasses + concrete clients + re-exported params -models listed here are the supported integration points. +details and may be reorganised in future releases without notice. """ from nemo_retriever.models.llm.types import ( AnswerJudge, AnswerRequest, AnswerResult, - CompletionClient, GeneratedTextResult, GenerationRequest, GenerationResult, @@ -47,6 +46,7 @@ LLMClient, RetrievalResult, RetrieverStrategy, + TextCompletionClient, ) from nemo_retriever.common.params.models import ( LLMInferenceParams, @@ -56,6 +56,7 @@ ) from nemo_retriever.models.llm.tasks import ( GenerationTask, + GenerationTaskError, GenericPromptTask, RagAnswerTask, SummarizeTask, @@ -79,10 +80,11 @@ def __getattr__(name: str): __all__ = [ # Protocols "AnswerJudge", - "CompletionClient", "GenerationTask", + "GenerationTaskError", "LLMClient", "RetrieverStrategy", + "TextCompletionClient", # Request/result models "AnswerRequest", "AnswerResult", diff --git a/nemo_retriever/src/nemo_retriever/models/llm/clients/litellm.py b/nemo_retriever/src/nemo_retriever/models/llm/clients/litellm.py index 42a10b18d..e498fe464 100644 --- a/nemo_retriever/src/nemo_retriever/models/llm/clients/litellm.py +++ b/nemo_retriever/src/nemo_retriever/models/llm/clients/litellm.py @@ -16,18 +16,33 @@ import time from typing import Any, Optional +from nemo_retriever.common.params.models import ( + LLMInferenceParams, + LLMRemoteClientParams, + resolve_environment_reference, + validate_llm_extra_params, +) from nemo_retriever.models.llm.tasks.rag_answer import ( RagAnswerTask, _build_rag_prompt as _task_build_rag_prompt, _deep_merge_dicts, - _format_rag_system_prompt, ) -from nemo_retriever.models.llm.types import GenerationResult -from nemo_retriever.common.params.models import LLMInferenceParams, LLMRemoteClientParams +from nemo_retriever.models.llm.types import ( + GenerationResult, + UnsupportedTextResponseError, +) logger = logging.getLogger(__name__) # Backwards-compatible helper export retained for existing callers. _build_rag_prompt = _task_build_rag_prompt +_NO_AUTH_API_KEY = "no-api-key" + + +def _field(value: object, name: str, default: object = None) -> object: + """Read a provider response field from either a mapping or an object.""" + if isinstance(value, dict): + return value.get(name, default) + return getattr(value, name, default) class LiteLLMClient: @@ -53,6 +68,7 @@ class LiteLLMClient: """ _DEFAULT_MODEL: str = "nvidia_nim/nvidia/llama-3.3-nemotron-super-49b-v1.5" + supports_concurrent_calls: bool = True def __init__( self, @@ -66,10 +82,6 @@ def __init__( # ``max_tokens=1024`` for captioning/summarization workloads; RAG # answers routinely exceed that, so the client overrides it. self.sampling = sampling if sampling is not None else LLMInferenceParams(temperature=0.0, max_tokens=4096) - self._formatted_rag_system_prompt = _format_rag_system_prompt( - rag_system_prompt=transport.rag_system_prompt, - rag_system_prompt_prefix=transport.rag_system_prompt_prefix, - ) @property def model(self) -> str: @@ -124,6 +136,8 @@ def complete( extra_params: Optional[dict[str, Any]] = None, ) -> tuple[str, float]: """Raw litellm completion call. Returns (content_text, latency_s).""" + validate_llm_extra_params(self.transport.extra_params, source="LLMRemoteClientParams.extra_params") + validate_llm_extra_params(extra_params, source="GenerationRequest.extra_params") import litellm sampling_kwargs = self.sampling.to_sampling_kwargs() @@ -139,8 +153,12 @@ def complete( } if self.transport.api_base: call_kwargs["api_base"] = self.transport.api_base - if self.transport.api_key: - call_kwargs["api_key"] = self.transport.api_key + if self.transport._uses_no_api_key("api_key"): + call_kwargs["api_key"] = _NO_AUTH_API_KEY + elif self.transport.api_key is not None: + resolved_api_key = resolve_environment_reference(self.transport.api_key) + if resolved_api_key: + call_kwargs["api_key"] = resolved_api_key call_kwargs.update(_deep_merge_dicts(self.transport.extra_params, extra_params or {})) t0 = time.monotonic() @@ -161,7 +179,26 @@ def complete( ) raise latency = time.monotonic() - t0 - content = (response.choices[0].message.content or "").strip() + + choices = _field(response, "choices") + if not isinstance(choices, (list, tuple)) or len(choices) != 1: + raise UnsupportedTextResponseError("provider response must contain exactly one choice") + choice = choices[0] + message = _field(choice, "message") + if message is None: + raise UnsupportedTextResponseError("provider response choice has no message") + if _field(choice, "finish_reason") in {"tool_calls", "function_call"}: + raise UnsupportedTextResponseError("tool-call responses are unsupported by the text completion contract") + if _field(message, "tool_calls") or _field(message, "function_call"): + raise UnsupportedTextResponseError("tool-call responses are unsupported by the text completion contract") + if _field(message, "refusal"): + raise UnsupportedTextResponseError("refusal responses are unsupported by the text completion contract") + if any(_field(message, field_name) is not None for field_name in ("audio", "images", "videos")): + raise UnsupportedTextResponseError("non-text response modalities are unsupported") + content = _field(message, "content") + if not isinstance(content, str): + raise UnsupportedTextResponseError("provider response content must be plain text") + content = content.strip() return content, latency def generate( @@ -176,7 +213,8 @@ def generate( self.transport.reasoning_enabled if reasoning_enabled is None else reasoning_enabled ) task = RagAnswerTask( - system_prompt=self._formatted_rag_system_prompt, + system_prompt=self.transport.rag_system_prompt, + system_prompt_prefix=self.transport.rag_system_prompt_prefix, reasoning_enabled=effective_reasoning_enabled, ) result = task.execute(self, query=query, chunks=chunks) diff --git a/nemo_retriever/src/nemo_retriever/models/llm/tasks/__init__.py b/nemo_retriever/src/nemo_retriever/models/llm/tasks/__init__.py index 775ec4ef9..76abf7e49 100644 --- a/nemo_retriever/src/nemo_retriever/models/llm/tasks/__init__.py +++ b/nemo_retriever/src/nemo_retriever/models/llm/tasks/__init__.py @@ -4,13 +4,14 @@ """Reusable text-generation task strategies.""" -from nemo_retriever.models.llm.tasks.base import GenerationTask +from nemo_retriever.models.llm.tasks.base import GenerationTask, GenerationTaskError from nemo_retriever.models.llm.tasks.generic import GenericPromptTask from nemo_retriever.models.llm.tasks.rag_answer import RagAnswerTask from nemo_retriever.models.llm.tasks.summarize import SummarizeTask __all__ = [ "GenerationTask", + "GenerationTaskError", "GenericPromptTask", "RagAnswerTask", "SummarizeTask", diff --git a/nemo_retriever/src/nemo_retriever/models/llm/tasks/base.py b/nemo_retriever/src/nemo_retriever/models/llm/tasks/base.py index 4d7782d83..517b390ee 100644 --- a/nemo_retriever/src/nemo_retriever/models/llm/tasks/base.py +++ b/nemo_retriever/src/nemo_retriever/models/llm/tasks/base.py @@ -7,16 +7,38 @@ from __future__ import annotations from abc import ABC, abstractmethod -from typing import Any, ClassVar, Optional +import time +from typing import Any, ClassVar, Literal, Optional from nemo_retriever.common.params.models import LLMInferenceParams from nemo_retriever.models.llm.types import ( - CompletionClient, GeneratedTextResult, GenerationRequest, + TextCompletionClient, + UnsupportedTextResponseError, ) +class GenerationTaskError(RuntimeError): + """Sanitized failure raised by the strict generation-task lifecycle.""" + + def __init__( + self, + *, + code: str, + phase: Literal["request", "transport", "response", "parse"], + retryable: bool, + public_message: str, + latency_s: float, + ) -> None: + super().__init__(public_message) + self.code = code + self.phase = phase + self.retryable = retryable + self.public_message = public_message + self.latency_s = latency_s + + class GenerationTask(ABC): """Stateless strategy that turns logical inputs into one completion call.""" @@ -45,36 +67,140 @@ def _preflight_error(self, **inputs: object) -> Optional[str]: """Return an error code when no provider request should be made.""" return None - def execute(self, client: CompletionClient, **inputs: object) -> GeneratedTextResult: - """Build, execute, and parse one request without leaking row failures.""" - latency_s = 0.0 + @staticmethod + def _client_model(client: object) -> str: + """Read a client model identifier without making error handling fail.""" + try: + model = getattr(client, "model", "") + except Exception: + return "" + return model if isinstance(model, str) else "" + + @staticmethod + def _elapsed(started_at: float) -> float: + return time.monotonic() - started_at + + def invoke(self, client: TextCompletionClient, **inputs: object) -> GeneratedTextResult: + """Strictly build, execute, and parse one text request. + + Failures are raised as :class:`GenerationTaskError` with stable codes + and sanitized messages. Batch operators collect those errors at the row + boundary; callers that need the historical collecting behavior can use + :meth:`execute`. + """ + started_at = time.monotonic() + + preflight_failure: Optional[GenerationTaskError] = None try: preflight_error = self._preflight_error(**inputs) - if preflight_error is not None: - return GeneratedTextResult( - text="", - latency_s=0.0, - model=client.model, - error=preflight_error, - ) + except Exception: + preflight_failure = GenerationTaskError( + code="request_error", + phase="request", + retryable=False, + public_message="generation request validation failed", + latency_s=self._elapsed(started_at), + ) + if preflight_failure is not None: + raise preflight_failure + if preflight_error is not None: + raise GenerationTaskError( + code=preflight_error, + phase="request", + retryable=False, + public_message="generation request was skipped", + latency_s=self._elapsed(started_at), + ) + request_failure: Optional[GenerationTaskError] = None + try: request = self.build_request(**inputs) + if not isinstance(request, GenerationRequest): + raise TypeError("build_request must return GenerationRequest") + request = GenerationRequest( + messages=request.messages, + max_tokens=request.max_tokens, + extra_params=request.extra_params, + ) + except Exception: + request_failure = GenerationTaskError( + code="request_error", + phase="request", + retryable=False, + public_message="generation request construction failed", + latency_s=self._elapsed(started_at), + ) + if request_failure is not None: + raise request_failure + + transport_failure: Optional[GenerationTaskError] = None + try: raw_text, latency_s = client.complete( request.messages, max_tokens=request.max_tokens, extra_params=request.extra_params, ) + except UnsupportedTextResponseError: + transport_failure = GenerationTaskError( + code="unsupported_response", + phase="response", + retryable=False, + public_message="provider response is not representable as a text completion", + latency_s=self._elapsed(started_at), + ) + except Exception as exc: + try: + retryable = bool(getattr(exc, "retryable", False)) + except Exception: + retryable = False + transport_failure = GenerationTaskError( + code="transport_error", + phase="transport", + retryable=retryable, + public_message="text completion request failed", + latency_s=self._elapsed(started_at), + ) + if transport_failure is not None: + raise transport_failure + + parse_failure: Optional[GenerationTaskError] = None + try: text = self.parse(raw_text) - return GeneratedTextResult( - text=text, + if not isinstance(text, str): + raise TypeError("parse must return text") + except Exception: + parse_failure = GenerationTaskError( + code="parse_error", + phase="parse", + retryable=False, + public_message="generation response parsing failed", + latency_s=self._elapsed(started_at), + ) + if parse_failure is not None: + raise parse_failure + if not text: + raise GenerationTaskError( + code=self.empty_output_error, + phase="response", + retryable=False, + public_message="generation produced no usable text", latency_s=latency_s, - model=client.model, - error=None if text else self.empty_output_error, ) - except Exception as exc: + return GeneratedTextResult( + text=text, + latency_s=latency_s, + model=self._client_model(client), + error=None, + ) + + def execute(self, client: TextCompletionClient, **inputs: object) -> GeneratedTextResult: + """Compatibility facade that collects strict failures into a result.""" + try: + return self.invoke(client, **inputs) + except GenerationTaskError as exc: return GeneratedTextResult( text="", - latency_s=latency_s, - model=client.model, - error=str(exc), + latency_s=exc.latency_s, + model=self._client_model(client), + error=exc.code, ) diff --git a/nemo_retriever/src/nemo_retriever/models/llm/tasks/generic.py b/nemo_retriever/src/nemo_retriever/models/llm/tasks/generic.py index 5bc27234b..1a26c18b9 100644 --- a/nemo_retriever/src/nemo_retriever/models/llm/tasks/generic.py +++ b/nemo_retriever/src/nemo_retriever/models/llm/tasks/generic.py @@ -6,6 +6,7 @@ from __future__ import annotations +from dataclasses import dataclass from string import Formatter from typing import Any, ClassVar, Optional, Sequence @@ -58,9 +59,15 @@ def _validate_prompt_template(prompt: str, input_names: tuple[str, ...]) -> None raise ValueError(f"prompt contains undeclared placeholders: {sorted(undeclared)}") +@dataclass(frozen=True, init=False) class GenericPromptTask(GenerationTask): """Render declared row inputs into a validated prompt template.""" + prompt: str + required_inputs: tuple[str, ...] + system_prompt: Optional[str] + reasoning_enabled: Optional[bool] + _default_sampling: ClassVar[dict[str, Any]] = { "temperature": 1.0, "top_p": None, @@ -79,10 +86,10 @@ def __init__( raise TypeError("input_names must be a sequence of names, not a string") names = tuple(input_names) _validate_prompt_template(prompt, names) - self.prompt = prompt - self.required_inputs = names - self.system_prompt = system_prompt - self.reasoning_enabled = reasoning_enabled + object.__setattr__(self, "prompt", prompt) + object.__setattr__(self, "required_inputs", names) + object.__setattr__(self, "system_prompt", system_prompt) + object.__setattr__(self, "reasoning_enabled", reasoning_enabled) def build_request(self, **inputs: object) -> GenerationRequest: """Render declared inputs and build one completion request.""" diff --git a/nemo_retriever/src/nemo_retriever/models/llm/types.py b/nemo_retriever/src/nemo_retriever/models/llm/types.py index 9093eeacf..7da26bf10 100644 --- a/nemo_retriever/src/nemo_retriever/models/llm/types.py +++ b/nemo_retriever/src/nemo_retriever/models/llm/types.py @@ -13,6 +13,7 @@ from __future__ import annotations +from copy import deepcopy from dataclasses import dataclass, field from typing import Any, Optional, Protocol, runtime_checkable @@ -40,8 +41,13 @@ def generate( @runtime_checkable -class CompletionClient(Protocol): - """Minimal client contract consumed by reusable generation tasks.""" +class TextCompletionClient(Protocol): + """Provisional synchronous, thread-safe, single-turn text client contract. + + Implementations return exactly one text completion. Tools, streaming, + multiple choices, and structured domain responses are intentionally + outside this contract. + """ @property def model(self) -> str: @@ -58,6 +64,10 @@ def complete( ... +class UnsupportedTextResponseError(RuntimeError): + """Raised when a provider response cannot be represented as plain text.""" + + @runtime_checkable class AnswerJudge(Protocol): """Pluggable answer scoring interface.""" @@ -83,16 +93,38 @@ class GenerationResult: error: Optional[str] = None -@dataclass +@dataclass(frozen=True) class GenerationRequest: - """Provider-neutral request produced by a generation task.""" + """One text-only request produced by a generation task. + + Tools, streaming, multiple choices, and structured domain results are not + supported by this provisional contract. + """ messages: list[dict[str, Any]] max_tokens: Optional[int] = None extra_params: Optional[dict[str, Any]] = None - -@dataclass + def __post_init__(self) -> None: + """Snapshot mutable inputs and reject non-text or protected state.""" + # Keep this types module lightweight and avoid a package import cycle. + from nemo_retriever.common.params.models import validate_llm_extra_params + + messages = deepcopy(self.messages) + extra_params = deepcopy(self.extra_params) + if not isinstance(messages, list) or not all(isinstance(message, dict) for message in messages): + raise TypeError("GenerationRequest.messages must be a list of message dictionaries") + for message in messages: + if not isinstance(message.get("role"), str) or not isinstance(message.get("content"), str): + raise TypeError("GenerationRequest messages require string role and content fields") + if {"tool_calls", "function_call", "tool_call_id"}.intersection(message): + raise ValueError("GenerationRequest does not support tool messages or tool calls") + validate_llm_extra_params(extra_params or {}, source="GenerationRequest.extra_params") + object.__setattr__(self, "messages", messages) + object.__setattr__(self, "extra_params", extra_params) + + +@dataclass(frozen=True) class GeneratedTextResult: """Task-neutral result from a single text-generation request.""" diff --git a/nemo_retriever/src/nemo_retriever/operators/generation/base.py b/nemo_retriever/src/nemo_retriever/operators/generation/base.py index d63454bd2..5c367187f 100644 --- a/nemo_retriever/src/nemo_retriever/operators/generation/base.py +++ b/nemo_retriever/src/nemo_retriever/operators/generation/base.py @@ -8,18 +8,20 @@ import inspect import logging +import time from abc import abstractmethod from collections.abc import Mapping from concurrent.futures import Future, ThreadPoolExecutor, as_completed +from copy import deepcopy from typing import Any, ClassVar import pandas as pd from pydantic import BaseModel -from nemo_retriever.common.params import TextGenerationParams +from nemo_retriever.common.params import LLMInferenceParams, TextGenerationParams from nemo_retriever.models.llm.clients import LiteLLMClient -from nemo_retriever.models.llm.tasks import GenerationTask -from nemo_retriever.models.llm.types import CompletionClient, GeneratedTextResult +from nemo_retriever.models.llm.tasks import GenerationTask, GenerationTaskError +from nemo_retriever.models.llm.types import GeneratedTextResult, TextCompletionClient from nemo_retriever.operators.abstract_operator import AbstractOperator from nemo_retriever.operators.cpu_operator import CPUOperator @@ -29,10 +31,11 @@ class TextGenerationOperator(AbstractOperator, CPUOperator): """Base operator for one text-generation request per DataFrame row. - Subclasses bind a concrete :class:`GenerationTask` through - :meth:`_create_task`. The base owns the common DataFrame concerns: column - validation, bounded threaded execution, positional result ordering, and - standard output metadata. + Concrete operators construct an immutable :class:`GenerationTask` before + calling this base. The task and client are runtime-only state; graph + reconstruction uses only defensive constructor state. The base owns + validation, safe bounded execution, positional ordering, and stable + output metadata. ``input_columns`` maps each task-level input name to a physical DataFrame column. Results are tracked by row position rather than index label so @@ -46,14 +49,16 @@ def __init__( self, params: TextGenerationParams, *, + task: GenerationTask, input_columns: Mapping[str, str], output_column: str, latency_column: str | None = None, model_column: str | None = None, error_column: str | None = None, overwrite: bool = False, - client: CompletionClient | None = None, + client: TextCompletionClient | None = None, ) -> None: + copied_params = params.model_copy(deep=True) logical_columns = dict(input_columns) self._validate_input_mapping(logical_columns) @@ -70,7 +75,7 @@ def __init__( super().__init__() - self._params = params + self._params = copied_params self._input_columns = logical_columns self._output_column = output_column self._latency_column = resolved_latency_column @@ -80,31 +85,30 @@ def __init__( self._model_column_arg = model_column self._error_column_arg = error_column self._overwrite = overwrite - self._max_workers = params.max_workers - self._configured_model = params.transport.model + self._max_workers = copied_params.max_workers + self._configured_model = copied_params.transport.model self.required_columns = tuple(dict.fromkeys(logical_columns.values())) self.output_columns = output_columns - self._task = self._create_task(params, tuple(logical_columns)) + self._task = task missing_inputs = [name for name in self._task.required_inputs if name not in logical_columns] if missing_inputs: raise ValueError(f"{type(self).__name__} is missing task input mappings: {missing_inputs}") if client is None: - sampling = params.resolve_sampling(self._task.default_sampling) - self._client: CompletionClient = LiteLLMClient(transport=params.transport, sampling=sampling) + sampling = copied_params.resolve_sampling(self._task.default_sampling) + self._client: TextCompletionClient = self._create_client(copied_params, sampling) else: self._client = client - @abstractmethod - def _create_task( + def _create_client( self, params: TextGenerationParams, - logical_inputs: tuple[str, ...], - ) -> GenerationTask: - """Create the stateless task bound to this concrete operator.""" - ... + sampling: LLMInferenceParams, + ) -> TextCompletionClient: + """Create the default client without introducing a global registry.""" + return LiteLLMClient(transport=params.transport, sampling=sampling) @abstractmethod def _get_generation_constructor_kwargs(self) -> dict[str, Any]: @@ -156,7 +160,10 @@ def get_constructor_kwargs(self) -> dict[str, Any]: signature.bind(None, **kwargs) except TypeError as exc: raise TypeError(f"{type(self).__name__} returned invalid graph constructor kwargs: {exc}") from exc - return kwargs + try: + return deepcopy(kwargs) + except Exception as exc: + raise TypeError(f"{type(self).__name__} graph constructor kwargs could not be copied safely") from exc @staticmethod def _validate_input_mapping(input_columns: Mapping[str, str]) -> None: @@ -230,10 +237,24 @@ def preprocess(self, data: Any, **kwargs: Any) -> pd.DataFrame: def _execute_task(self, inputs: dict[str, Any]) -> GeneratedTextResult: """Execute the configured task; subclasses may adapt legacy clients.""" - return self._task.execute(self._client, **inputs) + return self._task.invoke(self._client, **inputs) def _execute_row(self, position: int, inputs: dict[str, Any]) -> tuple[int, GeneratedTextResult]: - return position, self._execute_task(inputs) + started_at = time.monotonic() + try: + result = self._execute_task(inputs) + except GenerationTaskError as exc: + # Keep the strict task's measured lifecycle while covering custom + # adapters that report a shorter or zero failure duration. + elapsed = max(exc.latency_s, time.monotonic() - started_at) + result = self._failure_result(exc.code, elapsed) + logger.warning("Row %d generation failed (%s)", position, exc.code) + except Exception: + # Unexpected adapter/client failures remain isolated by row. Raw + # provider exception text is never persisted or logged. + result = self._failure_result("request_error", time.monotonic() - started_at) + logger.warning("Row %d generation failed (request_error)", position) + return position, result def _failure_model(self) -> str: try: @@ -242,21 +263,30 @@ def _failure_model(self) -> str: return self._configured_model return model if isinstance(model, str) and model else self._configured_model - def _failure_result(self, exc: Exception) -> GeneratedTextResult: + def _failure_result(self, error: str, latency_s: float) -> GeneratedTextResult: return GeneratedTextResult( text="", - latency_s=0.0, + latency_s=latency_s, model=self._failure_model(), - error=str(exc), + error=error, ) + def _effective_max_workers(self, row_count: int) -> int: + """Return safe concurrency for the current runtime client.""" + try: + supports_concurrency = getattr(self._client, "supports_concurrent_calls", False) is True + except Exception: + supports_concurrency = False + configured_workers = self._max_workers if supports_concurrency else 1 + return min(configured_workers, row_count) + def process(self, data: Any, **kwargs: Any) -> pd.DataFrame: df, input_positions = self._validate_and_resolve_dataframe(data) results: list[GeneratedTextResult | None] = [None] * len(df) if len(df): futures: dict[Future[tuple[int, GeneratedTextResult]], int] = {} - with ThreadPoolExecutor(max_workers=min(self._max_workers, len(df))) as pool: + with ThreadPoolExecutor(max_workers=self._effective_max_workers(len(df))) as pool: for position in range(len(df)): inputs = { name: df.iat[position, column_position] for name, column_position in input_positions.items() @@ -274,9 +304,9 @@ def process(self, data: Any, **kwargs: Any) -> pd.DataFrame: f"submitted position {position}" ) results[position] = result - except Exception as exc: - logger.warning("Row %d generation failed: %s", position, exc) - results[position] = self._failure_result(exc) + except Exception: + logger.warning("Row %d generation failed (request_error)", position) + results[position] = self._failure_result("request_error", 0.0) # Every non-empty row is assigned either a task result or a failure # result above. The cast-free local assertion catches future changes to @@ -286,10 +316,12 @@ def process(self, data: Any, **kwargs: Any) -> pd.DataFrame: completed_results = [result for result in results if result is not None] out = df.copy() - out[self._output_column] = [result.text for result in completed_results] - out[self._latency_column] = [result.latency_s for result in completed_results] - out[self._model_column] = [result.model for result in completed_results] - out[self._error_column] = [result.error for result in completed_results] + # Explicit dtypes keep empty and non-empty Ray/Arrow batches + # schema-compatible while retaining positional duplicate-index writes. + out[self._output_column] = pd.array([result.text for result in completed_results], dtype="object") + out[self._latency_column] = pd.array([result.latency_s for result in completed_results], dtype="float64") + out[self._model_column] = pd.array([result.model for result in completed_results], dtype="object") + out[self._error_column] = pd.array([result.error for result in completed_results], dtype="object") return out def postprocess(self, data: Any, **kwargs: Any) -> Any: diff --git a/nemo_retriever/src/nemo_retriever/operators/generation/generic.py b/nemo_retriever/src/nemo_retriever/operators/generation/generic.py index 69ff2ff96..4a35f81f9 100644 --- a/nemo_retriever/src/nemo_retriever/operators/generation/generic.py +++ b/nemo_retriever/src/nemo_retriever/operators/generation/generic.py @@ -9,8 +9,8 @@ from collections.abc import Mapping from nemo_retriever.common.params import TextGenerationParams -from nemo_retriever.models.llm.tasks import GenerationTask, GenericPromptTask -from nemo_retriever.models.llm.types import CompletionClient +from nemo_retriever.models.llm.tasks import GenericPromptTask +from nemo_retriever.models.llm.types import TextCompletionClient from nemo_retriever.operators.generation.base import TextGenerationOperator @@ -27,11 +27,23 @@ def __init__( model_column: str | None = None, error_column: str | None = None, overwrite: bool = False, - client: CompletionClient | None = None, + client: TextCompletionClient | None = None, ) -> None: normalized_input_columns = dict(input_columns) + if params.prompt is None: + raise ValueError("GenericGenerationOperator requires params.prompt") + reasoning_enabled = ( + params.reasoning_enabled if params.reasoning_enabled is not None else params.transport.reasoning_enabled + ) + task = GenericPromptTask( + prompt=params.prompt, + input_names=tuple(normalized_input_columns), + system_prompt=params.system_prompt, + reasoning_enabled=reasoning_enabled, + ) super().__init__( params, + task=task, input_columns=normalized_input_columns, output_column=output_column, latency_column=latency_column, @@ -43,28 +55,11 @@ def __init__( def _get_generation_constructor_kwargs(self) -> dict[str, object]: return { - "params": self._params, - "input_columns": dict(self._input_columns), + "params": self._params.model_copy(deep=True), + "input_columns": self._input_columns.copy(), "output_column": self._output_column, "latency_column": self._latency_column_arg, "model_column": self._model_column_arg, "error_column": self._error_column_arg, "overwrite": self._overwrite, } - - def _create_task( - self, - params: TextGenerationParams, - logical_inputs: tuple[str, ...], - ) -> GenerationTask: - if params.prompt is None: - raise ValueError("GenericGenerationOperator requires params.prompt") - reasoning_enabled = ( - params.reasoning_enabled if params.reasoning_enabled is not None else params.transport.reasoning_enabled - ) - return GenericPromptTask( - prompt=params.prompt, - input_names=logical_inputs, - system_prompt=params.system_prompt, - reasoning_enabled=reasoning_enabled, - ) diff --git a/nemo_retriever/src/nemo_retriever/operators/generation/summarization.py b/nemo_retriever/src/nemo_retriever/operators/generation/summarization.py index a18d14f0c..67c22b0c3 100644 --- a/nemo_retriever/src/nemo_retriever/operators/generation/summarization.py +++ b/nemo_retriever/src/nemo_retriever/operators/generation/summarization.py @@ -7,8 +7,8 @@ from __future__ import annotations from nemo_retriever.common.params import TextGenerationParams -from nemo_retriever.models.llm.tasks import GenerationTask, SummarizeTask -from nemo_retriever.models.llm.types import CompletionClient +from nemo_retriever.models.llm.tasks import SummarizeTask +from nemo_retriever.models.llm.types import TextCompletionClient from nemo_retriever.operators.generation.base import TextGenerationOperator @@ -25,10 +25,19 @@ def __init__( model_column: str | None = None, error_column: str | None = None, overwrite: bool = False, - client: CompletionClient | None = None, + client: TextCompletionClient | None = None, ) -> None: + reasoning_enabled = ( + params.reasoning_enabled if params.reasoning_enabled is not None else params.transport.reasoning_enabled + ) + task = SummarizeTask( + prompt=params.prompt, + system_prompt=params.system_prompt, + reasoning_enabled=reasoning_enabled, + ) super().__init__( params, + task=task, input_columns={"text": input_column}, output_column=output_column, latency_column=latency_column, @@ -40,7 +49,7 @@ def __init__( def _get_generation_constructor_kwargs(self) -> dict[str, object]: return { - "params": self._params, + "params": self._params.model_copy(deep=True), "input_column": self._input_columns["text"], "output_column": self._output_column, "latency_column": self._latency_column_arg, @@ -48,18 +57,3 @@ def _get_generation_constructor_kwargs(self) -> dict[str, object]: "error_column": self._error_column_arg, "overwrite": self._overwrite, } - - def _create_task( - self, - params: TextGenerationParams, - logical_inputs: tuple[str, ...], - ) -> GenerationTask: - del logical_inputs - reasoning_enabled = ( - params.reasoning_enabled if params.reasoning_enabled is not None else params.transport.reasoning_enabled - ) - return SummarizeTask( - prompt=params.prompt, - system_prompt=params.system_prompt, - reasoning_enabled=reasoning_enabled, - ) diff --git a/nemo_retriever/src/nemo_retriever/tools/evaluation/generation.py b/nemo_retriever/src/nemo_retriever/tools/evaluation/generation.py index d85fdc9bd..a6bc3ee24 100644 --- a/nemo_retriever/src/nemo_retriever/tools/evaluation/generation.py +++ b/nemo_retriever/src/nemo_retriever/tools/evaluation/generation.py @@ -6,11 +6,17 @@ from __future__ import annotations +from copy import deepcopy +import time from typing import Any, ClassVar, Optional from nemo_retriever.common.params import TextGenerationParams -from nemo_retriever.models.llm.tasks import GenerationTask, RagAnswerTask -from nemo_retriever.models.llm.types import CompletionClient, GeneratedTextResult, LLMClient +from nemo_retriever.models.llm.tasks import GenerationTaskError, RagAnswerTask +from nemo_retriever.models.llm.types import ( + GeneratedTextResult, + LLMClient, + TextCompletionClient, +) from nemo_retriever.operators.generation import TextGenerationOperator @@ -24,7 +30,12 @@ class QAGenerationOperator(TextGenerationOperator): """ required_columns: ClassVar[tuple[str, ...]] = ("query", "context") - output_columns: ClassVar[tuple[str, ...]] = ("answer", "latency_s", "model", "gen_error") + output_columns: ClassVar[tuple[str, ...]] = ( + "answer", + "latency_s", + "model", + "gen_error", + ) def __init__( self, @@ -58,8 +69,18 @@ def __init__( rag_system_prompt_prefix=rag_system_prompt_prefix, reasoning_enabled=reasoning_enabled, ) + task_reasoning_enabled = ( + params.reasoning_enabled if params.reasoning_enabled is not None else params.transport.reasoning_enabled + ) + task = RagAnswerTask( + prompt=params.prompt, + system_prompt=params.transport.rag_system_prompt, + system_prompt_prefix=params.transport.rag_system_prompt_prefix, + reasoning_enabled=task_reasoning_enabled, + ) super().__init__( params, + task=task, input_columns={"query": "query", "chunks": "context"}, output_column="answer", latency_column="latency_s", @@ -67,53 +88,58 @@ def __init__( error_column="gen_error", overwrite=True, ) - self._qa_constructor_kwargs = { - "model": model, - "api_base": api_base, - "api_key": api_key, - "temperature": temperature, - "top_p": top_p, - "max_tokens": max_tokens, - "extra_params": extra_params, - "num_retries": num_retries, - "timeout": timeout, - "max_workers": max_workers, - "rag_system_prompt": rag_system_prompt, - "rag_system_prompt_prefix": rag_system_prompt_prefix, - "reasoning_enabled": reasoning_enabled, - } + self._qa_constructor_kwargs = deepcopy( + { + "model": model, + "api_base": api_base, + "api_key": api_key, + "temperature": temperature, + "top_p": top_p, + "max_tokens": max_tokens, + "extra_params": extra_params, + "num_retries": num_retries, + "timeout": timeout, + "max_workers": max_workers, + "rag_system_prompt": rag_system_prompt, + "rag_system_prompt_prefix": rag_system_prompt_prefix, + "reasoning_enabled": reasoning_enabled, + } + ) def _get_generation_constructor_kwargs(self) -> dict[str, Any]: """Preserve the legacy flat QA constructor contract for graph workers.""" - return dict(self._qa_constructor_kwargs) + return deepcopy(self._qa_constructor_kwargs) def _execute_task(self, inputs: dict[str, Any]) -> GeneratedTextResult: """Prefer completion tasks while adapting legacy generate-only clients.""" client = self._client - if isinstance(client, CompletionClient): + if isinstance(client, TextCompletionClient): return super()._execute_task(inputs) if isinstance(client, LLMClient): - result = client.generate(inputs["query"], inputs["chunks"]) + started_at = time.monotonic() + failure: GenerationTaskError | None = None + try: + result = client.generate(inputs["query"], inputs["chunks"]) + except Exception: + failure = GenerationTaskError( + code="transport_error", + phase="transport", + retryable=False, + public_message="Legacy text generation request failed.", + latency_s=time.monotonic() - started_at, + ) + if failure is not None: + raise failure return GeneratedTextResult( text=result.answer, latency_s=result.latency_s, model=result.model, error=result.error, ) - raise TypeError("QAGenerationOperator client must implement CompletionClient or LLMClient") - - def _create_task( - self, - params: TextGenerationParams, - logical_inputs: tuple[str, ...], - ) -> GenerationTask: - del logical_inputs - reasoning_enabled = ( - params.reasoning_enabled if params.reasoning_enabled is not None else params.transport.reasoning_enabled - ) - return RagAnswerTask( - prompt=params.prompt, - system_prompt=params.transport.rag_system_prompt, - system_prompt_prefix=params.transport.rag_system_prompt_prefix, - reasoning_enabled=reasoning_enabled, + raise GenerationTaskError( + code="request_error", + phase="request", + retryable=False, + public_message=("QAGenerationOperator client must implement " "TextCompletionClient or LLMClient."), + latency_s=0.0, ) diff --git a/nemo_retriever/tests/test_generation_hardening.py b/nemo_retriever/tests/test_generation_hardening.py new file mode 100644 index 000000000..11e6a24f0 --- /dev/null +++ b/nemo_retriever/tests/test_generation_hardening.py @@ -0,0 +1,325 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Security and strict-contract tests for provisional text generation.""" + +from __future__ import annotations + +from dataclasses import FrozenInstanceError +from types import SimpleNamespace +import time +from unittest.mock import patch + +import pytest + + +class _Client: + model = "test/model" + + def __init__(self, handler): + self._handler = handler + self.calls = [] + + def complete(self, messages, max_tokens=None, extra_params=None): + self.calls.append((messages, max_tokens, extra_params)) + return self._handler(messages, max_tokens, extra_params) + + +def _response( + content="text", + *, + tool_calls=None, + function_call=None, + refusal=None, + audio=None, + finish_reason="stop", +): + message = SimpleNamespace( + content=content, + tool_calls=tool_calls, + function_call=function_call, + refusal=refusal, + audio=audio, + images=None, + videos=None, + ) + return SimpleNamespace(choices=[SimpleNamespace(message=message, finish_reason=finish_reason)]) + + +class TestStrictGenerationTask: + def test_invoke_raises_typed_sanitized_transport_error_and_execute_collects(self): + from nemo_retriever.models.llm.tasks import ( + GenerationTaskError, + SummarizeTask, + ) + + secret = "sk-TRANSPORT-MUST-NOT-LEAK" + + class RetryableFailure(RuntimeError): + retryable = True + + def fail(messages, max_tokens, extra_params): + time.sleep(0.001) + raise RetryableFailure(secret) + + client = _Client(fail) + task = SummarizeTask() + + with pytest.raises(GenerationTaskError) as exc_info: + task.invoke(client, text="source") + error = exc_info.value + assert error.code == "transport_error" + assert error.phase == "transport" + assert error.retryable is True + assert error.latency_s > 0.0 + assert secret not in str(error) + assert error.__cause__ is None + assert error.__context__ is None + + collected = task.execute(client, text="source") + assert collected.text == "" + assert collected.error == "transport_error" + assert collected.latency_s > 0.0 + assert secret not in repr(collected) + + def test_request_build_failures_are_request_errors(self): + from nemo_retriever.models.llm.tasks import GenerationTask, GenerationTaskError + + class BrokenTask(GenerationTask): + def build_request(self, **inputs): + raise RuntimeError("sk-REQUEST-MUST-NOT-LEAK") + + class WrongTypeTask(GenerationTask): + def build_request(self, **inputs): + return {"messages": []} + + client = _Client(lambda *args: ("unused", 0.0)) + for task in (BrokenTask(), WrongTypeTask()): + with pytest.raises(GenerationTaskError) as exc_info: + task.invoke(client) + assert exc_info.value.code == "request_error" + assert exc_info.value.phase == "request" + assert "sk-REQUEST-MUST-NOT-LEAK" not in str(exc_info.value) + assert client.calls == [] + + @pytest.mark.parametrize("parsed", [{"structured": True}, ["not", "text"], 7]) + def test_non_text_parse_results_are_parse_errors(self, parsed): + from nemo_retriever.models.llm.tasks import ( + GenerationTask, + GenerationTaskError, + ) + from nemo_retriever.models.llm.types import GenerationRequest + + class NonTextTask(GenerationTask): + def build_request(self, **inputs): + return GenerationRequest(messages=[{"role": "user", "content": "hello"}]) + + def parse(self, raw_text): + return parsed + + with pytest.raises(GenerationTaskError) as exc_info: + NonTextTask().invoke(_Client(lambda *args: ("raw", 0.1))) + assert exc_info.value.code == "parse_error" + assert exc_info.value.phase == "parse" + + def test_parser_exception_is_sanitized(self): + from nemo_retriever.models.llm.tasks import ( + GenerationTask, + GenerationTaskError, + ) + from nemo_retriever.models.llm.types import GenerationRequest + + class BrokenParser(GenerationTask): + def build_request(self, **inputs): + return GenerationRequest(messages=[{"role": "user", "content": "hello"}]) + + def parse(self, raw_text): + raise ValueError("sk-PARSE-MUST-NOT-LEAK") + + with pytest.raises(GenerationTaskError) as exc_info: + BrokenParser().invoke(_Client(lambda *args: ("raw", 0.1))) + assert exc_info.value.code == "parse_error" + assert exc_info.value.phase == "parse" + assert "sk-PARSE-MUST-NOT-LEAK" not in str(exc_info.value) + + +class TestImmutableTextContracts: + def test_request_snapshots_inputs_and_rejects_non_text_messages(self): + from nemo_retriever.models.llm.types import GenerationRequest + + messages = [{"role": "user", "content": "original"}] + extras = {"provider": {"seed": 1}} + request = GenerationRequest(messages=messages, extra_params=extras) + + messages[0]["content"] = "mutated" + extras["provider"]["seed"] = 2 + assert request.messages == [{"role": "user", "content": "original"}] + assert request.extra_params == {"provider": {"seed": 1}} + + with pytest.raises(FrozenInstanceError): + request.max_tokens = 10 + with pytest.raises(TypeError, match="string role and content"): + GenerationRequest(messages=[{"role": "user", "content": [{"type": "image"}]}]) + with pytest.raises(ValueError, match="tool"): + GenerationRequest( + messages=[ + { + "role": "assistant", + "content": "", + "tool_calls": [{"id": "call-1"}], + } + ] + ) + + def test_builtin_tasks_and_results_are_frozen(self): + from nemo_retriever.models.llm.tasks import ( + GenericPromptTask, + RagAnswerTask, + SummarizeTask, + ) + from nemo_retriever.models.llm.types import GeneratedTextResult + + values = [ + GenericPromptTask(prompt="{value}", input_names=("value",)), + RagAnswerTask(), + SummarizeTask(), + GeneratedTextResult(text="ok", latency_s=0.1, model="m"), + ] + for value in values: + with pytest.raises(FrozenInstanceError): + value.unexpected = True + + +class TestLiteLLMTextOnlyResponse: + @pytest.mark.parametrize( + "response", + [ + _response(content=None), + _response(content=[{"type": "text", "text": "hello"}]), + _response(content="ignored", tool_calls=[{"id": "call-1"}]), + _response(content="ignored", function_call={"name": "legacy"}), + _response(content="ignored", refusal="policy refusal"), + _response(content="ignored", audio={"id": "audio-1"}), + _response(content="ignored", finish_reason="tool_calls"), + SimpleNamespace(choices=[]), + SimpleNamespace( + choices=[ + SimpleNamespace(message=SimpleNamespace(content="one")), + SimpleNamespace(message=SimpleNamespace(content="two")), + ] + ), + ], + ) + @patch("litellm.completion") + def test_unsupported_provider_shapes_collect_stable_code( + self, + mock_completion, + response, + ): + from nemo_retriever.models.llm.clients import LiteLLMClient + from nemo_retriever.models.llm.tasks import SummarizeTask + + mock_completion.return_value = response + result = SummarizeTask().execute( + LiteLLMClient.from_kwargs(model="openai/mock"), + text="source", + ) + + assert result.text == "" + assert result.error == "unsupported_response" + assert result.latency_s >= 0.0 + + @patch("litellm.completion") + def test_plain_text_and_empty_text_remain_distinct(self, mock_completion): + from nemo_retriever.models.llm.clients import LiteLLMClient + from nemo_retriever.models.llm.tasks import SummarizeTask + + client = LiteLLMClient.from_kwargs(model="openai/mock") + task = SummarizeTask() + + mock_completion.return_value = _response(content=" plain text ") + success = task.execute(client, text="source") + assert success.text == "plain text" + assert success.error is None + + mock_completion.return_value = _response(content=" ") + empty = task.execute(client, text="source") + assert empty.text == "" + assert empty.error == "empty_output" + + +def _reconstruct_summary_operator(constructor_kwargs): + """Importable Ray target proving constructor-only worker reconstruction.""" + from nemo_retriever.operators.generation import SummarizationOperator + + operator = SummarizationOperator(**constructor_kwargs) + return { + "operator": type(operator).__name__, + "task": type(operator._task).__name__, + "client": type(operator._client).__name__, + "model": operator._params.transport.model, + "constructor_keys": sorted(operator.get_constructor_kwargs()), + } + + +class TestWorkerReconstruction: + def test_real_constructor_reconstructs_in_process_and_on_ray_when_available(self): + import importlib.util + + from nemo_retriever.common.params import TextGenerationParams + from nemo_retriever.operators.generation import SummarizationOperator + + params = TextGenerationParams.from_kwargs( + model="openai/mock", + api_key="", + prompt="Summarize: {text}", + ) + source = SummarizationOperator(params) + constructor_kwargs = source.get_constructor_kwargs() + expected = { + "operator": "SummarizationOperator", + "task": "SummarizeTask", + "client": "LiteLLMClient", + "model": "openai/mock", + "constructor_keys": sorted(constructor_kwargs), + } + + assert _reconstruct_summary_operator(constructor_kwargs) == expected + if importlib.util.find_spec("ray") is None: + return + + import ray + + def reconstruct_on_worker(values): + from nemo_retriever.operators.generation import SummarizationOperator + + operator = SummarizationOperator(**values) + return { + "operator": type(operator).__name__, + "task": type(operator._task).__name__, + "client": type(operator._client).__name__, + "model": operator._params.transport.model, + "constructor_keys": sorted(operator.get_constructor_kwargs()), + } + + owned_runtime = not ray.is_initialized() + if owned_runtime: + ray.init( + address="local", + num_cpus=1, + include_dashboard=False, + log_to_driver=False, + ) + try: + remote_reconstruct = ray.remote(reconstruct_on_worker) + assert ( + ray.get( + remote_reconstruct.remote(constructor_kwargs), + timeout=60, + ) + == expected + ) + finally: + if owned_runtime: + ray.shutdown() diff --git a/nemo_retriever/tests/test_generation_tasks.py b/nemo_retriever/tests/test_generation_tasks.py index eff2d91da..273d38f52 100644 --- a/nemo_retriever/tests/test_generation_tasks.py +++ b/nemo_retriever/tests/test_generation_tasks.py @@ -8,6 +8,8 @@ from collections.abc import Callable import json +import os +from pathlib import Path import threading import time from typing import Any @@ -96,15 +98,16 @@ def test_sampling_overrides_only_explicit_fields(self): assert overridden.top_p == 0.35 assert overridden.max_tokens == 123 - def test_api_key_is_resolved_but_redacted_from_display(self, monkeypatch): + def test_generic_api_key_resolution_is_deferred_and_explicit_keys_are_redacted(self, monkeypatch): from nemo_retriever.common.params import models as params_models monkeypatch.setattr(params_models, "resolve_remote_api_key", lambda: "resolved-secret") params = params_models.TextGenerationParams.from_kwargs(model="m") + explicit = params_models.TextGenerationParams.from_kwargs(model="m", api_key="explicit-secret") - assert params.transport.api_key == "resolved-secret" - assert "resolved-secret" not in repr(params) - assert "resolved-secret" not in str(params) + assert params.transport.api_key is None + assert "explicit-secret" not in repr(explicit) + assert "explicit-secret" not in str(explicit) def test_no_api_key_survives_nested_validation(self, monkeypatch): from nemo_retriever.common.params import models as params_models @@ -161,15 +164,18 @@ def fail(messages, max_tokens, extra_params): result = SummarizeTask().execute(FakeCompletionClient(fail), text="source") assert result.text == "" - assert result.latency_s == 0.0 + assert result.latency_s > 0.0 assert result.model == "fake/model" - assert result.error == "service unavailable" + assert result.error == "transport_error" def test_rag_request_applies_no_reasoning_controls_and_think_cleanup(self): from nemo_retriever.models.llm.tasks import RagAnswerTask client = FakeCompletionClient( - lambda messages, max_tokens, extra_params: ("private final answer ", 0.2) + lambda messages, max_tokens, extra_params: ( + "private final answer ", + 0.2, + ) ) task = RagAnswerTask(reasoning_enabled=False) @@ -185,7 +191,12 @@ def test_rag_request_applies_no_reasoning_controls_and_think_cleanup(self): def test_rag_think_only_output_uses_compatibility_error(self): from nemo_retriever.models.llm.tasks import RagAnswerTask - client = FakeCompletionClient(lambda messages, max_tokens, extra_params: ("unfinished", 0.3)) + client = FakeCompletionClient( + lambda messages, max_tokens, extra_params: ( + "unfinished", + 0.3, + ) + ) result = RagAnswerTask().execute(client, query="q", chunks=[]) assert result.text == "" @@ -238,7 +249,9 @@ def test_generic_missing_runtime_input_is_an_error_result(self): class TestTextGenerationOperators: def test_base_is_exported_from_canonical_operator_package(self): from nemo_retriever.operators import TextGenerationOperator - from nemo_retriever.operators.generation import TextGenerationOperator as DirectExport + from nemo_retriever.operators.generation import ( + TextGenerationOperator as DirectExport, + ) assert TextGenerationOperator is DirectExport @@ -300,7 +313,8 @@ def respond(messages, max_tokens, extra_params): assert out["text"].tolist() == ["slow", "boom", "fast"] assert out["summary"].tolist() == ["SLOW", "", "FAST"] assert out["summary_error"].tolist()[0] is None - assert out["summary_error"].tolist()[1] == "offline for boom" + assert out["summary_error"].tolist()[1] == "transport_error" + assert out["summary_latency_s"].tolist()[1] > 0.0 assert out["summary_error"].tolist()[2] is None assert out["summary_model"].tolist() == ["fake/model"] * 3 @@ -319,6 +333,13 @@ def test_empty_batch_adds_schema_without_calling_client(self): ] assert out.empty assert client.calls == [] + assert out.dtypes.astype(str).to_dict() == { + "text": "object", + "summary": "object", + "summary_latency_s": "float64", + "summary_model": "object", + "summary_error": "object", + } def test_generic_operator_maps_multiple_logical_inputs(self): from nemo_retriever.operators.generation import GenericGenerationOperator @@ -354,6 +375,96 @@ def test_injected_and_built_clients_are_not_graph_constructor_state(self): assert "task" not in kwargs assert client not in kwargs.values() + def test_params_and_graph_kwargs_are_defensive_snapshots(self): + from nemo_retriever.operators.generation import SummarizationOperator + + client = FakeCompletionClient() + params = _params( + prompt="Original {text}", + extra_params={"provider": {"seed": 1}}, + ) + operator = SummarizationOperator(params, client=client) + + params.prompt = "Changed {text}" + params.transport.model = "changed/model" + params.transport.extra_params["provider"]["seed"] = 99 + + first = operator.get_constructor_kwargs() + first["params"].prompt = "Mutated graph copy {text}" + first["params"].transport.extra_params["provider"]["seed"] = 7 + second = operator.get_constructor_kwargs() + + assert second["params"].prompt == "Original {text}" + assert second["params"].transport.model == "fake/model" + assert second["params"].transport.extra_params == {"provider": {"seed": 1}} + + operator.run(pd.DataFrame({"text": ["source"]})) + assert client.calls[0][0][-1]["content"] == "Original source" + + def test_client_factory_hook_receives_task_and_copied_params(self): + from nemo_retriever.models.llm.tasks import SummarizeTask + from nemo_retriever.operators.generation import SummarizationOperator + + created_client = FakeCompletionClient() + + class FactoryOperator(SummarizationOperator): + def _create_client(self, params, sampling): + self.factory_params = params + self.factory_sampling = sampling + self.task_at_factory = self._task + return created_client + + params = _params(temperature=0.3) + operator = FactoryOperator(params) + + assert operator._client is created_client + assert isinstance(operator.task_at_factory, SummarizeTask) + assert operator.factory_params is operator._params + assert operator.factory_params is not params + assert operator.factory_sampling.temperature == 0.3 + + def test_unknown_clients_serialize_and_opted_in_clients_run_concurrently(self): + from nemo_retriever.operators.generation import SummarizationOperator + + def peak_for(client): + state = {"active": 0, "peak": 0} + lock = threading.Lock() + + def respond(messages, max_tokens, extra_params): + with lock: + state["active"] += 1 + state["peak"] = max(state["peak"], state["active"]) + time.sleep(0.02) + with lock: + state["active"] -= 1 + return "done", 0.02 + + client._handler = respond + SummarizationOperator(_params(max_workers=4), client=client).run( + pd.DataFrame({"text": ["a", "b", "c", "d"]}) + ) + return state["peak"] + + unknown_client = FakeCompletionClient() + concurrent_client = FakeCompletionClient() + concurrent_client.supports_concurrent_calls = True + + assert peak_for(unknown_client) == 1 + assert peak_for(concurrent_client) > 1 + + def test_nonempty_batch_uses_the_same_output_dtypes_as_empty_batch(self): + from nemo_retriever.operators.generation import SummarizationOperator + + out = SummarizationOperator(_params(), client=FakeCompletionClient()).run(pd.DataFrame({"text": ["source"]})) + + assert out.dtypes.astype(str).to_dict() == { + "text": "object", + "summary": "object", + "summary_latency_s": "float64", + "summary_model": "object", + "summary_error": "object", + } + class TestQACompatibility: def test_qa_run_preserves_schema_order_and_overwrites_legacy_outputs(self): @@ -435,13 +546,22 @@ def test_qa_graph_kwargs_remain_flat_and_reconstructible(self): } assert not ({"params", "input_columns", "output_column", "task", "client"} & kwargs.keys()) + kwargs["extra_params"]["seed"] = 99 + assert operator.get_constructor_kwargs()["extra_params"] == {"seed": 4} + kwargs["extra_params"]["seed"] = 4 + reconstructed = QAGenerationOperator(**kwargs) assert reconstructed._client.transport.model == "fake/model" assert reconstructed._client.sampling.temperature == 0.2 assert reconstructed._client.sampling.top_p == 0.7 assert reconstructed._client.sampling.max_tokens == 321 assert reconstructed.required_columns == ("query", "context") - assert reconstructed.output_columns == ("answer", "latency_s", "model", "gen_error") + assert reconstructed.output_columns == ( + "answer", + "latency_s", + "model", + "gen_error", + ) def test_litellm_generate_preserves_legacy_result_sentinels(self): from nemo_retriever.models.llm.clients import LiteLLMClient @@ -467,12 +587,10 @@ def fail(messages, max_tokens=None, extra_params=None): client.complete = fail failed = client.generate("question", ["context"]) - assert failed == GenerationResult( - answer="", - latency_s=0.0, - model="fake/model", - error="transport unavailable", - ) + assert failed.answer == "" + assert failed.latency_s > 0.0 + assert failed.model == "fake/model" + assert failed.error == "transport_error" def test_rag_prompt_helper_reexport_preserves_exact_messages(self): from nemo_retriever.models.llm.clients import _build_rag_prompt @@ -496,7 +614,10 @@ def test_rag_prompt_helper_reexport_preserves_exact_messages(self): class TestSamplingRemediation: def test_omitted_and_explicit_null_overrides_survive_round_trip(self): - from nemo_retriever.common.params import LLMInferenceParams, LLMSamplingOverrides + from nemo_retriever.common.params import ( + LLMInferenceParams, + LLMSamplingOverrides, + ) defaults = LLMInferenceParams(temperature=0.4, top_p=0.7, max_tokens=99) omitted = LLMSamplingOverrides() @@ -623,6 +744,7 @@ def generate(self, query, chunks, *, reasoning_enabled=None): operator = QAGenerationOperator(model="configured/model", api_key="") operator._client = client + assert operator._effective_max_workers(4) == 1 out = operator.run(pd.DataFrame({"query": ["q"], "context": [["c"]]})) assert client.calls == [("q", ["c"])] @@ -644,7 +766,31 @@ def generate(self, query, chunks, *, reasoning_enabled=None): assert out.loc[0, "answer"] == "" assert out.loc[0, "model"] == "configured/model" - assert out.loc[0, "gen_error"] == "legacy unavailable" + assert out.loc[0, "gen_error"] == "transport_error" + assert out.loc[0, "latency_s"] > 0.0 + + def test_generate_only_result_preserves_legacy_error_sentinel(self): + from nemo_retriever.models.llm.types import GenerationResult + from nemo_retriever.tools.evaluation.generation import QAGenerationOperator + + class LegacyClient: + def generate(self, query, chunks, *, reasoning_enabled=None): + return GenerationResult( + answer="", + latency_s=0.4, + model="legacy/model", + error="legacy_sentinel", + ) + + operator = QAGenerationOperator(model="configured/model", api_key="") + operator._client = LegacyClient() + + out = operator.run(pd.DataFrame({"query": ["q"], "context": [["c"]]})) + + assert out.loc[0, "answer"] == "" + assert out.loc[0, "latency_s"] == 0.4 + assert out.loc[0, "model"] == "legacy/model" + assert out.loc[0, "gen_error"] == "legacy_sentinel" def test_dual_protocol_client_prefers_completion_contract(self): from nemo_retriever.models.llm.types import GenerationResult @@ -652,7 +798,12 @@ def test_dual_protocol_client_prefers_completion_contract(self): class DualClient(FakeCompletionClient): def __init__(self): - super().__init__(lambda messages, max_tokens, extra_params: ("completion answer", 0.2)) + super().__init__( + lambda messages, max_tokens, extra_params: ( + "completion answer", + 0.2, + ) + ) self.generate_calls = 0 def generate(self, query, chunks, *, reasoning_enabled=None): @@ -672,7 +823,12 @@ def generate(self, query, chunks, *, reasoning_enabled=None): class TestPublicLLMExports: def test_direct_and_star_imports_resolve_canonical_clients(self): - from nemo_retriever.models.llm import LLMJudge, LiteLLMClient + import nemo_retriever.models.llm as llm_module + from nemo_retriever.models.llm import ( + LLMJudge, + LiteLLMClient, + TextCompletionClient, + ) namespace: dict[str, Any] = {} exec("from nemo_retriever.models.llm import *", namespace) @@ -681,6 +837,9 @@ def test_direct_and_star_imports_resolve_canonical_clients(self): assert LLMJudge.__module__ == "nemo_retriever.models.llm.clients.judge" assert namespace["LiteLLMClient"] is LiteLLMClient assert namespace["LLMJudge"] is LLMJudge + assert namespace["TextCompletionClient"] is TextCompletionClient + assert "CompletionClient" not in namespace + assert not hasattr(llm_module, "CompletionClient") def test_type_contract_import_does_not_eagerly_load_clients(self): import subprocess @@ -688,7 +847,7 @@ def test_type_contract_import_does_not_eagerly_load_clients(self): code = ( "import sys; " - "from nemo_retriever.models.llm import CompletionClient; " + "from nemo_retriever.models.llm import TextCompletionClient; " "assert 'nemo_retriever.models.llm.clients.litellm' not in sys.modules; " "assert 'litellm' not in sys.modules; " "from nemo_retriever.models.llm import LiteLLMClient; " @@ -696,9 +855,14 @@ def test_type_contract_import_does_not_eagerly_load_clients(self): "assert 'litellm' not in sys.modules" ) + source_path = str(Path(__file__).resolve().parents[1] / "src") + env = os.environ.copy() + env["PYTHONPATH"] = os.pathsep.join([source_path, env.get("PYTHONPATH", "")]) + subprocess.run( [sys.executable, "-c", code], check=True, + env=env, ) @@ -725,9 +889,10 @@ def test_typed_params_are_secret_free_and_rehydrate_from_environment( ) from nemo_retriever.operators.generation import SummarizationOperator + api_key_reference = "os.environ/OPENAI_API_KEY" params = TextGenerationParams.from_kwargs( - model="fake/model", - api_key="constructor-secret", + model="openai/fake-model", + api_key=api_key_reference, temperature=None, top_p=0.7, ) @@ -737,16 +902,16 @@ def test_typed_params_are_secret_free_and_rehydrate_from_environment( encoded = json.dumps(payload) assert payload["format_version"] == 2 - assert "constructor-secret" not in encoded assert "__pydantic_model__" in encoded - assert "__secret_env__" in encoded + assert api_key_reference in encoded - monkeypatch.setenv("NVIDIA_API_KEY", "worker-secret") + monkeypatch.setenv("NVIDIA_API_KEY", "wrong-provider-secret") + monkeypatch.setenv("OPENAI_API_KEY", "worker-secret") restored = deserialize_graph(payload) restored_operator = restored.roots[0].operator assert isinstance(restored_operator, SummarizationOperator) - assert restored_operator._params.transport.api_key == "worker-secret" + assert restored_operator._params.transport.api_key == api_key_reference assert restored_operator._params.sampling.model_fields_set == { "temperature", "top_p", @@ -756,8 +921,10 @@ def test_typed_params_are_secret_free_and_rehydrate_from_environment( path = tmp_path / "generation.json" save_graph(graph, path) - assert "constructor-secret" not in path.read_text() - assert isinstance(load_graph(path).roots[0].operator, SummarizationOperator) + assert "worker-secret" not in path.read_text() + loaded_operator = load_graph(path).roots[0].operator + assert isinstance(loaded_operator, SummarizationOperator) + assert loaded_operator._params.transport.api_key == api_key_reference def test_no_auth_survives_clone_with_environment_key_present(self, monkeypatch): from nemo_retriever.graph.graph_pipeline_registry import clone_graph @@ -788,17 +955,19 @@ def test_qa_and_registry_round_trip_to_concrete_operators( ) from nemo_retriever.tools.evaluation.generation import QAGenerationOperator + api_key_reference = "os.environ/OPENAI_API_KEY" operator = QAGenerationOperator( - model="fake/model", - api_key="qa-constructor-secret", + model="openai/fake-model", + api_key=api_key_reference, ) graph = self._single_root_graph(operator) payload = serialize_graph(graph) - monkeypatch.setenv("NVIDIA_API_KEY", "qa-worker-secret") + monkeypatch.setenv("NVIDIA_API_KEY", "wrong-provider-secret") + monkeypatch.setenv("OPENAI_API_KEY", "qa-worker-secret") restored = deserialize_graph(payload) assert isinstance(restored.roots[0].operator, QAGenerationOperator) - assert restored.roots[0].operator._client.transport.api_key == "qa-worker-secret" + assert restored.roots[0].operator._client.transport.api_key == api_key_reference registry = GraphPipelineRegistry() registry.register_graph( @@ -813,7 +982,9 @@ def test_qa_and_registry_round_trip_to_concrete_operators( assert raw["format_version"] == 2 assert set(raw["graphs"]) == {"qa"} - assert "qa-constructor-secret" not in path.read_text() + assert api_key_reference in path.read_text() + assert "qa-worker-secret" not in path.read_text() + assert "wrong-provider-secret" not in path.read_text() loaded_registry = GraphPipelineRegistry() assert loaded_registry.load_all(path) == ["qa"] @@ -916,7 +1087,8 @@ def test_v2_constructor_failures_raise_instead_of_using_placeholder(self): graph = self._single_root_graph(SummarizationOperator(_params())) payload = serialize_graph(graph) - payload["roots"][0]["operator_kwargs"]["unexpected"] = True + root_id = payload["roots"][0] + payload["nodes"][root_id]["operator_kwargs"]["unexpected"] = True with pytest.raises( GraphSerializationError, diff --git a/nemo_retriever/tests/test_graph_pipeline_registry.py b/nemo_retriever/tests/test_graph_pipeline_registry.py index e73761ab3..4f253cdc9 100644 --- a/nemo_retriever/tests/test_graph_pipeline_registry.py +++ b/nemo_retriever/tests/test_graph_pipeline_registry.py @@ -11,8 +11,9 @@ from typing import Any import pytest -from pydantic import BaseModel -from nemo_retriever.common.params import LLMRemoteClientParams +from pydantic import BaseModel, SecretStr, field_validator + +from nemo_retriever.common.params import LLMRemoteClientParams, RemoteInvokeParams, StoreParams from nemo_retriever.graph.graph_pipeline_registry import ( GraphBlueprint, GraphDiff, @@ -101,6 +102,141 @@ def postprocess(self, data: Any, **kw: Any) -> Any: return data +class PathOp(AbstractOperator): + def __init__(self, some_path: Path) -> None: + super().__init__(some_path=some_path) + self.some_path = some_path + + def preprocess(self, data: Any, **kw: Any) -> Any: + return data + + def process(self, data: Any, **kw: Any) -> Any: + return data + + def postprocess(self, data: Any, **kw: Any) -> Any: + return data + + +class ApiKeyOp(AbstractOperator): + def __init__(self, api_key: str | None = None) -> None: + super().__init__(api_key=api_key) + self.api_key = api_key + + def preprocess(self, data: Any, **kw: Any) -> Any: + return data + + def process(self, data: Any, **kw: Any) -> Any: + return data + + def postprocess(self, data: Any, **kw: Any) -> Any: + return data + + +class StoreParamsOp(AbstractOperator): + def __init__(self, params: StoreParams) -> None: + super().__init__(params=params) + self.params = params + + def preprocess(self, data: Any, **kw: Any) -> Any: + return data + + def process(self, data: Any, **kw: Any) -> Any: + return data + + def postprocess(self, data: Any, **kw: Any) -> Any: + return data + + +class CredentialReferenceParams(BaseModel): + api_key: str | None = None + + def _api_key_env_reference(self, field_name: str) -> str | None: + return "os.environ/OPENAI_API_KEY" if field_name == "api_key" else None + + +class CredentialReferenceOp(AbstractOperator): + def __init__(self, params: CredentialReferenceParams) -> None: + super().__init__(params=params) + self.params = params + + def preprocess(self, data: Any, **kw: Any) -> Any: + return data + + def process(self, data: Any, **kw: Any) -> Any: + return data + + def postprocess(self, data: Any, **kw: Any) -> Any: + return data + + +class RemoteParamsOp(AbstractOperator): + def __init__(self, params: RemoteInvokeParams) -> None: + super().__init__(params=params) + self.params = params + + def preprocess(self, data: Any, **kw: Any) -> Any: + return data + + def process(self, data: Any, **kw: Any) -> Any: + return data + + def postprocess(self, data: Any, **kw: Any) -> Any: + return data + + +class TypedValuesOp(AbstractOperator): + def __init__(self, values: dict[str, Any]) -> None: + super().__init__(values=values) + self.values = values + + def preprocess(self, data: Any, **kw: Any) -> Any: + return data + + def process(self, data: Any, **kw: Any) -> Any: + return data + + def postprocess(self, data: Any, **kw: Any) -> Any: + return data + + +class LeakyConstructorOp(AbstractOperator): + def __init__(self, value: str) -> None: + raise RuntimeError(f"constructor rejected {value}") + + def preprocess(self, data: Any, **kw: Any) -> Any: + return data + + def process(self, data: Any, **kw: Any) -> Any: + return data + + def postprocess(self, data: Any, **kw: Any) -> Any: + return data + + +class LeakyValidationParams(BaseModel): + value: str + + @field_validator("value") + @classmethod + def reject_value(cls, value: str) -> str: + raise ValueError(f"validation rejected {value}") + + +class ValidationParamsOp(AbstractOperator): + def __init__(self, params: LeakyValidationParams) -> None: + super().__init__(params=params) + self.params = params + + def preprocess(self, data: Any, **kw: Any) -> Any: + return data + + def process(self, data: Any, **kw: Any) -> Any: + return data + + def postprocess(self, data: Any, **kw: Any) -> Any: + return data + + class ParamsContainer(BaseModel): children: dict[str, LLMRemoteClientParams] @@ -149,6 +285,35 @@ def _multi_root_graph() -> Graph: return g +def _single_node_v2_payload(operator_class: type, operator_kwargs: dict[str, Any]) -> dict[str, Any]: + return { + "format_version": 2, + "roots": ["node_0"], + "nodes": { + "node_0": { + "name": "Test", + "operator_class": _qualified_name(operator_class), + "operator_kwargs": operator_kwargs, + "children": [], + } + }, + } + + +def _shared_child_graph() -> Graph: + root = Node(AddOp(1), name="Root") + left = Node(MulOp(2), name="Left") + right = Node(MulOp(3), name="Right") + shared = Node(AppendOp("_shared"), name="Shared") + root.add_child(left) + root.add_child(right) + left.add_child(shared) + right.add_child(shared) + graph = Graph() + graph.add_root(root) + return graph + + # ===================================================================== # Helper function tests # ===================================================================== @@ -664,7 +829,7 @@ def test_metadata_present(self): assert "metadata" in data assert data["metadata"]["node_count"] == 3 assert data["metadata"]["max_depth"] == 2 - assert "serialized_at" in data["metadata"] + assert "serialized_at" not in data["metadata"] def test_kwargs_preserved(self): original = _linear_graph() @@ -690,6 +855,155 @@ def test_diff_after_round_trip(self): assert result.identical is True +class TestNormalizedDagFormat: + def test_uses_node_ids_and_preserves_shared_identity(self): + graph = _shared_child_graph() + payload = serialize_graph(graph) + + assert payload["roots"] == ["node_0"] + assert list(payload["nodes"]) == ["node_0", "node_1", "node_2", "node_3"] + assert payload["nodes"]["node_1"]["children"] == ["node_2"] + assert payload["nodes"]["node_3"]["children"] == ["node_2"] + + restored = deserialize_graph(payload) + left, right = restored.roots[0].children + assert left.children[0] is right.children[0] + assert node_count(restored) == 4 + + def test_serialization_is_deterministic(self): + graph = _shared_child_graph() + first = serialize_graph(graph) + second = serialize_graph(graph) + + assert first == second + assert json.dumps(first, separators=(",", ":")) == json.dumps(second, separators=(",", ":")) + + def test_saved_json_bytes_are_deterministic(self, tmp_path): + graph = _shared_child_graph() + first = save_graph(graph, tmp_path / "first.json") + second = save_graph(graph, tmp_path / "second.json") + + assert first.read_bytes() == second.read_bytes() + + def test_multiple_roots_round_trip(self): + restored = deserialize_graph(serialize_graph(_multi_root_graph())) + assert [root.name for root in restored.roots] == ["R1", "R2"] + + def test_serialize_rejects_cycle(self): + first = Node(AddOp(1), name="First") + second = Node(MulOp(2), name="Second") + first.add_child(second) + second.add_child(first) + graph = Graph() + graph.add_root(first) + + with pytest.raises(GraphSerializationError, match="cycle detected"): + serialize_graph(graph) + + def test_deserialize_rejects_cycle(self): + payload = serialize_graph(_linear_graph()) + payload["nodes"]["node_2"]["children"] = ["node_0"] + with pytest.raises(GraphSerializationError, match="cycle detected"): + deserialize_graph(payload) + + def test_deserialize_rejects_missing_reference(self): + payload = serialize_graph(_linear_graph()) + payload["nodes"]["node_0"]["children"] = ["missing"] + with pytest.raises(GraphSerializationError, match="unknown node ID 'missing'"): + deserialize_graph(payload) + + def test_deserialize_rejects_unreachable_record(self): + payload = serialize_graph(_linear_graph()) + payload["nodes"]["orphan"] = { + "name": "Orphan", + "operator_class": _qualified_name(AddOp), + "operator_kwargs": {"value": 1}, + "children": [], + } + with pytest.raises(GraphSerializationError, match="unreachable node IDs"): + deserialize_graph(payload) + + def test_serialize_rejects_duplicate_edges_and_roots(self): + root = Node(AddOp(1), name="Root") + child = Node(MulOp(2), name="Child") + root.add_child(child) + root.add_child(child) + graph = Graph() + graph.add_root(root) + with pytest.raises(GraphSerializationError, match="duplicate child"): + serialize_graph(graph) + + fresh_root = Node(AddOp(1), name="Root") + graph = Graph() + graph.roots.extend([fresh_root, fresh_root]) + with pytest.raises(GraphSerializationError, match="duplicate root"): + serialize_graph(graph) + + def test_serialize_rejects_non_operator_class_override(self): + node = Node(AddOp(1), operator_class=dict) + graph = Graph() + graph.add_root(node) + with pytest.raises(GraphSerializationError, match="AbstractOperator class"): + serialize_graph(graph) + + def test_deserialize_rejects_duplicate_edges_and_roots(self): + payload = serialize_graph(_linear_graph()) + payload["nodes"]["node_0"]["children"] = ["node_1", "node_1"] + with pytest.raises(GraphSerializationError, match="duplicate child"): + deserialize_graph(payload) + + payload = serialize_graph(_multi_root_graph()) + payload["roots"].append(payload["roots"][0]) + with pytest.raises(GraphSerializationError, match="duplicate root"): + deserialize_graph(payload) + + def test_load_rejects_duplicate_node_ids_in_json(self, tmp_path): + path = tmp_path / "duplicates.json" + path.write_text('{"format_version":2,"roots":[],"nodes":{"node_0":{},"node_0":{}}}') + with pytest.raises(GraphSerializationError, match="duplicate JSON object key"): + load_graph(path) + + def test_draft_nested_v2_is_rejected(self): + payload = { + "format_version": 2, + "roots": [ + { + "name": "Test", + "operator_class": _qualified_name(AddOp), + "operator_kwargs": {"value": 1}, + "children": [], + } + ], + } + with pytest.raises(GraphSerializationError, match="missing required top-level"): + deserialize_graph(payload) + + def test_rejects_unknown_top_level_fields_and_inconsistent_metadata(self): + payload = serialize_graph(_linear_graph()) + payload["unexpected"] = True + with pytest.raises(GraphSerializationError, match="unknown top-level"): + deserialize_graph(payload) + + payload = serialize_graph(_linear_graph()) + payload["metadata"]["node_count"] = 999 + with pytest.raises(GraphSerializationError, match="metadata.node_count"): + deserialize_graph(payload) + + def test_rejects_malformed_typed_envelope(self): + payload = _single_node_v2_payload( + ApiKeyOp, + { + "api_key": { + "__secret_no_auth__": "", + "unexpected": "MUST-NOT-LEAK", + } + }, + ) + with pytest.raises(GraphSerializationError, match="malformed no-auth") as exc_info: + deserialize_graph(payload) + assert "MUST-NOT-LEAK" not in str(exc_info.value) + + class TestSaveLoadGraph: def test_save_and_load(self, tmp_path): graph = _linear_graph() @@ -753,23 +1067,38 @@ def test_process_raises(self): op.process("data") +class TestTypedValueSerialization: + def test_v2_round_trip_preserves_supported_value_types(self): + values = { + "path": Path("/tmp/example"), + "tuple": ("a", 1), + "set": {"b", "a"}, + "frozenset": frozenset({"c", "d"}), + "type": MulOp, + "callable": len, + } + graph = Graph() + graph.add_root(TypedValuesOp(values)) + + restored = deserialize_graph(serialize_graph(graph)).roots[0].operator.values + + assert restored["path"] == Path("/tmp/example") + assert restored["tuple"] == ("a", 1) + assert restored["set"] == {"a", "b"} + assert restored["frozenset"] == frozenset({"c", "d"}) + assert restored["type"] is MulOp + assert restored["callable"] is len + + class TestSpecialValueSerialization: def test_path_round_trip(self, tmp_path): - node = Node(AddOp(1), name="Test") - node.operator_kwargs["some_path"] = Path("/usr/local/bin") - data = serialize_graph(Graph()) - data["roots"] = [ - { - "name": "Test", - "operator_class": _qualified_name(AddOp), - "operator_kwargs": {"value": 1, "some_path": {"__path__": "/usr/local/bin"}}, - "children": [], - } - ] - restored = deserialize_graph(data) + graph = Graph() + graph.add_root(PathOp(Path("/usr/local/bin"))) + restored = deserialize_graph(serialize_graph(graph)) kw = restored.roots[0].operator_kwargs assert isinstance(kw["some_path"], Path) assert str(kw["some_path"]) == "/usr/local/bin" + assert restored.roots[0].operator.some_path == Path("/usr/local/bin") def test_set_round_trip(self): data = { @@ -777,7 +1106,10 @@ def test_set_round_trip(self): { "name": "Test", "operator_class": _qualified_name(AddOp), - "operator_kwargs": {"value": 1, "tags": {"__set__": ["a", "b", "c"]}}, + "operator_kwargs": { + "value": 1, + "tags": {"__set__": ["a", "b", "c"]}, + }, "children": [], } ] @@ -791,7 +1123,10 @@ def test_type_ref_round_trip(self): { "name": "Test", "operator_class": _qualified_name(AddOp), - "operator_kwargs": {"value": 1, "cls": {"__type_ref__": _qualified_name(MulOp)}}, + "operator_kwargs": { + "value": 1, + "cls": {"__type_ref__": _qualified_name(MulOp)}, + }, "children": [], } ] @@ -1091,6 +1426,9 @@ def test_save_graph_single(self, tmp_path): assert "blueprint" in data assert data["blueprint"]["description"] == "single test" + restored = load_graph(path) + assert node_count(restored) == 3 + def test_load_graph_single(self, tmp_path): reg = GraphPipelineRegistry() reg.register_graph("orig", _linear_graph, description="for load") @@ -1130,6 +1468,43 @@ def test_load_graph_no_overwrite_raises(self, tmp_path): with pytest.raises(ValueError, match="already registered"): reg.load_graph(path, name="x", overwrite=False) + def test_load_graph_validates_constructor_before_registration(self, tmp_path): + payload = _single_node_v2_payload( + LeakyConstructorOp, + {"value": "MUST-NOT-LEAK"}, + ) + payload["blueprint"] = {"name": "invalid"} + path = tmp_path / "invalid.json" + path.write_text(json.dumps(payload)) + + registry = GraphPipelineRegistry() + with pytest.raises(GraphSerializationError, match="failed to construct") as exc_info: + registry.load_graph(path) + + assert len(registry) == 0 + assert "MUST-NOT-LEAK" not in str(exc_info.value) + assert exc_info.value.__cause__ is None + assert exc_info.value.__context__ is None + + def test_load_all_validates_atomically_before_registration(self, tmp_path): + source = GraphPipelineRegistry() + source.register_graph("valid", _linear_graph) + source.register_graph("invalid", _fan_out_graph) + path = tmp_path / "all.json" + source.save_all(path) + + payload = json.loads(path.read_text()) + invalid_graph = payload["graphs"]["invalid"] + root_id = invalid_graph["roots"][0] + invalid_graph["nodes"][root_id]["children"] = ["missing"] + path.write_text(json.dumps(payload)) + + registry = GraphPipelineRegistry() + with pytest.raises(GraphSerializationError, match="unknown node ID"): + registry.load_all(path) + + assert len(registry) == 0 + def test_blueprint_metadata_preserved(self, tmp_path): reg = GraphPipelineRegistry() reg.register_graph( @@ -1222,6 +1597,29 @@ def my_factory(): assert isinstance(my_factory(), Graph) +class TestVersionlessV1Compatibility: + @pytest.mark.parametrize("explicit_version", [False, True]) + def test_concrete_operator_reconstructs_and_executes(self, explicit_version): + payload = { + "roots": [ + { + "name": "LegacyAdd", + "operator_class": _qualified_name(AddOp), + "operator_kwargs": {"value": 7}, + "children": [], + } + ] + } + if explicit_version: + payload["format_version"] = 1 + + restored = deserialize_graph(payload) + + assert type(restored.roots[0].operator) is AddOp + assert restored.roots[0].operator.value == 7 + assert restored.execute(5) == [12] + + class TestLegacyRegistryCompatibility: def test_v1_graph_may_be_named_format_version(self, tmp_path): payload = { @@ -1244,6 +1642,285 @@ def test_v1_graph_may_be_named_format_version(self, tmp_path): assert registry.build("format_version").roots == [] +class TestGraphSerializationSecurity: + def test_auto_resolved_params_preserve_exact_environment_provenance(self, monkeypatch): + monkeypatch.delenv("NVIDIA_API_KEY", raising=False) + monkeypatch.setenv("NGC_API_KEY", "DRIVER-KEY") + params = RemoteInvokeParams(api_key=None) + graph = Graph() + graph.add_root(RemoteParamsOp(params)) + + payload = serialize_graph(graph) + encoded = json.dumps(payload) + assert "DRIVER-KEY" not in encoded + assert "os.environ/NGC_API_KEY" in encoded + + monkeypatch.setenv("NGC_API_KEY", "WORKER-KEY") + restored = deserialize_graph(payload).roots[0].operator.params + + assert restored.api_key == "WORKER-KEY" + assert restored._api_key_env_reference("api_key") == "os.environ/NGC_API_KEY" + + def test_model_credential_provenance_preserves_provider_reference(self, monkeypatch): + monkeypatch.setenv("NVIDIA_API_KEY", "NVIDIA-WORKER-KEY") + params = CredentialReferenceParams(api_key="OPENAI-DRIVER-KEY") + graph = Graph() + graph.add_root(CredentialReferenceOp(params)) + + payload = serialize_graph(graph) + encoded = json.dumps(payload) + restored = deserialize_graph(payload).roots[0].operator.params + + assert "OPENAI-DRIVER-KEY" not in encoded + assert "NVIDIA-WORKER-KEY" not in encoded + assert "os.environ/OPENAI_API_KEY" in encoded + assert restored.api_key == "os.environ/OPENAI_API_KEY" + + @pytest.mark.parametrize( + "reference", + [ + "os.environ/", + "os.environ/1INVALID", + "os.environ/INVALID-NAME", + "os.environ/INVALID NAME", + ], + ) + def test_invalid_environment_reference_is_rejected(self, reference): + graph = Graph() + graph.add_root(ApiKeyOp(reference)) + with pytest.raises(GraphSerializationError, match="invalid environment reference"): + serialize_graph(graph) + + def test_no_auth_marker_round_trips(self): + graph = Graph() + graph.add_root(ApiKeyOp("")) + payload = serialize_graph(graph) + restored = deserialize_graph(payload).roots[0].operator + + assert "__secret_no_auth__" in json.dumps(payload) + assert restored.api_key == "" + + def test_v2_payload_rejects_literal_api_key_without_echoing_it(self): + payload = _single_node_v2_payload( + ApiKeyOp, + {"api_key": "MUST-NOT-LEAK"}, + ) + with pytest.raises(GraphSerializationError, match="cannot be persisted literally") as exc_info: + deserialize_graph(payload) + assert "MUST-NOT-LEAK" not in str(exc_info.value) + + def test_v2_reader_rejects_non_string_api_key_without_echoing_it(self): + payload = _single_node_v2_payload( + ApiKeyOp, + {"api_key": ["MUST-NOT-LEAK"]}, + ) + with pytest.raises(GraphSerializationError, match="invalid encoded API-key") as exc_info: + deserialize_graph(payload) + assert "MUST-NOT-LEAK" not in str(exc_info.value) + + @pytest.mark.parametrize( + "secret_field", + [ + "authorization", + "authorizationHeader", + "cookieHeader", + "credentialFile", + "key", + ], + ) + def test_v2_reader_rejects_nested_storage_secrets(self, secret_field): + payload = _single_node_v2_payload( + TypedValuesOp, + { + "values": { + "__mapping__": { + "storage_options": { + "__mapping__": { + secret_field: "MUST-NOT-LEAK", + } + } + } + } + }, + ) + with pytest.raises(GraphSerializationError, match="serialized secret fields must be empty") as exc_info: + deserialize_graph(payload) + assert "MUST-NOT-LEAK" not in str(exc_info.value) + + def test_safe_storage_options_round_trip(self): + params = StoreParams( + storage_uri="s3://bucket/prefix", + storage_options={ + "anon": True, + "region_name": "us-west-2", + "endpoint_url": "https://storage.example.test", + "timeout": 30, + }, + ) + graph = Graph() + graph.add_root(StoreParamsOp(params)) + restored = deserialize_graph(serialize_graph(graph)).roots[0].operator.params + + assert restored.storage_options == params.storage_options + + @pytest.mark.parametrize( + "secret_field", + [ + "authorization", + "authorizationHeader", + "password", + "refresh_token", + "secret", + "key", + "access-key", + "aws_access_key_id", + "account_key", + "cookie", + "cookie_header", + "cookieHeader", + "bearer", + "credential", + "credential_file", + "credentialFile", + ], + ) + def test_nested_storage_option_secrets_are_rejected(self, secret_field): + graph = Graph() + graph.add_root( + StoreParamsOp( + StoreParams( + storage_uri="s3://bucket/prefix", + storage_options={"inner": {secret_field: "MUST-NOT-LEAK"}}, + ) + ) + ) + with pytest.raises(GraphSerializationError, match="non-rehydratable secret field") as exc_info: + serialize_graph(graph) + assert "MUST-NOT-LEAK" not in str(exc_info.value) + assert f"storage_options.inner.{secret_field}" in str(exc_info.value) + + def test_storage_diagnostics_redact_contextual_key(self): + node = Node( + AddOp(1), + name="Diagnostic", + operator_kwargs={ + "storage_options": { + "anon": True, + "key": "MUST-NOT-LEAK", + } + }, + ) + rendered = format_node_details(node) + + assert "MUST-NOT-LEAK" not in rendered + assert "'anon': True" in rendered + assert "'key': '***'" in rendered + + @pytest.mark.parametrize( + "empty_container", + [{}, [], (), set(), frozenset()], + ) + def test_empty_secret_containers_are_rejected_before_json_encoding(self, empty_container): + with pytest.raises(GraphSerializationError, match="non-rehydratable secret field"): + _safe_serialize_value({"password": empty_container}) + + def test_unknown_mapping_key_error_never_calls_repr(self): + class DangerousKey: + def __hash__(self) -> int: + return 1 + + def __repr__(self) -> str: + raise AssertionError("repr must not be called") + + with pytest.raises(GraphSerializationError, match="DangerousKey"): + _safe_serialize_value({DangerousKey(): "value"}) + + def test_secret_emptiness_check_never_invokes_custom_equality(self): + class DangerousSecret: + def __eq__(self, other: object) -> bool: + raise RuntimeError("MUST-NOT-LEAK") + + def __repr__(self) -> str: + return "MUST-NOT-LEAK" + + with pytest.raises(GraphSerializationError, match="non-rehydratable secret field") as exc_info: + _safe_serialize_value({"password": DangerousSecret()}) + assert "MUST-NOT-LEAK" not in str(exc_info.value) + + def test_dataframe_is_rejected_with_a_type_only_error(self): + import pandas as pd + + with pytest.raises(GraphSerializationError, match="pandas.core.frame.DataFrame") as exc_info: + _safe_serialize_value(pd.DataFrame({"secret": ["MUST-NOT-LEAK"]})) + assert "MUST-NOT-LEAK" not in str(exc_info.value) + + def test_secret_str_is_rejected_without_revealing_value(self): + with pytest.raises(GraphSerializationError, match="unsupported value") as exc_info: + _safe_serialize_value({"label": SecretStr("MUST-NOT-LEAK")}) + assert "MUST-NOT-LEAK" not in str(exc_info.value) + + def test_diagnostics_use_type_only_for_unknown_objects(self): + class DangerousValue: + def __repr__(self) -> str: + return "MUST-NOT-LEAK" + + node = Node( + AddOp(1), + name="Diagnostic", + operator_kwargs={"value": 1, "opaque": DangerousValue()}, + ) + rendered = format_node_details(node) + + assert "MUST-NOT-LEAK" not in rendered + assert "DangerousValue" in rendered + + def test_diff_never_calls_repr_when_equality_raises(self): + class DangerousValue: + def __eq__(self, other: object) -> bool: + raise RuntimeError("comparison failed") + + def __repr__(self) -> str: + raise AssertionError("repr must not be called") + + graph_a = Graph() + graph_a.add_root(Node(AddOp(1), operator_kwargs={"opaque": DangerousValue()})) + graph_b = Graph() + graph_b.add_root(Node(AddOp(1), operator_kwargs={"opaque": DangerousValue()})) + + result = diff_graphs(graph_a, graph_b) + rendered = result.format() + + assert result.identical is False + assert "DangerousValue" in rendered + + def test_constructor_and_model_validation_errors_are_sanitized(self): + constructor_payload = _single_node_v2_payload( + LeakyConstructorOp, + {"value": "MUST-NOT-LEAK"}, + ) + with pytest.raises(GraphSerializationError, match="builtins.RuntimeError") as exc_info: + deserialize_graph(constructor_payload) + assert "MUST-NOT-LEAK" not in str(exc_info.value) + assert exc_info.value.__cause__ is None + assert exc_info.value.__context__ is None + + model_payload = _single_node_v2_payload( + ValidationParamsOp, + { + "params": { + "__pydantic_model__": _qualified_name(LeakyValidationParams), + "fields": {"value": "MUST-NOT-LEAK"}, + "fields_set": ["value"], + } + }, + ) + with pytest.raises(GraphSerializationError, match="pydantic_core.*ValidationError") as exc_info: + deserialize_graph(model_payload) + assert "MUST-NOT-LEAK" not in str(exc_info.value) + assert exc_info.value.__cause__ is None + assert exc_info.value.__context__ is None + + class TestTypedCodecRegressions: def test_nested_model_container_preserves_no_auth(self, monkeypatch): monkeypatch.setenv("NVIDIA_API_KEY", "ENV-SECRET") @@ -1282,33 +1959,32 @@ class LocalModel(BaseModel): def test_common_token_and_camel_case_secrets_are_rejected(self, field_name): with pytest.raises( GraphSerializationError, - match="non-rehydratable secret field", + match="non-rehydratable secret field|opaque mapping", ): _safe_serialize_value( {field_name: "MUST-NOT-LEAK"}, ) - def test_flat_operator_api_key_rehydrates_on_worker(self, monkeypatch): + def test_flat_operator_api_key_requires_reference_and_resolves_at_use(self, monkeypatch): from nemo_retriever.operators.graph_ops.subquery_operator import ( SubQueryGeneratorOperator, ) - original = SubQueryGeneratorOperator(llm_model="model", api_key="ORIGINAL") - graph = Graph() - graph.add_root(original) - payload = serialize_graph(graph) - - assert "ORIGINAL" not in json.dumps(payload) - monkeypatch.setenv("NVIDIA_API_KEY", "WORKER-KEY") - restored = deserialize_graph(payload).roots[0].operator - assert restored._api_key == "WORKER-KEY" - assert restored._resolve_api_key() == "WORKER-KEY" + literal_graph = Graph() + literal_graph.add_root(SubQueryGeneratorOperator(llm_model="model", api_key="ORIGINAL")) + with pytest.raises(GraphSerializationError, match="cannot be persisted literally") as exc_info: + serialize_graph(literal_graph) + assert "ORIGINAL" not in str(exc_info.value) monkeypatch.setenv("CUSTOM_LLM_KEY", "CUSTOM-WORKER-KEY") reference = "os.environ/CUSTOM_LLM_KEY" - referenced = SubQueryGeneratorOperator(llm_model="model", api_key=reference) referenced_graph = Graph() - referenced_graph.add_root(referenced) - restored_reference = deserialize_graph(serialize_graph(referenced_graph)).roots[0].operator - assert restored_reference._api_key == reference - assert restored_reference._resolve_api_key() == "CUSTOM-WORKER-KEY" + referenced_graph.add_root(SubQueryGeneratorOperator(llm_model="model", api_key=reference)) + payload = serialize_graph(referenced_graph) + encoded = json.dumps(payload) + restored = deserialize_graph(payload).roots[0].operator + + assert reference in encoded + assert "CUSTOM-WORKER-KEY" not in encoded + assert restored._api_key == reference + assert restored._resolve_api_key() == "CUSTOM-WORKER-KEY" diff --git a/nemo_retriever/tests/test_llm_params.py b/nemo_retriever/tests/test_llm_params.py index 0f7bc1169..b3893e758 100644 --- a/nemo_retriever/tests/test_llm_params.py +++ b/nemo_retriever/tests/test_llm_params.py @@ -70,17 +70,23 @@ def test_extra_forbid(self): with pytest.raises(ValueError): LLMRemoteClientParams(model="m", unknown_field=123) # type: ignore[call-arg] - def test_api_key_auto_resolved_from_env(self, monkeypatch): - """api_key=None should resolve from the remote-auth helper.""" + def test_api_key_none_defers_to_provider_native_lookup(self, monkeypatch): + """Generic LiteLLM transport must not substitute NVIDIA credentials.""" from nemo_retriever.common.params import models as params_models - monkeypatch.setattr(params_models, "resolve_remote_api_key", lambda: "resolved-secret") + def unexpected_resolution(): + raise AssertionError("generic transport must not resolve NVIDIA credentials") + + monkeypatch.setattr(params_models, "resolve_remote_api_key", unexpected_resolution) p = params_models.LLMRemoteClientParams(model="m") - assert p.api_key == "resolved-secret" + assert p.api_key is None def test_api_key_no_api_key_sentinel_yields_none(self): """Explicit NO_API_KEY sentinel suppresses auto-resolution.""" - from nemo_retriever.common.params.models import NO_API_KEY, LLMRemoteClientParams + from nemo_retriever.common.params.models import ( + NO_API_KEY, + LLMRemoteClientParams, + ) p = LLMRemoteClientParams(model="m", api_key=NO_API_KEY) assert p.api_key is None @@ -91,7 +97,10 @@ class TestLiteLLMClientConstruction: def test_structured_construction(self): from nemo_retriever.models.llm.clients import LiteLLMClient - from nemo_retriever.common.params.models import LLMInferenceParams, LLMRemoteClientParams + from nemo_retriever.common.params.models import ( + LLMInferenceParams, + LLMRemoteClientParams, + ) transport = LLMRemoteClientParams(model="openai/gpt-4o-mini", api_key="k") sampling = LLMInferenceParams(temperature=0.2, top_p=0.9, max_tokens=512) @@ -111,7 +120,10 @@ def test_default_sampling_matches_from_kwargs_for_rag_determinism(self): :meth:`LiteLLMClient.from_kwargs`. """ from nemo_retriever.models.llm.clients import LiteLLMClient - from nemo_retriever.common.params.models import LLMInferenceParams, LLMRemoteClientParams + from nemo_retriever.common.params.models import ( + LLMInferenceParams, + LLMRemoteClientParams, + ) client = LiteLLMClient(transport=LLMRemoteClientParams(model="m")) assert isinstance(client.sampling, LLMInferenceParams) @@ -213,20 +225,194 @@ def test_api_key_and_api_base_forwarded(self, mock_completion): assert kwargs["api_key"] == "secret" @patch("litellm.completion") - def test_extra_params_merged_last(self, mock_completion): - """extra_params should win over keys it overlaps with.""" + def test_allowed_nested_extra_params_merge_recursively(self, mock_completion): + """Per-request extensions win without discarding sibling values.""" from nemo_retriever.models.llm.clients import LiteLLMClient mock_completion.return_value = _fake_litellm_response("hi") client = LiteLLMClient.from_kwargs( model="m", - extra_params={"user": "tester", "num_retries": 99}, + extra_params={ + "user": "tester", + "provider": {"seed": 1, "mode": "stable"}, + }, + ) + client.complete( + [{"role": "user", "content": "hi"}], + extra_params={ + "provider": {"seed": 2, "extension": True}, + "stop": ["END"], + }, ) - client.complete([{"role": "user", "content": "hi"}]) kwargs = mock_completion.call_args.kwargs assert kwargs["user"] == "tester" - assert kwargs["num_retries"] == 99 + assert kwargs["provider"] == { + "seed": 2, + "mode": "stable", + "extension": True, + } + assert kwargs["stop"] == ["END"] + + +class TestLiteLLMHardening: + """Credential, protected-field, and text-only response contracts.""" + + @pytest.mark.parametrize( + "key", + [ + "model", + "messages", + "api_key", + "api_base", + "timeout", + "num_retries", + "temperature", + "top_p", + "max_tokens", + "tools", + "tool_choice", + "parallel_tool_calls", + "functions", + "function_call", + "stream", + "n", + ], + ) + def test_transport_rejects_every_protected_extra(self, key): + from nemo_retriever.common.params.models import LLMRemoteClientParams + + with pytest.raises(ValueError, match=key): + LLMRemoteClientParams(model="m", extra_params={key: "forbidden"}) + + @pytest.mark.parametrize( + "key", + [ + "model", + "messages", + "api_key", + "api_base", + "timeout", + "num_retries", + "temperature", + "top_p", + "max_tokens", + "tools", + "tool_choice", + "parallel_tool_calls", + "functions", + "function_call", + "stream", + "n", + ], + ) + @patch("litellm.completion") + def test_per_request_rejects_every_protected_extra(self, mock_completion, key): + from nemo_retriever.models.llm.clients import LiteLLMClient + + client = LiteLLMClient.from_kwargs(model="m") + with pytest.raises(ValueError, match=key): + client.complete( + [{"role": "user", "content": "hi"}], + extra_params={key: "forbidden"}, + ) + mock_completion.assert_not_called() + + @pytest.mark.parametrize( + ("model", "environment_name"), + [ + ("nvidia_nim/model", "NVIDIA_API_KEY"), + ("openai/model", "OPENAI_API_KEY"), + ("huggingface/model", "HUGGINGFACE_API_KEY"), + ("openai/custom", "MY_CUSTOM_PROVIDER_KEY"), + ], + ) + @patch("litellm.completion") + def test_explicit_environment_reference_resolves_at_call_time( + self, + mock_completion, + monkeypatch, + model, + environment_name, + ): + from nemo_retriever.models.llm.clients import LiteLLMClient + + expected = f"value-for-{environment_name}" + monkeypatch.setenv(environment_name, expected) + mock_completion.return_value = _fake_litellm_response("ok") + client = LiteLLMClient.from_kwargs( + model=model, + api_key=f"os.environ/{environment_name}", + ) + + assert client.transport.api_key == f"os.environ/{environment_name}" + client.complete([{"role": "user", "content": "hi"}]) + assert mock_completion.call_args.kwargs["api_key"] == expected + + @pytest.mark.parametrize("value", [None, ""]) + @patch("litellm.completion") + def test_missing_or_blank_explicit_environment_reference_fails_before_call( + self, + mock_completion, + monkeypatch, + value, + ): + from nemo_retriever.models.llm.clients import LiteLLMClient + + if value is None: + monkeypatch.delenv("MISSING_PROVIDER_KEY", raising=False) + else: + monkeypatch.setenv("MISSING_PROVIDER_KEY", value) + client = LiteLLMClient.from_kwargs( + model="openai/model", + api_key="os.environ/MISSING_PROVIDER_KEY", + ) + + with pytest.raises(ValueError, match="MISSING_PROVIDER_KEY"): + client.complete([{"role": "user", "content": "hi"}]) + mock_completion.assert_not_called() + + @patch("litellm.completion") + def test_none_omits_api_key_and_no_auth_forwards_inert_nonempty_key( + self, + mock_completion, + monkeypatch, + ): + from nemo_retriever.common.params.models import NO_API_KEY + from nemo_retriever.models.llm.clients import LiteLLMClient + + monkeypatch.setenv("NVIDIA_API_KEY", "must-not-be-substituted") + monkeypatch.setenv("OPENAI_API_KEY", "must-not-be-substituted-either") + mock_completion.return_value = _fake_litellm_response("ok") + + provider_native = LiteLLMClient.from_kwargs(model="openai/model") + provider_native.complete([{"role": "user", "content": "hi"}]) + assert "api_key" not in mock_completion.call_args.kwargs + + no_auth = LiteLLMClient.from_kwargs(model="openai/local", api_key=NO_API_KEY) + no_auth.complete([{"role": "user", "content": "hi"}]) + forwarded = mock_completion.call_args.kwargs["api_key"] + assert forwarded + assert forwarded not in { + "must-not-be-substituted", + "must-not-be-substituted-either", + } + + @pytest.mark.parametrize( + "reference", + [ + "os.environ/", + "os.environ/NOT-VALID", + "os.environ/9INVALID", + "os.environ/SPACE KEY", + " os.environ/OPENAI_API_KEY ", + ], + ) + def test_invalid_environment_reference_names_are_rejected(self, reference): + from nemo_retriever.common.params.models import LLMRemoteClientParams + + with pytest.raises(ValueError, match="environment references"): + LLMRemoteClientParams(model="m", api_key=reference) class TestLiteLLMRAGPrompt: @@ -249,7 +435,10 @@ def test_generate_disables_reasoning_with_portable_request_controls(self, mock_c assert messages[0]["role"] == "system" assert messages[0]["content"].startswith("/no_think\n") assert "precise question-answering assistant" in messages[0]["content"] - assert kwargs["chat_template_kwargs"] == {"reasoning_budget": 32, "enable_thinking": False} + assert kwargs["chat_template_kwargs"] == { + "reasoning_budget": 32, + "enable_thinking": False, + } assert result.answer == "answer" @patch("litellm.completion") @@ -308,7 +497,10 @@ def test_structured_construction_uses_defaults(self): def test_custom_sampling_override(self): from nemo_retriever.models.llm.clients import LLMJudge - from nemo_retriever.common.params.models import LLMInferenceParams, LLMRemoteClientParams + from nemo_retriever.common.params.models import ( + LLMInferenceParams, + LLMRemoteClientParams, + ) transport = LLMRemoteClientParams(model="m") sampling = LLMInferenceParams(temperature=0.4, max_tokens=1024) @@ -450,7 +642,9 @@ def test_pipeline_builder_generate_forwards_reasoning_enabled(self): assert len(generation_ops) == 1 assert generation_ops[0]._client.transport.reasoning_enabled is True - def test_pipeline_builder_generate_defaults_reasoning_enabled_for_legacy_client(self): + def test_pipeline_builder_generate_defaults_reasoning_enabled_for_legacy_client( + self, + ): from types import SimpleNamespace from unittest.mock import MagicMock @@ -581,7 +775,10 @@ def test_api_key_attribute_is_plain_str(self): def test_empty_api_key_not_masked(self): """Redaction only fires when a key is actually present.""" - from nemo_retriever.common.params.models import NO_API_KEY, LLMRemoteClientParams + from nemo_retriever.common.params.models import ( + NO_API_KEY, + LLMRemoteClientParams, + ) p = LLMRemoteClientParams(model="m", api_key=NO_API_KEY) assert p.api_key is None @@ -615,6 +812,46 @@ def test_nested_api_key_fields_also_masked(self): assert "page_elements_api_key=***" in rendered assert "ocr_api_key=***" in rendered + def test_nested_extra_param_secrets_are_redacted_from_repr(self): + from nemo_retriever.common.params.models import LLMRemoteClientParams + + secret = "sk-NESTED-MUST-NOT-LEAK" + params = LLMRemoteClientParams( + model="m", + extra_params={ + "provider": {"api_key": secret}, + "authorization": f"Bearer {secret}", + "authorizationHeader": f"Bearer {secret}", + }, + ) + + rendered = repr(params) + assert secret not in rendered + assert "***" in rendered + + def test_protected_extra_validation_error_hides_raw_input(self): + from nemo_retriever.common.params.models import LLMRemoteClientParams + + secret = "sk-VALIDATION-MUST-NOT-LEAK" + with pytest.raises(ValueError, match="protected request fields") as exc_info: + LLMRemoteClientParams( + model="m", + extra_params={"api_key": secret}, + ) + + assert secret not in str(exc_info.value) + + def test_mutated_literal_invalidates_environment_provenance(self): + from nemo_retriever.common.params.models import LLMRemoteClientParams + + params = LLMRemoteClientParams( + model="m", + api_key="os.environ/OPENAI_API_KEY", + ) + params.api_key = "sk-LITERAL-MUTATION" + + assert params._api_key_env_reference("api_key") is None + class TestLiteLLMDefaultModel: """Mirror of LLMJudge._DEFAULT_MODEL coverage for LiteLLMClient.""" @@ -663,7 +900,10 @@ def test_structured_and_flat_paths_agree_on_defaults(self): def test_explicit_sampling_is_not_overridden(self): """Passing an explicit ``LLMInferenceParams`` must win over the default.""" from nemo_retriever.models.llm.clients import LiteLLMClient - from nemo_retriever.common.params import LLMInferenceParams, LLMRemoteClientParams + from nemo_retriever.common.params import ( + LLMInferenceParams, + LLMRemoteClientParams, + ) client = LiteLLMClient( transport=LLMRemoteClientParams(model="m"), From 7d14bcf9c1a60d373df4b3cb4948807dbe895e07 Mon Sep 17 00:00:00 2001 From: Kyle Zheng <126034466+KyleZheng1284@users.noreply.github.com> Date: Mon, 6 Jul 2026 17:21:37 +0000 Subject: [PATCH 3/3] Address generation review feedback Signed-off-by: Kyle Zheng <126034466+KyleZheng1284@users.noreply.github.com> --- .../operators/generation/base.py | 31 ++++++----- .../tools/evaluation/generation.py | 6 ++- nemo_retriever/tests/test_generation_tasks.py | 53 +++++++++++++++++++ 3 files changed, 77 insertions(+), 13 deletions(-) diff --git a/nemo_retriever/src/nemo_retriever/operators/generation/base.py b/nemo_retriever/src/nemo_retriever/operators/generation/base.py index 5c367187f..b8eecfc90 100644 --- a/nemo_retriever/src/nemo_retriever/operators/generation/base.py +++ b/nemo_retriever/src/nemo_retriever/operators/generation/base.py @@ -259,7 +259,15 @@ def _execute_row(self, position: int, inputs: dict[str, Any]) -> tuple[int, Gene def _failure_model(self) -> str: try: model = self._client.model - except Exception: + except Exception as exc: + # Failure reporting must remain best-effort, but a broken client + # property should still be diagnosable. Log only the exception + # type: provider messages may contain request data or credentials. + exc_type = f"{type(exc).__module__}.{type(exc).__qualname__}" + logger.debug( + "Unable to read generation client model metadata (%s); using configured model", + exc_type, + ) return self._configured_model return model if isinstance(model, str) and model else self._configured_model @@ -296,17 +304,16 @@ def process(self, data: Any, **kwargs: Any) -> pd.DataFrame: for future in as_completed(futures): position = futures[future] - try: - result_position, result = future.result() - if result_position != position: - raise RuntimeError( - f"generation result position {result_position} does not match " - f"submitted position {position}" - ) - results[position] = result - except Exception: - logger.warning("Row %d generation failed (request_error)", position) - results[position] = self._failure_result("request_error", 0.0) + # _execute_row owns per-row failure collection. Exceptions + # or position mismatches here are executor/programming + # failures and must not be silently converted into row data. + result_position, result = future.result() + if result_position != position: + raise RuntimeError( + f"generation result position {result_position} does not match " + f"submitted position {position}" + ) + results[position] = result # Every non-empty row is assigned either a task result or a failure # result above. The cast-free local assertion catches future changes to diff --git a/nemo_retriever/src/nemo_retriever/tools/evaluation/generation.py b/nemo_retriever/src/nemo_retriever/tools/evaluation/generation.py index a6bc3ee24..c5f6ee6a1 100644 --- a/nemo_retriever/src/nemo_retriever/tools/evaluation/generation.py +++ b/nemo_retriever/src/nemo_retriever/tools/evaluation/generation.py @@ -119,7 +119,11 @@ def _execute_task(self, inputs: dict[str, Any]) -> GeneratedTextResult: started_at = time.monotonic() failure: GenerationTaskError | None = None try: - result = client.generate(inputs["query"], inputs["chunks"]) + result = client.generate( + inputs["query"], + inputs["chunks"], + reasoning_enabled=self._task.reasoning_enabled, + ) except Exception: failure = GenerationTaskError( code="transport_error", diff --git a/nemo_retriever/tests/test_generation_tasks.py b/nemo_retriever/tests/test_generation_tasks.py index 273d38f52..27c8b29ad 100644 --- a/nemo_retriever/tests/test_generation_tasks.py +++ b/nemo_retriever/tests/test_generation_tasks.py @@ -465,6 +465,35 @@ def test_nonempty_batch_uses_the_same_output_dtypes_as_empty_batch(self): "summary_error": "object", } + def test_failure_model_fallback_logs_only_sanitized_exception_type(self, caplog): + from nemo_retriever.operators.generation import SummarizationOperator + + class BrokenModelClient(FakeCompletionClient): + @property + def model(self): + raise RuntimeError("must-not-appear: secret-provider-message") + + operator = SummarizationOperator(_params(), client=BrokenModelClient()) + + with caplog.at_level("DEBUG", logger="nemo_retriever.operators.generation.base"): + assert operator._failure_model() == "fake/model" + + assert "builtins.RuntimeError" in caplog.text + assert "secret-provider-message" not in caplog.text + + def test_result_position_mismatch_is_an_operator_failure(self): + from nemo_retriever.operators.generation import SummarizationOperator + + class MisroutedResultOperator(SummarizationOperator): + def _execute_row(self, position, inputs): + _, result = super()._execute_row(position, inputs) + return position + 1, result + + operator = MisroutedResultOperator(_params(), client=FakeCompletionClient()) + + with pytest.raises(RuntimeError, match="result position 1 does not match submitted position 0"): + operator.run(pd.DataFrame({"text": ["source"]})) + class TestQACompatibility: def test_qa_run_preserves_schema_order_and_overwrites_legacy_outputs(self): @@ -752,6 +781,30 @@ def generate(self, query, chunks, *, reasoning_enabled=None): assert out.loc[0, "model"] == "legacy/model" assert out.loc[0, "gen_error"] is None + def test_generate_only_client_receives_operator_reasoning_setting(self): + from nemo_retriever.models.llm.types import GenerationResult + from nemo_retriever.tools.evaluation.generation import QAGenerationOperator + + class LegacyClient: + def __init__(self): + self.reasoning_values: list[bool | None] = [] + + def generate(self, query, chunks, *, reasoning_enabled=None): + self.reasoning_values.append(reasoning_enabled) + return GenerationResult("legacy answer", 0.2, "legacy/model") + + client = LegacyClient() + operator = QAGenerationOperator( + model="configured/model", + api_key="", + reasoning_enabled=False, + ) + operator._client = client + + operator.run(pd.DataFrame({"query": ["q"], "context": [["c"]]})) + + assert client.reasoning_values == [False] + def test_generate_only_failure_uses_configured_model(self): from nemo_retriever.tools.evaluation.generation import QAGenerationOperator