diff --git a/docs/index.yml b/docs/index.yml index 85eb68156c..5561684a47 100644 --- a/docs/index.yml +++ b/docs/index.yml @@ -605,6 +605,9 @@ navigation: - endpoint: GET / title: Get Server Health or Chat UI slug: get-root + - page: Migrating to 0.23 + path: migration/0.23.mdx + slug: 0-23 - page: Migrating to 0.22 path: migration/0.22.mdx slug: 0-22 diff --git a/docs/migration/0.23.mdx b/docs/migration/0.23.mdx new file mode 100644 index 0000000000..2e69ec5f24 --- /dev/null +++ b/docs/migration/0.23.mdx @@ -0,0 +1,14 @@ +--- +title: "Migrating to 0.23" +sidebar-title: "Migrating to 0.23" +description: Behavior changes to account for when updating a working NeMo Guardrails configuration to 0.23. +keywords: ["nemoguardrails 0.23", "migration", "LLMRails", "explain"] +content: + type: "how_to" +--- + +# Migrating to v0.23.0 + +## Behavior Changes + +- `LLMRails.explain().llm_calls` now reflects only the most recent `generate`/`generate_async` call instead of accumulating across calls on a reused `LLMRails` instance. This matches the documented "latest call" contract and is a side effect of per-request context isolation in the decomposed generation pipeline. diff --git a/nemoguardrails/rails/llm/checks/__init__.py b/nemoguardrails/rails/llm/checks/__init__.py new file mode 100644 index 0000000000..dc33922794 --- /dev/null +++ b/nemoguardrails/rails/llm/checks/__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/checks/rails_check.py b/nemoguardrails/rails/llm/checks/rails_check.py new file mode 100644 index 0000000000..cbf5362fb9 --- /dev/null +++ b/nemoguardrails/rails/llm/checks/rails_check.py @@ -0,0 +1,168 @@ +# 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. + +"""Message rail check helpers.""" + +import logging +from copy import deepcopy +from dataclasses import dataclass +from typing import Any, List, Optional, Protocol + +from nemoguardrails.rails.llm.options import ( + GenerationResponse, + RailsResult, + RailStatus, + RailType, +) + +log = logging.getLogger(__name__) + +__all__ = ["RailsCheckRuntime", "check_messages"] + + +class RailsCheckRuntime(Protocol): + config: Any + runtime: Any + + async def generate_async( + self, + *, + messages: List[dict], + options: dict, + ) -> object: ... + + +@dataclass +class RailsCheckPlan: + messages: List[dict] + options: dict + original_content: str + + +async def check_messages( + rails: RailsCheckRuntime, + messages: List[dict], + rail_types: Optional[List[RailType]] = None, +) -> RailsResult: + """Run input/output rails for a message list and return rail check status.""" + plan = _plan_rails_check(messages, rail_types) + if plan is None: + last_content = messages[-1].get("content", "") if messages else "" + return RailsResult(status=RailStatus.PASSED, content=last_content) + + response = await rails.generate_async(messages=plan.messages, options=plan.options) + + if not isinstance(response, GenerationResponse): + raise RuntimeError(f"Expected GenerationResponse, got {type(response).__name__}") + + return _classify_rails_response(response, plan.original_content) + + +def _plan_rails_check( + messages: List[dict], + rail_types: Optional[List[RailType]] = None, +) -> Optional[RailsCheckPlan]: + if rail_types is not None: + options: Optional[dict] = {"rails": [r.value for r in rail_types]} + else: + options = _determine_rails_from_messages(messages) + + if options is None: + return None + + rails_to_run = options["rails"] + if "output" in rails_to_run: + original_content = _get_last_content_by_role(messages, "assistant") + else: + original_content = _get_last_content_by_role(messages, "user") + + messages = _normalize_messages_for_rails(deepcopy(messages), rails_to_run) + options["log"] = {"activated_rails": True} + + return RailsCheckPlan( + messages=messages, + options=options, + original_content=original_content, + ) + + +def _classify_rails_response( + response: GenerationResponse, + original_content: str, +) -> RailsResult: + blocking_rail = _get_blocking_rail(response) + result_content = _get_last_response_content(response) + + if blocking_rail: + return RailsResult(status=RailStatus.BLOCKED, content=result_content, rail=blocking_rail) + + if result_content != original_content: + return RailsResult(status=RailStatus.MODIFIED, content=result_content) + return RailsResult(status=RailStatus.PASSED, content=result_content) + + +def _determine_rails_from_messages(messages: List[dict]) -> Optional[dict]: + roles = {msg.get("role") for msg in reversed(messages)} + has_user = "user" in roles + has_assistant = "assistant" in roles + + if not has_user and not has_assistant: + log.warning( + "check() called with no user or assistant messages. " + "Only system, context, or tool messages found. " + "Returning passing result without running rails." + ) + return None + + if has_user and has_assistant: + return {"rails": ["input", "output"]} + if has_user: + return {"rails": ["input"]} + return {"rails": ["output"]} + + +def _normalize_messages_for_rails( + messages: List[dict], + rails: List[str], +) -> List[dict]: + if rails == ["output"]: + has_user = any(msg.get("role") == "user" for msg in messages) + if not has_user: + return [{"role": "user", "content": ""}] + messages + + return messages + + +def _get_last_content_by_role(messages: List[dict], role: str) -> str: + for msg in reversed(messages): + if msg.get("role") == role: + return msg.get("content", "") + return "" + + +def _get_blocking_rail(response: GenerationResponse) -> Optional[str]: + if response.log and response.log.activated_rails: + for rail in response.log.activated_rails: + if rail.stop: + return rail.name + return None + + +def _get_last_response_content(response: GenerationResponse) -> str: + if isinstance(response.response, list) and response.response: + return response.response[-1].get("content", "") + if isinstance(response.response, str): + return response.response + return "" diff --git a/nemoguardrails/rails/llm/llmrails.py b/nemoguardrails/rails/llm/llmrails.py index d1f60a3a2a..3f551f12c2 100644 --- a/nemoguardrails/rails/llm/llmrails.py +++ b/nemoguardrails/rails/llm/llmrails.py @@ -42,11 +42,8 @@ from nemoguardrails.logging.explain import ExplainInfo from nemoguardrails.logging.verbose import set_verbose from nemoguardrails.patch_asyncio import check_sync_call_from_async_loop -from nemoguardrails.rails.llm.config import ( - OutputRailsStreamingConfig, - RailsConfig, -) -from nemoguardrails.rails.llm.conversation.conversation_events import events_for_messages +from nemoguardrails.rails.llm.checks import rails_check +from nemoguardrails.rails.llm.config import OutputRailsStreamingConfig, RailsConfig from nemoguardrails.rails.llm.generation.generation_context import ( ensure_explain_info, explain_info_for_current_context, @@ -57,13 +54,7 @@ validate_public_state, ) from nemoguardrails.rails.llm.generation.generation_workflow import generate_standard_async -from nemoguardrails.rails.llm.options import ( - GenerationOptions, - GenerationResponse, - RailsResult, - RailStatus, - RailType, -) +from nemoguardrails.rails.llm.options import GenerationOptions, GenerationResponse, RailsResult, RailType from nemoguardrails.rails.llm.runtime.colang_runtime import runtime_for_colang_version from nemoguardrails.rails.llm.runtime.colang_turns import ( generate_colang_events, @@ -231,17 +222,10 @@ def llm_generation_actions(self): @property def passthrough_fn(self): - """The optional passthrough function that bypasses LLM generation. - - When set, the rails pipeline calls this function instead of the main LLM - for generating responses. LLMGenerationActions is private, expose only - `passthrough_fn` as a public API - """ return self._llm_generation_actions._passthrough_fn @passthrough_fn.setter def passthrough_fn(self, fn): - """LLMGenerationActions is private, set passthrough_fn directly""" self._llm_generation_actions._passthrough_fn = fn @property @@ -367,12 +351,6 @@ def _init_llms(self): load_llm_action_models(self, init_llm=init_llm_model) initialize_llm_action_caches(self) - 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 events_for_messages(self, messages, state) - @staticmethod def _ensure_explain_info() -> ExplainInfo: """Ensure that the ExplainInfo variable is present in the current context @@ -382,9 +360,8 @@ def _ensure_explain_info() -> ExplainInfo: """ return ensure_explain_info() - def _validate_public_state(self, state: Optional[Union[dict, State]]) -> None: - """Validate public dict state passed through generate/generate_async.""" - validate_public_state(self.config, state) + def _get_embeddings_search_provider_instance(self, esp_config=None): + return self.embedding_search.get_provider_instance(esp_config) async def generate_async( self, @@ -662,38 +639,7 @@ async def check_async( result = await rails.check_async(messages, rail_types=[RailType.INPUT]) """ - if rail_types is not None: - options: Optional[dict] = {"rails": [r.value for r in rail_types]} - else: - options = _determine_rails_from_messages(messages) - - if options is None: - last_content = messages[-1].get("content", "") if messages else "" - return RailsResult(status=RailStatus.PASSED, content=last_content) - - rails_to_run = options["rails"] - if "output" in rails_to_run: - original_content = _get_last_content_by_role(messages, "assistant") - else: - original_content = _get_last_content_by_role(messages, "user") - - messages = _normalize_messages_for_rails(messages, rails_to_run) - options["log"] = {"activated_rails": True} - - response = await self.generate_async(messages=messages, options=options) - - if not isinstance(response, GenerationResponse): - raise RuntimeError(f"Expected GenerationResponse, got {type(response).__name__}") - - blocking_rail = _get_blocking_rail(response) - result_content = _get_last_response_content(response) - - if blocking_rail: - return RailsResult(status=RailStatus.BLOCKED, content=result_content, rail=blocking_rail) - - if result_content != original_content: - return RailsResult(status=RailStatus.MODIFIED, content=result_content) - return RailsResult(status=RailStatus.PASSED, content=result_content) + return await rails_check.check_messages(self, messages, rail_types=rail_types) def check( self, @@ -720,22 +666,38 @@ def check( return loop.run_until_complete(self.check_async(messages, rail_types=rail_types)) def register_action(self, action: Callable, name: Optional[str] = None) -> Self: - """Register a custom action for the rails configuration.""" + """Register a custom action for the rails configuration. + + This mutates the runtime action registry and is intended to be called + during application startup, before concurrent generation requests begin. + """ self.runtime.register_action(action, name) return self def register_action_param(self, name: str, value: Any) -> Self: - """Registers a custom action parameter.""" + """Register a custom action parameter. + + This mutates runtime action dependencies and is intended to be called + during application startup, before concurrent generation requests begin. + """ self.runtime.register_action_param(name, value) return self def register_filter(self, filter_fn: Callable, name: Optional[str] = None) -> Self: - """Register a custom filter for the rails configuration.""" + """Register a custom filter for the rails configuration. + + This mutates the runtime task manager and is intended to be called + during application startup, before concurrent generation requests begin. + """ self.runtime.llm_task_manager.register_filter(filter_fn, name) return self def register_output_parser(self, output_parser: Callable, name: str) -> Self: - """Register a custom output parser for the rails configuration.""" + """Register a custom output parser for the rails configuration. + + This mutates the runtime task manager and is intended to be called + during application startup, before concurrent generation requests begin. + """ self.runtime.llm_task_manager.register_output_parser(output_parser, name) return self @@ -744,6 +706,9 @@ def register_prompt_context(self, name: str, value_or_fn: Any) -> Self: :name: The name of the variable or function that will be used. :value_or_fn: The value or function that will be used to generate the value. + + This mutates the runtime task manager and is intended to be called + during application startup, before concurrent generation requests begin. """ self.runtime.llm_task_manager.register_prompt_context(name, value_or_fn) return self @@ -754,6 +719,9 @@ def register_embedding_search_provider(self, name: str, cls: Type[EmbeddingsInde Args: name: The name of the embedding search provider that will be used. cls: The class that will be used to generate and search embedding + + This updates the instance provider registry. Register providers during + startup so knowledge-base and generation-action setup can see them. """ self.embedding_search.register_provider(name, cls) @@ -769,6 +737,9 @@ def register_embedding_provider(self, cls: Type[EmbeddingModel], name: Optional[ Raises: ValueError: If the engine name is not provided and the model does not have an engine name. ValueError: If the model does not have 'encode' or 'encode_async' methods. + + This updates the process-global embedding provider registry and is + intended to be called during application startup. """ register_embedding_provider(engine_name=name, model=cls) return self @@ -810,58 +781,3 @@ async def _run_output_rails_in_streaming( stream_first=stream_first, ): yield chunk - - -def _determine_rails_from_messages(messages: List[dict]) -> Optional[dict]: - roles = {msg.get("role") for msg in reversed(messages)} - has_user = "user" in roles - has_assistant = "assistant" in roles - - if not has_user and not has_assistant: - log.warning( - "check() called with no user or assistant messages. " - "Only system, context, or tool messages found. " - "Returning passing result without running rails." - ) - return None - - if has_user and has_assistant: - return {"rails": ["input", "output"]} - if has_user: - return {"rails": ["input"]} - return {"rails": ["output"]} - - -def _normalize_messages_for_rails( - messages: List[dict], - rails: List[str], -) -> List[dict]: - if rails == ["output"]: - has_user = any(msg.get("role") == "user" for msg in messages) - if not has_user: - return [{"role": "user", "content": ""}] + messages - - return messages - - -def _get_last_content_by_role(messages: List[dict], role: str) -> str: - for msg in reversed(messages): - if msg.get("role") == role: - return msg.get("content", "") - return "" - - -def _get_blocking_rail(response: "GenerationResponse") -> Optional[str]: - if response.log and response.log.activated_rails: - for rail in response.log.activated_rails: - if rail.stop: - return rail.name - return None - - -def _get_last_response_content(response: "GenerationResponse") -> str: - if isinstance(response.response, list) and response.response: - return response.response[-1].get("content", "") - if isinstance(response.response, str): - return response.response - return "" diff --git a/tests/rails/llm/test_llmrails_api_contract.py b/tests/rails/llm/test_llmrails_api_contract.py new file mode 100644 index 0000000000..dd4399da35 --- /dev/null +++ b/tests/rails/llm/test_llmrails_api_contract.py @@ -0,0 +1,107 @@ +# 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 pickle + +import pytest + +from nemoguardrails import LLMRails, RailsConfig +from nemoguardrails.context import explain_info_var +from nemoguardrails.logging.explain import ExplainInfo +from tests.utils import FakeLLMModel + + +def _config() -> RailsConfig: + return RailsConfig.from_content(config={"models": []}) + + +def test_constructor_keeps_public_compatibility_state_visible(): + config = _config() + llm = FakeLLMModel(responses=[]) + + rails = LLMRails(config=config, llm=llm) + + assert rails.config is config + assert rails.llm is llm + assert rails.runtime is not None + assert rails.llm_generation_actions is not None + assert rails.embedding_search.providers == {} + assert rails.events_history_cache == {} + assert rails.explain_info is None + + +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"]) + + rails.update_llm(new_llm) + + assert rails.llm is new_llm + assert rails.llm_generation_actions.llm is new_llm + assert rails.runtime.registered_action_params["llm"] is new_llm + + +@pytest.mark.asyncio +async def test_sync_wrappers_raise_when_called_from_async_loop(): + rails = LLMRails(config=_config(), llm=FakeLLMModel(responses=[])) + + with pytest.raises(RuntimeError, match="sync `generate` inside async code"): + rails.generate(prompt="hi") + + with pytest.raises(RuntimeError, match="sync `generate_events` inside async code"): + rails.generate_events([]) + + with pytest.raises(RuntimeError, match="sync `generate_events` inside async code"): + rails.process_events([]) + + with pytest.raises(RuntimeError, match="sync `check` inside async code"): + rails.check([{"role": "user", "content": "hi"}]) + + +def test_getstate_serializes_config_only(): + rails = LLMRails(config=_config(), llm=FakeLLMModel(responses=[])) + rails.events_history_cache["cached"] = [{"type": "CachedEvent"}] + rails.explain_info = ExplainInfo() + rails.register_action_param("custom_param", object()) + + state = rails.__getstate__() + + assert state == {"config": rails.config} + + +def test_explain_prefers_active_request_context_over_latest_instance_info(): + rails = LLMRails(config=_config(), llm=FakeLLMModel(responses=[])) + latest_explain_info = ExplainInfo() + request_explain_info = ExplainInfo() + rails.explain_info = latest_explain_info + token = explain_info_var.set(request_explain_info) + try: + assert rails.explain() is request_explain_info + assert rails.explain_info is request_explain_info + finally: + explain_info_var.reset(token) + + +def test_pickle_round_trip_reinitializes_from_config_without_runtime_state(): + rails = LLMRails(config=_config(), llm=FakeLLMModel(responses=[])) + rails.events_history_cache["cached"] = [{"type": "CachedEvent"}] + rails.register_action_param("custom_param", "value") + + restored = pickle.loads(pickle.dumps(rails)) + + assert isinstance(restored, LLMRails) + assert restored.llm is None + assert restored.events_history_cache == {} + assert "custom_param" not in restored.runtime.registered_action_params diff --git a/tests/rails/llm/test_llmrails_hardening.py b/tests/rails/llm/test_llmrails_hardening.py new file mode 100644 index 0000000000..eee6744f32 --- /dev/null +++ b/tests/rails/llm/test_llmrails_hardening.py @@ -0,0 +1,451 @@ +# 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 asyncio +import logging +from copy import deepcopy +from typing import Any, cast +from unittest.mock import patch + +import pytest + +from nemoguardrails import LLMRails, RailsConfig +from nemoguardrails.context import ( + explain_info_var, + generation_options_var, + llm_stats_var, + raw_llm_request, + streaming_handler_var, +) +from nemoguardrails.logging.explain import ExplainInfo +from nemoguardrails.logging.stats import LLMStats +from nemoguardrails.rails.llm.config import Model +from nemoguardrails.rails.llm.options import GenerationOptions, GenerationResponse +from nemoguardrails.streaming import END_OF_STREAM +from tests.utils import FakeLLMModel + +COLANG = """ +define user express greeting + "hi" + +define flow + user express greeting + $user_greeted = True + bot express greeting + +define bot express greeting + "Hello there!" +""" + + +def _config(yaml_content: str | None = None) -> RailsConfig: + return RailsConfig.from_content(colang_content=COLANG, yaml_content=yaml_content) + + +def _rails(config: RailsConfig | None = None) -> LLMRails: + return LLMRails(config=config or _config(), llm=FakeLLMModel(responses=[" express greeting"])) + + +@pytest.mark.asyncio +async def test_generate_standard_info_logs_keep_llmrails_logger_name(caplog): + with caplog.at_level(logging.INFO): + await _rails().generate_async(messages=[{"role": "user", "content": "hi"}]) + + total_processing_logs = [record for record in caplog.records if "--- :: Total processing took" in record.message] + + assert [record.name for record in total_processing_logs] == ["nemoguardrails.rails.llm.llmrails"] + + +@pytest.mark.asyncio +async def test_generate_prompt_and_messages_return_shapes_without_options(): + prompt_result = await _rails().generate_async(prompt="hi") + assert prompt_result == "Hello there!" + + messages_result = await _rails().generate_async(messages=[{"role": "user", "content": "hi"}]) + assert messages_result == {"role": "assistant", "content": "Hello there!"} + + +@pytest.mark.asyncio +async def test_generate_options_dict_and_object_return_generation_response(): + dict_result = await _rails().generate_async( + prompt="hi", + options={"output_vars": ["user_greeted"]}, + ) + + assert isinstance(dict_result, GenerationResponse) + assert dict_result.response == "Hello there!" + assert dict_result.output_data == {"user_greeted": True} + + options = GenerationOptions(output_vars=["user_greeted"]) + object_result = await _rails().generate_async( + messages=[{"role": "user", "content": "hi"}], + options=options, + ) + + assert isinstance(object_result, GenerationResponse) + assert object_result.response == [{"role": "assistant", "content": "Hello there!"}] + assert object_result.output_data == {"user_greeted": True} + + +@pytest.mark.asyncio +async def test_generate_state_forces_generation_response_without_options(): + state = {"events": []} + + result = await _rails().generate_async( + messages=[{"role": "user", "content": "hi"}], + state=state, + ) + + assert isinstance(result, GenerationResponse) + assert result.response == [{"role": "assistant", "content": "Hello there!"}] + assert result.state is not None + assert "events" in result.state + assert state == {"events": []} + + +@pytest.mark.asyncio +async def test_generate_tracing_enabled_returns_generation_response_without_requested_log(): + config = _config( + yaml_content=""" +tracing: + enabled: true + adapters: [] +""" + ) + + result = await _rails(config).generate_async(prompt="hi") + + assert isinstance(result, GenerationResponse) + assert result.response == "Hello there!" + assert result.log is None + + +@pytest.mark.asyncio +async def test_generate_tracing_does_not_mutate_generation_options_object(): + config = _config( + yaml_content=""" +tracing: + enabled: true + adapters: [] +""" + ) + options = GenerationOptions() + + result = await _rails(config).generate_async(prompt="hi", options=options) + + assert isinstance(result, GenerationResponse) + assert result.response == "Hello there!" + assert result.log is None + assert options.log.activated_rails is False + assert options.log.llm_calls is False + assert options.log.internal_events is False + + +@pytest.mark.asyncio +async def test_generate_does_not_mutate_input_messages_or_options_dict(): + messages = [{"role": "user", "content": "hi"}] + options = { + "rails": ["input", "output", "retrieval", "dialog", "tool_input", "tool_output"], + "output_vars": ["user_greeted"], + } + original_messages = deepcopy(messages) + original_options = deepcopy(options) + + result = await _rails().generate_async(messages=messages, options=options) + + assert isinstance(result, GenerationResponse) + assert result.response == [{"role": "assistant", "content": "Hello there!"}] + assert messages == original_messages + assert options == original_options + + +@pytest.mark.asyncio +async def test_concurrent_generate_requests_keep_context_isolated(): + token_options = generation_options_var.set(None) + token_request = raw_llm_request.set(None) + snapshots = {} + both_requests_started = asyncio.Event() + + async def fake_run_colang_turn(rails, events, state, processing_log): + request = cast(list[dict[str, Any]], deepcopy(raw_llm_request.get())) + options = cast(GenerationOptions, generation_options_var.get()) + prompt = request[0]["content"] + snapshots[prompt] = { + "request": request, + "llm_params": deepcopy(options.llm_params), + } + + if len(snapshots) == 2: + both_requests_started.set() + + await both_requests_started.wait() + processing_log.extend( + [ + {"type": "noop", "timestamp": 0.0}, + {"type": "noop", "timestamp": 0.1}, + ] + ) + return [ + { + "type": "StartUtteranceBotAction", + "script": f"reply {prompt}", + } + ] + + try: + rails = _rails() + + with patch( + "nemoguardrails.rails.llm.generation.generation_workflow.run_colang_turn", + fake_run_colang_turn, + ): + first, second = await asyncio.gather( + rails.generate_async( + prompt="first", + options={"llm_params": {"request_id": "first"}}, + ), + rails.generate_async( + prompt="second", + options={"llm_params": {"request_id": "second"}}, + ), + ) + + first = cast(GenerationResponse, first) + second = cast(GenerationResponse, second) + assert first.response == "reply first" + assert second.response == "reply second" + assert snapshots == { + "first": { + "request": [{"role": "user", "content": "first"}], + "llm_params": {"request_id": "first"}, + }, + "second": { + "request": [{"role": "user", "content": "second"}], + "llm_params": {"request_id": "second"}, + }, + } + assert generation_options_var.get() is None + assert raw_llm_request.get() is None + + finally: + generation_options_var.reset(token_options) + raw_llm_request.reset(token_request) + + +@pytest.mark.asyncio +async def test_generate_resets_request_context_after_success_in_same_task(): + outer_options = GenerationOptions(llm_params={"request_id": "outer"}) + outer_request = [{"role": "user", "content": "outer"}] + outer_explain_info = ExplainInfo() + outer_stats = LLMStats() + token_options = generation_options_var.set(outer_options) + token_request = raw_llm_request.set(outer_request) + token_explain_info = explain_info_var.set(outer_explain_info) + token_stats = llm_stats_var.set(outer_stats) + token_streaming_handler = streaming_handler_var.set(None) + + try: + result = await _rails().generate_async( + prompt="hi", + options={"llm_params": {"request_id": "inner"}}, + ) + + assert isinstance(result, GenerationResponse) + assert result.response == "Hello there!" + assert generation_options_var.get() is outer_options + assert raw_llm_request.get() is outer_request + assert explain_info_var.get() is outer_explain_info + request_stats = llm_stats_var.get() + assert isinstance(request_stats, LLMStats) + assert request_stats is not outer_stats + assert streaming_handler_var.get() is None + finally: + streaming_handler_var.reset(token_streaming_handler) + llm_stats_var.reset(token_stats) + explain_info_var.reset(token_explain_info) + raw_llm_request.reset(token_request) + generation_options_var.reset(token_options) + + +@pytest.mark.asyncio +async def test_generate_resets_request_context_after_failure_in_same_task(): + outer_options = GenerationOptions(llm_params={"request_id": "outer"}) + outer_request = [{"role": "user", "content": "outer"}] + outer_explain_info = ExplainInfo() + outer_stats = LLMStats() + token_options = generation_options_var.set(outer_options) + token_request = raw_llm_request.set(outer_request) + token_explain_info = explain_info_var.set(outer_explain_info) + token_stats = llm_stats_var.set(outer_stats) + token_streaming_handler = streaming_handler_var.set(None) + + async def failing_run_colang_turn(rails, events, state, processing_log): + raise RuntimeError("generation failed") + + try: + rails = _rails() + with patch( + "nemoguardrails.rails.llm.generation.generation_workflow.run_colang_turn", + failing_run_colang_turn, + ): + with pytest.raises(RuntimeError, match="generation failed"): + await rails.generate_async( + prompt="hi", + options={"llm_params": {"request_id": "inner"}}, + ) + + assert generation_options_var.get() is outer_options + assert raw_llm_request.get() is outer_request + assert explain_info_var.get() is outer_explain_info + request_stats = llm_stats_var.get() + assert isinstance(request_stats, LLMStats) + assert request_stats is not outer_stats + assert streaming_handler_var.get() is None + finally: + streaming_handler_var.reset(token_streaming_handler) + llm_stats_var.reset(token_stats) + explain_info_var.reset(token_explain_info) + raw_llm_request.reset(token_request) + generation_options_var.reset(token_options) + + +@pytest.mark.asyncio +async def test_generate_closes_request_streaming_handler_once_on_success(): + class StreamingHandler: + def __init__(self): + self.chunks = [] + + async def push_chunk(self, chunk): + self.chunks.append(chunk) + + streaming_handler = StreamingHandler() + + result = await _rails().generate_async( + prompt="hi", + streaming_handler=cast(Any, streaming_handler), + ) + + assert result == "Hello there!" + assert streaming_handler.chunks.count(END_OF_STREAM) == 1 + assert streaming_handler.chunks[-1] is END_OF_STREAM + + +@pytest.mark.asyncio +async def test_generate_closes_request_streaming_handler_once_on_failure(): + class StreamingHandler: + def __init__(self): + self.chunks = [] + + async def push_chunk(self, chunk): + self.chunks.append(chunk) + + async def failing_run_colang_turn(rails, events, state, processing_log): + raise RuntimeError("generation failed") + + streaming_handler = StreamingHandler() + + with patch( + "nemoguardrails.rails.llm.generation.generation_workflow.run_colang_turn", + failing_run_colang_turn, + ): + with pytest.raises(RuntimeError, match="generation failed"): + await _rails().generate_async( + prompt="hi", + streaming_handler=cast(Any, streaming_handler), + ) + + assert streaming_handler.chunks == [END_OF_STREAM] + + +@pytest.mark.asyncio +async def test_generate_closes_request_streaming_handler_before_trace_export(): + order = [] + + class StreamingHandler: + async def push_chunk(self, chunk): + if chunk is END_OF_STREAM: + order.append("stream-closed") + + async def export_trace(**kwargs): + order.append("trace-exported") + + config = _config( + yaml_content=""" +tracing: + enabled: true + adapters: [] +""" + ) + + with patch( + "nemoguardrails.rails.llm.generation.generation_workflow.export_generation_trace", + export_trace, + ): + result = await _rails(config).generate_async( + prompt="hi", + streaming_handler=cast(Any, StreamingHandler()), + ) + + assert isinstance(result, GenerationResponse) + assert result.response == "Hello there!" + assert order == ["stream-closed", "trace-exported"] + + +@pytest.mark.asyncio +async def test_generate_colang_1_writes_implicit_history_cache_when_state_is_none(): + rails = _rails() + + result = await rails.generate_async(messages=[{"role": "user", "content": "hi"}]) + + assert result == {"role": "assistant", "content": "Hello there!"} + assert len(rails.events_history_cache) == 1 + + +@pytest.mark.asyncio +async def test_generate_colang_1_does_not_write_implicit_history_cache_with_explicit_state(): + rails = _rails() + + result = await rails.generate_async( + messages=[{"role": "user", "content": "hi"}], + state={"events": []}, + ) + + assert isinstance(result, GenerationResponse) + assert result.response == [{"role": "assistant", "content": "Hello there!"}] + assert rails.events_history_cache == {} + + +def test_env_api_key_model_kwargs_do_not_write_back_to_config(monkeypatch): + monkeypatch.setenv("NGR_TEST_API_KEY", "env-secret") + config = RailsConfig( + models=[ + Model( + type="main", + engine="fake", + model="fake", + api_key_env_var="NGR_TEST_API_KEY", + parameters={"base_url": "https://example.test"}, + ) + ] + ) + + with patch("nemoguardrails.rails.llm.llmrails.init_llm_model") as init_llm: + init_llm.return_value = FakeLLMModel(responses=[]) + LLMRails(config=config) + + init_kwargs = init_llm.call_args.kwargs["kwargs"] + assert init_kwargs["api_key"] == "env-secret" + assert init_kwargs["base_url"] == "https://example.test" + assert config.models[0].parameters == {"base_url": "https://example.test"} diff --git a/tests/rails/llm/test_llmrails_public_contract.py b/tests/rails/llm/test_llmrails_public_contract.py index 9f5ee48fa3..6a23ac8228 100644 --- a/tests/rails/llm/test_llmrails_public_contract.py +++ b/tests/rails/llm/test_llmrails_public_contract.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import pickle from unittest.mock import patch import pytest @@ -114,3 +115,16 @@ def test_getstate_serializes_config_only(): state = rails.__getstate__() assert state == {"config": rails.config} + + +def test_pickle_round_trip_reinitializes_from_config_without_runtime_state(): + rails = LLMRails(config=_config(), llm=FakeLLMModel(responses=[])) + rails.events_history_cache["cached"] = [{"type": "CachedEvent"}] + rails.register_action_param("custom_param", "value") + + restored = pickle.loads(pickle.dumps(rails)) + + assert isinstance(restored, LLMRails) + assert restored.llm is None + assert restored.events_history_cache == {} + assert "custom_param" not in restored.runtime.registered_action_params diff --git a/tests/rails/llm/test_llmrails_startup_boundary.py b/tests/rails/llm/test_llmrails_startup_boundary.py new file mode 100644 index 0000000000..2e0649abcc --- /dev/null +++ b/tests/rails/llm/test_llmrails_startup_boundary.py @@ -0,0 +1,246 @@ +# 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 types import ModuleType, SimpleNamespace + +from nemoguardrails.embeddings.index import EmbeddingsIndex +from nemoguardrails.rails.llm import llmrails +from nemoguardrails.rails.llm.config import RailsConfig, TracingConfig +from nemoguardrails.rails.llm.startup.embedding_search import DEFAULT_EMBEDDING_ENGINE, DEFAULT_EMBEDDING_MODEL +from nemoguardrails.rails.llm.startup.tracing import create_startup_tracing_adapters + + +class RuntimeWithRegistries: + def __init__(self): + self.llm_task_manager = object() + self.registered_action_params = {} + + def register_action_param(self, name, value): + self.registered_action_params[name] = value + + +def _prepared_config(config: RailsConfig) -> RailsConfig: + setattr(config, "_prepared_for_startup_test", True) + return config + + +def test_constructor_startup_order_from_config_preparation_through_kb(monkeypatch): + sequence = [] + + def prepare_llmrails_config(**kwargs): + sequence.append("prepare config") + return _prepared_config(kwargs["config"]) + + def load_config_py_modules(config): + sequence.append("load config.py modules") + assert getattr(config, "_prepared_for_startup_test") is True + return [ModuleType("config")] + + def runtime_for_colang_version(config, verbose): + sequence.append("runtime creation") + assert getattr(config, "_prepared_for_startup_test") is True + assert verbose is True + return RuntimeWithRegistries() + + def run_config_py_init_hooks(rails, config_modules): + sequence.append("config.py init hooks") + assert len(config_modules) == 1 + assert getattr(rails.config, "_prepared_for_startup_test") is True + assert isinstance(rails.runtime, RuntimeWithRegistries) + + def apply_embedding_model_config( + config, + default_embedding_model, + default_embedding_engine, + default_embedding_params, + ): + sequence.append("embedding config") + assert getattr(config, "_prepared_for_startup_test") is True + assert default_embedding_model == DEFAULT_EMBEDDING_MODEL + assert default_embedding_engine == DEFAULT_EMBEDDING_ENGINE + assert default_embedding_params == {} + return "prepared-model", "prepared-engine", {"prepared": True} + + def create_log_adapters(tracing_config): + sequence.append("tracing adapter creation") + assert tracing_config.enabled is True + return ["adapter"] + + def validate_config(config): + sequence.append("validation") + + def init_llms(rails): + sequence.append("LLM/model-cache initialization") + + def register_llm_generation_actions(rails, verbose): + sequence.append("generation-action registration") + assert verbose is True + rails._llm_generation_actions = object() + + def init_knowledge_base(rails): + sequence.append("KB initialization") + rails._kb = None + + monkeypatch.setattr(llmrails, "prepare_llmrails_config", prepare_llmrails_config) + monkeypatch.setattr(llmrails, "load_config_py_modules", load_config_py_modules) + monkeypatch.setattr(llmrails, "runtime_for_colang_version", runtime_for_colang_version) + monkeypatch.setattr(llmrails, "run_config_py_init_hooks", run_config_py_init_hooks) + monkeypatch.setattr(llmrails, "apply_embedding_model_config", apply_embedding_model_config) + monkeypatch.setattr("nemoguardrails.tracing.create_log_adapters", create_log_adapters) + monkeypatch.setattr(llmrails, "validate_llmrails_config", validate_config) + monkeypatch.setattr(llmrails.LLMRails, "_init_llms", init_llms) + monkeypatch.setattr(llmrails, "register_llm_generation_actions", register_llm_generation_actions) + monkeypatch.setattr(llmrails, "init_knowledge_base", init_knowledge_base) + + config = RailsConfig(models=[], tracing=TracingConfig(enabled=True)) + + rails = llmrails.LLMRails(config, verbose=True) + + assert sequence == [ + "prepare config", + "load config.py modules", + "runtime creation", + "config.py init hooks", + "embedding config", + "tracing adapter creation", + "validation", + "LLM/model-cache initialization", + "generation-action registration", + "KB initialization", + ] + assert rails._log_adapters == ["adapter"] + assert rails.embedding_search.default_model == "prepared-model" + assert rails.embedding_search.default_engine == "prepared-engine" + assert rails.embedding_search.default_params == {"prepared": True} + assert rails.explain_info is None + + +def test_config_py_hook_sees_prepared_state_and_provider_registration_reaches_later_startup(monkeypatch): + provider_seen_by_generation_actions = None + provider_seen_by_kb = None + + class CustomProvider(EmbeddingsIndex): + pass + + config_module = ModuleType("config") + + def init(app): + assert getattr(app.config, "_prepared_for_startup_test") is True + assert isinstance(app.runtime, RuntimeWithRegistries) + assert app.embedding_search.default_model == DEFAULT_EMBEDDING_MODEL + assert app.embedding_search.default_engine == DEFAULT_EMBEDDING_ENGINE + assert app.embedding_search.default_params == {} + assert app.embedding_search.providers == {} + app.register_embedding_search_provider("custom", CustomProvider) + + setattr(config_module, "init", init) + + def register_llm_generation_actions(rails, verbose): + nonlocal provider_seen_by_generation_actions + del verbose + provider_seen_by_generation_actions = rails.embedding_search.providers["custom"] + rails._llm_generation_actions = SimpleNamespace() + + def init_knowledge_base(rails): + nonlocal provider_seen_by_kb + provider_seen_by_kb = rails.embedding_search.providers["custom"] + rails._kb = None + + monkeypatch.setattr(llmrails, "prepare_llmrails_config", lambda **kwargs: _prepared_config(kwargs["config"])) + monkeypatch.setattr(llmrails, "load_config_py_modules", lambda config: [config_module]) + monkeypatch.setattr(llmrails, "runtime_for_colang_version", lambda config, verbose: RuntimeWithRegistries()) + monkeypatch.setattr(llmrails, "validate_llmrails_config", lambda config: None) + monkeypatch.setattr(llmrails.LLMRails, "_init_llms", lambda rails: None) + monkeypatch.setattr(llmrails, "register_llm_generation_actions", register_llm_generation_actions) + monkeypatch.setattr(llmrails, "init_knowledge_base", init_knowledge_base) + + config = RailsConfig(models=[]) + object.__setattr__(config, "tracing", None) + + rails = llmrails.LLMRails(config) + + assert rails.embedding_search.providers["custom"] is CustomProvider + assert provider_seen_by_generation_actions is CustomProvider + assert provider_seen_by_kb is CustomProvider + assert rails._log_adapters is None + + +def test_tracing_adapters_are_created_after_config_py_init_hooks(monkeypatch): + events = [] + config_module = ModuleType("config") + + def init(app): + app.tracing_marker = "hook-ran" + events.append("hook") + + setattr(config_module, "init", init) + + def create_log_adapters(tracing_config): + assert tracing_config.enabled is True + assert events == ["hook"] + events.append("tracing") + return [{"tracing": tracing_config}] + + monkeypatch.setattr(llmrails, "prepare_llmrails_config", lambda **kwargs: _prepared_config(kwargs["config"])) + monkeypatch.setattr(llmrails, "load_config_py_modules", lambda config: [config_module]) + monkeypatch.setattr(llmrails, "runtime_for_colang_version", lambda config, verbose: RuntimeWithRegistries()) + monkeypatch.setattr("nemoguardrails.tracing.create_log_adapters", create_log_adapters) + monkeypatch.setattr(llmrails, "validate_llmrails_config", lambda config: None) + monkeypatch.setattr(llmrails.LLMRails, "_init_llms", lambda rails: None) + monkeypatch.setattr( + llmrails, + "register_llm_generation_actions", + lambda rails, verbose: setattr(rails, "_llm_generation_actions", object()), + ) + monkeypatch.setattr(llmrails, "init_knowledge_base", lambda rails: setattr(rails, "_kb", None)) + + config = RailsConfig(models=[], tracing=TracingConfig(enabled=True)) + + rails = llmrails.LLMRails(config) + + assert events == ["hook", "tracing"] + assert getattr(rails, "tracing_marker") == "hook-ran" + assert rails._log_adapters == [{"tracing": config.tracing}] + + +def test_startup_tracing_helper_returns_none_without_tracing_config(monkeypatch): + calls = [] + + def create_log_adapters(tracing_config): + calls.append(tracing_config) + return ["adapter"] + + monkeypatch.setattr("nemoguardrails.tracing.create_log_adapters", create_log_adapters) + + config = RailsConfig(models=[]) + object.__setattr__(config, "tracing", None) + + assert create_startup_tracing_adapters(config) is None + assert calls == [] + + +def test_startup_tracing_helper_creates_adapters_when_tracing_config_is_present(monkeypatch): + calls = [] + + def create_log_adapters(tracing_config): + calls.append(tracing_config) + return ["adapter"] + + monkeypatch.setattr("nemoguardrails.tracing.create_log_adapters", create_log_adapters) + + config = RailsConfig(models=[], tracing=TracingConfig(enabled=True)) + + assert create_startup_tracing_adapters(config) == ["adapter"] + assert calls == [config.tracing] diff --git a/tests/rails/llm/test_rails_check.py b/tests/rails/llm/test_rails_check.py new file mode 100644 index 0000000000..eace3aa7ba --- /dev/null +++ b/tests/rails/llm/test_rails_check.py @@ -0,0 +1,141 @@ +# 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 + +import pytest + +from nemoguardrails.rails.llm.checks.rails_check import check_messages +from nemoguardrails.rails.llm.options import ( + ActivatedRail, + GenerationLog, + GenerationResponse, + RailStatus, + RailType, +) + + +class FakeRails: + def __init__(self, response): + self.config = SimpleNamespace(colang_version="1.0") + self.runtime = SimpleNamespace() + self.response = response + self.generate_calls = [] + + async def generate_async(self, *, messages, options): + self.generate_calls.append({"messages": messages, "options": options}) + return self.response + + +class MutatingRails(FakeRails): + async def generate_async(self, *, messages, options): + messages[0]["content"]["nested"] = "changed" + return await super().generate_async(messages=messages, options=options) + + +@pytest.mark.asyncio +async def test_check_messages_detects_input_rails_and_marks_passed(): + rails = FakeRails(GenerationResponse(response=[{"role": "user", "content": "hello"}])) + + result = await check_messages(rails, [{"role": "user", "content": "hello"}]) + + assert result.status == RailStatus.PASSED + assert result.content == "hello" + assert rails.generate_calls == [ + { + "messages": [{"role": "user", "content": "hello"}], + "options": { + "rails": ["input"], + "log": {"activated_rails": True}, + }, + } + ] + + +@pytest.mark.asyncio +async def test_check_messages_normalizes_output_only_messages(): + rails = FakeRails(GenerationResponse(response=[{"role": "assistant", "content": "hello"}])) + + result = await check_messages(rails, [{"role": "assistant", "content": "hello"}]) + + assert result.status == RailStatus.PASSED + assert rails.generate_calls[0]["messages"] == [ + {"role": "user", "content": ""}, + {"role": "assistant", "content": "hello"}, + ] + assert rails.generate_calls[0]["options"]["rails"] == ["output"] + + +@pytest.mark.asyncio +async def test_check_messages_honors_explicit_rail_types_and_reports_modified_content(): + rails = FakeRails(GenerationResponse(response=[{"role": "assistant", "content": "updated"}])) + messages = [ + {"role": "user", "content": "hello"}, + {"role": "assistant", "content": "original"}, + ] + + result = await check_messages(rails, messages, rail_types=[RailType.OUTPUT]) + + assert result.status == RailStatus.MODIFIED + assert result.content == "updated" + assert rails.generate_calls[0]["messages"] == messages + assert rails.generate_calls[0]["options"]["rails"] == ["output"] + + +@pytest.mark.asyncio +async def test_check_messages_does_not_expose_caller_messages_to_generation_mutation(): + rails = MutatingRails(GenerationResponse(response=[{"role": "user", "content": "hello"}])) + messages = [ + {"role": "context", "content": {"nested": "original"}}, + {"role": "user", "content": "hello"}, + ] + + result = await check_messages(rails, messages) + + assert result.status == RailStatus.PASSED + assert messages == [ + {"role": "context", "content": {"nested": "original"}}, + {"role": "user", "content": "hello"}, + ] + assert rails.generate_calls[0]["messages"][0]["content"] == {"nested": "changed"} + + +@pytest.mark.asyncio +async def test_check_messages_reports_first_blocking_rail(): + rails = FakeRails( + GenerationResponse( + response="blocked", + log=GenerationLog( + activated_rails=[ + ActivatedRail(type="input", name="first rail", stop=False), + ActivatedRail(type="output", name="blocking rail", stop=True), + ] + ), + ) + ) + + result = await check_messages(rails, [{"role": "user", "content": "hello"}]) + + assert result.status == RailStatus.BLOCKED + assert result.content == "blocked" + assert result.rail == "blocking rail" + + +@pytest.mark.asyncio +async def test_check_messages_rejects_unexpected_generation_response_type(): + rails = FakeRails("not a generation response") + + with pytest.raises(RuntimeError, match="Expected GenerationResponse, got str"): + await check_messages(rails, [{"role": "user", "content": "hello"}]) diff --git a/tests/test_llmrails_check_async.py b/tests/test_llmrails_check_async.py index bf0b527ef4..30e720d8c1 100644 --- a/tests/test_llmrails_check_async.py +++ b/tests/test_llmrails_check_async.py @@ -18,7 +18,7 @@ import pytest from nemoguardrails import LLMRails, RailsConfig -from nemoguardrails.rails.llm.llmrails import ( +from nemoguardrails.rails.llm.checks.rails_check import ( _determine_rails_from_messages, _get_blocking_rail, _get_last_content_by_role, diff --git a/tests/test_system_message_conversion.py b/tests/test_system_message_conversion.py index 928d285499..a21d70e5fc 100644 --- a/tests/test_system_message_conversion.py +++ b/tests/test_system_message_conversion.py @@ -16,6 +16,7 @@ import pytest from nemoguardrails import LLMRails, RailsConfig +from nemoguardrails.rails.llm.conversation.conversation_events import events_for_messages from tests.utils import FakeLLMModel @@ -44,7 +45,7 @@ async def test_system_message_conversion_v1(): {"role": "user", "content": "Hello!"}, ] - events = llm_rails._get_events_for_messages(messages, None) + events = events_for_messages(llm_rails, messages, None) system_messages = [event for event in events if event["type"] == "SystemMessage"] assert len(system_messages) == 1 @@ -76,7 +77,7 @@ async def test_system_message_conversion_v2x(): {"role": "user", "content": "Hello!"}, ] - events = llm_rails._get_events_for_messages(messages, None) + events = events_for_messages(llm_rails, messages, None) system_messages = [event for event in events if event["type"] == "SystemMessage"] assert len(system_messages) == 1 @@ -108,7 +109,7 @@ async def test_system_message_conversion_multiple(): {"role": "user", "content": "Hello!"}, ] - events = llm_rails._get_events_for_messages(messages, None) + events = events_for_messages(llm_rails, messages, None) system_messages = [event for event in events if event["type"] == "SystemMessage"] assert len(system_messages) == 2 diff --git a/tests/test_token_usage_integration.py b/tests/test_token_usage_integration.py index e184ac7211..b8e64bdfa6 100644 --- a/tests/test_token_usage_integration.py +++ b/tests/test_token_usage_integration.py @@ -25,7 +25,6 @@ import pytest from nemoguardrails import RailsConfig -from nemoguardrails.context import llm_stats_var from nemoguardrails.rails.llm.options import GenerationOptions, GenerationResponse from tests.utils import TestChat @@ -284,12 +283,15 @@ async def test_token_usage_not_set_for_unsupported_provider(): token_usage=token_usage_data, ) - result = await chat.app.generate_async(messages=[{"role": "user", "content": "Hi!"}]) - - assert result["content"] == "Hello there!" + # Read the per-call stats from the response log rather than the request-scoped + # llm_stats_var contextvar, which generation_context now resets on request close. + result = await chat.app.generate_async( + messages=[{"role": "user", "content": "Hi!"}], + options={"log": {"llm_calls": True}}, + ) - llm_stats = llm_stats_var.get() + assert result.response[-1]["content"] == "Hello there!" - assert llm_stats is not None - assert llm_stats.get_stat("total_tokens") == 0 - assert llm_stats.get_stat("total_calls") == 1 + assert result.log is not None + assert result.log.stats.llm_calls_total_tokens == 0 + assert result.log.stats.llm_calls_count == 1