From 5587b97fe39a6f81ec04db9b60041e88cbf85a7c Mon Sep 17 00:00:00 2001 From: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> Date: Mon, 1 Jun 2026 14:34:46 +0200 Subject: [PATCH 1/3] refactor(llmrails): extract startup config loading Move config preparation, config.py hooks, Colang flow loading, embedding search defaults, and flow validation into startup helpers. --- nemoguardrails/rails/llm/startup/__init__.py | 16 ++ .../rails/llm/startup/colang_flows.py | 136 ++++++++++++++ .../rails/llm/startup/config_preparation.py | 76 ++++++++ nemoguardrails/rails/llm/startup/config_py.py | 55 ++++++ .../rails/llm/startup/config_validation.py | 32 ++++ .../rails/llm/startup/embedding_search.py | 147 +++++++++++++++ tests/rails/llm/test_colang_flows.py | 177 ++++++++++++++++++ tests/rails/llm/test_config_preparation.py | 121 ++++++++++++ tests/rails/llm/test_config_py.py | 103 ++++++++++ tests/rails/llm/test_config_validation.py | 52 +++++ tests/rails/llm/test_embedding_search.py | 151 +++++++++++++++ 11 files changed, 1066 insertions(+) create mode 100644 nemoguardrails/rails/llm/startup/__init__.py create mode 100644 nemoguardrails/rails/llm/startup/colang_flows.py create mode 100644 nemoguardrails/rails/llm/startup/config_preparation.py create mode 100644 nemoguardrails/rails/llm/startup/config_py.py create mode 100644 nemoguardrails/rails/llm/startup/config_validation.py create mode 100644 nemoguardrails/rails/llm/startup/embedding_search.py create mode 100644 tests/rails/llm/test_colang_flows.py create mode 100644 tests/rails/llm/test_config_preparation.py create mode 100644 tests/rails/llm/test_config_py.py create mode 100644 tests/rails/llm/test_config_validation.py create mode 100644 tests/rails/llm/test_embedding_search.py diff --git a/nemoguardrails/rails/llm/startup/__init__.py b/nemoguardrails/rails/llm/startup/__init__.py new file mode 100644 index 0000000000..dc33922794 --- /dev/null +++ b/nemoguardrails/rails/llm/startup/__init__.py @@ -0,0 +1,16 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +__all__ = [] diff --git a/nemoguardrails/rails/llm/startup/colang_flows.py b/nemoguardrails/rails/llm/startup/colang_flows.py new file mode 100644 index 0000000000..e4c577841a --- /dev/null +++ b/nemoguardrails/rails/llm/startup/colang_flows.py @@ -0,0 +1,136 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Colang flow loading, rail flow marking, and rail config validation.""" + +import importlib.resources as resources +import logging +from collections.abc import Iterator +from importlib.abc import Traversable + +from nemoguardrails.colang import parse_colang_file +from nemoguardrails.colang.v1_0.runtime.flows import _normalize_flow_id +from nemoguardrails.exceptions import InvalidRailsConfigurationError +from nemoguardrails.rails.llm.config import RailsConfig + +log = logging.getLogger("nemoguardrails.rails.llm.llmrails") + +_NEMOGUARDRAILS_PACKAGE = "nemoguardrails" + +__all__ = [ + "load_default_colang_1_flows", + "load_guardrails_library_flows_and_bot_messages", + "mark_rail_flows_as_system_subflows", + "validate_rail_flow_names", +] + + +def _package_resource(*parts: str) -> Traversable: + resource = resources.files(_NEMOGUARDRAILS_PACKAGE) + for part in parts: + resource = resource.joinpath(part) + return resource + + +def _iter_colang_resources(resource: Traversable) -> Iterator[Traversable]: + if resource.is_file(): + if resource.name.endswith(".co"): + yield resource + return + + # Sort children by name so the library load order is deterministic and independent + # of the filesystem's directory-iteration order. ``bot_messages`` are inserted on a + # first-seen-wins basis, so a non-deterministic order can change which utterances win. + children = sorted(resource.iterdir(), key=lambda child: child.name) + for child in children: + if child.is_file() and child.name.endswith(".co"): + yield child + + for child in children: + if child.is_dir(): + yield from _iter_colang_resources(child) + + +def load_default_colang_1_flows(config: RailsConfig) -> None: + """Load the built-in Colang 1.0 LLM flows into the rails config.""" + if config.colang_version != "1.0": + return + + default_flows_resource = _package_resource("rails", "llm", "llm_flows.co") + default_flows_content = default_flows_resource.read_text(encoding="utf-8") + default_flows = parse_colang_file(default_flows_resource.name, default_flows_content)["flows"] + + for flow_config in default_flows: + flow_config["is_system_flow"] = True + + config.flows.extend(default_flows) + + +def load_guardrails_library_flows_and_bot_messages(config: RailsConfig) -> None: + """Load Colang 1.0 flows and bot messages from the guardrails library.""" + if config.colang_version != "1.0": + return + + library_resource = _package_resource("library") + for colang_resource in _iter_colang_resources(library_resource): + log.debug(f"Loading file: {colang_resource}") + content = parse_colang_file( + colang_resource.name, + content=colang_resource.read_text(encoding="utf-8"), + version=config.colang_version, + ) + if not content: + continue + + for flow_config in content["flows"]: + flow_config["is_system_flow"] = True + + config.flows.extend(content["flows"]) + + for message_id, utterances in content.get("bot_messages", {}).items(): + if message_id not in config.bot_messages: + config.bot_messages[message_id] = utterances + + +def mark_rail_flows_as_system_subflows(config: RailsConfig) -> None: + """Mark configured input, output, and retrieval rail flows as system subflows.""" + rail_flow_ids = config.rails.input.flows + config.rails.output.flows + config.rails.retrieval.flows + + for flow_config in config.flows: + if flow_config.get("id") in rail_flow_ids: + flow_config["is_system_flow"] = True + flow_config["is_subflow"] = True + + +def validate_rail_flow_names(config: RailsConfig) -> None: + """Validate that configured rail flows exist in the config.""" + if config.colang_version == "1.0": + existing_flows_names = set([flow.get("id") for flow in config.flows]) + else: + existing_flows_names = set([flow.get("name") for flow in config.flows]) + + for flow_name in config.rails.input.flows: + flow_name = _normalize_flow_id(flow_name) + if flow_name not in existing_flows_names: + raise InvalidRailsConfigurationError(f"The provided input rail flow `{flow_name}` does not exist") + + for flow_name in config.rails.output.flows: + flow_name = _normalize_flow_id(flow_name) + if flow_name not in existing_flows_names: + raise InvalidRailsConfigurationError(f"The provided output rail flow `{flow_name}` does not exist") + + for flow_name in config.rails.retrieval.flows: + if flow_name not in existing_flows_names: + raise InvalidRailsConfigurationError(f"The provided retrieval rail flow `{flow_name}` does not exist") diff --git a/nemoguardrails/rails/llm/startup/config_preparation.py b/nemoguardrails/rails/llm/startup/config_preparation.py new file mode 100644 index 0000000000..65b877e86e --- /dev/null +++ b/nemoguardrails/rails/llm/startup/config_preparation.py @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""LLMRails config preparation.""" + +from dataclasses import dataclass +from typing import Any, Dict, Optional + +from nemoguardrails.rails.llm.config import RailsConfig +from nemoguardrails.rails.llm.startup.colang_flows import ( + load_default_colang_1_flows, + load_guardrails_library_flows_and_bot_messages, + mark_rail_flows_as_system_subflows, +) +from nemoguardrails.rails.llm.startup.embedding_search import apply_embedding_model_config + +_COLANG_FLOWS_PREPARED_ATTR = "_llmrails_colang_flows_prepared" + +__all__ = ["PreparedLLMRailsConfig", "prepare_llmrails_config"] + + +@dataclass +class PreparedLLMRailsConfig: + config: RailsConfig + default_embedding_model: Optional[str] + default_embedding_engine: Optional[str] + default_embedding_params: Dict[str, Any] + + +def prepare_llmrails_config( + *, + config: RailsConfig, + default_embedding_model: Optional[str], + default_embedding_engine: Optional[str], + default_embedding_params: Dict[str, Any], + in_place: bool = True, +) -> PreparedLLMRailsConfig: + """Prepare a RailsConfig for the standard LLMRails runtime.""" + prepared_config = config if in_place else config.model_copy(deep=True) + + if not getattr(prepared_config, _COLANG_FLOWS_PREPARED_ATTR, False): + load_default_colang_1_flows(prepared_config) + load_guardrails_library_flows_and_bot_messages(prepared_config) + setattr(prepared_config, _COLANG_FLOWS_PREPARED_ATTR, True) + + mark_rail_flows_as_system_subflows(prepared_config) + + ( + default_embedding_model, + default_embedding_engine, + default_embedding_params, + ) = apply_embedding_model_config( + config=prepared_config, + default_embedding_model=default_embedding_model, + default_embedding_engine=default_embedding_engine, + default_embedding_params=default_embedding_params, + ) + + return PreparedLLMRailsConfig( + config=prepared_config, + default_embedding_model=default_embedding_model, + default_embedding_engine=default_embedding_engine, + default_embedding_params=default_embedding_params, + ) diff --git a/nemoguardrails/rails/llm/startup/config_py.py b/nemoguardrails/rails/llm/startup/config_py.py new file mode 100644 index 0000000000..3ae315fef5 --- /dev/null +++ b/nemoguardrails/rails/llm/startup/config_py.py @@ -0,0 +1,55 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Config.py module hooks.""" + +import importlib.util +import os +from types import ModuleType +from typing import Any, List + +from nemoguardrails.rails.llm.config import RailsConfig + +__all__ = ["load_config_py_modules", "run_config_py_init_hooks"] + + +def load_config_py_modules(config: RailsConfig) -> List[ModuleType]: + """Load config.py modules from imported config paths and the main config path.""" + config_modules = [] + config_paths = list(config.imported_paths.values() if config.imported_paths else []) + [config.config_path] + + for config_path in config_paths: + if not config_path: + continue + + filepath = os.path.join(config_path, "config.py") + if not os.path.exists(filepath): + continue + + filename = os.path.basename(filepath) + spec = importlib.util.spec_from_file_location(filename, filepath) + if spec and spec.loader: + config_module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(config_module) + config_modules.append(config_module) + + return config_modules + + +def run_config_py_init_hooks(rails: Any, config_modules: List[ModuleType]) -> None: + """Run config.py init hooks with the LLMRails instance.""" + for config_module in config_modules: + if hasattr(config_module, "init"): + config_module.init(rails) diff --git a/nemoguardrails/rails/llm/startup/config_validation.py b/nemoguardrails/rails/llm/startup/config_validation.py new file mode 100644 index 0000000000..bb5a15d9f8 --- /dev/null +++ b/nemoguardrails/rails/llm/startup/config_validation.py @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""LLMRails startup config validation.""" + +from nemoguardrails.exceptions import InvalidRailsConfigurationError +from nemoguardrails.rails.llm.config import RailsConfig +from nemoguardrails.rails.llm.startup.colang_flows import validate_rail_flow_names + +__all__ = ["validate_llmrails_config"] + + +def validate_llmrails_config(config: RailsConfig) -> None: + validate_rail_flow_names(config) + + if config.passthrough and config.rails.dialog.single_call.enabled: + raise InvalidRailsConfigurationError( + "The passthrough mode and the single call dialog rails mode can't be used at the same time. " + "The single call mode needs to use an altered prompt when prompting the LLM. " + ) diff --git a/nemoguardrails/rails/llm/startup/embedding_search.py b/nemoguardrails/rails/llm/startup/embedding_search.py new file mode 100644 index 0000000000..01a63bcc71 --- /dev/null +++ b/nemoguardrails/rails/llm/startup/embedding_search.py @@ -0,0 +1,147 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Embedding search provider setup.""" + +from dataclasses import dataclass +from typing import Any, Dict, Optional, Tuple, Type + +from nemoguardrails.embeddings.index import EmbeddingsIndex +from nemoguardrails.rails.llm.config import EmbeddingSearchProvider, RailsConfig + +DEFAULT_EMBEDDING_MODEL = "all-MiniLM-L6-v2" +DEFAULT_EMBEDDING_ENGINE = "FastEmbed" + +__all__ = [ + "DEFAULT_EMBEDDING_ENGINE", + "DEFAULT_EMBEDDING_MODEL", + "EmbeddingSearchState", + "apply_embedding_model_config", + "get_embedding_search_provider_instance", +] + + +@dataclass +class EmbeddingSearchState: + providers: Dict[str, Type[EmbeddingsIndex]] + default_model: Optional[str] + default_engine: Optional[str] + default_params: Dict[str, Any] + + @classmethod + def default(cls) -> "EmbeddingSearchState": + return cls( + providers={}, + default_model=DEFAULT_EMBEDDING_MODEL, + default_engine=DEFAULT_EMBEDDING_ENGINE, + default_params={}, + ) + + def update_defaults( + self, + *, + default_model: Optional[str], + default_engine: Optional[str], + default_params: Dict[str, Any], + ) -> None: + self.default_model = default_model + self.default_engine = default_engine + self.default_params = default_params + + def register_provider(self, name: str, cls: Type[EmbeddingsIndex]) -> None: + self.providers[name] = cls + + def get_provider_instance( + self, + esp_config: Optional[EmbeddingSearchProvider] = None, + ) -> EmbeddingsIndex: + return get_embedding_search_provider_instance( + embedding_search_providers=self.providers, + default_embedding_model=self.default_model, + default_embedding_engine=self.default_engine, + default_embedding_params=self.default_params, + esp_config=esp_config, + ) + + +def apply_embedding_model_config( + config: RailsConfig, + default_embedding_model: Optional[str], + default_embedding_engine: Optional[str], + default_embedding_params: Dict[str, Any], +) -> Tuple[Optional[str], Optional[str], Dict[str, Any]]: + """Apply an embeddings model config to the default embedding search settings.""" + for model in config.models: + if model.type != "embeddings": + continue + + default_embedding_model = model.model + default_embedding_engine = model.engine + default_embedding_params = model.parameters or {} + + for esp in [ + config.core.embedding_search_provider, + config.knowledge_base.embedding_search_provider, + ]: + if esp.name != "default": + continue + if "embedding_model" not in esp.parameters and model.model is not None: + esp.parameters["embedding_model"] = model.model + if "embedding_engine" not in esp.parameters and model.engine is not None: + esp.parameters["embedding_engine"] = model.engine + + break + + return default_embedding_model, default_embedding_engine, default_embedding_params + + +def get_embedding_search_provider_instance( + embedding_search_providers: Dict[str, Type[EmbeddingsIndex]], + default_embedding_model: Optional[str], + default_embedding_engine: Optional[str], + default_embedding_params: Dict[str, Any], + esp_config: Optional[EmbeddingSearchProvider] = None, +) -> EmbeddingsIndex: + """Create the embedding search provider selected by a provider config.""" + if esp_config is None: + esp_config = EmbeddingSearchProvider() + + if esp_config.name == "default": + from nemoguardrails.embeddings.basic import BasicEmbeddingsIndex + + return BasicEmbeddingsIndex( + embedding_model=esp_config.parameters.get("embedding_model", default_embedding_model), + embedding_engine=esp_config.parameters.get("embedding_engine", default_embedding_engine), + embedding_params=esp_config.parameters.get("embedding_parameters", default_embedding_params), + cache_config=esp_config.cache, + **{ + k: v + for k, v in esp_config.parameters.items() + if k + in [ + "use_batching", + "max_batch_size", + "matx_batch_hold", + "search_threshold", + ] + and v is not None + }, + ) + + if esp_config.name not in embedding_search_providers: + raise Exception(f"Unknown embedding search provider: {esp_config.name}") + + kwargs = esp_config.parameters + return embedding_search_providers[esp_config.name](**kwargs) diff --git a/tests/rails/llm/test_colang_flows.py b/tests/rails/llm/test_colang_flows.py new file mode 100644 index 0000000000..8a5857c93a --- /dev/null +++ b/tests/rails/llm/test_colang_flows.py @@ -0,0 +1,177 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest + +import nemoguardrails.rails.llm.startup.colang_flows as colang_flows +from nemoguardrails.exceptions import InvalidRailsConfigurationError +from nemoguardrails.rails.llm.config import RailsConfig +from nemoguardrails.rails.llm.startup.colang_flows import ( + load_default_colang_1_flows, + load_guardrails_library_flows_and_bot_messages, + mark_rail_flows_as_system_subflows, + validate_rail_flow_names, +) + + +def test_load_default_colang_1_flows_marks_loaded_flows_as_system(): + config = RailsConfig(models=[]) + + load_default_colang_1_flows(config) + + assert config.flows + assert "process user input" in {flow.get("id") for flow in config.flows} + assert all(flow.get("is_system_flow") is True for flow in config.flows) + + +def test_load_guardrails_library_flows_and_bot_messages_loads_package_resources(): + config = RailsConfig(models=[]) + + load_guardrails_library_flows_and_bot_messages(config) + + flow_ids = {flow.get("id") for flow in config.flows} + assert "content safety check input" in flow_ids + assert "regex check input" in flow_ids + assert all(flow.get("is_system_flow") is True for flow in config.flows) + assert config.bot_messages["refuse to respond"] == ["I'm sorry, I can't respond to that."] + assert config.bot_messages["inform cannot engage with sensitive content"] == [ + "I will not engage with sensitive content." + ] + + +def test_load_guardrails_library_flows_and_bot_messages_keeps_existing_bot_messages(): + config = RailsConfig( + models=[], + bot_messages={"refuse to respond": ["Custom refusal."]}, + ) + + load_guardrails_library_flows_and_bot_messages(config) + + assert config.bot_messages["refuse to respond"] == ["Custom refusal."] + assert config.bot_messages["inform cannot engage with sensitive content"] == [ + "I will not engage with sensitive content." + ] + + +def test_colang_flow_loading_does_not_depend_on_colang_flows_file_location(monkeypatch, tmp_path): + monkeypatch.setattr(colang_flows, "__file__", str(tmp_path / "moved" / "colang_flows.py")) + + default_config = RailsConfig(models=[]) + colang_flows.load_default_colang_1_flows(default_config) + + library_config = RailsConfig(models=[]) + colang_flows.load_guardrails_library_flows_and_bot_messages(library_config) + + assert "process user input" in {flow.get("id") for flow in default_config.flows} + assert "content safety check input" in {flow.get("id") for flow in library_config.flows} + assert library_config.bot_messages["inform cannot engage with sensitive content"] == [ + "I will not engage with sensitive content." + ] + + +def test_mark_rail_flows_as_system_subflows_mutates_matching_flow_configs(): + config = RailsConfig( + models=[], + flows=[ + {"id": "check input", "is_system_flow": False}, + {"id": "regular flow", "is_system_flow": False}, + ], + ) + config.rails.input.flows = ["check input"] + + mark_rail_flows_as_system_subflows(config) + + assert config.flows[0]["is_system_flow"] is True + assert config.flows[0]["is_subflow"] is True + assert config.flows[1]["is_system_flow"] is False + assert "is_subflow" not in config.flows[1] + + +class _FakeResource: + """Minimal Traversable stand-in whose ``iterdir`` order we control.""" + + def __init__(self, name, *, is_dir=False, children=None): + self._name = name + self._is_dir = is_dir + self._children = children or [] + + @property + def name(self): + return self._name + + def is_file(self): + return not self._is_dir + + def is_dir(self): + return self._is_dir + + def iterdir(self): + return iter(self._children) + + +def test_iter_colang_resources_yields_sorted_filesystem_independent_order(): + # iterdir() returns children in an arbitrary (here: unsorted) order; the traversal + # must still yield .co files in a deterministic, sorted order so the library load + # order does not depend on the filesystem (the dropped os.walk + sort behavior). + resource = _FakeResource( + "library", + is_dir=True, + children=[ + _FakeResource("zebra.co"), + _FakeResource("alpha.co"), + _FakeResource("not_colang.txt"), + _FakeResource( + "subpack", + is_dir=True, + children=[_FakeResource("yray.co"), _FakeResource("beta.co")], + ), + _FakeResource("middle.co"), + ], + ) + + names = [item.name for item in colang_flows._iter_colang_resources(resource)] + + # Files in the directory come first (sorted), then sorted subdirectories are + # descended into (their files sorted); non-.co files are skipped. + assert names == ["alpha.co", "middle.co", "zebra.co", "beta.co", "yray.co"] + + +def test_iter_colang_resources_order_does_not_depend_on_iterdir_order(): + # The same children in two different iterdir() orders must yield the same sequence. + files_forward = [_FakeResource(f"{name}.co") for name in ("a", "b", "c")] + files_reversed = list(reversed(files_forward)) + + forward = [ + r.name for r in colang_flows._iter_colang_resources(_FakeResource("d", is_dir=True, children=files_forward)) + ] + backward = [ + r.name for r in colang_flows._iter_colang_resources(_FakeResource("d", is_dir=True, children=files_reversed)) + ] + + assert forward == backward == ["a.co", "b.co", "c.co"] + + +def test_validate_rail_flow_names_keeps_existing_flow_validation_behavior(): + config = RailsConfig( + models=[], + flows=[{"id": "content safety check input"}], + ) + config.rails.input.flows = ["content safety check input $model=content_safety"] + + validate_rail_flow_names(config) + + config.rails.output.flows = ["missing output rail"] + with pytest.raises(InvalidRailsConfigurationError, match="`missing output rail` does not exist"): + validate_rail_flow_names(config) diff --git a/tests/rails/llm/test_config_preparation.py b/tests/rails/llm/test_config_preparation.py new file mode 100644 index 0000000000..5cbd0067b0 --- /dev/null +++ b/tests/rails/llm/test_config_preparation.py @@ -0,0 +1,121 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from nemoguardrails.rails.llm.config import Model, RailsConfig +from nemoguardrails.rails.llm.startup import config_preparation as startup_config_preparation +from nemoguardrails.rails.llm.startup.config_preparation import ( + prepare_llmrails_config, +) +from nemoguardrails.rails.llm.startup.embedding_search import ( + DEFAULT_EMBEDDING_ENGINE, + DEFAULT_EMBEDDING_MODEL, +) + + +def test_config_preparation_startup_module_exports_public_helpers(): + assert startup_config_preparation.__all__ == [ + "PreparedLLMRailsConfig", + "prepare_llmrails_config", + ] + + +def test_prepare_llmrails_config_loads_colang_flows_once_in_place(): + config = RailsConfig(models=[]) + + first = prepare_llmrails_config( + config=config, + default_embedding_model=DEFAULT_EMBEDDING_MODEL, + default_embedding_engine=DEFAULT_EMBEDDING_ENGINE, + default_embedding_params={}, + ) + first_flow_count = len(config.flows) + + second = prepare_llmrails_config( + config=config, + default_embedding_model=DEFAULT_EMBEDDING_MODEL, + default_embedding_engine=DEFAULT_EMBEDDING_ENGINE, + default_embedding_params={}, + ) + + assert first.config is config + assert second.config is config + assert first_flow_count > 0 + assert len(config.flows) == first_flow_count + + +def test_prepare_llmrails_config_marks_rail_flows_each_time(): + config = RailsConfig( + models=[], + flows=[{"id": "check input", "is_system_flow": False}], + ) + + prepare_llmrails_config( + config=config, + default_embedding_model=DEFAULT_EMBEDDING_MODEL, + default_embedding_engine=DEFAULT_EMBEDDING_ENGINE, + default_embedding_params={}, + ) + + config.rails.input.flows = ["check input"] + prepare_llmrails_config( + config=config, + default_embedding_model=DEFAULT_EMBEDDING_MODEL, + default_embedding_engine=DEFAULT_EMBEDDING_ENGINE, + default_embedding_params={}, + ) + + assert config.flows[0]["is_system_flow"] is True + assert config.flows[0]["is_subflow"] is True + + +def test_prepare_llmrails_config_can_prepare_a_copy(): + config = RailsConfig(models=[]) + + prepared = prepare_llmrails_config( + config=config, + default_embedding_model=DEFAULT_EMBEDDING_MODEL, + default_embedding_engine=DEFAULT_EMBEDDING_ENGINE, + default_embedding_params={}, + in_place=False, + ) + + assert prepared.config is not config + assert len(config.flows) == 0 + assert len(prepared.config.flows) > 0 + + +def test_prepare_llmrails_config_returns_embedding_defaults(): + config = RailsConfig( + models=[ + Model( + type="embeddings", + engine="SentenceTransformers", + model="intfloat/e5-large-v2", + parameters={"device": "cpu"}, + ) + ], + ) + + prepared = prepare_llmrails_config( + config=config, + default_embedding_model=DEFAULT_EMBEDDING_MODEL, + default_embedding_engine=DEFAULT_EMBEDDING_ENGINE, + default_embedding_params={}, + ) + + assert prepared.default_embedding_model == "intfloat/e5-large-v2" + assert prepared.default_embedding_engine == "SentenceTransformers" + assert prepared.default_embedding_params == {"device": "cpu"} + assert config.core.embedding_search_provider.parameters["embedding_model"] == "intfloat/e5-large-v2" diff --git a/tests/rails/llm/test_config_py.py b/tests/rails/llm/test_config_py.py new file mode 100644 index 0000000000..45b4822794 --- /dev/null +++ b/tests/rails/llm/test_config_py.py @@ -0,0 +1,103 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from types import ModuleType + +from nemoguardrails.rails.llm.config import RailsConfig +from nemoguardrails.rails.llm.llmrails import LLMRails +from nemoguardrails.rails.llm.startup.config_py import ( + load_config_py_modules, + run_config_py_init_hooks, +) + + +def test_load_config_py_modules_preserves_imported_then_main_order(tmp_path): + imported_path = tmp_path / "imported" + main_path = tmp_path / "main" + imported_path.mkdir() + main_path.mkdir() + (imported_path / "config.py").write_text("SOURCE = 'imported'\n") + (main_path / "config.py").write_text("SOURCE = 'main'\n") + + config = RailsConfig( + config_path=str(main_path), + imported_paths={"imported": str(imported_path)}, + models=[], + ) + + config_modules = load_config_py_modules(config) + + assert [config_module.SOURCE for config_module in config_modules] == [ + "imported", + "main", + ] + + +def test_load_config_py_modules_skips_missing_paths(tmp_path): + main_path = tmp_path / "main" + imported_path = tmp_path / "imported" + main_path.mkdir() + imported_path.mkdir() + (main_path / "config.py").write_text("SOURCE = 'main'\n") + + config = RailsConfig( + config_path=str(main_path), + imported_paths={"imported": str(imported_path)}, + models=[], + ) + + config_modules = load_config_py_modules(config) + + assert [config_module.SOURCE for config_module in config_modules] == ["main"] + + +def test_llmrails_runs_imported_config_py_init_before_main(tmp_path): + imported_path = tmp_path / "imported" + main_path = tmp_path / "main" + imported_path.mkdir() + main_path.mkdir() + (imported_path / "config.py").write_text("def init(app):\n app.hook_order = ['imported']\n") + (main_path / "config.py").write_text("def init(app):\n app.hook_order.append('main')\n") + + config = RailsConfig( + config_path=str(main_path), + imported_paths={"imported": str(imported_path)}, + models=[], + ) + + rails = LLMRails(config) + + assert getattr(rails, "hook_order") == ["imported", "main"] + + +def test_run_config_py_init_hooks_passes_rails_instance(): + rails = object() + config_module = ModuleType("config") + + def init(received_rails): + setattr(config_module, "received_rails", received_rails) + + setattr(config_module, "init", init) + + run_config_py_init_hooks(rails, [config_module]) + + assert getattr(config_module, "received_rails") is rails + + +def test_run_config_py_init_hooks_ignores_modules_without_init(): + rails = object() + config_module = ModuleType("config") + + run_config_py_init_hooks(rails, [config_module]) diff --git a/tests/rails/llm/test_config_validation.py b/tests/rails/llm/test_config_validation.py new file mode 100644 index 0000000000..9d50875da3 --- /dev/null +++ b/tests/rails/llm/test_config_validation.py @@ -0,0 +1,52 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest + +from nemoguardrails.exceptions import InvalidRailsConfigurationError +from nemoguardrails.rails.llm.config import RailsConfig +from nemoguardrails.rails.llm.startup import config_validation +from nemoguardrails.rails.llm.startup.config_validation import validate_llmrails_config + + +def test_config_validation_startup_module_exports_public_helpers(): + assert config_validation.__all__ == ["validate_llmrails_config"] + + +def test_validate_llmrails_config_delegates_rail_flow_validation(monkeypatch): + config = RailsConfig(models=[]) + calls = [] + + def validate_rail_flow_names(config_to_validate): + calls.append(config_to_validate) + + monkeypatch.setattr(config_validation, "validate_rail_flow_names", validate_rail_flow_names) + + validate_llmrails_config(config) + + assert calls == [config] + + +def test_validate_llmrails_config_rejects_passthrough_with_single_call(): + config = RailsConfig.from_content( + config={ + "models": [], + "passthrough": True, + "rails": {"dialog": {"single_call": {"enabled": True}}}, + } + ) + + with pytest.raises(InvalidRailsConfigurationError, match="passthrough mode"): + validate_llmrails_config(config) diff --git a/tests/rails/llm/test_embedding_search.py b/tests/rails/llm/test_embedding_search.py new file mode 100644 index 0000000000..4676b12651 --- /dev/null +++ b/tests/rails/llm/test_embedding_search.py @@ -0,0 +1,151 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest + +from nemoguardrails.embeddings.basic import BasicEmbeddingsIndex +from nemoguardrails.embeddings.index import EmbeddingsIndex +from nemoguardrails.rails.llm.config import EmbeddingSearchProvider, Model, RailsConfig +from nemoguardrails.rails.llm.startup import embedding_search as startup_embedding_search +from nemoguardrails.rails.llm.startup.embedding_search import ( + DEFAULT_EMBEDDING_ENGINE, + DEFAULT_EMBEDDING_MODEL, + EmbeddingSearchState, + apply_embedding_model_config, + get_embedding_search_provider_instance, +) + + +def test_embedding_search_startup_module_exports_public_helpers(): + assert startup_embedding_search.__all__ == [ + "DEFAULT_EMBEDDING_ENGINE", + "DEFAULT_EMBEDDING_MODEL", + "EmbeddingSearchState", + "apply_embedding_model_config", + "get_embedding_search_provider_instance", + ] + + +def test_embedding_search_state_owns_providers_defaults_and_provider_creation(): + class CustomEmbeddingsIndex(EmbeddingsIndex): + def __init__(self, prefix: str): + self.prefix = prefix + + state = EmbeddingSearchState.default() + + assert state.providers == {} + assert state.default_model == DEFAULT_EMBEDDING_MODEL + assert state.default_engine == DEFAULT_EMBEDDING_ENGINE + assert state.default_params == {} + + state.update_defaults( + default_model="prepared-model", + default_engine="PreparedEngine", + default_params={"device": "cpu"}, + ) + state.register_provider("custom", CustomEmbeddingsIndex) + + index = state.get_provider_instance(EmbeddingSearchProvider(name="custom", parameters={"prefix": "docs"})) + + assert isinstance(index, CustomEmbeddingsIndex) + assert index.prefix == "docs" + assert state.default_model == "prepared-model" + assert state.default_engine == "PreparedEngine" + assert state.default_params == {"device": "cpu"} + + +def test_apply_embedding_model_config_backfills_default_search_providers(): + config = RailsConfig( + models=[ + Model( + type="embeddings", + engine="SentenceTransformers", + model="intfloat/e5-large-v2", + parameters={"device": "cpu"}, + ) + ], + ) + + default_model, default_engine, default_params = apply_embedding_model_config( + config=config, + default_embedding_model=DEFAULT_EMBEDDING_MODEL, + default_embedding_engine=DEFAULT_EMBEDDING_ENGINE, + default_embedding_params={}, + ) + + assert default_model == "intfloat/e5-large-v2" + assert default_engine == "SentenceTransformers" + assert default_params == {"device": "cpu"} + assert config.core.embedding_search_provider.parameters["embedding_model"] == "intfloat/e5-large-v2" + assert config.core.embedding_search_provider.parameters["embedding_engine"] == "SentenceTransformers" + assert config.knowledge_base.embedding_search_provider.parameters["embedding_model"] == "intfloat/e5-large-v2" + assert config.knowledge_base.embedding_search_provider.parameters["embedding_engine"] == "SentenceTransformers" + + +def test_get_embedding_search_provider_instance_uses_default_provider_parameters(): + esp_config = EmbeddingSearchProvider( + parameters={ + "embedding_model": "custom-model", + "embedding_engine": "CustomEngine", + "embedding_parameters": {"device": "cpu"}, + "use_batching": True, + "max_batch_size": 3, + "search_threshold": 0.7, + } + ) + + index = get_embedding_search_provider_instance( + embedding_search_providers={}, + default_embedding_model=DEFAULT_EMBEDDING_MODEL, + default_embedding_engine=DEFAULT_EMBEDDING_ENGINE, + default_embedding_params={}, + esp_config=esp_config, + ) + + assert isinstance(index, BasicEmbeddingsIndex) + assert index.embedding_model == "custom-model" + assert index.embedding_engine == "CustomEngine" + assert index.embedding_params == {"device": "cpu"} + assert index.use_batching is True + assert index.max_batch_size == 3 + assert index.search_threshold == 0.7 + + +def test_get_embedding_search_provider_instance_uses_registered_custom_provider(): + class CustomEmbeddingsIndex(EmbeddingsIndex): + def __init__(self, prefix: str): + self.prefix = prefix + + index = get_embedding_search_provider_instance( + embedding_search_providers={"custom": CustomEmbeddingsIndex}, + default_embedding_model=DEFAULT_EMBEDDING_MODEL, + default_embedding_engine=DEFAULT_EMBEDDING_ENGINE, + default_embedding_params={}, + esp_config=EmbeddingSearchProvider(name="custom", parameters={"prefix": "docs"}), + ) + + assert isinstance(index, CustomEmbeddingsIndex) + assert index.prefix == "docs" + + +def test_get_embedding_search_provider_instance_rejects_unknown_provider(): + with pytest.raises(Exception, match="Unknown embedding search provider: missing"): + get_embedding_search_provider_instance( + embedding_search_providers={}, + default_embedding_model=DEFAULT_EMBEDDING_MODEL, + default_embedding_engine=DEFAULT_EMBEDDING_ENGINE, + default_embedding_params={}, + esp_config=EmbeddingSearchProvider(name="missing"), + ) From f818a6067acd06c6f17c40da1f0a6bbac4ad184a Mon Sep 17 00:00:00 2001 From: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> Date: Mon, 1 Jun 2026 14:35:07 +0200 Subject: [PATCH 2/3] refactor(llmrails): extract startup model wiring Move LLM and action model setup, action caches, generation action registration, knowledge base setup, and tracing adapter creation into startup helpers. --- .../rails/llm/startup/generation_actions.py | 72 ++++++++ .../rails/llm/startup/knowledge_base.py | 70 ++++++++ .../rails/llm/startup/llm_action_caches.py | 94 ++++++++++ .../rails/llm/startup/llm_action_models.py | 154 ++++++++++++++++ nemoguardrails/rails/llm/startup/tracing.py | 28 +++ tests/rails/llm/test_generation_actions.py | 107 +++++++++++ tests/rails/llm/test_knowledge_base.py | 144 +++++++++++++++ tests/rails/llm/test_llm_action_caches.py | 111 ++++++++++++ tests/rails/llm/test_llm_action_models.py | 166 ++++++++++++++++++ 9 files changed, 946 insertions(+) create mode 100644 nemoguardrails/rails/llm/startup/generation_actions.py create mode 100644 nemoguardrails/rails/llm/startup/knowledge_base.py create mode 100644 nemoguardrails/rails/llm/startup/llm_action_caches.py create mode 100644 nemoguardrails/rails/llm/startup/llm_action_models.py create mode 100644 nemoguardrails/rails/llm/startup/tracing.py create mode 100644 tests/rails/llm/test_generation_actions.py create mode 100644 tests/rails/llm/test_knowledge_base.py create mode 100644 tests/rails/llm/test_llm_action_caches.py create mode 100644 tests/rails/llm/test_llm_action_models.py diff --git a/nemoguardrails/rails/llm/startup/generation_actions.py b/nemoguardrails/rails/llm/startup/generation_actions.py new file mode 100644 index 0000000000..13e9c20f25 --- /dev/null +++ b/nemoguardrails/rails/llm/startup/generation_actions.py @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""LLM generation action registration.""" + +from typing import Any, Protocol, Type + +from nemoguardrails.actions.llm.generation import LLMGenerationActions +from nemoguardrails.actions.v2_x.generation import LLMGenerationActionsV2dotx + +__all__ = [ + "GenerationActionRails", + "GenerationActionRuntime", + "LLMGenerationActions", + "LLMGenerationActionsV2dotx", + "generation_actions_class_for_colang_version", + "register_llm_generation_actions", +] + + +class GenerationActionRuntime(Protocol): + @property + def llm_task_manager(self) -> Any: ... + + def register_actions(self, actions_obj: Any, /, override: bool = True) -> None: ... + + +class GenerationActionRails(Protocol): + llm_generation_actions: Any + + @property + def config(self) -> Any: ... + + @property + def llm(self) -> Any: ... + + @property + def runtime(self) -> GenerationActionRuntime: ... + + @property + def embedding_search(self) -> Any: ... + + +def generation_actions_class_for_colang_version(colang_version: str) -> Type[Any]: + """Return the LLM generation actions class selected by the Colang version.""" + return LLMGenerationActions if colang_version == "1.0" else LLMGenerationActionsV2dotx + + +def register_llm_generation_actions(rails: GenerationActionRails, verbose: bool) -> None: + """Create and register the LLM generation actions for an LLMRails instance.""" + llm_generation_actions_class = generation_actions_class_for_colang_version(rails.config.colang_version) + rails.llm_generation_actions = llm_generation_actions_class( + config=rails.config, + llm=rails.llm, + llm_task_manager=rails.runtime.llm_task_manager, + get_embedding_search_provider_instance=rails.embedding_search.get_provider_instance, + verbose=verbose, + ) + + rails.runtime.register_actions(rails.llm_generation_actions, override=False) diff --git a/nemoguardrails/rails/llm/startup/knowledge_base.py b/nemoguardrails/rails/llm/startup/knowledge_base.py new file mode 100644 index 0000000000..4334501990 --- /dev/null +++ b/nemoguardrails/rails/llm/startup/knowledge_base.py @@ -0,0 +1,70 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Knowledge base setup for LLMRails.""" + +import asyncio +import threading +from typing import Callable, Optional + +from nemoguardrails.embeddings.index import EmbeddingsIndex +from nemoguardrails.kb.kb import KnowledgeBase +from nemoguardrails.patch_asyncio import check_sync_call_from_async_loop +from nemoguardrails.rails.llm.config import EmbeddingSearchProvider, RailsConfig +from nemoguardrails.utils import get_or_create_event_loop + +__all__ = ["build_knowledge_base_for_docs", "init_knowledge_base"] + + +async def build_knowledge_base_for_docs( + config: RailsConfig, + get_embedding_search_provider_instance: Callable[[Optional[EmbeddingSearchProvider]], EmbeddingsIndex], +) -> Optional[KnowledgeBase]: + """Build the docs-backed knowledge base configured for LLMRails.""" + if not config.docs: + return None + + documents = [doc.content for doc in config.docs] + kb = KnowledgeBase( + documents=documents, + config=config.knowledge_base, + get_embedding_search_provider_instance=get_embedding_search_provider_instance, + ) + kb.init() + await kb.build() + return kb + + +def init_knowledge_base(rails) -> None: + """Initialize and register the LLMRails knowledge base.""" + rails.kb = None + + async def _init_kb(): + rails.kb = await build_knowledge_base_for_docs( + config=rails.config, + get_embedding_search_provider_instance=rails.embedding_search.get_provider_instance, + ) + + # There are still some edge cases not covered by nest_asyncio. + # Using a separate thread always for now. + loop = get_or_create_event_loop() + if True or check_sync_call_from_async_loop(): + t = threading.Thread(target=asyncio.run, args=(_init_kb(),)) + t.start() + t.join() + else: + loop.run_until_complete(_init_kb()) + + rails.runtime.register_action_param("kb", rails.kb) diff --git a/nemoguardrails/rails/llm/startup/llm_action_caches.py b/nemoguardrails/rails/llm/startup/llm_action_caches.py new file mode 100644 index 0000000000..7fb02b87ac --- /dev/null +++ b/nemoguardrails/rails/llm/startup/llm_action_caches.py @@ -0,0 +1,94 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""LLM action model caches.""" + +import logging +from typing import Dict, Protocol + +from nemoguardrails.llm.cache import CacheInterface, LFUCache +from nemoguardrails.rails.llm.config import Model, RailsConfig + +log = logging.getLogger("nemoguardrails.rails.llm.llmrails") + +__all__ = [ + "build_llm_action_cache", + "build_llm_action_caches", + "initialize_llm_action_caches", +] + + +class LLMActionCacheRuntime(Protocol): + def register_action_param(self, name: str, value: object) -> None: ... + + +class LLMActionCacheRails(Protocol): + @property + def config(self) -> RailsConfig: ... + + @property + def runtime(self) -> LLMActionCacheRuntime: ... + + +def build_llm_action_cache(model: Model) -> LFUCache: + """Build the cache configured for one LLM action model.""" + if model.cache is None: + raise ValueError(f"Missing cache configuration for model '{model.type}'.") + + if model.cache.maxsize <= 0: + raise ValueError( + f"Invalid cache maxsize for model '{model.type}': {model.cache.maxsize}. " + "Capacity must be greater than 0. Skipping cache creation." + ) + + stats_logging_interval = None + if model.cache.stats.enabled and model.cache.stats.log_interval is not None: + stats_logging_interval = model.cache.stats.log_interval + + cache = LFUCache( + maxsize=model.cache.maxsize, + track_stats=model.cache.stats.enabled, + stats_logging_interval=stats_logging_interval, + ) + + log.info(f"Created cache for model '{model.type}' with maxsize {model.cache.maxsize}") + + return cache + + +def build_llm_action_caches(config: RailsConfig) -> Dict[str, CacheInterface]: + """Build caches for configured action models.""" + model_caches: Dict[str, CacheInterface] = dict() + for model in config.models: + if model.type in ["main", "embeddings"]: + continue + + if model.cache and model.cache.enabled: + cache = build_llm_action_cache(model) + model_caches[model.type] = cache + + log.info( + f"Initialized model '{model.type}' with cache %s", + "enabled" if cache else "disabled", + ) + + return model_caches + + +def initialize_llm_action_caches(rails: LLMActionCacheRails) -> None: + """Initialize configured action model caches for an LLMRails instance.""" + model_caches = build_llm_action_caches(rails.config) + if model_caches: + rails.runtime.register_action_param("model_caches", model_caches) diff --git a/nemoguardrails/rails/llm/startup/llm_action_models.py b/nemoguardrails/rails/llm/startup/llm_action_models.py new file mode 100644 index 0000000000..539099e336 --- /dev/null +++ b/nemoguardrails/rails/llm/startup/llm_action_models.py @@ -0,0 +1,154 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""LLM action model loading.""" + +import logging +import os +from typing import Any, Callable, Dict, Protocol + +from nemoguardrails.exceptions import InvalidModelConfigurationError +from nemoguardrails.llm.models.initializer import ModelInitializationError +from nemoguardrails.types import LLMModel + +log = logging.getLogger("nemoguardrails.rails.llm.llmrails") + +InitLLM = Callable[..., LLMModel] + +__all__ = [ + "InitLLM", + "LLMActionRails", + "LLMActionRuntime", + "load_llm_action_models", + "model_kwargs_from_config", + "sync_update_llm_bindings", +] + + +class LLMActionRuntime(Protocol): + def register_action_param(self, name: str, value: Any) -> None: ... + + +class LLMActionRails(Protocol): + llm: Any + llm_generation_actions: Any + + @property + def config(self) -> Any: ... + + @property + def runtime(self) -> LLMActionRuntime: ... + + +def model_kwargs_from_config(model_config: Any) -> Dict[str, Any]: + """Prepare model kwargs, including optional API key environment variables.""" + kwargs = dict(model_config.parameters or {}) + + if model_config.api_key_env_var: + api_key = os.environ.get(model_config.api_key_env_var) + if api_key: + kwargs["api_key"] = api_key + + return kwargs + + +def sync_update_llm_bindings(rails: LLMActionRails, llm: LLMModel) -> None: + """Synchronize the main LLM bindings after a public update.""" + rails.llm = llm + rails.llm_generation_actions.llm = llm + rails.runtime.register_action_param("llm", llm) + + +def load_llm_action_models(rails: LLMActionRails, init_llm: InitLLM) -> None: + """Load the main and action LLMs configured for an LLMRails instance.""" + from nemoguardrails._compat.langchain_kwargs import check_langchain_kwargs + from nemoguardrails.llm.frameworks import get_default_framework + + prepare_model_kwargs = getattr(rails, "_prepare_model_kwargs", model_kwargs_from_config) + models_to_check = ( + [model for model in rails.config.models if model.type != "main"] if rails.llm else rails.config.models + ) + check_langchain_kwargs(models_to_check, get_default_framework()) + + if rails.llm: + if any(model.type == "main" for model in rails.config.models): + log.warning( + "Both an LLM was provided via constructor and a main LLM is specified in the config. " + "The LLM provided via constructor will be used and the main LLM from config will be ignored." + ) + rails.runtime.register_action_param("llm", rails.llm) + + else: + main_model = next((model for model in rails.config.models if model.type == "main"), None) + + if main_model and main_model.model: + kwargs = prepare_model_kwargs(main_model) + rails.llm = init_llm( + model_name=main_model.model, + provider_name=main_model.engine, + mode="chat", + kwargs=kwargs, + ) + rails.runtime.register_action_param("llm", rails.llm) + + else: + log.info("No main LLM specified in the config and no LLM provided via constructor.") + + llms = dict() + + for llm_config in rails.config.models: + if llm_config.type in ["embeddings", "jailbreak_detection"]: + continue + + if rails.llm and llm_config.type == "main": + continue + + try: + model_name = llm_config.model + if not model_name: + raise InvalidModelConfigurationError( + f"`model` field must be set in model configuration: {llm_config.model_dump_json()}" + ) + + provider_name = llm_config.engine + kwargs = prepare_model_kwargs(llm_config) + mode = llm_config.mode + + llm_model = init_llm( + model_name=model_name, + provider_name=provider_name, + mode=mode, + kwargs=kwargs, + ) + + if llm_config.type == "main": + if not rails.llm: + rails.llm = llm_model + rails.runtime.register_action_param("llm", rails.llm) + else: + model_attr = f"{llm_config.type}_llm" + if not hasattr(rails, model_attr): + setattr(rails, model_attr, llm_model) + rails.runtime.register_action_param(model_attr, getattr(rails, model_attr)) + llms[llm_config.type] = getattr(rails, model_attr) + + except ModelInitializationError as e: + log.error("Failed to initialize model: %s", str(e)) + raise + except Exception as e: + log.error("Unexpected error initializing model: %s", str(e)) + raise + + rails.runtime.register_action_param("llms", llms) diff --git a/nemoguardrails/rails/llm/startup/tracing.py b/nemoguardrails/rails/llm/startup/tracing.py new file mode 100644 index 0000000000..b5ca19987a --- /dev/null +++ b/nemoguardrails/rails/llm/startup/tracing.py @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from nemoguardrails.rails.llm.config import RailsConfig +from nemoguardrails.tracing.adapters.base import InteractionLogAdapter + +__all__ = ["create_startup_tracing_adapters"] + + +def create_startup_tracing_adapters(config: RailsConfig) -> list[InteractionLogAdapter] | None: + if not config.tracing: + return None + + from nemoguardrails.tracing import create_log_adapters + + return create_log_adapters(config.tracing) diff --git a/tests/rails/llm/test_generation_actions.py b/tests/rails/llm/test_generation_actions.py new file mode 100644 index 0000000000..db6f929b9e --- /dev/null +++ b/tests/rails/llm/test_generation_actions.py @@ -0,0 +1,107 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from types import SimpleNamespace + +from nemoguardrails.rails.llm.startup import generation_actions +from nemoguardrails.rails.llm.startup.generation_actions import ( + generation_actions_class_for_colang_version, + register_llm_generation_actions, +) + + +class RuntimeWithActionRegistration: + def __init__(self): + self.llm_task_manager = object() + self.registered_actions = None + self.override = None + + def register_actions(self, actions, override=True): + self.registered_actions = actions + self.override = override + + +class FakeGenerationActions: + def __init__( + self, + config, + llm, + llm_task_manager, + get_embedding_search_provider_instance, + verbose, + ): + self.config = config + self.llm = llm + self.llm_task_manager = llm_task_manager + self.get_embedding_search_provider_instance = get_embedding_search_provider_instance + self.verbose = verbose + + +class FakeGenerationActionsV2(FakeGenerationActions): + pass + + +class FakeRails: + def __init__(self, config, llm, runtime, get_embedding_search_provider_instance): + self.config = config + self.llm = llm + self.runtime = runtime + self.embedding_search = SimpleNamespace(get_provider_instance=get_embedding_search_provider_instance) + self.llm_generation_actions = None + + +def test_generation_actions_class_for_colang_version_uses_v1_actions(monkeypatch): + monkeypatch.setattr(generation_actions, "LLMGenerationActions", FakeGenerationActions) + + actions_class = generation_actions_class_for_colang_version("1.0") + + assert actions_class is FakeGenerationActions + + +def test_generation_actions_class_for_colang_version_uses_v2_actions(monkeypatch): + monkeypatch.setattr(generation_actions, "LLMGenerationActionsV2dotx", FakeGenerationActionsV2) + + actions_class = generation_actions_class_for_colang_version("2.x") + + assert actions_class is FakeGenerationActionsV2 + + +def test_register_llm_generation_actions_registers_without_overriding(monkeypatch): + monkeypatch.setattr(generation_actions, "LLMGenerationActions", FakeGenerationActions) + config = SimpleNamespace(colang_version="1.0") + llm = object() + runtime = RuntimeWithActionRegistration() + + def get_embedding_search_provider_instance(esp_config=None): + del esp_config + return None + + rails = FakeRails( + config=config, + llm=llm, + runtime=runtime, + get_embedding_search_provider_instance=get_embedding_search_provider_instance, + ) + + register_llm_generation_actions(rails, verbose=True) + + assert isinstance(rails.llm_generation_actions, FakeGenerationActions) + assert rails.llm_generation_actions.config is config + assert rails.llm_generation_actions.llm is llm + assert rails.llm_generation_actions.llm_task_manager is runtime.llm_task_manager + assert rails.llm_generation_actions.get_embedding_search_provider_instance is get_embedding_search_provider_instance + assert rails.llm_generation_actions.verbose is True + assert runtime.registered_actions is rails.llm_generation_actions + assert runtime.override is False diff --git a/tests/rails/llm/test_knowledge_base.py b/tests/rails/llm/test_knowledge_base.py new file mode 100644 index 0000000000..9f7dc26fbb --- /dev/null +++ b/tests/rails/llm/test_knowledge_base.py @@ -0,0 +1,144 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from types import SimpleNamespace +from typing import Any, cast + +import pytest + +from nemoguardrails.embeddings.index import EmbeddingsIndex +from nemoguardrails.rails.llm.config import Document, RailsConfig +from nemoguardrails.rails.llm.startup import knowledge_base +from nemoguardrails.rails.llm.startup.knowledge_base import ( + build_knowledge_base_for_docs, + init_knowledge_base, +) + + +class RuntimeWithActionParams: + def __init__(self): + self.registered_action_params = {} + + def register_action_param(self, name, value): + self.registered_action_params[name] = value + + +@pytest.mark.asyncio +async def test_build_knowledge_base_for_docs_returns_none_without_docs(monkeypatch): + def fail_if_constructed(*args, **kwargs): + raise AssertionError("KnowledgeBase should not be constructed without docs.") + + monkeypatch.setattr(knowledge_base, "KnowledgeBase", fail_if_constructed) + + kb = await build_knowledge_base_for_docs( + config=RailsConfig(models=[]), + get_embedding_search_provider_instance=lambda _: cast(EmbeddingsIndex, None), + ) + + assert kb is None + + +@pytest.mark.asyncio +async def test_build_knowledge_base_for_docs_builds_docs_backed_kb(monkeypatch): + provider = cast(EmbeddingsIndex, object()) + created = {} + + class FakeKnowledgeBase: + def __init__(self, documents, config, get_embedding_search_provider_instance): + self.documents = documents + self.config = config + self.get_embedding_search_provider_instance = get_embedding_search_provider_instance + self.init_called = False + self.build_called = False + created["kb"] = self + + def init(self): + self.init_called = True + + async def build(self): + self.build_called = True + + monkeypatch.setattr(knowledge_base, "KnowledgeBase", FakeKnowledgeBase) + config = RailsConfig( + models=[], + docs=[ + Document(format="md", content="# Alpha\n\nA"), + Document(format="md", content="# Beta\n\nB"), + ], + ) + + kb = cast( + Any, + await build_knowledge_base_for_docs( + config=config, + get_embedding_search_provider_instance=lambda _: provider, + ), + ) + + assert kb is created["kb"] + assert kb.documents == ["# Alpha\n\nA", "# Beta\n\nB"] + assert kb.config is config.knowledge_base + assert kb.get_embedding_search_provider_instance(None) is provider + assert kb.init_called is True + assert kb.build_called is True + + +def test_init_knowledge_base_uses_thread_path_and_registers_kb(monkeypatch): + built_kb = object() + calls = {"thread_started": False, "thread_joined": False} + + async def fake_build_knowledge_base_for_docs(config, get_embedding_search_provider_instance): + assert rails.kb is None + assert config is rails.config + assert get_embedding_search_provider_instance is get_provider_instance + return built_kb + + class FakeThread: + def __init__(self, target, args): + self.target = target + self.args = args + + def start(self): + calls["thread_started"] = True + self.target(*self.args) + + def join(self): + calls["thread_joined"] = True + + def get_provider_instance(_=None): + return None + + rails = SimpleNamespace( + config=RailsConfig( + models=[], + docs=[Document(format="md", content="# Alpha\n\nA")], + ), + kb="existing", + runtime=RuntimeWithActionParams(), + embedding_search=SimpleNamespace(get_provider_instance=get_provider_instance), + ) + monkeypatch.setattr( + knowledge_base, + "build_knowledge_base_for_docs", + fake_build_knowledge_base_for_docs, + ) + monkeypatch.setattr(knowledge_base.threading, "Thread", FakeThread) + monkeypatch.setattr(knowledge_base, "get_or_create_event_loop", lambda: object()) + + init_knowledge_base(rails) + + assert rails.kb is built_kb + assert rails.runtime.registered_action_params["kb"] is built_kb + assert calls == {"thread_started": True, "thread_joined": True} diff --git a/tests/rails/llm/test_llm_action_caches.py b/tests/rails/llm/test_llm_action_caches.py new file mode 100644 index 0000000000..a122e99df2 --- /dev/null +++ b/tests/rails/llm/test_llm_action_caches.py @@ -0,0 +1,111 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest + +from nemoguardrails.rails.llm.config import ( + CacheStatsConfig, + Model, + ModelCacheConfig, + RailsConfig, +) +from nemoguardrails.rails.llm.startup.llm_action_caches import ( + build_llm_action_cache, + build_llm_action_caches, +) + + +def test_build_llm_action_caches_skips_main_embeddings_and_disabled_caches(): + config = RailsConfig( + models=[ + Model( + type="main", + engine="fake", + model="fake", + cache=ModelCacheConfig(enabled=True), + ), + Model( + type="embeddings", + engine="fake", + model="fake", + cache=ModelCacheConfig(enabled=True), + ), + Model( + type="content_safety", + engine="fake", + model="fake", + cache=ModelCacheConfig(enabled=False), + ), + ] + ) + + assert build_llm_action_caches(config) == {} + + +def test_build_llm_action_caches_creates_enabled_action_model_caches(): + config = RailsConfig( + models=[ + Model( + type="content_safety", + engine="fake", + model="fake", + cache=ModelCacheConfig(enabled=True, maxsize=1000), + ), + Model( + type="jailbreak_detection", + engine="fake", + model="fake", + cache=ModelCacheConfig(enabled=True, maxsize=2000), + ), + ] + ) + + model_caches = build_llm_action_caches(config) + + assert set(model_caches.keys()) == {"content_safety", "jailbreak_detection"} + assert model_caches["content_safety"].maxsize == 1000 + assert model_caches["jailbreak_detection"].maxsize == 2000 + + +def test_build_llm_action_cache_preserves_stats_config(): + model = Model( + type="content_safety", + engine="fake", + model="fake", + cache=ModelCacheConfig( + enabled=True, + maxsize=5000, + stats=CacheStatsConfig(enabled=True, log_interval=60.0), + ), + ) + + cache = build_llm_action_cache(model) + + assert cache.maxsize == 5000 + assert cache.track_stats is True + assert cache.stats_logging_interval == 60.0 + assert cache.supports_stats_logging() is True + + +def test_build_llm_action_cache_rejects_invalid_maxsize(): + model = Model( + type="content_safety", + engine="fake", + model="fake", + cache=ModelCacheConfig(enabled=True, maxsize=0), + ) + + with pytest.raises(ValueError, match="Invalid cache maxsize"): + build_llm_action_cache(model) diff --git a/tests/rails/llm/test_llm_action_models.py b/tests/rails/llm/test_llm_action_models.py new file mode 100644 index 0000000000..573f373dd8 --- /dev/null +++ b/tests/rails/llm/test_llm_action_models.py @@ -0,0 +1,166 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Any, cast +from unittest.mock import MagicMock, patch + +from nemoguardrails.rails.llm.config import Model, RailsConfig +from nemoguardrails.rails.llm.startup.llm_action_models import ( + load_llm_action_models, + model_kwargs_from_config, + sync_update_llm_bindings, +) +from nemoguardrails.types import LLMModel + + +class RuntimeWithActionParams: + def __init__(self): + self.registered_action_params = {} + + def register_action_param(self, name, value): + self.registered_action_params[name] = value + + +class FakeRails: + content_safety_llm: Any + + def __init__(self, config: RailsConfig, llm: Any, runtime: RuntimeWithActionParams): + self.config = config + self.llm = llm + self.runtime = runtime + self.llm_generation_actions = MagicMock() + + +def test_model_kwargs_from_config_adds_api_key_from_environment(): + model = Model( + type="main", + engine="openai", + model="gpt-3.5-turbo", + api_key_env_var="TEST_OPENAI_KEY", + parameters={"temperature": 0.7}, + ) + + with patch.dict("os.environ", {"TEST_OPENAI_KEY": "secret-api-key-from-env"}): + kwargs = model_kwargs_from_config(model) + + assert kwargs["api_key"] == "secret-api-key-from-env" + assert kwargs["temperature"] == 0.7 + + +def test_model_kwargs_from_config_omits_missing_api_key_environment_variable(): + model = Model( + type="main", + engine="openai", + model="gpt-3.5-turbo", + api_key_env_var="MISSING_OPENAI_KEY", + parameters={"temperature": 0.5}, + ) + + with patch.dict("os.environ", {}, clear=True): + kwargs = model_kwargs_from_config(model) + + assert "api_key" not in kwargs + assert kwargs["temperature"] == 0.5 + + +def test_model_kwargs_from_config_preserves_direct_api_key_parameter(): + model = Model( + type="main", + engine="openai", + model="gpt-3.5-turbo", + parameters={"api_key": "direct-key", "temperature": 0.3}, + ) + + kwargs = model_kwargs_from_config(model) + + assert kwargs["api_key"] == "direct-key" + assert kwargs["temperature"] == 0.3 + + +def test_load_llm_action_models_uses_constructor_llm_and_loads_action_models(): + injected_llm = object() + content_safety_llm = object() + init_llm = MagicMock(return_value=content_safety_llm) + rails = FakeRails( + config=RailsConfig( + models=[ + Model(type="main", engine="fake", model="main-model"), + Model(type="content_safety", engine="fake", model="content-safety-model"), + ] + ), + llm=injected_llm, + runtime=RuntimeWithActionParams(), + ) + + load_llm_action_models(rails, init_llm=init_llm) + + init_llm.assert_called_once_with( + model_name="content-safety-model", + provider_name="fake", + mode="chat", + kwargs={}, + ) + assert rails.llm is injected_llm + assert rails.content_safety_llm is content_safety_llm + assert rails.runtime.registered_action_params["llm"] is injected_llm + assert rails.runtime.registered_action_params["content_safety_llm"] is content_safety_llm + assert rails.runtime.registered_action_params["llms"] == {"content_safety": content_safety_llm} + + +def test_load_llm_action_models_initializes_main_llm_from_config(): + main_llm = object() + init_llm = MagicMock(return_value=main_llm) + rails = FakeRails( + config=RailsConfig( + models=[ + Model( + type="main", + engine="fake", + model="main-model", + parameters={"temperature": 0.2}, + ) + ] + ), + llm=None, + runtime=RuntimeWithActionParams(), + ) + + load_llm_action_models(rails, init_llm=init_llm) + + init_llm.assert_called_once_with( + model_name="main-model", + provider_name="fake", + mode="chat", + kwargs={"temperature": 0.2}, + ) + assert rails.llm is main_llm + assert rails.runtime.registered_action_params["llm"] is main_llm + assert rails.runtime.registered_action_params["llms"] == {} + + +def test_sync_update_llm_bindings_updates_llm_generation_actions_and_runtime_param(): + initial_llm = cast(LLMModel, object()) + updated_llm = cast(LLMModel, object()) + rails = FakeRails( + config=RailsConfig(models=[]), + llm=initial_llm, + runtime=RuntimeWithActionParams(), + ) + + sync_update_llm_bindings(rails, updated_llm) + + assert rails.llm is updated_llm + assert rails.llm_generation_actions.llm is updated_llm + assert rails.runtime.registered_action_params["llm"] is updated_llm From a3cec09ed0fcee5bd4c4d3e7846b9079641cefad Mon Sep 17 00:00:00 2001 From: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> Date: Mon, 8 Jun 2026 21:36:50 +0200 Subject: [PATCH 3/3] refactor(llmrails): delegate startup wiring --- nemoguardrails/rails/llm/llmrails.py | 477 +++--------------- nemoguardrails/rails/llm/runtime/__init__.py | 16 + .../rails/llm/runtime/colang_runtime.py | 40 ++ .../rails/llm/startup/colang_flows.py | 6 +- .../rails/llm/startup/config_preparation.py | 45 +- .../rails/llm/startup/generation_actions.py | 11 +- .../rails/llm/startup/knowledge_base.py | 8 +- .../rails/llm/startup/llm_action_models.py | 4 +- .../test_public_api_deprecations.py | 50 +- tests/rails/llm/test_colang_runtime.py | 95 ++++ tests/rails/llm/test_config_preparation.py | 59 +-- tests/rails/llm/test_generation_actions.py | 20 +- tests/rails/llm/test_knowledge_base.py | 8 +- tests/rails/llm/test_llm_action_models.py | 4 +- .../llm/test_llmrails_public_contract.py | 41 ++ tests/test_llmrails.py | 27 - 16 files changed, 333 insertions(+), 578 deletions(-) create mode 100644 nemoguardrails/rails/llm/runtime/__init__.py create mode 100644 nemoguardrails/rails/llm/runtime/colang_runtime.py create mode 100644 tests/rails/llm/test_colang_runtime.py diff --git a/nemoguardrails/rails/llm/llmrails.py b/nemoguardrails/rails/llm/llmrails.py index c777d5ae3a..66001fbe65 100644 --- a/nemoguardrails/rails/llm/llmrails.py +++ b/nemoguardrails/rails/llm/llmrails.py @@ -16,12 +16,9 @@ """LLM Rails entry point.""" import asyncio -import importlib.util import json import logging -import os import re -import threading import time import warnings from functools import partial @@ -42,7 +39,6 @@ from typing_extensions import Self -from nemoguardrails.actions.llm.generation import LLMGenerationActions from nemoguardrails.actions.llm.utils import ( extract_bot_thinking_from_events, extract_tool_calls_from_events, @@ -50,11 +46,9 @@ get_colang_history, ) from nemoguardrails.actions.output_mapping import is_output_blocked -from nemoguardrails.actions.v2_x.generation import LLMGenerationActionsV2dotx from nemoguardrails.base_guardrails import BaseGuardrails -from nemoguardrails.colang import parse_colang_file -from nemoguardrails.colang.v1_0.runtime.flows import _normalize_flow_id, compute_context -from nemoguardrails.colang.v1_0.runtime.runtime import Runtime, RuntimeV1_0 +from nemoguardrails.colang.v1_0.runtime.flows import compute_context +from nemoguardrails.colang.v1_0.runtime.runtime import Runtime from nemoguardrails.colang.v2_x.runtime.flows import Action, State from nemoguardrails.colang.v2_x.runtime.runtime import RuntimeV2_x from nemoguardrails.context import ( @@ -68,17 +62,10 @@ from nemoguardrails.embeddings.providers import register_embedding_provider from nemoguardrails.embeddings.providers.base import EmbeddingModel from nemoguardrails.exceptions import ( - InvalidModelConfigurationError, - InvalidRailsConfigurationError, InvalidStateError, StreamingNotSupportedError, ) -from nemoguardrails.kb.kb import KnowledgeBase -from nemoguardrails.llm.cache import CacheInterface, LFUCache -from nemoguardrails.llm.models.initializer import ( - ModelInitializationError, - init_llm_model, -) +from nemoguardrails.llm.models.initializer import init_llm_model from nemoguardrails.logging.explain import ExplainInfo from nemoguardrails.logging.processing_log import compute_generation_log from nemoguardrails.logging.stats import LLMStats @@ -86,7 +73,6 @@ from nemoguardrails.patch_asyncio import check_sync_call_from_async_loop from nemoguardrails.rails.llm.buffer import get_buffer_strategy from nemoguardrails.rails.llm.config import ( - EmbeddingSearchProvider, OutputRailsStreamingConfig, RailsConfig, ) @@ -98,6 +84,20 @@ RailStatus, RailType, ) +from nemoguardrails.rails.llm.runtime.colang_runtime import runtime_for_colang_version +from nemoguardrails.rails.llm.startup.config_preparation import prepare_llmrails_config +from nemoguardrails.rails.llm.startup.config_py import load_config_py_modules, run_config_py_init_hooks +from nemoguardrails.rails.llm.startup.config_validation import validate_llmrails_config +from nemoguardrails.rails.llm.startup.embedding_search import EmbeddingSearchState, apply_embedding_model_config +from nemoguardrails.rails.llm.startup.generation_actions import register_llm_generation_actions +from nemoguardrails.rails.llm.startup.knowledge_base import init_knowledge_base +from nemoguardrails.rails.llm.startup.llm_action_caches import initialize_llm_action_caches +from nemoguardrails.rails.llm.startup.llm_action_models import ( + load_llm_action_models, + model_kwargs_from_config, + sync_update_llm_bindings, +) +from nemoguardrails.rails.llm.startup.tracing import create_startup_tracing_adapters from nemoguardrails.rails.llm.utils import ( get_action_details_from_flow_id, get_history_cache_key, @@ -137,6 +137,11 @@ class LLMRails(BaseGuardrails): """Rails based on a given configuration.""" config: RailsConfig + embedding_search: EmbeddingSearchState + _explain_info: Optional[ExplainInfo] + _kb: Any + _llm_generation_actions: Any + _verbose: bool llm: Optional[LLMModel] runtime: Runtime @@ -158,7 +163,7 @@ def embedding_search_providers(self): DeprecationWarning, stacklevel=2, ) - return self._embedding_search_providers + return self.embedding_search.providers @property def default_embedding_model(self): @@ -167,7 +172,7 @@ def default_embedding_model(self): DeprecationWarning, stacklevel=2, ) - return self._default_embedding_model + return self.embedding_search.default_model @default_embedding_model.setter def default_embedding_model(self, value): @@ -176,7 +181,7 @@ def default_embedding_model(self, value): DeprecationWarning, stacklevel=2, ) - self._default_embedding_model = value + self.embedding_search.default_model = value @property def default_embedding_engine(self): @@ -185,7 +190,7 @@ def default_embedding_engine(self): DeprecationWarning, stacklevel=2, ) - return self._default_embedding_engine + return self.embedding_search.default_engine @default_embedding_engine.setter def default_embedding_engine(self, value): @@ -194,7 +199,7 @@ def default_embedding_engine(self, value): DeprecationWarning, stacklevel=2, ) - self._default_embedding_engine = value + self.embedding_search.default_engine = value @property def default_embedding_params(self): @@ -203,7 +208,7 @@ def default_embedding_params(self): DeprecationWarning, stacklevel=2, ) - return self._default_embedding_params + return self.embedding_search.default_params @default_embedding_params.setter def default_embedding_params(self, value): @@ -212,7 +217,7 @@ def default_embedding_params(self, value): DeprecationWarning, stacklevel=2, ) - self._default_embedding_params = value + self.embedding_search.default_params = value @property def explain_info(self): @@ -260,6 +265,14 @@ def passthrough_fn(self, fn): """LLMGenerationActions is private, set passthrough_fn directly""" self._llm_generation_actions._passthrough_fn = fn + @property + def verbose(self) -> bool: + return self._verbose + + @verbose.setter + def verbose(self, verbose: bool) -> None: + self._verbose = verbose + def __init__( self, config: RailsConfig, @@ -284,184 +297,58 @@ def __init__( if self.verbose: set_verbose(True, llm_calls=True) - # We allow the user to register additional embedding search providers, so we keep - # an index of them. - self._embedding_search_providers = {} - - # The default embeddings model is using FastEmbed - self._default_embedding_model = "all-MiniLM-L6-v2" - self._default_embedding_engine = "FastEmbed" - self._default_embedding_params = {} + self.embedding_search = EmbeddingSearchState.default() # We keep a cache of the events history associated with a sequence of user messages. # TODO: when we update the interface to allow to return a "state object", this # should be removed self.events_history_cache = {} - # We also load the default flows from the `default_flows.yml` file in the current folder. - # But only for version 1.0. - # TODO: decide on the default flows for 2.x. - if config.colang_version == "1.0": - # We also load the default flows from the `llm_flows.co` file in the current folder. - current_folder = os.path.dirname(__file__) - default_flows_file = "llm_flows.co" - default_flows_path = os.path.join(current_folder, default_flows_file) - with open(default_flows_path, "r") as f: - default_flows_content = f.read() - default_flows = parse_colang_file(default_flows_file, default_flows_content)["flows"] - - # We mark all the default flows as system flows. - for flow_config in default_flows: - flow_config["is_system_flow"] = True - - # We add the default flows to the config. - self.config.flows.extend(default_flows) - - # We also need to load the content from the components library. - # Sort entries so the traversal order is filesystem-independent; - # otherwise the order in which library bot_messages are inserted - # (and which definition wins on collisions) varies between platforms. - library_path = os.path.join(os.path.dirname(__file__), "../../library") - for root, dirs, files in os.walk(library_path): - dirs.sort() - for file in sorted(files): - # Extract the full path for the file - full_path = os.path.join(root, file) - if file.endswith(".co"): - log.debug(f"Loading file: {full_path}") - with open(full_path, "r", encoding="utf-8") as f: - content = parse_colang_file(file, content=f.read(), version=config.colang_version) - if not content: - continue - - # We mark all the flows coming from the guardrails library as system flows. - for flow_config in content["flows"]: - flow_config["is_system_flow"] = True - - # We load all the flows - self.config.flows.extend(content["flows"]) - - # And all the messages as well, if they have not been overwritten - for message_id, utterances in content.get("bot_messages", {}).items(): - if message_id not in self.config.bot_messages: - self.config.bot_messages[message_id] = utterances - - # Last but not least, we mark all the flows that are used in any of the rails - # as system flows (so they don't end up in the prompt). - - rail_flow_ids = config.rails.input.flows + config.rails.output.flows + config.rails.retrieval.flows - - for flow_config in self.config.flows: - if flow_config.get("id") in rail_flow_ids: - flow_config["is_system_flow"] = True - - # We also mark them as subflows by default, to simplify the syntax - flow_config["is_subflow"] = True + self.config = prepare_llmrails_config(config=self.config) # We check if the configuration or any of the imported ones have config.py modules. - config_modules = [] - for _path in list(self.config.imported_paths.values() if self.config.imported_paths else []) + [ - self.config.config_path - ]: - if _path: - filepath = os.path.join(_path, "config.py") - if os.path.exists(filepath): - filename = os.path.basename(filepath) - spec = importlib.util.spec_from_file_location(filename, filepath) - if spec and spec.loader: - config_module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(config_module) - config_modules.append(config_module) - - colang_version_to_runtime: Dict[str, Type[Runtime]] = { - "1.0": RuntimeV1_0, - "2.x": RuntimeV2_x, - } - if config.colang_version not in colang_version_to_runtime: - raise InvalidRailsConfigurationError( - f"Unsupported colang version: {config.colang_version}. Supported versions: {list(colang_version_to_runtime.keys())}" - ) + config_modules = load_config_py_modules(self.config) # First, we initialize the runtime. - self.runtime = colang_version_to_runtime[config.colang_version](config=config, verbose=verbose) + self.runtime = runtime_for_colang_version(config=self.config, verbose=verbose) # If we have a config_modules with an `init` function, we call it. # We need to call this here because the `init` might register additional # LLM providers. - for config_module in config_modules: - if hasattr(config_module, "init"): - config_module.init(self) - - # If we have a customized embedding model, we'll use it. - for model in self.config.models: - if model.type == "embeddings": - self._default_embedding_model = model.model - self._default_embedding_engine = model.engine - self._default_embedding_params = model.parameters or {} - - for esp in [ - self.config.core.embedding_search_provider, - self.config.knowledge_base.embedding_search_provider, - ]: - if esp.name != "default": - continue - if "embedding_model" not in esp.parameters and model.model is not None: - esp.parameters["embedding_model"] = model.model - if "embedding_engine" not in esp.parameters and model.engine is not None: - esp.parameters["embedding_engine"] = model.engine - - break - - # InteractionLogAdapters used for tracing - # We ensure that it is used after config.py is loaded - if config.tracing: - from nemoguardrails.tracing import create_log_adapters - - self._log_adapters = create_log_adapters(config.tracing) - else: - self._log_adapters = None + run_config_py_init_hooks(self, config_modules) + + default_embedding_model, default_embedding_engine, default_embedding_params = apply_embedding_model_config( + config=self.config, + default_embedding_model=self.embedding_search.default_model, + default_embedding_engine=self.embedding_search.default_engine, + default_embedding_params=self.embedding_search.default_params, + ) + self.embedding_search.update_defaults( + default_model=default_embedding_model, + default_engine=default_embedding_engine, + default_params=default_embedding_params, + ) + + self._log_adapters = create_startup_tracing_adapters(self.config) # We run some additional checks on the config - self._validate_config() + validate_llmrails_config(self.config) # Next, we initialize the LLM engines (main engine and action engines if specified). self._init_llms() # Next, we initialize the LLM Generate actions and register them. - llm_generation_actions_class = ( - LLMGenerationActions if config.colang_version == "1.0" else LLMGenerationActionsV2dotx - ) - self._llm_generation_actions = llm_generation_actions_class( - config=config, - llm=self.llm, - llm_task_manager=self.runtime.llm_task_manager, - get_embedding_search_provider_instance=self._get_embeddings_search_provider_instance, - verbose=verbose, - ) + register_llm_generation_actions(self, verbose=verbose) - # If there's already an action registered, we don't override. - self.runtime.register_actions(self._llm_generation_actions, override=False) - - # Next, we initialize the Knowledge Base - # There are still some edge cases not covered by nest_asyncio. - # Using a separate thread always for now. - loop = get_or_create_event_loop() - if True or check_sync_call_from_async_loop(): - t = threading.Thread(target=asyncio.run, args=(self._init_kb(),)) - t.start() - t.join() - else: - loop.run_until_complete(self._init_kb()) - - # We also register the kb as a parameter that can be passed to actions. - self.runtime.register_action_param("kb", self._kb) + # Next, we initialize the Knowledge Base. + init_knowledge_base(self) # Reference to the general ExplainInfo object. self._explain_info = None from nemoguardrails.telemetry import report_usage - report_usage(config, deployment_type="library", rails_engine="LLMRails") + report_usage(self.config, deployment_type="library", rails_engine="LLMRails") def update_llm(self, llm: LLMModel): """Replace the main LLM with the provided one. @@ -471,55 +358,7 @@ def update_llm(self, llm: LLMModel): """ if not isinstance(llm, LLMModel): llm = _wrap_legacy_llm(llm) - self.llm = llm - self._llm_generation_actions.llm = llm - self.runtime.register_action_param("llm", llm) - - def _validate_config(self): - """Runs additional validation checks on the config.""" - - if self.config.colang_version == "1.0": - existing_flows_names = set([flow.get("id") for flow in self.config.flows]) - else: - existing_flows_names = set([flow.get("name") for flow in self.config.flows]) - - for flow_name in self.config.rails.input.flows: - # content safety check input/output flows are special as they have parameters - flow_name = _normalize_flow_id(flow_name) - if flow_name not in existing_flows_names: - raise InvalidRailsConfigurationError(f"The provided input rail flow `{flow_name}` does not exist") - - for flow_name in self.config.rails.output.flows: - flow_name = _normalize_flow_id(flow_name) - if flow_name not in existing_flows_names: - raise InvalidRailsConfigurationError(f"The provided output rail flow `{flow_name}` does not exist") - - for flow_name in self.config.rails.retrieval.flows: - if flow_name not in existing_flows_names: - raise InvalidRailsConfigurationError(f"The provided retrieval rail flow `{flow_name}` does not exist") - - # If both passthrough mode and single call mode are specified, we raise an exception. - if self.config.passthrough and self.config.rails.dialog.single_call.enabled: - raise InvalidRailsConfigurationError( - "The passthrough mode and the single call dialog rails mode can't be used at the same time. " - "The single call mode needs to use an altered prompt when prompting the LLM. " - ) - - async def _init_kb(self): - """Initializes the knowledge base.""" - self._kb = None - - if not self.config.docs: - return - - documents = [doc.content for doc in self.config.docs] - self._kb = KnowledgeBase( - documents=documents, - config=self.config.knowledge_base, - get_embedding_search_provider_instance=self._get_embeddings_search_provider_instance, - ) - self._kb.init() - await self._kb.build() + sync_update_llm_bindings(self, llm) def _prepare_model_kwargs(self, model_config): """ @@ -531,15 +370,7 @@ def _prepare_model_kwargs(self, model_config): Returns: dict: The prepared kwargs for model initialization """ - kwargs = model_config.parameters or {} - - # If the optional API Key Environment Variable is set, add it to kwargs - if model_config.api_key_env_var: - api_key = os.environ.get(model_config.api_key_env_var) - if api_key: - kwargs["api_key"] = api_key - - return kwargs + return model_kwargs_from_config(model_config) def _init_llms(self): """ @@ -554,183 +385,11 @@ def _init_llms(self): Raises: ModelInitializationError: If any model initialization fails """ - from nemoguardrails._compat.langchain_kwargs import check_langchain_kwargs - from nemoguardrails.llm.frameworks import get_default_framework - - models_to_check = ( - [model for model in self.config.models if model.type != "main"] if self.llm else self.config.models - ) - check_langchain_kwargs(models_to_check, get_default_framework()) - - # If the user supplied an already-constructed LLM via the constructor we - # treat it as the *main* model, but **still** iterate through the - # configuration to load any additional models (e.g. `content_safety`). - - if self.llm: - # If an LLM was provided via constructor, use it as the main LLM - # Log a warning if a main LLM is also specified in the config - if any(model.type == "main" for model in self.config.models): - log.warning( - "Both an LLM was provided via constructor and a main LLM is specified in the config. " - "The LLM provided via constructor will be used and the main LLM from config will be ignored." - ) - self.runtime.register_action_param("llm", self.llm) - - else: - # Otherwise, initialize the main LLM from the config - main_model = next((model for model in self.config.models if model.type == "main"), None) - - if main_model and main_model.model: - kwargs = self._prepare_model_kwargs(main_model) - self.llm = init_llm_model( - model_name=main_model.model, - provider_name=main_model.engine, - mode="chat", - kwargs=kwargs, - ) - self.runtime.register_action_param("llm", self.llm) - - else: - log.info("No main LLM specified in the config and no LLM provided via constructor.") + load_llm_action_models(self, init_llm=init_llm_model) + initialize_llm_action_caches(self) - llms = dict() - - for llm_config in self.config.models: - if llm_config.type in ["embeddings", "jailbreak_detection"]: - continue - - # If a constructor LLM is provided, skip initializing any 'main' model from config - if self.llm and llm_config.type == "main": - continue - - try: - model_name = llm_config.model - if not model_name: - raise InvalidModelConfigurationError( - f"`model` field must be set in model configuration: {llm_config.model_dump_json()}" - ) - - provider_name = llm_config.engine - kwargs = self._prepare_model_kwargs(llm_config) - mode = llm_config.mode - - llm_model = init_llm_model( - model_name=model_name, - provider_name=provider_name, - mode=mode, - kwargs=kwargs, - ) - - # Configure the model based on its type - if llm_config.type == "main": - # If a main LLM was already injected, skip creating another - # one. Otherwise, create and register it. - if not self.llm: - self.llm = llm_model - self.runtime.register_action_param("llm", self.llm) - else: - model_name = f"{llm_config.type}_llm" - if not hasattr(self, model_name): - setattr(self, model_name, llm_model) - self.runtime.register_action_param(model_name, getattr(self, model_name)) - # this is used for content safety and topic control - llms[llm_config.type] = getattr(self, model_name) - - except ModelInitializationError as e: - log.error("Failed to initialize model: %s", str(e)) - raise - except Exception as e: - log.error("Unexpected error initializing model: %s", str(e)) - raise - - self.runtime.register_action_param("llms", llms) - - self._initialize_model_caches() - - def _create_model_cache(self, model) -> LFUCache: - """ - Create cache instance for a model based on its configuration. - - Args: - model: The model configuration object - - Returns: - LFUCache: The cache instance - """ - - if model.cache.maxsize <= 0: - raise ValueError( - f"Invalid cache maxsize for model '{model.type}': {model.cache.maxsize}. " - "Capacity must be greater than 0. Skipping cache creation." - ) - - stats_logging_interval = None - if model.cache.stats.enabled and model.cache.stats.log_interval is not None: - stats_logging_interval = model.cache.stats.log_interval - - cache = LFUCache( - maxsize=model.cache.maxsize, - track_stats=model.cache.stats.enabled, - stats_logging_interval=stats_logging_interval, - ) - - log.info(f"Created cache for model '{model.type}' with maxsize {model.cache.maxsize}") - - return cache - - def _initialize_model_caches(self) -> None: - """Initialize caches for configured models.""" - model_caches: Optional[Dict[str, CacheInterface]] = dict() - for model in self.config.models: - if model.type in ["main", "embeddings"]: - continue - - if model.cache and model.cache.enabled: - cache = self._create_model_cache(model) - model_caches[model.type] = cache - - log.info( - f"Initialized model '{model.type}' with cache %s", - "enabled" if cache else "disabled", - ) - - if model_caches: - self.runtime.register_action_param("model_caches", model_caches) - - def _get_embeddings_search_provider_instance( - self, esp_config: Optional[EmbeddingSearchProvider] = None - ) -> EmbeddingsIndex: - if esp_config is None: - esp_config = EmbeddingSearchProvider() - - if esp_config.name == "default": - from nemoguardrails.embeddings.basic import BasicEmbeddingsIndex - - return BasicEmbeddingsIndex( - embedding_model=esp_config.parameters.get("embedding_model", self._default_embedding_model), - embedding_engine=esp_config.parameters.get("embedding_engine", self._default_embedding_engine), - embedding_params=esp_config.parameters.get("embedding_parameters", self._default_embedding_params), - cache_config=esp_config.cache, - # We make sure we also pass additional relevant params. - **{ - k: v - for k, v in esp_config.parameters.items() - if k - in [ - "use_batching", - "max_batch_size", - "matx_batch_hold", - "search_threshold", - ] - and v is not None - }, - ) - else: - if esp_config.name not in self._embedding_search_providers: - raise Exception(f"Unknown embedding search provider: {esp_config.name}") - else: - kwargs = esp_config.parameters - return self._embedding_search_providers[esp_config.name](**kwargs) + def _get_embeddings_search_provider_instance(self, esp_config=None): + return self.embedding_search.get_provider_instance(esp_config) def _get_events_for_messages(self, messages: List[dict], state: Any): """Return the list of events corresponding to the provided messages. @@ -1759,7 +1418,7 @@ def register_embedding_search_provider(self, name: str, cls: Type[EmbeddingsInde cls: The class that will be used to generate and search embedding """ - self._embedding_search_providers[name] = cls + self.embedding_search.register_provider(name, cls) return self def register_embedding_provider(self, cls: Type[EmbeddingModel], name: Optional[str] = None) -> Self: diff --git a/nemoguardrails/rails/llm/runtime/__init__.py b/nemoguardrails/rails/llm/runtime/__init__.py new file mode 100644 index 0000000000..dc33922794 --- /dev/null +++ b/nemoguardrails/rails/llm/runtime/__init__.py @@ -0,0 +1,16 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +__all__ = [] diff --git a/nemoguardrails/rails/llm/runtime/colang_runtime.py b/nemoguardrails/rails/llm/runtime/colang_runtime.py new file mode 100644 index 0000000000..7d1c665557 --- /dev/null +++ b/nemoguardrails/rails/llm/runtime/colang_runtime.py @@ -0,0 +1,40 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Colang runtime selection.""" + +from typing import Dict, Type + +from nemoguardrails.colang.v1_0.runtime.runtime import Runtime, RuntimeV1_0 +from nemoguardrails.colang.v2_x.runtime.runtime import RuntimeV2_x +from nemoguardrails.exceptions import InvalidRailsConfigurationError +from nemoguardrails.rails.llm.config import RailsConfig + +__all__ = ["runtime_for_colang_version"] + + +def runtime_for_colang_version(config: RailsConfig, verbose: bool) -> Runtime: + """Create the Colang runtime selected by the config version.""" + colang_version_to_runtime: Dict[str, Type[Runtime]] = { + "1.0": RuntimeV1_0, + "2.x": RuntimeV2_x, + } + if config.colang_version not in colang_version_to_runtime: + raise InvalidRailsConfigurationError( + f"Unsupported colang version: {config.colang_version}. " + f"Supported versions: {list(colang_version_to_runtime.keys())}" + ) + + return colang_version_to_runtime[config.colang_version](config=config, verbose=verbose) diff --git a/nemoguardrails/rails/llm/startup/colang_flows.py b/nemoguardrails/rails/llm/startup/colang_flows.py index e4c577841a..49e81e0267 100644 --- a/nemoguardrails/rails/llm/startup/colang_flows.py +++ b/nemoguardrails/rails/llm/startup/colang_flows.py @@ -18,7 +18,11 @@ import importlib.resources as resources import logging from collections.abc import Iterator -from importlib.abc import Traversable + +try: + from importlib.resources.abc import Traversable +except ImportError: + from importlib.abc import Traversable from nemoguardrails.colang import parse_colang_file from nemoguardrails.colang.v1_0.runtime.flows import _normalize_flow_id diff --git a/nemoguardrails/rails/llm/startup/config_preparation.py b/nemoguardrails/rails/llm/startup/config_preparation.py index 65b877e86e..9fe2fdb820 100644 --- a/nemoguardrails/rails/llm/startup/config_preparation.py +++ b/nemoguardrails/rails/llm/startup/config_preparation.py @@ -15,62 +15,27 @@ """LLMRails config preparation.""" -from dataclasses import dataclass -from typing import Any, Dict, Optional - from nemoguardrails.rails.llm.config import RailsConfig from nemoguardrails.rails.llm.startup.colang_flows import ( load_default_colang_1_flows, load_guardrails_library_flows_and_bot_messages, mark_rail_flows_as_system_subflows, ) -from nemoguardrails.rails.llm.startup.embedding_search import apply_embedding_model_config - -_COLANG_FLOWS_PREPARED_ATTR = "_llmrails_colang_flows_prepared" - -__all__ = ["PreparedLLMRailsConfig", "prepare_llmrails_config"] - -@dataclass -class PreparedLLMRailsConfig: - config: RailsConfig - default_embedding_model: Optional[str] - default_embedding_engine: Optional[str] - default_embedding_params: Dict[str, Any] +__all__ = ["prepare_llmrails_config"] def prepare_llmrails_config( *, config: RailsConfig, - default_embedding_model: Optional[str], - default_embedding_engine: Optional[str], - default_embedding_params: Dict[str, Any], in_place: bool = True, -) -> PreparedLLMRailsConfig: +) -> RailsConfig: """Prepare a RailsConfig for the standard LLMRails runtime.""" prepared_config = config if in_place else config.model_copy(deep=True) - if not getattr(prepared_config, _COLANG_FLOWS_PREPARED_ATTR, False): - load_default_colang_1_flows(prepared_config) - load_guardrails_library_flows_and_bot_messages(prepared_config) - setattr(prepared_config, _COLANG_FLOWS_PREPARED_ATTR, True) + load_default_colang_1_flows(prepared_config) + load_guardrails_library_flows_and_bot_messages(prepared_config) mark_rail_flows_as_system_subflows(prepared_config) - ( - default_embedding_model, - default_embedding_engine, - default_embedding_params, - ) = apply_embedding_model_config( - config=prepared_config, - default_embedding_model=default_embedding_model, - default_embedding_engine=default_embedding_engine, - default_embedding_params=default_embedding_params, - ) - - return PreparedLLMRailsConfig( - config=prepared_config, - default_embedding_model=default_embedding_model, - default_embedding_engine=default_embedding_engine, - default_embedding_params=default_embedding_params, - ) + return prepared_config diff --git a/nemoguardrails/rails/llm/startup/generation_actions.py b/nemoguardrails/rails/llm/startup/generation_actions.py index 13e9c20f25..16b41e7299 100644 --- a/nemoguardrails/rails/llm/startup/generation_actions.py +++ b/nemoguardrails/rails/llm/startup/generation_actions.py @@ -38,7 +38,7 @@ def register_actions(self, actions_obj: Any, /, override: bool = True) -> None: class GenerationActionRails(Protocol): - llm_generation_actions: Any + _llm_generation_actions: Any @property def config(self) -> Any: ... @@ -49,8 +49,7 @@ def llm(self) -> Any: ... @property def runtime(self) -> GenerationActionRuntime: ... - @property - def embedding_search(self) -> Any: ... + def _get_embeddings_search_provider_instance(self, esp_config: Any = None) -> Any: ... def generation_actions_class_for_colang_version(colang_version: str) -> Type[Any]: @@ -61,12 +60,12 @@ def generation_actions_class_for_colang_version(colang_version: str) -> Type[Any def register_llm_generation_actions(rails: GenerationActionRails, verbose: bool) -> None: """Create and register the LLM generation actions for an LLMRails instance.""" llm_generation_actions_class = generation_actions_class_for_colang_version(rails.config.colang_version) - rails.llm_generation_actions = llm_generation_actions_class( + rails._llm_generation_actions = llm_generation_actions_class( config=rails.config, llm=rails.llm, llm_task_manager=rails.runtime.llm_task_manager, - get_embedding_search_provider_instance=rails.embedding_search.get_provider_instance, + get_embedding_search_provider_instance=rails._get_embeddings_search_provider_instance, verbose=verbose, ) - rails.runtime.register_actions(rails.llm_generation_actions, override=False) + rails.runtime.register_actions(rails._llm_generation_actions, override=False) diff --git a/nemoguardrails/rails/llm/startup/knowledge_base.py b/nemoguardrails/rails/llm/startup/knowledge_base.py index 4334501990..9fb3603a32 100644 --- a/nemoguardrails/rails/llm/startup/knowledge_base.py +++ b/nemoguardrails/rails/llm/startup/knowledge_base.py @@ -49,12 +49,12 @@ async def build_knowledge_base_for_docs( def init_knowledge_base(rails) -> None: """Initialize and register the LLMRails knowledge base.""" - rails.kb = None + rails._kb = None async def _init_kb(): - rails.kb = await build_knowledge_base_for_docs( + rails._kb = await build_knowledge_base_for_docs( config=rails.config, - get_embedding_search_provider_instance=rails.embedding_search.get_provider_instance, + get_embedding_search_provider_instance=rails._get_embeddings_search_provider_instance, ) # There are still some edge cases not covered by nest_asyncio. @@ -67,4 +67,4 @@ async def _init_kb(): else: loop.run_until_complete(_init_kb()) - rails.runtime.register_action_param("kb", rails.kb) + rails.runtime.register_action_param("kb", rails._kb) diff --git a/nemoguardrails/rails/llm/startup/llm_action_models.py b/nemoguardrails/rails/llm/startup/llm_action_models.py index 539099e336..01f17cbc37 100644 --- a/nemoguardrails/rails/llm/startup/llm_action_models.py +++ b/nemoguardrails/rails/llm/startup/llm_action_models.py @@ -43,7 +43,7 @@ def register_action_param(self, name: str, value: Any) -> None: ... class LLMActionRails(Protocol): llm: Any - llm_generation_actions: Any + _llm_generation_actions: Any @property def config(self) -> Any: ... @@ -67,7 +67,7 @@ def model_kwargs_from_config(model_config: Any) -> Dict[str, Any]: def sync_update_llm_bindings(rails: LLMActionRails, llm: LLMModel) -> None: """Synchronize the main LLM bindings after a public update.""" rails.llm = llm - rails.llm_generation_actions.llm = llm + rails._llm_generation_actions.llm = llm rails.runtime.register_action_param("llm", llm) diff --git a/tests/guardrails/test_public_api_deprecations.py b/tests/guardrails/test_public_api_deprecations.py index 1f913817f4..0495fdb3d6 100644 --- a/tests/guardrails/test_public_api_deprecations.py +++ b/tests/guardrails/test_public_api_deprecations.py @@ -36,6 +36,7 @@ from nemoguardrails.guardrails.iorails import IORails from nemoguardrails.rails.llm.config import RailsConfig from nemoguardrails.rails.llm.llmrails import LLMRails +from nemoguardrails.rails.llm.startup.embedding_search import EmbeddingSearchState from tests.guardrails.test_data import CONTENT_SAFETY_CONFIG @@ -69,21 +70,12 @@ def iorails_guardrails(_content_safety_config): class TestLLMRailsDeprecatedAliases: - """Deprecated property aliases on LLMRails. - - Each public name is a thin getter (and, for default_embedding_*, also a - setter) that warns and forwards to the underscored attribute. Tests - verify both the warning and the forwarded value/assignment. - """ + """Deprecated property aliases on LLMRails.""" @pytest.mark.parametrize( "public_name, private_name", [ ("kb", "_kb"), - ("embedding_search_providers", "_embedding_search_providers"), - ("default_embedding_model", "_default_embedding_model"), - ("default_embedding_engine", "_default_embedding_engine"), - ("default_embedding_params", "_default_embedding_params"), ("llm_generation_actions", "_llm_generation_actions"), ], ) @@ -99,26 +91,40 @@ def test_read_emits_deprecation_warning(self, public_name, private_name): assert value is sentinel @pytest.mark.parametrize( - "public_name, private_name", + "public_name, state_name, sentinel", [ - ("default_embedding_model", "_default_embedding_model"), - ("default_embedding_engine", "_default_embedding_engine"), - ("default_embedding_params", "_default_embedding_params"), + ("embedding_search_providers", "providers", {}), + ("default_embedding_model", "default_model", "model"), + ("default_embedding_engine", "default_engine", "engine"), + ("default_embedding_params", "default_params", {"device": "cpu"}), ], ) - def test_write_emits_deprecation_warning(self, public_name, private_name): - """Writing a deprecated default_embedding_* alias warns and forwards to the underscored attribute. + def test_embedding_read_emits_deprecation_warning(self, public_name, state_name, sentinel): + rails = LLMRails.__new__(LLMRails) + rails.embedding_search = EmbeddingSearchState.default() + setattr(rails.embedding_search, state_name, sentinel) - These have setters (unlike kb / embedding_search_providers / llm_generation_actions) - because they were previously plain instance attributes that downstream code may have - written to. The setter preserves that write path during deprecation. - """ - # Bypass __init__: testing only the @setter descriptor. + with pytest.warns(DeprecationWarning, match=rf"LLMRails\.{public_name}.*deprecated"): + value = getattr(rails, public_name) + + assert value is sentinel + + @pytest.mark.parametrize( + "public_name, state_name", + [ + ("default_embedding_model", "default_model"), + ("default_embedding_engine", "default_engine"), + ("default_embedding_params", "default_params"), + ], + ) + def test_write_emits_deprecation_warning(self, public_name, state_name): + """Writing a deprecated default_embedding_* alias warns and updates embedding_search state.""" rails = LLMRails.__new__(LLMRails) + rails.embedding_search = EmbeddingSearchState.default() sentinel = object() with pytest.warns(DeprecationWarning, match=rf"Setting LLMRails\.{public_name}.*deprecated"): setattr(rails, public_name, sentinel) - assert getattr(rails, private_name) is sentinel + assert getattr(rails.embedding_search, state_name) is sentinel class TestLLMRailsExplainInfoDeprecation: diff --git a/tests/rails/llm/test_colang_runtime.py b/tests/rails/llm/test_colang_runtime.py new file mode 100644 index 0000000000..9f6f0077de --- /dev/null +++ b/tests/rails/llm/test_colang_runtime.py @@ -0,0 +1,95 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest + +from nemoguardrails.colang.v1_0.runtime.runtime import RuntimeV1_0 +from nemoguardrails.colang.v2_x.runtime.runtime import RuntimeV2_x +from nemoguardrails.exceptions import InvalidRailsConfigurationError +from nemoguardrails.rails.llm import runtime as runtime_package +from nemoguardrails.rails.llm.config import RailsConfig +from nemoguardrails.rails.llm.llmrails import LLMRails +from nemoguardrails.rails.llm.runtime import colang_runtime +from nemoguardrails.rails.llm.runtime.colang_runtime import runtime_for_colang_version + + +def test_runtime_package_has_no_star_exports(): + assert runtime_package.__all__ == [] + + +def test_runtime_for_colang_version_uses_v1_runtime(): + config = RailsConfig(models=[], colang_version="1.0") + + runtime = runtime_for_colang_version(config=config, verbose=False) + + assert isinstance(runtime, RuntimeV1_0) + + +def test_runtime_for_colang_version_uses_v2_runtime(): + config = RailsConfig(models=[], colang_version="2.x") + + runtime = runtime_for_colang_version(config=config, verbose=True) + + assert isinstance(runtime, RuntimeV2_x) + + +def test_runtime_for_colang_version_rejects_unsupported_version(): + config = RailsConfig(models=[], colang_version="3.x") + + with pytest.raises(InvalidRailsConfigurationError) as exc_info: + runtime_for_colang_version(config=config, verbose=False) + + assert str(exc_info.value) == "Unsupported colang version: 3.x. Supported versions: ['1.0', '2.x']" + + +def test_colang_runtime_module_exports_public_helpers(): + assert colang_runtime.__all__ == ["runtime_for_colang_version"] + + +def test_llmrails_initializes_v1_runtime(): + rails = LLMRails(RailsConfig(models=[], colang_version="1.0")) + + assert isinstance(rails.runtime, RuntimeV1_0) + + +def test_llmrails_initializes_v2_runtime(): + rails = LLMRails(RailsConfig(models=[], colang_version="2.x")) + + assert isinstance(rails.runtime, RuntimeV2_x) + + +def test_llmrails_rejects_unsupported_runtime_version(): + config = RailsConfig(models=[], colang_version="3.x") + + with pytest.raises(InvalidRailsConfigurationError) as exc_info: + LLMRails(config) + + assert str(exc_info.value) == "Unsupported colang version: 3.x. Supported versions: ['1.0', '2.x']" + + +def test_config_py_init_hook_sees_initialized_runtime(tmp_path): + config_path = tmp_path / "config" + config_path.mkdir() + (config_path / "config.py").write_text( + "def init(app):\n app.runtime_seen_by_init = type(app.runtime).__name__\n" + ) + config = RailsConfig( + config_path=str(config_path), + models=[], + ) + + rails = LLMRails(config) + + assert getattr(rails, "runtime_seen_by_init") == "RuntimeV1_0" diff --git a/tests/rails/llm/test_config_preparation.py b/tests/rails/llm/test_config_preparation.py index 5cbd0067b0..dfe7a991f9 100644 --- a/tests/rails/llm/test_config_preparation.py +++ b/tests/rails/llm/test_config_preparation.py @@ -13,46 +13,35 @@ # See the License for the specific language governing permissions and # limitations under the License. -from nemoguardrails.rails.llm.config import Model, RailsConfig +from nemoguardrails.rails.llm.config import RailsConfig from nemoguardrails.rails.llm.startup import config_preparation as startup_config_preparation from nemoguardrails.rails.llm.startup.config_preparation import ( prepare_llmrails_config, ) -from nemoguardrails.rails.llm.startup.embedding_search import ( - DEFAULT_EMBEDDING_ENGINE, - DEFAULT_EMBEDDING_MODEL, -) def test_config_preparation_startup_module_exports_public_helpers(): assert startup_config_preparation.__all__ == [ - "PreparedLLMRailsConfig", "prepare_llmrails_config", ] -def test_prepare_llmrails_config_loads_colang_flows_once_in_place(): +def test_prepare_llmrails_config_loads_colang_flows_each_time_in_place(): config = RailsConfig(models=[]) first = prepare_llmrails_config( config=config, - default_embedding_model=DEFAULT_EMBEDDING_MODEL, - default_embedding_engine=DEFAULT_EMBEDDING_ENGINE, - default_embedding_params={}, ) first_flow_count = len(config.flows) second = prepare_llmrails_config( config=config, - default_embedding_model=DEFAULT_EMBEDDING_MODEL, - default_embedding_engine=DEFAULT_EMBEDDING_ENGINE, - default_embedding_params={}, ) - assert first.config is config - assert second.config is config + assert first is config + assert second is config assert first_flow_count > 0 - assert len(config.flows) == first_flow_count + assert len(config.flows) > first_flow_count def test_prepare_llmrails_config_marks_rail_flows_each_time(): @@ -63,17 +52,11 @@ def test_prepare_llmrails_config_marks_rail_flows_each_time(): prepare_llmrails_config( config=config, - default_embedding_model=DEFAULT_EMBEDDING_MODEL, - default_embedding_engine=DEFAULT_EMBEDDING_ENGINE, - default_embedding_params={}, ) config.rails.input.flows = ["check input"] prepare_llmrails_config( config=config, - default_embedding_model=DEFAULT_EMBEDDING_MODEL, - default_embedding_engine=DEFAULT_EMBEDDING_ENGINE, - default_embedding_params={}, ) assert config.flows[0]["is_system_flow"] is True @@ -85,37 +68,9 @@ def test_prepare_llmrails_config_can_prepare_a_copy(): prepared = prepare_llmrails_config( config=config, - default_embedding_model=DEFAULT_EMBEDDING_MODEL, - default_embedding_engine=DEFAULT_EMBEDDING_ENGINE, - default_embedding_params={}, in_place=False, ) - assert prepared.config is not config + assert prepared is not config assert len(config.flows) == 0 - assert len(prepared.config.flows) > 0 - - -def test_prepare_llmrails_config_returns_embedding_defaults(): - config = RailsConfig( - models=[ - Model( - type="embeddings", - engine="SentenceTransformers", - model="intfloat/e5-large-v2", - parameters={"device": "cpu"}, - ) - ], - ) - - prepared = prepare_llmrails_config( - config=config, - default_embedding_model=DEFAULT_EMBEDDING_MODEL, - default_embedding_engine=DEFAULT_EMBEDDING_ENGINE, - default_embedding_params={}, - ) - - assert prepared.default_embedding_model == "intfloat/e5-large-v2" - assert prepared.default_embedding_engine == "SentenceTransformers" - assert prepared.default_embedding_params == {"device": "cpu"} - assert config.core.embedding_search_provider.parameters["embedding_model"] == "intfloat/e5-large-v2" + assert len(prepared.flows) > 0 diff --git a/tests/rails/llm/test_generation_actions.py b/tests/rails/llm/test_generation_actions.py index db6f929b9e..a14af6b83d 100644 --- a/tests/rails/llm/test_generation_actions.py +++ b/tests/rails/llm/test_generation_actions.py @@ -58,8 +58,8 @@ def __init__(self, config, llm, runtime, get_embedding_search_provider_instance) self.config = config self.llm = llm self.runtime = runtime - self.embedding_search = SimpleNamespace(get_provider_instance=get_embedding_search_provider_instance) - self.llm_generation_actions = None + self._get_embeddings_search_provider_instance = get_embedding_search_provider_instance + self._llm_generation_actions = None def test_generation_actions_class_for_colang_version_uses_v1_actions(monkeypatch): @@ -97,11 +97,13 @@ def get_embedding_search_provider_instance(esp_config=None): register_llm_generation_actions(rails, verbose=True) - assert isinstance(rails.llm_generation_actions, FakeGenerationActions) - assert rails.llm_generation_actions.config is config - assert rails.llm_generation_actions.llm is llm - assert rails.llm_generation_actions.llm_task_manager is runtime.llm_task_manager - assert rails.llm_generation_actions.get_embedding_search_provider_instance is get_embedding_search_provider_instance - assert rails.llm_generation_actions.verbose is True - assert runtime.registered_actions is rails.llm_generation_actions + assert isinstance(rails._llm_generation_actions, FakeGenerationActions) + assert rails._llm_generation_actions.config is config + assert rails._llm_generation_actions.llm is llm + assert rails._llm_generation_actions.llm_task_manager is runtime.llm_task_manager + assert ( + rails._llm_generation_actions.get_embedding_search_provider_instance is get_embedding_search_provider_instance + ) + assert rails._llm_generation_actions.verbose is True + assert runtime.registered_actions is rails._llm_generation_actions assert runtime.override is False diff --git a/tests/rails/llm/test_knowledge_base.py b/tests/rails/llm/test_knowledge_base.py index 9f7dc26fbb..d09939ee9a 100644 --- a/tests/rails/llm/test_knowledge_base.py +++ b/tests/rails/llm/test_knowledge_base.py @@ -100,7 +100,7 @@ def test_init_knowledge_base_uses_thread_path_and_registers_kb(monkeypatch): calls = {"thread_started": False, "thread_joined": False} async def fake_build_knowledge_base_for_docs(config, get_embedding_search_provider_instance): - assert rails.kb is None + assert rails._kb is None assert config is rails.config assert get_embedding_search_provider_instance is get_provider_instance return built_kb @@ -125,9 +125,9 @@ def get_provider_instance(_=None): models=[], docs=[Document(format="md", content="# Alpha\n\nA")], ), - kb="existing", + _kb="existing", runtime=RuntimeWithActionParams(), - embedding_search=SimpleNamespace(get_provider_instance=get_provider_instance), + _get_embeddings_search_provider_instance=get_provider_instance, ) monkeypatch.setattr( knowledge_base, @@ -139,6 +139,6 @@ def get_provider_instance(_=None): init_knowledge_base(rails) - assert rails.kb is built_kb + assert rails._kb is built_kb assert rails.runtime.registered_action_params["kb"] is built_kb assert calls == {"thread_started": True, "thread_joined": True} diff --git a/tests/rails/llm/test_llm_action_models.py b/tests/rails/llm/test_llm_action_models.py index 573f373dd8..37bd906403 100644 --- a/tests/rails/llm/test_llm_action_models.py +++ b/tests/rails/llm/test_llm_action_models.py @@ -40,7 +40,7 @@ def __init__(self, config: RailsConfig, llm: Any, runtime: RuntimeWithActionPara self.config = config self.llm = llm self.runtime = runtime - self.llm_generation_actions = MagicMock() + self._llm_generation_actions = MagicMock() def test_model_kwargs_from_config_adds_api_key_from_environment(): @@ -162,5 +162,5 @@ def test_sync_update_llm_bindings_updates_llm_generation_actions_and_runtime_par sync_update_llm_bindings(rails, updated_llm) assert rails.llm is updated_llm - assert rails.llm_generation_actions.llm is updated_llm + assert rails._llm_generation_actions.llm is updated_llm assert rails.runtime.registered_action_params["llm"] is updated_llm diff --git a/tests/rails/llm/test_llmrails_public_contract.py b/tests/rails/llm/test_llmrails_public_contract.py index 3c040c1a5f..9f5ee48fa3 100644 --- a/tests/rails/llm/test_llmrails_public_contract.py +++ b/tests/rails/llm/test_llmrails_public_contract.py @@ -13,6 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +from unittest.mock import patch + import pytest from nemoguardrails import LLMRails, RailsConfig @@ -37,6 +39,45 @@ def test_constructor_keeps_public_state_visible(): assert rails.explain_info is None +def test_constructor_reports_library_usage(): + config = _config() + + with patch("nemoguardrails.telemetry.report_usage") as report_usage: + LLMRails(config=config, llm=FakeLLMModel(responses=[])) + + report_usage.assert_called_once_with(config, deployment_type="library", rails_engine="LLMRails") + + +def test_config_py_init_embeddings_model_updates_default_search_config(tmp_path): + config = _config() + config.config_path = str(tmp_path) + (tmp_path / "config.py").write_text( + """ +from nemoguardrails.rails.llm.config import Model + + +def init(rails): + rails.config.models.append( + Model( + type="embeddings", + engine="SentenceTransformers", + model="intfloat/e5-large-v2", + parameters={"device": "cpu"}, + ) + ) +""", + encoding="utf-8", + ) + + rails = LLMRails(config=config, llm=FakeLLMModel(responses=[])) + + assert rails.embedding_search.default_model == "intfloat/e5-large-v2" + assert rails.embedding_search.default_engine == "SentenceTransformers" + assert rails.embedding_search.default_params == {"device": "cpu"} + assert config.core.embedding_search_provider.parameters["embedding_model"] == "intfloat/e5-large-v2" + assert config.core.embedding_search_provider.parameters["embedding_engine"] == "SentenceTransformers" + + def test_update_llm_keeps_runtime_generation_actions_and_public_attr_in_sync(): rails = LLMRails(config=_config(), llm=FakeLLMModel(responses=[])) new_llm = FakeLLMModel(responses=["updated"]) diff --git a/tests/test_llmrails.py b/tests/test_llmrails.py index ea65bbfc5d..ffb4d156d9 100644 --- a/tests/test_llmrails.py +++ b/tests/test_llmrails.py @@ -1586,30 +1586,3 @@ async def test_warning_behavior(self, no_main_llm_config, caplog, options, has_l else: await rails.generate_async(messages=messages, options=options) assert _count_no_llm_warnings(caplog) == expected_warnings - - -def test_load_library_sorts_files_for_deterministic_overrides(tmp_path, monkeypatch): - """Library files are traversed in sorted order so collisions resolve identically - on every filesystem. - - The library loader in ``LLMRails.__init__`` inserts each ``.co`` file's - ``bot_messages`` first-wins, so an unsorted ``os.walk`` would let the winner of a - message-id collision depend on filesystem ordering. This is the library-traversal - sibling of - ``test_prompt_override.py::test_load_prompts_sorts_files_for_deterministic_overrides``. - """ - library_dir = tmp_path / "library" - library_dir.mkdir() - (library_dir / "z.co").write_text('define bot test det msg\n "from_z"\n', encoding="utf-8") - (library_dir / "a.co").write_text('define bot test det msg\n "from_a"\n', encoding="utf-8") - - def walk(_path): - # Yield in non-sorted order; the loader must sort so "a.co" wins the collision. - yield str(library_dir), [], ["z.co", "a.co"] - - monkeypatch.setattr("nemoguardrails.rails.llm.llmrails.os.walk", walk) - - config = RailsConfig.from_content(yaml_content="models: []\n") - rails = LLMRails(config, llm=FakeLLMModel(responses=["unused"])) - - assert rails.config.bot_messages["test det msg"] == ["from_a"]