From d858c02ff8862c49fe14ce9c34b7178dcf04ac09 Mon Sep 17 00:00:00 2001 From: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> Date: Mon, 8 Jun 2026 21:42:27 +0200 Subject: [PATCH] refactor(llmrails): extract runtime conversation flow --- .../rails/llm/conversation/__init__.py | 16 ++ .../llm/conversation/conversation_events.py | 177 ++++++++++++++ nemoguardrails/rails/llm/llmrails.py | 208 ++-------------- .../rails/llm/runtime/colang_turns.py | 100 ++++++++ tests/rails/llm/test_colang_turns.py | 181 ++++++++++++++ tests/rails/llm/test_conversation_events.py | 225 ++++++++++++++++++ 6 files changed, 715 insertions(+), 192 deletions(-) create mode 100644 nemoguardrails/rails/llm/conversation/__init__.py create mode 100644 nemoguardrails/rails/llm/conversation/conversation_events.py create mode 100644 nemoguardrails/rails/llm/runtime/colang_turns.py create mode 100644 tests/rails/llm/test_colang_turns.py create mode 100644 tests/rails/llm/test_conversation_events.py diff --git a/nemoguardrails/rails/llm/conversation/__init__.py b/nemoguardrails/rails/llm/conversation/__init__.py new file mode 100644 index 0000000000..f2e0a46a5c --- /dev/null +++ b/nemoguardrails/rails/llm/conversation/__init__.py @@ -0,0 +1,16 @@ +# 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. + +__all__ = [] diff --git a/nemoguardrails/rails/llm/conversation/conversation_events.py b/nemoguardrails/rails/llm/conversation/conversation_events.py new file mode 100644 index 0000000000..2b26d318e4 --- /dev/null +++ b/nemoguardrails/rails/llm/conversation/conversation_events.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. + +"""Conversation message conversion to Colang events.""" + +from collections.abc import MutableMapping +from typing import Any, List, Protocol + +from nemoguardrails.colang.v2_x.runtime.flows import Action +from nemoguardrails.rails.llm.utils import get_history_cache_key +from nemoguardrails.utils import new_event_dict, new_uuid + +__all__ = [ + "ConversationEventRails", + "events_for_messages", + "events_history_cache_prefix", +] + + +class ConversationEventRails(Protocol): + @property + def config(self) -> Any: ... + + @property + def events_history_cache(self) -> MutableMapping[str, list[dict]]: ... + + +def events_history_cache_prefix( + events_history_cache: MutableMapping[str, list[dict]], + messages: List[dict], +) -> tuple[int, list[dict]]: + """Return the longest cached event prefix for the provided messages. + + The cache stores generated Colang events for prior complete message + histories. Reusing the longest cached prefix avoids replaying old messages + through intent detection on every turn. + """ + prefix_len = len(messages) - 1 + while prefix_len > 0: + cache_key = get_history_cache_key(messages[0:prefix_len]) + if cache_key in events_history_cache: + return prefix_len, events_history_cache[cache_key].copy() + + prefix_len -= 1 + + return 0, [] + + +def events_for_messages(rails: ConversationEventRails, messages: List[dict], state: Any) -> List[dict]: + """Return Colang events corresponding to OpenAI-style messages.""" + events = [] + + if rails.config.colang_version == "1.0": + prefix_len, events = events_history_cache_prefix(rails.events_history_cache, messages) + + for idx in range(prefix_len, len(messages)): + msg = messages[idx] + if msg["role"] == "user": + events.append( + { + "type": "UtteranceUserActionFinished", + "final_transcript": msg["content"], + } + ) + + if idx != len(messages) - 1: + events.append( + { + "type": "UserMessage", + "text": msg["content"], + } + ) + + elif msg["role"] == "assistant": + if msg.get("tool_calls"): + events.append({"type": "BotToolCalls", "tool_calls": msg["tool_calls"]}) + else: + action_uid = new_uuid() + start_event = new_event_dict( + "StartUtteranceBotAction", + script=msg["content"], + action_uid=action_uid, + ) + finished_event = new_event_dict( + "UtteranceBotActionFinished", + final_script=msg["content"], + is_success=True, + action_uid=action_uid, + ) + events.extend([start_event, finished_event]) + elif msg["role"] == "context": + events.append({"type": "ContextUpdate", "data": msg["content"]}) + elif msg["role"] == "event": + events.append(msg["event"]) + elif msg["role"] == "system": + events.append({"type": "SystemMessage", "content": msg["content"]}) + elif msg["role"] == "tool": + if idx == len(messages) - 1: + user_message = None + for prev_msg in reversed(messages[:idx]): + if prev_msg["role"] == "user": + user_message = prev_msg["content"] + break + + if user_message: + if rails.config.rails.tool_input.flows: + tool_messages = [] + for tool_idx in range(len(messages)): + if messages[tool_idx]["role"] == "tool": + tool_messages.append( + { + "content": messages[tool_idx]["content"], + "name": messages[tool_idx].get("name", "unknown"), + "tool_call_id": messages[tool_idx].get("tool_call_id", ""), + } + ) + + events.append( + { + "type": "UserToolMessages", + "tool_messages": tool_messages, + } + ) + + else: + events.append({"type": "UserMessage", "text": user_message}) + + else: + for idx in range(len(messages)): + msg = messages[idx] + if msg["role"] == "user": + events.append( + { + "type": "UtteranceUserActionFinished", + "final_transcript": msg["content"], + } + ) + + elif msg["role"] == "assistant": + raise ValueError( + "Providing `assistant` messages as input is not supported for Colang 2.0 configurations." + ) + elif msg["role"] == "context": + events.append({"type": "ContextUpdate", "data": msg["content"]}) + elif msg["role"] == "event": + events.append(msg["event"]) + elif msg["role"] == "system": + events.append({"type": "SystemMessage", "content": msg["content"]}) + elif msg["role"] == "tool": + action_uid = msg["tool_call_id"] + return_value = msg["content"] + action: Action = state.actions[action_uid] + events.append( + new_event_dict( + f"{action.name}Finished", + action_uid=action_uid, + action_name=action.name, + status="success", + is_success=True, + return_value=return_value, + events=[], + ) + ) + + return events diff --git a/nemoguardrails/rails/llm/llmrails.py b/nemoguardrails/rails/llm/llmrails.py index 66001fbe65..6c2406ca00 100644 --- a/nemoguardrails/rails/llm/llmrails.py +++ b/nemoguardrails/rails/llm/llmrails.py @@ -49,7 +49,7 @@ from nemoguardrails.base_guardrails import BaseGuardrails 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.flows import State from nemoguardrails.colang.v2_x.runtime.runtime import RuntimeV2_x from nemoguardrails.context import ( explain_info_var, @@ -76,6 +76,7 @@ OutputRailsStreamingConfig, RailsConfig, ) +from nemoguardrails.rails.llm.conversation.conversation_events import events_for_messages from nemoguardrails.rails.llm.options import ( GenerationLog, GenerationOptions, @@ -85,6 +86,11 @@ RailType, ) from nemoguardrails.rails.llm.runtime.colang_runtime import runtime_for_colang_version +from nemoguardrails.rails.llm.runtime.colang_turns import ( + generate_colang_events, + process_colang_events, + process_events_semaphore, +) 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 @@ -107,14 +113,10 @@ from nemoguardrails.utils import ( extract_error_json, get_or_create_event_loop, - new_event_dict, - new_uuid, ) log = logging.getLogger(__name__) -process_events_semaphore = asyncio.Semaphore(1) - def _wrap_legacy_llm(llm): try: @@ -392,156 +394,7 @@ 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. - - Tries to find a prefix of messages for which we have already a list of events - in the cache. For the rest, they are converted as is. - - The reason this cache exists is that we want to benefit from events generated in - previous turns, which can't be computed again because it would be expensive (e.g., - involving multiple LLM calls). - - When an explicit state object will be added, this mechanism can be removed. - - Args: - messages: The list of messages. - - Returns: - A list of events. - """ - events = [] - - if self.config.colang_version == "1.0": - # We try to find the longest prefix of messages for which we have a cache - # of events. - p = len(messages) - 1 - while p > 0: - cache_key = get_history_cache_key(messages[0:p]) - if cache_key in self.events_history_cache: - events = self.events_history_cache[cache_key].copy() - break - - p -= 1 - - # For the rest of the messages, we transform them directly into events. - # TODO: Move this to separate function once more types of messages are supported. - for idx in range(p, len(messages)): - msg = messages[idx] - if msg["role"] == "user": - events.append( - { - "type": "UtteranceUserActionFinished", - "final_transcript": msg["content"], - } - ) - - # If it's not the last message, we also need to add the `UserMessage` event - if idx != len(messages) - 1: - events.append( - { - "type": "UserMessage", - "text": msg["content"], - } - ) - - elif msg["role"] == "assistant": - if msg.get("tool_calls"): - events.append({"type": "BotToolCalls", "tool_calls": msg["tool_calls"]}) - else: - action_uid = new_uuid() - start_event = new_event_dict( - "StartUtteranceBotAction", - script=msg["content"], - action_uid=action_uid, - ) - finished_event = new_event_dict( - "UtteranceBotActionFinished", - final_script=msg["content"], - is_success=True, - action_uid=action_uid, - ) - events.extend([start_event, finished_event]) - elif msg["role"] == "context": - events.append({"type": "ContextUpdate", "data": msg["content"]}) - elif msg["role"] == "event": - events.append(msg["event"]) - elif msg["role"] == "system": - # Handle system messages - convert them to SystemMessage events - events.append({"type": "SystemMessage", "content": msg["content"]}) - elif msg["role"] == "tool": - # For the last tool message, create grouped tool event and synthetic UserMessage - if idx == len(messages) - 1: - # Find the original user message for response generation - user_message = None - for prev_msg in reversed(messages[:idx]): - if prev_msg["role"] == "user": - user_message = prev_msg["content"] - break - - if user_message: - # If tool input rails are configured, group all tool messages - if self.config.rails.tool_input.flows: - # Collect all tool messages for grouped processing - tool_messages = [] - for tool_idx in range(len(messages)): - if messages[tool_idx]["role"] == "tool": - tool_messages.append( - { - "content": messages[tool_idx]["content"], - "name": messages[tool_idx].get("name", "unknown"), - "tool_call_id": messages[tool_idx].get("tool_call_id", ""), - } - ) - - events.append( - { - "type": "UserToolMessages", - "tool_messages": tool_messages, - } - ) - - else: - events.append({"type": "UserMessage", "text": user_message}) - - else: - for idx in range(len(messages)): - msg = messages[idx] - if msg["role"] == "user": - events.append( - { - "type": "UtteranceUserActionFinished", - "final_transcript": msg["content"], - } - ) - - elif msg["role"] == "assistant": - raise ValueError( - "Providing `assistant` messages as input is not supported for Colang 2.0 configurations." - ) - elif msg["role"] == "context": - events.append({"type": "ContextUpdate", "data": msg["content"]}) - elif msg["role"] == "event": - events.append(msg["event"]) - elif msg["role"] == "system": - # Handle system messages - convert them to SystemMessage events - events.append({"type": "SystemMessage", "content": msg["content"]}) - elif msg["role"] == "tool": - action_uid = msg["tool_call_id"] - return_value = msg["content"] - action: Action = state.actions[action_uid] - events.append( - new_event_dict( - f"{action.name}Finished", - action_uid=action_uid, - action_name=action.name, - status="success", - is_success=True, - return_value=return_value, - events=[], - ) - ) - - return events + return events_for_messages(self, messages, state) @staticmethod def _ensure_explain_info() -> ExplainInfo: @@ -1185,26 +1038,7 @@ async def generate_events_async( The newly generate event(s). """ - t0 = time.time() - - # Initialize the LLM stats - llm_stats = LLMStats() - llm_stats_var.set(llm_stats) - - # Compute the new events. - processing_log = [] - new_events = await self.runtime.generate_events(events, processing_log=processing_log) - - # If logging is enabled, we log the conversation - # TODO: add support for logging flag - if self.verbose: - history = get_colang_history(events) - log.info(f"Conversation history so far: \n{history}") - - log.info("--- :: Total processing took %.2f seconds." % (time.time() - t0)) - log.info("--- :: Stats: %s" % llm_stats) - - return new_events + return await generate_colang_events(self, events) def generate_events( self, @@ -1240,23 +1074,13 @@ async def process_events_async( (output_events, output_state) Returns a sequence of output events and an output state. """ - t0 = time.time() - llm_stats = LLMStats() - llm_stats_var.set(llm_stats) - - # Compute the new events. - # We need to protect 'process_events' to be called only once at a time - # TODO (cschueller): Why is this? - async with process_events_semaphore: - output_events, output_state = await self.runtime.process_events(events, state, blocking) - - took = time.time() - t0 - # Small tweak, disable this when there were no events (or it was just too fast). - if took > 0.1: - log.info("--- :: Total processing took %.2f seconds." % took) - log.info("--- :: Stats: %s" % llm_stats) - - return output_events, output_state + return await process_colang_events( + self, + events, + state, + blocking, + semaphore=process_events_semaphore, + ) def process_events( self, diff --git a/nemoguardrails/rails/llm/runtime/colang_turns.py b/nemoguardrails/rails/llm/runtime/colang_turns.py new file mode 100644 index 0000000000..86bf26ff8a --- /dev/null +++ b/nemoguardrails/rails/llm/runtime/colang_turns.py @@ -0,0 +1,100 @@ +# 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 event helpers.""" + +import asyncio +import logging +import time +from typing import Any, List, Optional, Protocol, Tuple, Union + +from nemoguardrails.actions.llm.utils import get_colang_history +from nemoguardrails.colang.v2_x.runtime.flows import State +from nemoguardrails.context import llm_stats_var +from nemoguardrails.logging.stats import LLMStats + +log = logging.getLogger(__name__) + +process_events_semaphore = asyncio.Semaphore(1) + +__all__ = [ + "ColangTurnRails", + "generate_colang_events", + "process_colang_events", + "process_events_semaphore", +] + + +class ColangTurnRails(Protocol): + @property + def config(self) -> Any: ... + + @property + def runtime(self) -> Any: ... + + @property + def verbose(self) -> bool: ... + + +async def generate_colang_events(rails: ColangTurnRails, events: List[dict]) -> List[dict]: + """Generate the next Colang 1.0 events for an event history.""" + t0 = time.time() + + # Initialize the LLM stats + llm_stats = LLMStats() + llm_stats_var.set(llm_stats) + + # Compute the new events. + processing_log = [] + new_events = await rails.runtime.generate_events(events, processing_log=processing_log) + + # If logging is enabled, we log the conversation + # TODO: add support for logging flag + if rails.verbose: + history = get_colang_history(events) + log.info(f"Conversation history so far: \n{history}") + + log.info("--- :: Total processing took %.2f seconds." % (time.time() - t0)) + log.info("--- :: Stats: %s" % llm_stats) + + return new_events + + +async def process_colang_events( + rails: ColangTurnRails, + events: List[dict], + state: Union[Optional[dict], State] = None, + blocking: bool = False, + *, + semaphore: asyncio.Semaphore, +) -> Tuple[List[dict], Union[dict, State]]: + """Process Colang events in order and return output events plus state.""" + t0 = time.time() + llm_stats = LLMStats() + llm_stats_var.set(llm_stats) + + # Compute the new events. + # We need to protect 'process_events' to be called only once at a time + # TODO (cschueller): Why is this? + async with semaphore: + output_events, output_state = await rails.runtime.process_events(events, state, blocking) + + took = time.time() - t0 + # Small tweak, disable this when there were no events (or it was just too fast). + if took > 0.1: + log.info("--- :: Total processing took %.2f seconds." % took) + log.info("--- :: Stats: %s" % llm_stats) + + return output_events, output_state diff --git a/tests/rails/llm/test_colang_turns.py b/tests/rails/llm/test_colang_turns.py new file mode 100644 index 0000000000..b25cd816fb --- /dev/null +++ b/tests/rails/llm/test_colang_turns.py @@ -0,0 +1,181 @@ +# 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 asyncio +from types import SimpleNamespace +from typing import Any, cast + +import pytest + +from nemoguardrails.context import llm_stats_var +from nemoguardrails.logging.stats import LLMStats +from nemoguardrails.rails.llm import llmrails as llmrails_module +from nemoguardrails.rails.llm.llmrails import LLMRails +from nemoguardrails.rails.llm.runtime.colang_turns import ( + generate_colang_events, + process_colang_events, +) + + +class FakeRails: + def __init__(self, colang_version: str, runtime, instant_actions=None): + self.config = SimpleNamespace( + colang_version=colang_version, + rails=SimpleNamespace( + actions=SimpleNamespace(instant_actions=instant_actions), + ), + ) + self.runtime = runtime + self.verbose = False + + +def make_rails(colang_version: str, runtime, instant_actions=None): + return FakeRails(colang_version=colang_version, runtime=runtime, instant_actions=instant_actions) + + +@pytest.fixture(autouse=True) +def reset_llm_stats_context(): + token = llm_stats_var.set(None) + try: + yield + finally: + llm_stats_var.reset(token) + + +@pytest.mark.asyncio +async def test_generate_colang_events_sets_stats_and_passes_processing_log(): + class Runtime: + def __init__(self): + self.processing_log = None + + async def generate_events(self, events, processing_log): + self.processing_log = processing_log + processing_log.append({"event": "generated"}) + return [{"type": "BotIntent", "intent": "express greeting"}] + + runtime = Runtime() + new_events = await generate_colang_events( + make_rails("1.0", runtime), + events=[{"type": "UserMessage", "text": "Hi"}], + ) + + assert new_events == [{"type": "BotIntent", "intent": "express greeting"}] + assert runtime.processing_log == [{"event": "generated"}] + assert isinstance(llm_stats_var.get(), LLMStats) + + +@pytest.mark.asyncio +async def test_process_colang_events_sets_stats_and_forwards_state_and_blocking(): + class Runtime: + def __init__(self): + self.calls = [] + + async def process_events(self, events, state, blocking): + self.calls.append((events, state, blocking)) + return [{"type": "OutputEvent"}], {"state": "updated"} + + runtime = Runtime() + output_events, output_state = await process_colang_events( + make_rails("2.x", runtime), + events=[{"type": "InputEvent"}], + state={"state": "input"}, + blocking=True, + semaphore=asyncio.Semaphore(1), + ) + + assert runtime.calls == [([{"type": "InputEvent"}], {"state": "input"}, True)] + assert output_events == [{"type": "OutputEvent"}] + assert output_state == {"state": "updated"} + assert isinstance(llm_stats_var.get(), LLMStats) + + +@pytest.mark.asyncio +async def test_process_colang_events_uses_explicit_semaphore_only(): + class Runtime: + def __init__(self): + self.active_calls = 0 + self.max_active_calls = 0 + + async def process_events(self, events, state, blocking): + self.active_calls += 1 + self.max_active_calls = max(self.max_active_calls, self.active_calls) + await asyncio.sleep(0) + self.active_calls -= 1 + return events, state + + runtime = Runtime() + rails = make_rails("2.x", runtime) + + await asyncio.gather( + process_colang_events(rails, [{"type": "First"}], semaphore=asyncio.Semaphore(1)), + process_colang_events(rails, [{"type": "Second"}], semaphore=asyncio.Semaphore(1)), + ) + + assert runtime.max_active_calls == 2 + + +@pytest.mark.asyncio +async def test_process_colang_events_serializes_runtime_calls_with_shared_semaphore(): + class Runtime: + def __init__(self): + self.active_calls = 0 + self.max_active_calls = 0 + + async def process_events(self, events, state, blocking): + self.active_calls += 1 + self.max_active_calls = max(self.max_active_calls, self.active_calls) + await asyncio.sleep(0) + self.active_calls -= 1 + return events, state + + runtime = Runtime() + rails = make_rails("2.x", runtime) + semaphore = asyncio.Semaphore(1) + + await asyncio.gather( + process_colang_events(rails, [{"type": "First"}], semaphore=semaphore), + process_colang_events(rails, [{"type": "Second"}], semaphore=semaphore), + ) + + assert runtime.max_active_calls == 1 + + +@pytest.mark.asyncio +async def test_llmrails_process_events_uses_module_semaphore_patch_point(): + class Runtime: + def __init__(self): + self.calls = 0 + + async def process_events(self, events, state, blocking): + self.calls += 1 + return events, state + + runtime = Runtime() + rails = cast(Any, object.__new__(LLMRails)) + rails.runtime = runtime + + original_semaphore = llmrails_module.process_events_semaphore + llmrails_module.process_events_semaphore = asyncio.Semaphore(0) + try: + task = asyncio.create_task(rails.process_events_async([{"type": "InputEvent"}])) + await asyncio.sleep(0) + assert runtime.calls == 0 + + llmrails_module.process_events_semaphore.release() + await task + finally: + llmrails_module.process_events_semaphore = original_semaphore + + assert runtime.calls == 1 diff --git a/tests/rails/llm/test_conversation_events.py b/tests/rails/llm/test_conversation_events.py new file mode 100644 index 0000000000..7579cb25d0 --- /dev/null +++ b/tests/rails/llm/test_conversation_events.py @@ -0,0 +1,225 @@ +# 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.config import RailsConfig +from nemoguardrails.rails.llm.conversation import conversation_events +from nemoguardrails.rails.llm.conversation.conversation_events import ( + events_for_messages, + events_history_cache_prefix, +) +from nemoguardrails.rails.llm.llmrails import LLMRails +from nemoguardrails.rails.llm.utils import get_history_cache_key + + +class FakeRails: + def __init__(self, config: RailsConfig): + self.config = config + self.events_history_cache = {} + + +def make_rails(config: RailsConfig): + return FakeRails(config=config) + + +def test_conversation_events_public_exports(): + assert conversation_events.__all__ == [ + "ConversationEventRails", + "events_for_messages", + "events_history_cache_prefix", + ] + + +def test_events_for_messages_v1_converts_supported_message_roles(monkeypatch): + monkeypatch.setattr(conversation_events, "new_uuid", lambda: "action-1") + tool_calls = [{"id": "call-1", "function": {"name": "lookup", "arguments": "{}"}}] + passthrough_event = {"type": "UserSilent"} + rails = make_rails(RailsConfig(models=[], colang_version="1.0")) + + events = events_for_messages( + rails, + [ + {"role": "context", "content": {"user_name": "Ada"}}, + {"role": "system", "content": "Be concise."}, + {"role": "user", "content": "Hi"}, + {"role": "assistant", "content": "Hello"}, + {"role": "assistant", "content": "", "tool_calls": tool_calls}, + {"role": "event", "event": passthrough_event}, + ], + state=None, + ) + + assert events[0] == {"type": "ContextUpdate", "data": {"user_name": "Ada"}} + assert events[1] == {"type": "SystemMessage", "content": "Be concise."} + assert events[2] == {"type": "UtteranceUserActionFinished", "final_transcript": "Hi"} + assert events[3] == {"type": "UserMessage", "text": "Hi"} + assert events[4]["type"] == "StartUtteranceBotAction" + assert events[4]["script"] == "Hello" + assert events[4]["action_uid"] == "action-1" + assert events[5]["type"] == "UtteranceBotActionFinished" + assert events[5]["final_script"] == "Hello" + assert events[5]["action_uid"] == "action-1" + assert events[6] == {"type": "BotToolCalls", "tool_calls": tool_calls} + assert events[7] is passthrough_event + + +def test_events_for_messages_v1_reuses_longest_history_cache_prefix(): + config = RailsConfig(models=[], colang_version="1.0") + rails = make_rails(config) + cached_prefix = [{"type": "CachedEvent"}] + cached_messages = [ + {"role": "user", "content": "Hi"}, + {"role": "assistant", "content": "Hello"}, + ] + rails.events_history_cache[get_history_cache_key(cached_messages)] = cached_prefix + + events = events_for_messages( + rails, + cached_messages + [{"role": "user", "content": "Next"}], + state=None, + ) + + assert events == [ + {"type": "CachedEvent"}, + {"type": "UtteranceUserActionFinished", "final_transcript": "Next"}, + ] + assert events is not cached_prefix + assert cached_prefix == [{"type": "CachedEvent"}] + + +def test_events_history_cache_prefix_returns_copy_of_longest_match(): + events_history_cache = {} + short_messages = [{"role": "user", "content": "Hi"}] + long_messages = short_messages + [{"role": "assistant", "content": "Hello"}] + cached_events = [{"type": "LongCachedEvent"}] + events_history_cache[get_history_cache_key(short_messages)] = [{"type": "ShortCachedEvent"}] + events_history_cache[get_history_cache_key(long_messages)] = cached_events + + prefix_len, events = events_history_cache_prefix( + events_history_cache, + long_messages + [{"role": "user", "content": "Next"}], + ) + + assert prefix_len == len(long_messages) + assert events == [{"type": "LongCachedEvent"}] + assert events is not cached_events + + +def test_events_for_messages_v1_final_tool_message_creates_synthetic_user_message_without_tool_input_rails(): + rails = make_rails(RailsConfig(models=[], colang_version="1.0")) + + events = events_for_messages( + rails, + [ + {"role": "user", "content": "Use the tool"}, + {"role": "tool", "content": "tool result", "name": "lookup", "tool_call_id": "call-1"}, + ], + state=None, + ) + + assert events[-1] == {"type": "UserMessage", "text": "Use the tool"} + + +def test_events_for_messages_v1_final_tool_message_groups_tool_messages_when_tool_input_rails_exist(): + config = RailsConfig(models=[], colang_version="1.0") + config.rails.tool_input.flows = ["check tool input"] + rails = make_rails(config) + + events = events_for_messages( + rails, + [ + {"role": "user", "content": "Use tools"}, + {"role": "tool", "content": "first", "name": "lookup", "tool_call_id": "call-1"}, + {"role": "tool", "content": "second", "tool_call_id": "call-2"}, + ], + state=None, + ) + + assert events[-1] == { + "type": "UserToolMessages", + "tool_messages": [ + {"content": "first", "name": "lookup", "tool_call_id": "call-1"}, + {"content": "second", "name": "unknown", "tool_call_id": "call-2"}, + ], + } + + +def test_events_for_messages_v2_converts_supported_messages(): + rails = make_rails(RailsConfig(models=[], colang_version="2.x")) + passthrough_event = {"type": "UserSilent"} + + events = events_for_messages( + rails, + [ + {"role": "context", "content": {"generation_options": {}}}, + {"role": "system", "content": "Be concise."}, + {"role": "user", "content": "Hi"}, + {"role": "event", "event": passthrough_event}, + ], + state=SimpleNamespace(actions={}), + ) + + assert events == [ + {"type": "ContextUpdate", "data": {"generation_options": {}}}, + {"type": "SystemMessage", "content": "Be concise."}, + {"type": "UtteranceUserActionFinished", "final_transcript": "Hi"}, + passthrough_event, + ] + + +def test_events_for_messages_v2_rejects_assistant_messages(): + rails = make_rails(RailsConfig(models=[], colang_version="2.x")) + + with pytest.raises(ValueError, match="assistant.*not supported"): + events_for_messages( + rails, + [{"role": "assistant", "content": "Hello"}], + state=SimpleNamespace(actions={}), + ) + + +def test_events_for_messages_v2_converts_tool_message_to_action_finished_event(): + rails = make_rails(RailsConfig(models=[], colang_version="2.x")) + state = SimpleNamespace(actions={"call-1": SimpleNamespace(name="LookupAction")}) + + events = events_for_messages( + rails, + [{"role": "tool", "content": "tool result", "tool_call_id": "call-1"}], + state=state, + ) + + assert len(events) == 1 + assert events[0]["type"] == "LookupActionFinished" + assert events[0]["action_uid"] == "call-1" + assert events[0]["action_name"] == "LookupAction" + assert events[0]["status"] == "success" + assert events[0]["is_success"] is True + assert events[0]["return_value"] == "tool result" + assert events[0]["events"] == [] + + +def test_events_for_messages_accepts_llmrails_instance(): + rails = LLMRails(RailsConfig(models=[], colang_version="1.0")) + + events = events_for_messages( + rails, + [{"role": "user", "content": "Hi"}], + state=None, + ) + + assert events == [{"type": "UtteranceUserActionFinished", "final_transcript": "Hi"}]