diff --git a/nemoguardrails/rails/llm/generation/__init__.py b/nemoguardrails/rails/llm/generation/__init__.py new file mode 100644 index 0000000000..2ff3ef21a0 --- /dev/null +++ b/nemoguardrails/rails/llm/generation/__init__.py @@ -0,0 +1,18 @@ +# 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.generation.generation_workflow import generate_standard_async + +__all__ = ["generate_standard_async"] diff --git a/nemoguardrails/rails/llm/generation/bot_messages.py b/nemoguardrails/rails/llm/generation/bot_messages.py new file mode 100644 index 0000000000..fae0cee0be --- /dev/null +++ b/nemoguardrails/rails/llm/generation/bot_messages.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. + +"""Bot message extraction from Colang runtime events.""" + +import re +from dataclasses import dataclass, field +from typing import List + +__all__ = ["BotMessageFromEvents", "bot_message_from_colang_events"] + + +@dataclass +class BotMessageFromEvents: + message: dict + extra_events: List[dict] = field(default_factory=list) + + +def bot_message_from_colang_events( + colang_version: str, + events: List[dict], +) -> BotMessageFromEvents: + """Return the outward message represented by Colang runtime events.""" + if colang_version == "1.0": + return _bot_message_from_colang_1_events(events) + return _bot_message_from_colang_2_events(events) + + +def _bot_message_from_colang_1_events(events: List[dict]) -> BotMessageFromEvents: + responses = [] + exception = None + + for event in events: + if event["type"] == "StartUtteranceBotAction": + # Check if we need to remove a message + if event["script"] == "(remove last message)": + responses = responses[0:-1] + else: + responses.append(event["script"]) + elif event["type"].endswith("Exception"): + exception = event + + if exception: + return BotMessageFromEvents(message={"role": "exception", "content": exception}) + + return BotMessageFromEvents(message=_assistant_message_from_responses(responses)) + + +def _bot_message_from_colang_2_events(events: List[dict]) -> BotMessageFromEvents: + responses = [] + response_tool_calls = [] + response_events = [] + + for event in events: + start_action_match = re.match(r"Start(.*Action)", event["type"]) + + if start_action_match: + action_name = start_action_match[1] + # TODO: is there an elegant way to extract just the arguments? + arguments = { + k: v + for k, v in event.items() + if k != "type" and k != "uid" and k != "event_created_at" and k != "source_uid" and k != "action_uid" + } + response_tool_calls.append( + { + "id": event["action_uid"], + "type": "function", + "function": {"name": action_name, "arguments": arguments}, + } + ) + + elif event["type"] == "UtteranceBotActionFinished": + responses.append(event["final_script"]) + else: + # We just append the event + response_events.append(event) + + message = _assistant_message_from_responses(responses) + if response_tool_calls: + message["tool_calls"] = response_tool_calls + if response_events: + message["events"] = response_events + + return BotMessageFromEvents(message=message) + + +def _assistant_message_from_responses(responses: List) -> dict: + # Ensure all items in responses are strings + string_responses = [str(response) if not isinstance(response, str) else response for response in responses] + return {"role": "assistant", "content": "\n".join(string_responses)} diff --git a/nemoguardrails/rails/llm/generation/generation_context.py b/nemoguardrails/rails/llm/generation/generation_context.py new file mode 100644 index 0000000000..33a1f389f9 --- /dev/null +++ b/nemoguardrails/rails/llm/generation/generation_context.py @@ -0,0 +1,134 @@ +# 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. + +"""Generation context helpers.""" + +from contextvars import Token +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Tuple, cast + +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.options import GenerationOptions +from nemoguardrails.streaming import END_OF_STREAM, StreamingHandler + +__all__ = [ + "GenerationRequestContext", + "ensure_explain_info", + "explain_info_for_current_context", + "start_generation_request_context", + "start_generation_stats", +] + + +@dataclass +class GenerationRequestContext: + """Request-scoped generation context bindings.""" + + explain_info: ExplainInfo + _generation_options_token: Token + _raw_llm_request_token: Token + _explain_info_token: Token + _streaming_handler_token: Optional[Token] + _streaming_handler: Optional[StreamingHandler] + _streaming_handler_closed: bool = False + _closed: bool = False + + async def close_streaming_handler(self) -> None: + """Close the request streaming handler once.""" + if self._streaming_handler and not self._streaming_handler_closed: + self._streaming_handler_closed = True + await self._streaming_handler.push_chunk(cast(Any, END_OF_STREAM)) + + async def close(self) -> None: + """Close request resources and restore previous context bindings.""" + if self._closed: + return + + self._closed = True + try: + await self.close_streaming_handler() + finally: + if self._streaming_handler_token is not None: + streaming_handler_var.reset(self._streaming_handler_token) + explain_info_var.reset(self._explain_info_token) + raw_llm_request.reset(self._raw_llm_request_token) + generation_options_var.reset(self._generation_options_token) + + +def ensure_explain_info() -> ExplainInfo: + """Ensure that an ExplainInfo object is present in the current context.""" + explain_info = explain_info_var.get() + if explain_info is None: + explain_info = ExplainInfo() + explain_info_var.set(explain_info) + + return explain_info + + +def explain_info_for_current_context(fallback: Optional[ExplainInfo]) -> ExplainInfo: + """Return request-scoped explain info when present, otherwise use fallback.""" + explain_info = explain_info_var.get() + if explain_info is not None: + return explain_info + + if fallback is not None: + return fallback + + return ensure_explain_info() + + +def start_generation_request_context( + *, + gen_options: Optional[GenerationOptions], + messages: Optional[List[Dict[str, Any]]], + streaming_handler: Optional[StreamingHandler], +) -> GenerationRequestContext: + """Bind request-scoped generation context and return its cleanup scope.""" + generation_options_token = generation_options_var.set(gen_options) + + streaming_handler_token = None + if streaming_handler: + streaming_handler_token = streaming_handler_var.set(streaming_handler) + + explain_info = explain_info_var.get() + if explain_info is None: + explain_info = ExplainInfo() + explain_info_token = explain_info_var.set(explain_info) + + raw_llm_request_token = raw_llm_request.set(messages) + + return GenerationRequestContext( + explain_info=explain_info, + _generation_options_token=generation_options_token, + _raw_llm_request_token=raw_llm_request_token, + _explain_info_token=explain_info_token, + _streaming_handler_token=streaming_handler_token, + _streaming_handler=streaming_handler, + ) + + +def start_generation_stats() -> Tuple[LLMStats, List[dict]]: + """Initialize LLM stats and processing log for a generation request.""" + llm_stats = LLMStats() + llm_stats_var.set(llm_stats) + return llm_stats, [] diff --git a/nemoguardrails/rails/llm/generation/generation_request.py b/nemoguardrails/rails/llm/generation/generation_request.py new file mode 100644 index 0000000000..a984792a0c --- /dev/null +++ b/nemoguardrails/rails/llm/generation/generation_request.py @@ -0,0 +1,177 @@ +# 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. + +"""Generation request normalization.""" + +from dataclasses import dataclass +from typing import Any, List, Optional, Union + +from nemoguardrails.colang.v2_x.runtime.flows import State +from nemoguardrails.colang.v2_x.runtime.serialization import json_to_state +from nemoguardrails.exceptions import InvalidStateError +from nemoguardrails.rails.llm.options import GenerationOptions + +__all__ = [ + "GenerationRequest", + "PreparedGenerationRequest", + "normalize_generation_request", + "prepare_generation_request_for_runtime", +] + + +@dataclass +class GenerationRequest: + prompt: Optional[str] + messages: List[dict] + options: Optional[GenerationOptions] + state: Optional[Union[dict, State]] + + +@dataclass +class PreparedGenerationRequest: + prompt: Optional[str] + request_messages: List[dict] + runtime_messages: List[dict] + options: Optional[GenerationOptions] + state: Optional[Union[dict, State]] + needs_llm: bool + + +def normalize_generation_request( + *, + prompt: Optional[str], + messages: Optional[List[dict]], + options: Optional[Union[dict, GenerationOptions]], + state: Optional[Union[dict, State]], +) -> GenerationRequest: + """Normalize public generate inputs into the internal request shape.""" + if prompt is None and messages is None: + raise ValueError("Either prompt or messages must be provided.") + + if prompt is not None and messages is not None: + raise ValueError("Only one of prompt or messages can be provided.") + + if prompt is not None: + messages = [{"role": "user", "content": prompt}] + + normalized_state = _normalize_generation_state(state) + gen_options = _normalize_generation_options(options, state=normalized_state) + + return GenerationRequest( + prompt=prompt, + messages=messages or [], + options=gen_options, + state=normalized_state, + ) + + +def prepare_generation_request_for_runtime( + *, + prompt: Optional[str], + messages: Optional[List[dict]], + options: Optional[Union[dict, GenerationOptions]], + state: Optional[Union[dict, State]], +) -> PreparedGenerationRequest: + request = normalize_generation_request( + prompt=prompt, + messages=messages, + options=options, + state=state, + ) + gen_options = request.options + runtime_messages = request.messages + + if gen_options: + runtime_messages = [ + { + "role": "context", + "content": {"generation_options": gen_options.model_dump()}, + } + ] + (request.messages or []) + + if ( + runtime_messages + and runtime_messages[-1]["role"] == "assistant" + and not runtime_messages[-1].get("tool_calls") + and gen_options + and gen_options.rails.dialog is False + ): + runtime_messages[0]["content"]["bot_message"] = runtime_messages[-1]["content"] + runtime_messages = runtime_messages[0:-1] + + needs_llm = gen_options is None or gen_options.rails.dialog is not False + + return PreparedGenerationRequest( + prompt=request.prompt, + request_messages=request.messages, + runtime_messages=runtime_messages, + options=gen_options, + state=request.state, + needs_llm=needs_llm, + ) + + +def _normalize_generation_state( + state: Optional[Union[dict, State]], +) -> Optional[Union[dict, State]]: + if isinstance(state, dict) and state.get("version", "1.0") == "2.x": + return json_to_state(state["state"]) + return state + + +def validate_public_state(config: Any, state: Optional[Union[dict, State]]) -> None: + if not isinstance(state, dict) or not state: + return + + if config.colang_version == "1.0" and state.get("version") != "2.x": + if "state" in state: + raise InvalidStateError("Invalid Colang 1.0 state format: expected transcript state with an 'events' list.") + if "events" not in state: + raise InvalidStateError( + "Invalid Colang 1.0 state format: state must contain an 'events' key. " + "Use an empty dict {} to start a new conversation." + ) + if not isinstance(state["events"], list): + raise InvalidStateError("Invalid Colang 1.0 state format: 'events' must be a list.") + return + + raise InvalidStateError( + "Colang 2.0 dict state is not supported by generate/generate_async. " + "Use rails.process_events_async(events, state) with a live State object " + "for trusted in-process multi-turn execution. Public serialized Colang " + "2.0 runtime state is not accepted." + ) + + +def _normalize_generation_options( + options: Optional[Union[dict, GenerationOptions]], + *, + state: Optional[Any], +) -> Optional[GenerationOptions]: + if state is not None: + if options is None: + return GenerationOptions() + if isinstance(options, dict): + return GenerationOptions(**options) + return options + + if options and isinstance(options, dict): + return GenerationOptions(**options) + if isinstance(options, GenerationOptions): + return options + if options is None: + return None + + raise TypeError("options must be a dict or GenerationOptions") diff --git a/nemoguardrails/rails/llm/generation/generation_response.py b/nemoguardrails/rails/llm/generation/generation_response.py new file mode 100644 index 0000000000..cfecbc7ff0 --- /dev/null +++ b/nemoguardrails/rails/llm/generation/generation_response.py @@ -0,0 +1,174 @@ +# 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. + +"""Generation response envelope construction.""" + +from dataclasses import dataclass +from typing import Any, List, Optional + +from nemoguardrails.actions.llm.utils import ( + extract_bot_thinking_from_events, + extract_tool_calls_from_events, + get_and_clear_response_metadata_contextvar, + get_colang_history, +) +from nemoguardrails.colang.v1_0.runtime.flows import compute_context +from nemoguardrails.logging.processing_log import compute_generation_log +from nemoguardrails.rails.llm.options import ( + GenerationLog, + GenerationOptions, + GenerationResponse, +) + +__all__ = [ + "GenerationEventMetadata", + "generation_event_metadata", + "generation_response_from_colang_turn", +] + + +@dataclass +class GenerationEventMetadata: + tool_calls: Optional[List[dict]] + llm_metadata: Optional[dict] + reasoning_content: Optional[str] + + +def generation_event_metadata(events: List[dict]) -> GenerationEventMetadata: + """Extract response metadata carried by generated events/context vars.""" + return GenerationEventMetadata( + tool_calls=extract_tool_calls_from_events(events), + llm_metadata=get_and_clear_response_metadata_contextvar(), + reasoning_content=extract_bot_thinking_from_events(events), + ) + + +def generation_response_from_colang_turn( + *, + colang_version: str, + prompt: Optional[str], + new_message: dict, + events: List[dict], + new_events: List[dict], + processing_log: List[dict], + gen_options: GenerationOptions, + state: Any, + output_state: Any, + event_metadata: GenerationEventMetadata, +) -> GenerationResponse: + """Build a GenerationResponse from a completed Colang turn.""" + if prompt: + res = GenerationResponse(response=new_message["content"]) + else: + res = GenerationResponse(response=[new_message]) + + if event_metadata.reasoning_content: + res.reasoning_content = event_metadata.reasoning_content + + if event_metadata.tool_calls: + res.tool_calls = event_metadata.tool_calls + + if event_metadata.llm_metadata: + res.llm_metadata = event_metadata.llm_metadata + + if colang_version == "1.0": + _populate_colang_1_response( + res=res, + events=events, + new_events=new_events, + processing_log=processing_log, + gen_options=gen_options, + ) + else: + _validate_colang_2_response_options(gen_options) + + if state is not None: + res.state = output_state + + return res + + +def _populate_colang_1_response( + *, + res: GenerationResponse, + events: List[dict], + new_events: List[dict], + processing_log: List[dict], + gen_options: GenerationOptions, +) -> None: + if gen_options.output_vars: + context = compute_context(events) + output_vars = gen_options.output_vars + if isinstance(output_vars, list): + # If we have only a selection of keys, we filter to only that. + res.output_data = {k: context.get(k) for k in output_vars} + else: + # Otherwise, we return the full context + res.output_data = context + + generation_log = compute_generation_log(processing_log) + log_options = gen_options.log + + if log_options.activated_rails or log_options.llm_calls: + res.log = GenerationLog() + + # We always include the stats + res.log.stats = generation_log.stats + + if log_options.activated_rails: + res.log.activated_rails = generation_log.activated_rails + + if log_options.llm_calls: + res.log.llm_calls = [] + for activated_rail in generation_log.activated_rails: + for executed_action in activated_rail.executed_actions: + res.log.llm_calls.extend(executed_action.llm_calls) + + if log_options.internal_events: + if res.log is None: + res.log = GenerationLog() + + res.log.internal_events = new_events + + if log_options.colang_history: + if res.log is None: + res.log = GenerationLog() + + res.log.colang_history = get_colang_history(events) + + if gen_options.llm_output: + # Currently, we include the output from the generation LLM calls. + for activated_rail in generation_log.activated_rails: + if activated_rail.type == "generation": + for executed_action in activated_rail.executed_actions: + for llm_call in executed_action.llm_calls: + res.llm_output = llm_call.raw_response + + +def _validate_colang_2_response_options(gen_options: GenerationOptions) -> None: + if gen_options.output_vars: + raise ValueError("The `output_vars` option is not supported for Colang 2.0 configurations.") + + log_options = gen_options.log + if ( + log_options.activated_rails + or log_options.llm_calls + or log_options.internal_events + or log_options.colang_history + ): + raise ValueError("The `log` option is not supported for Colang 2.0 configurations.") + + if gen_options.llm_output: + raise ValueError("The `llm_output` option is not supported for Colang 2.0 configurations.") diff --git a/nemoguardrails/rails/llm/generation/generation_tracing.py b/nemoguardrails/rails/llm/generation/generation_tracing.py new file mode 100644 index 0000000000..b2a23db330 --- /dev/null +++ b/nemoguardrails/rails/llm/generation/generation_tracing.py @@ -0,0 +1,120 @@ +# 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. + +"""Generation tracing option preparation and log restoration.""" + +from dataclasses import dataclass +from typing import Any, List, Optional + +from nemoguardrails.rails.llm.options import ( + GenerationLogOptions, + GenerationOptions, + GenerationResponse, +) + +__all__ = [ + "GenerationTracingState", + "export_generation_trace", + "prepare_generation_tracing", + "restore_generation_trace_log", +] + + +@dataclass +class GenerationTracingState: + gen_options: Optional[GenerationOptions] + original_log_options: Optional[GenerationLogOptions] + + +def prepare_generation_tracing( + *, + tracing_enabled: bool, + gen_options: Optional[GenerationOptions], +) -> GenerationTracingState: + """Prepare generation options needed for tracing without mutating callers.""" + if not tracing_enabled: + return GenerationTracingState( + gen_options=gen_options, + original_log_options=None, + ) + + if gen_options is None: + tracing_options = GenerationOptions() + else: + tracing_options = gen_options.model_copy(deep=True) + + original_log_options = tracing_options.log.model_copy(deep=True) + tracing_options.log.activated_rails = True + tracing_options.log.llm_calls = True + tracing_options.log.internal_events = True + + return GenerationTracingState( + gen_options=tracing_options, + original_log_options=original_log_options, + ) + + +async def export_generation_trace( + *, + tracing_config: Any, + log_adapters: Optional[List[Any]], + messages: Optional[List[dict]], + response: GenerationResponse, +) -> None: + """Export tracing data for a completed generation response.""" + # Lazy import to avoid circular dependencies through eval/tracing modules. + from nemoguardrails.tracing import Tracer + + span_format = getattr(tracing_config, "span_format", "opentelemetry") + enable_content_capture = getattr(tracing_config, "enable_content_capture", False) + tracer = Tracer( + input=messages, + response=response, + adapters=log_adapters, + span_format=span_format, + enable_content_capture=enable_content_capture, + ) + await tracer.export_async() + + +def restore_generation_trace_log( + *, + response: GenerationResponse, + original_log_options: Optional[GenerationLogOptions], +) -> None: + """Restore the response log fields requested by the caller before tracing.""" + if original_log_options is None: + return + + if not any( + ( + original_log_options.internal_events, + original_log_options.activated_rails, + original_log_options.llm_calls, + original_log_options.colang_history, + ) + ): + response.log = None + return + + if response.log is None: + return + + if not original_log_options.internal_events: + response.log.internal_events = [] + if not original_log_options.activated_rails: + response.log.activated_rails = [] + if not original_log_options.llm_calls: + response.log.llm_calls = [] diff --git a/nemoguardrails/rails/llm/generation/generation_workflow.py b/nemoguardrails/rails/llm/generation/generation_workflow.py new file mode 100644 index 0000000000..1b0d1c7da6 --- /dev/null +++ b/nemoguardrails/rails/llm/generation/generation_workflow.py @@ -0,0 +1,131 @@ +# 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 logging +import time +from typing import Any, Optional, Union + +from nemoguardrails.actions.llm.utils import get_colang_history +from nemoguardrails.colang.v2_x.runtime.flows import State +from nemoguardrails.rails.llm.conversation.conversation_events import events_for_messages +from nemoguardrails.rails.llm.generation.bot_messages import bot_message_from_colang_events +from nemoguardrails.rails.llm.generation.generation_context import ( + GenerationRequestContext, + start_generation_stats, +) +from nemoguardrails.rails.llm.generation.generation_response import ( + generation_event_metadata, + generation_response_from_colang_turn, +) +from nemoguardrails.rails.llm.generation.generation_tracing import ( + export_generation_trace, + prepare_generation_tracing, + restore_generation_trace_log, +) +from nemoguardrails.rails.llm.options import GenerationOptions, GenerationResponse +from nemoguardrails.rails.llm.runtime.colang_turns import run_colang_turn +from nemoguardrails.rails.llm.utils import get_history_cache_key + +__all__ = ["generate_standard_async"] + +log = logging.getLogger("nemoguardrails.rails.llm.llmrails") + + +async def generate_standard_async( + rails: Any, + *, + prompt: Optional[str], + messages: list[dict], + gen_options: Optional[GenerationOptions], + state: Optional[Union[dict, State]], + request_context: GenerationRequestContext, +) -> Union[str, dict, GenerationResponse, tuple[dict, dict]]: + rails._explain_info = request_context.explain_info + + t0 = time.time() + llm_stats, processing_log = start_generation_stats() + + events = events_for_messages(rails, messages, state) + + new_events = await run_colang_turn(rails, events, state, processing_log) + output_state = None + + bot_message = bot_message_from_colang_events(rails.config.colang_version, new_events) + new_message = bot_message.message + + if rails.config.colang_version == "1.0": + events.extend(new_events) + events.extend(bot_message.extra_events) + + if state is None: + cache_key = get_history_cache_key((messages) + [new_message]) + rails.events_history_cache[cache_key] = events + else: + output_state = {"events": events} + + rails._explain_info.colang_history = get_colang_history(events) + if rails.verbose: + log.info(f"Conversation history so far: \n{rails._explain_info.colang_history}") + + total_time = time.time() - t0 + log.info("--- :: Total processing took %.2f seconds. LLM Stats: %s" % (total_time, llm_stats)) + + await request_context.close_streaming_handler() + + tracing_state = prepare_generation_tracing( + tracing_enabled=rails.config.tracing.enabled, + gen_options=gen_options, + ) + gen_options = tracing_state.gen_options + + response_event_metadata = generation_event_metadata(new_events) + if gen_options: + res = generation_response_from_colang_turn( + colang_version=rails.config.colang_version, + prompt=prompt, + new_message=new_message, + events=events, + new_events=new_events, + processing_log=processing_log, + gen_options=gen_options, + state=state, + output_state=output_state, + event_metadata=response_event_metadata, + ) + + if rails.config.tracing.enabled: + await export_generation_trace( + tracing_config=rails.config.tracing, + log_adapters=rails._log_adapters, + messages=messages, + response=res, + ) + restore_generation_trace_log( + response=res, + original_log_options=tracing_state.original_log_options, + ) + + return res + + if response_event_metadata.reasoning_content: + thinking_trace = f"{response_event_metadata.reasoning_content}\n" + new_message["content"] = thinking_trace + new_message["content"] + + if prompt: + return new_message["content"] + + if response_event_metadata.tool_calls: + new_message["tool_calls"] = response_event_metadata.tool_calls + return new_message diff --git a/nemoguardrails/rails/llm/llmrails.py b/nemoguardrails/rails/llm/llmrails.py index 6c2406ca00..3f9101c61b 100644 --- a/nemoguardrails/rails/llm/llmrails.py +++ b/nemoguardrails/rails/llm/llmrails.py @@ -18,8 +18,6 @@ import asyncio import json import logging -import re -import time import warnings from functools import partial from typing import ( @@ -33,42 +31,21 @@ Tuple, Type, Union, - cast, overload, ) from typing_extensions import Self -from nemoguardrails.actions.llm.utils import ( - extract_bot_thinking_from_events, - extract_tool_calls_from_events, - get_and_clear_response_metadata_contextvar, - get_colang_history, -) from nemoguardrails.actions.output_mapping import is_output_blocked 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 State -from nemoguardrails.colang.v2_x.runtime.runtime import RuntimeV2_x -from nemoguardrails.context import ( - explain_info_var, - generation_options_var, - llm_stats_var, - raw_llm_request, - streaming_handler_var, -) from nemoguardrails.embeddings.index import EmbeddingsIndex from nemoguardrails.embeddings.providers import register_embedding_provider from nemoguardrails.embeddings.providers.base import EmbeddingModel -from nemoguardrails.exceptions import ( - InvalidStateError, - StreamingNotSupportedError, -) +from nemoguardrails.exceptions import StreamingNotSupportedError 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 from nemoguardrails.logging.verbose import set_verbose from nemoguardrails.patch_asyncio import check_sync_call_from_async_loop from nemoguardrails.rails.llm.buffer import get_buffer_strategy @@ -77,8 +54,17 @@ RailsConfig, ) from nemoguardrails.rails.llm.conversation.conversation_events import events_for_messages +from nemoguardrails.rails.llm.generation.generation_context import ( + ensure_explain_info, + explain_info_for_current_context, + start_generation_request_context, +) +from nemoguardrails.rails.llm.generation.generation_request import ( + prepare_generation_request_for_runtime, + validate_public_state, +) +from nemoguardrails.rails.llm.generation.generation_workflow import generate_standard_async from nemoguardrails.rails.llm.options import ( - GenerationLog, GenerationOptions, GenerationResponse, RailsResult, @@ -106,7 +92,6 @@ 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, ) from nemoguardrails.streaming import END_OF_STREAM, StreamingHandler from nemoguardrails.types import LLMModel @@ -403,38 +388,11 @@ def _ensure_explain_info() -> ExplainInfo: Returns: A ExplainInfo class containing the llm calls' statistics """ - explain_info = explain_info_var.get() - if explain_info is None: - explain_info = ExplainInfo() - explain_info_var.set(explain_info) - - return explain_info + return ensure_explain_info() def _validate_public_state(self, state: Optional[Union[dict, State]]) -> None: """Validate public dict state passed through generate/generate_async.""" - if not isinstance(state, dict) or not state: - return - - if self.config.colang_version == "1.0" and state.get("version") != "2.x": - if "state" in state: - raise InvalidStateError( - "Invalid Colang 1.0 state format: expected transcript state with an 'events' list." - ) - if "events" not in state: - raise InvalidStateError( - "Invalid Colang 1.0 state format: state must contain an 'events' key. " - "Use an empty dict {} to start a new conversation." - ) - if not isinstance(state["events"], list): - raise InvalidStateError("Invalid Colang 1.0 state format: 'events' must be a list.") - return - - raise InvalidStateError( - "Colang 2.0 dict state is not supported by generate/generate_async. " - "Use rails.process_events_async(events, state) with a live State object " - "for trusted in-process multi-turn execution. Public serialized Colang " - "2.0 runtime state is not accepted." - ) + validate_public_state(self.config, state) async def generate_async( self, @@ -468,398 +426,44 @@ async def generate_async( The completion (when a prompt is provided) or the next message. System messages are not yet supported.""" - # convert options to gen_options of type GenerationOptions - gen_options: Optional[GenerationOptions] = None - if prompt is None and messages is None: raise ValueError("Either prompt or messages must be provided.") if prompt is not None and messages is not None: raise ValueError("Only one of prompt or messages can be provided.") - if prompt is not None: - # Currently, we transform the prompt request into a single turn conversation - messages = [{"role": "user", "content": prompt}] - - # If a state object is specified, then we switch to "generation options" mode. - # This is because we want the output to be a GenerationResponse which will contain - # the output state. - if state is not None: - self._validate_public_state(state) - - if options is None: - gen_options = GenerationOptions() - elif isinstance(options, dict): - gen_options = GenerationOptions(**options) - else: - gen_options = options - else: - # We allow options to be specified both as a dict and as an object. - if options and isinstance(options, dict): - gen_options = GenerationOptions(**options) - elif isinstance(options, GenerationOptions): - gen_options = options - elif options is None: - gen_options = None - else: - raise TypeError("options must be a dict or GenerationOptions") - - # Save the generation options in the current async context. - # At this point, gen_options is either None or GenerationOptions - generation_options_var.set(gen_options) - - needs_llm = gen_options is None or gen_options.rails.dialog is not False - if needs_llm and not self.llm: - log.warning("No main LLM specified in the config and no LLM provided via constructor.") - - if streaming_handler: - streaming_handler_var.set(streaming_handler) - - # Initialize the object with additional explanation information. - # We allow this to also be set externally. This is useful when multiple parallel - # requests are made. - self._explain_info = self._ensure_explain_info() - - raw_llm_request.set(messages) - - # If we have generation options, we also add them to the context - if gen_options: - messages = [ - { - "role": "context", - "content": {"generation_options": gen_options.model_dump()}, - } - ] + (messages or []) - - # If the last message is from the assistant, rather than the user, then - # we move that to the `$bot_message` variable. This is to enable a more - # convenient interface for text output rails. Tool-call assistant messages - # must remain in the history so they can be converted into BotToolCalls - # events and evaluated by tool output rails. - if ( - messages - and messages[-1]["role"] == "assistant" - and not messages[-1].get("tool_calls") - and gen_options - and gen_options.rails.dialog is False - ): - # We already have the first message with a context update, so we use that - messages[0]["content"]["bot_message"] = messages[-1]["content"] - messages = messages[0:-1] - - # TODO: Add support to load back history of events, next to history of messages - # This is important as without it, the LLM prediction is not as good. - - t0 = time.time() - - # Initialize the LLM stats - llm_stats = LLMStats() - llm_stats_var.set(llm_stats) - processing_log = [] - - # The array of events corresponding to the provided sequence of messages. - events = self._get_events_for_messages(messages, state) # type: ignore - - if self.config.colang_version == "1.0": - # If we had a state object, we also need to prepend the events from the state. - state_events = [] - if state: - assert isinstance(state, dict) - state_events = state["events"] - - new_events = [] - # Compute the new events. - try: - new_events = await self.runtime.generate_events(state_events + events, processing_log=processing_log) - output_state = None + validate_public_state(self.config, state) + prepared_request = prepare_generation_request_for_runtime( + prompt=prompt, + messages=messages, + options=options, + state=state, + ) + prompt = prepared_request.prompt + request_messages = prepared_request.request_messages + messages = prepared_request.runtime_messages + gen_options = prepared_request.options + state = prepared_request.state + + request_context = start_generation_request_context( + gen_options=gen_options, + messages=request_messages, + streaming_handler=streaming_handler, + ) + try: + if prepared_request.needs_llm and not self.llm: + log.warning("No main LLM specified in the config and no LLM provided via constructor.") - except Exception as e: - log.error("Error in generate_async: %s", e, exc_info=True) - streaming_handler = streaming_handler_var.get() - if streaming_handler: - # Push an error chunk instead of None. - error_message = str(e) - error_dict = extract_error_json(error_message) - error_payload: str = json.dumps(error_dict) - await streaming_handler.push_chunk(error_payload) - # push a termination signal - await streaming_handler.push_chunk(END_OF_STREAM) # type: ignore - # Re-raise the exact exception - raise - else: - # In generation mode, by default the bot response is an instant action. - instant_actions = ["UtteranceBotAction"] - if self.config.rails.actions.instant_actions is not None: - instant_actions = self.config.rails.actions.instant_actions - - # Cast this explicitly to avoid certain warnings - runtime: RuntimeV2_x = cast(RuntimeV2_x, self.runtime) - - # Compute the new events. - # In generation mode, the processing is always blocking, i.e., it waits for - # all local actions (sync and async). - new_events, _output_state = await runtime.process_events( - events, state=state, instant_actions=instant_actions, blocking=True + return await generate_standard_async( + self, + prompt=prompt, + messages=messages, + gen_options=gen_options, + state=state, + request_context=request_context, ) - # The runtime State for 2.x is not publicly exposed through generate_async. - # Callers that need stateful 2.x execution use process_events_async, which - # returns the live State object directly. - output_state = None - - # Extract and join all the messages from StartUtteranceBotAction events as the response. - responses = [] - response_tool_calls = [] - response_events = [] - new_extra_events = [] - exception = None - - # The processing is different for Colang 1.0 and 2.0 - if self.config.colang_version == "1.0": - for event in new_events: - if event["type"] == "StartUtteranceBotAction": - # Check if we need to remove a message - if event["script"] == "(remove last message)": - responses = responses[0:-1] - else: - responses.append(event["script"]) - elif event["type"].endswith("Exception"): - exception = event - - else: - for event in new_events: - start_action_match = re.match(r"Start(.*Action)", event["type"]) - - if start_action_match: - action_name = start_action_match[1] - # TODO: is there an elegant way to extract just the arguments? - arguments = { - k: v - for k, v in event.items() - if k != "type" - and k != "uid" - and k != "event_created_at" - and k != "source_uid" - and k != "action_uid" - } - response_tool_calls.append( - { - "id": event["action_uid"], - "type": "function", - "function": {"name": action_name, "arguments": arguments}, - } - ) - - elif event["type"] == "UtteranceBotActionFinished": - responses.append(event["final_script"]) - else: - # We just append the event - response_events.append(event) - - if exception: - new_message: dict = {"role": "exception", "content": exception} - - else: - # Ensure all items in responses are strings - responses = [str(response) if not isinstance(response, str) else response for response in responses] - new_message: dict = {"role": "assistant", "content": "\n".join(responses)} - if response_tool_calls: - new_message["tool_calls"] = response_tool_calls - if response_events: - new_message["events"] = response_events - - if self.config.colang_version == "1.0": - events.extend(new_events) - events.extend(new_extra_events) - - # If a state object is not used, then we use the implicit caching - if state is None: - # Save the new events in the history and update the cache - cache_key = get_history_cache_key((messages) + [new_message]) # type: ignore - self.events_history_cache[cache_key] = events - else: - output_state = {"events": events} - - # If logging is enabled, we log the conversation - # TODO: add support for logging flag - self._explain_info.colang_history = get_colang_history(events) - if self.verbose: - log.info(f"Conversation history so far: \n{self._explain_info.colang_history}") - - total_time = time.time() - t0 - log.info("--- :: Total processing took %.2f seconds. LLM Stats: %s" % (total_time, llm_stats)) - - # If there is a streaming handler, we make sure we close it now - streaming_handler = streaming_handler_var.get() - if streaming_handler: - # print("Closing the stream handler explicitly") - await streaming_handler.push_chunk(END_OF_STREAM) # type: ignore - - # IF tracing is enabled we need to set GenerationLog attrs - original_log_options = None - if self.config.tracing.enabled: - if gen_options is None: - gen_options = GenerationOptions() - else: - # create a copy of the gen_options to avoid modifying the original - gen_options = gen_options.model_copy(deep=True) - original_log_options = gen_options.log.model_copy(deep=True) - - # enable log options - # it is aggressive, but these are required for tracing - if ( - not gen_options.log.activated_rails - or not gen_options.log.llm_calls - or not gen_options.log.internal_events - ): - gen_options.log.activated_rails = True - gen_options.log.llm_calls = True - gen_options.log.internal_events = True - - tool_calls = extract_tool_calls_from_events(new_events) - llm_metadata = get_and_clear_response_metadata_contextvar() - reasoning_content = extract_bot_thinking_from_events(new_events) - # If we have generation options, we prepare a GenerationResponse instance. - if gen_options: - # If a prompt was used, we only need to return the content of the message. - if prompt: - res = GenerationResponse(response=new_message["content"]) - else: - res = GenerationResponse(response=[new_message]) - - if reasoning_content: - res.reasoning_content = reasoning_content - - if tool_calls: - res.tool_calls = tool_calls - - if llm_metadata: - res.llm_metadata = llm_metadata - - if self.config.colang_version == "1.0": - # If output variables are specified, we extract their values - if gen_options and gen_options.output_vars: - context = compute_context(events) - output_vars = gen_options.output_vars - if isinstance(output_vars, list): - # If we have only a selection of keys, we filter to only that. - res.output_data = {k: context.get(k) for k in output_vars} - else: - # Otherwise, we return the full context - res.output_data = context - - _log = compute_generation_log(processing_log) - - # Include information about activated rails and LLM calls if requested - log_options = gen_options.log if gen_options else None - if log_options and (log_options.activated_rails or log_options.llm_calls): - res.log = GenerationLog() - - # We always include the stats - res.log.stats = _log.stats - - if log_options.activated_rails: - res.log.activated_rails = _log.activated_rails - - if log_options.llm_calls: - res.log.llm_calls = [] - for activated_rail in _log.activated_rails: - for executed_action in activated_rail.executed_actions: - res.log.llm_calls.extend(executed_action.llm_calls) - - # Include internal events if requested - if log_options and log_options.internal_events: - if res.log is None: - res.log = GenerationLog() - - res.log.internal_events = new_events - - # Include the Colang history if requested - if log_options and log_options.colang_history: - if res.log is None: - res.log = GenerationLog() - - res.log.colang_history = get_colang_history(events) - - # Include the raw llm output if requested - if gen_options and gen_options.llm_output: - # Currently, we include the output from the generation LLM calls. - for activated_rail in _log.activated_rails: - if activated_rail.type == "generation": - for executed_action in activated_rail.executed_actions: - for llm_call in executed_action.llm_calls: - res.llm_output = llm_call.raw_response - else: - if gen_options and gen_options.output_vars: - raise ValueError("The `output_vars` option is not supported for Colang 2.0 configurations.") - - log_options = gen_options.log if gen_options else None - if log_options and ( - log_options.activated_rails - or log_options.llm_calls - or log_options.internal_events - or log_options.colang_history - ): - raise ValueError("The `log` option is not supported for Colang 2.0 configurations.") - - if gen_options and gen_options.llm_output: - raise ValueError("The `llm_output` option is not supported for Colang 2.0 configurations.") - - # Include the state - if state is not None: - res.state = output_state - - if self.config.tracing.enabled: - # TODO: move it to the top once resolved circular dependency of eval - # lazy import to avoid circular dependency - from nemoguardrails.tracing import Tracer - - span_format = getattr(self.config.tracing, "span_format", "opentelemetry") - enable_content_capture = getattr(self.config.tracing, "enable_content_capture", False) - # Create a Tracer instance with instantiated adapters and span configuration - tracer = Tracer( - input=messages, - response=res, - adapters=self._log_adapters, - span_format=span_format, - enable_content_capture=enable_content_capture, - ) - await tracer.export_async() - - # respect original log specification, if tracing added information to the output - if original_log_options: - if not any( - ( - original_log_options.internal_events, - original_log_options.activated_rails, - original_log_options.llm_calls, - original_log_options.colang_history, - ) - ): - res.log = None - else: - # Ensure res.log exists before setting attributes - if res.log is not None: - if not original_log_options.internal_events: - res.log.internal_events = [] - if not original_log_options.activated_rails: - res.log.activated_rails = [] - if not original_log_options.llm_calls: - res.log.llm_calls = [] - - return res - else: - # If a prompt is used, we only return the content of the message. - - if reasoning_content: - thinking_trace = f"{reasoning_content}\n" - new_message["content"] = thinking_trace + new_message["content"] - - if prompt: - return new_message["content"] - else: - if tool_calls: - new_message["tool_calls"] = tool_calls - return new_message + finally: + await request_context.close() def _validate_streaming_with_output_rails(self) -> None: if len(self.config.rails.output.flows) > 0 and ( @@ -1261,8 +865,7 @@ def register_embedding_provider(self, cls: Type[EmbeddingModel], name: Optional[ def explain(self) -> ExplainInfo: """Helper function to return the latest ExplainInfo object.""" - if self._explain_info is None: - self._explain_info = self._ensure_explain_info() + self._explain_info = explain_info_for_current_context(self._explain_info) return self._explain_info def __getstate__(self): diff --git a/nemoguardrails/rails/llm/runtime/colang_turns.py b/nemoguardrails/rails/llm/runtime/colang_turns.py index 86bf26ff8a..3cbe368c86 100644 --- a/nemoguardrails/rails/llm/runtime/colang_turns.py +++ b/nemoguardrails/rails/llm/runtime/colang_turns.py @@ -16,14 +16,18 @@ """Colang runtime event helpers.""" import asyncio +import json import logging import time -from typing import Any, List, Optional, Protocol, Tuple, Union +from typing import Any, List, Optional, Protocol, Tuple, Union, cast 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.colang.v2_x.runtime.runtime import RuntimeV2_x +from nemoguardrails.context import llm_stats_var, streaming_handler_var from nemoguardrails.logging.stats import LLMStats +from nemoguardrails.streaming import END_OF_STREAM +from nemoguardrails.utils import extract_error_json log = logging.getLogger(__name__) @@ -34,6 +38,7 @@ "generate_colang_events", "process_colang_events", "process_events_semaphore", + "run_colang_turn", ] @@ -48,6 +53,66 @@ def runtime(self) -> Any: ... def verbose(self) -> bool: ... +async def run_colang_turn( + rails: ColangTurnRails, + events: List[dict], + state: Any, + processing_log: List[dict], +) -> List[dict]: + """Run one Colang turn for events already converted from messages.""" + if rails.config.colang_version == "1.0": + return await _run_colang_1_turn(rails, events, state, processing_log) + return await _run_colang_2_turn(rails, events, state) + + +async def _run_colang_1_turn( + rails: ColangTurnRails, + events: List[dict], + state: Any, + processing_log: List[dict], +) -> List[dict]: + state_events = [] + if state: + assert isinstance(state, dict) + state_events = state["events"] + + try: + return await rails.runtime.generate_events( + state_events + events, + processing_log=processing_log, + ) + except Exception as e: + log.error("Error in generate_async: %s", e, exc_info=True) + streaming_handler = streaming_handler_var.get() + if streaming_handler: + error_message = str(e) + error_dict = extract_error_json(error_message) + error_payload: str = json.dumps(error_dict) + await streaming_handler.push_chunk(error_payload) + await streaming_handler.push_chunk(END_OF_STREAM) # type: ignore + raise + + +async def _run_colang_2_turn( + rails: ColangTurnRails, + events: List[dict], + state: Any, +) -> List[dict]: + instant_actions = ["UtteranceBotAction"] + if rails.config.rails.actions.instant_actions is not None: + instant_actions = rails.config.rails.actions.instant_actions + + runtime: RuntimeV2_x = cast(RuntimeV2_x, rails.runtime) + new_events, _output_state = await runtime.process_events( + events, + state=state, + instant_actions=instant_actions, + blocking=True, + ) + + return new_events + + 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() diff --git a/tests/rails/llm/test_bot_messages.py b/tests/rails/llm/test_bot_messages.py new file mode 100644 index 0000000000..fbee911802 --- /dev/null +++ b/tests/rails/llm/test_bot_messages.py @@ -0,0 +1,87 @@ +# 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 nemoguardrails.rails.llm.generation.bot_messages import bot_message_from_colang_events + + +def test_bot_message_from_colang_1_events_joins_utterances_and_removes_last_message(): + result = bot_message_from_colang_events( + "1.0", + [ + {"type": "StartUtteranceBotAction", "script": "First"}, + {"type": "StartUtteranceBotAction", "script": "Second"}, + {"type": "StartUtteranceBotAction", "script": "(remove last message)"}, + {"type": "StartUtteranceBotAction", "script": 3}, + ], + ) + + assert result.message == {"role": "assistant", "content": "First\n3"} + assert result.extra_events == [] + + +def test_bot_message_from_colang_1_events_returns_exception_message(): + exception_event = {"type": "CustomException", "message": "blocked"} + + result = bot_message_from_colang_events( + "1.0", + [ + {"type": "StartUtteranceBotAction", "script": "First"}, + exception_event, + ], + ) + + assert result.message == {"role": "exception", "content": exception_event} + + +def test_bot_message_from_colang_2_events_extracts_text_tool_calls_and_response_events(): + response_event = {"type": "CustomEvent", "value": 1} + + result = bot_message_from_colang_events( + "2.x", + [ + { + "type": "StartGestureBotAction", + "action_uid": "gesture-1", + "uid": "event-1", + "event_created_at": "now", + "source_uid": "source-1", + "gesture": "Wave", + }, + {"type": "UtteranceBotActionFinished", "final_script": "Hello"}, + response_event, + ], + ) + + assert result.message == { + "role": "assistant", + "content": "Hello", + "tool_calls": [ + { + "id": "gesture-1", + "type": "function", + "function": { + "name": "GestureBotAction", + "arguments": {"gesture": "Wave"}, + }, + } + ], + "events": [response_event], + } + + +def test_bot_message_from_colang_2_events_returns_empty_assistant_message_for_no_outputs(): + result = bot_message_from_colang_events("2.x", []) + + assert result.message == {"role": "assistant", "content": ""} diff --git a/tests/rails/llm/test_colang_turns.py b/tests/rails/llm/test_colang_turns.py index b25cd816fb..75f0c8453c 100644 --- a/tests/rails/llm/test_colang_turns.py +++ b/tests/rails/llm/test_colang_turns.py @@ -14,19 +14,22 @@ # limitations under the License. import asyncio +import json from types import SimpleNamespace from typing import Any, cast import pytest -from nemoguardrails.context import llm_stats_var +from nemoguardrails.context import llm_stats_var, streaming_handler_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, + run_colang_turn, ) +from nemoguardrails.streaming import END_OF_STREAM class FakeRails: @@ -54,6 +57,128 @@ def reset_llm_stats_context(): llm_stats_var.reset(token) +@pytest.mark.asyncio +async def test_run_colang_turn_v1_prepends_state_events_and_passes_processing_log(): + class Runtime: + def __init__(self): + self.events = None + self.processing_log = None + + async def generate_events(self, events, processing_log): + self.events = events + self.processing_log = processing_log + processing_log.append({"event": "generated"}) + return [{"type": "BotIntent", "intent": "express greeting"}] + + runtime = Runtime() + processing_log = [] + + new_events = await run_colang_turn( + make_rails("1.0", runtime), + events=[{"type": "UserMessage", "text": "Hi"}], + state={"events": [{"type": "ContextUpdate"}]}, + processing_log=processing_log, + ) + + assert runtime.events == [ + {"type": "ContextUpdate"}, + {"type": "UserMessage", "text": "Hi"}, + ] + assert runtime.processing_log is processing_log + assert processing_log == [{"event": "generated"}] + assert new_events == [{"type": "BotIntent", "intent": "express greeting"}] + + +@pytest.mark.asyncio +async def test_run_colang_turn_v1_streams_error_chunk_before_reraising(): + class Runtime: + async def generate_events(self, events, processing_log): + raise RuntimeError("boom") + + class StreamingHandler: + def __init__(self): + self.chunks = [] + + async def push_chunk(self, chunk): + self.chunks.append(chunk) + + streaming_handler = StreamingHandler() + token = streaming_handler_var.set(streaming_handler) + try: + with pytest.raises(RuntimeError, match="boom"): + await run_colang_turn( + make_rails("1.0", Runtime()), + events=[{"type": "UserMessage", "text": "Hi"}], + state=None, + processing_log=[], + ) + finally: + streaming_handler_var.reset(token) + + assert json.loads(streaming_handler.chunks[0]) == {"error": {"message": "boom"}} + assert streaming_handler.chunks[1] == END_OF_STREAM + + +@pytest.mark.asyncio +async def test_run_colang_turn_v2_processes_events_with_default_instant_actions(): + class Runtime: + def __init__(self): + self.calls = [] + + async def process_events(self, events, state, instant_actions, blocking): + self.calls.append( + { + "events": events, + "state": state, + "instant_actions": instant_actions, + "blocking": blocking, + } + ) + return [{"type": "StartUtteranceBotAction", "script": "Hi"}], object() + + runtime = Runtime() + events = [{"type": "UtteranceUserActionFinished", "final_transcript": "hi"}] + state = object() + + new_events = await run_colang_turn( + make_rails("2.x", runtime), + events=events, + state=state, + processing_log=[], + ) + + assert runtime.calls == [ + { + "events": events, + "state": state, + "instant_actions": ["UtteranceBotAction"], + "blocking": True, + } + ] + assert new_events == [{"type": "StartUtteranceBotAction", "script": "Hi"}] + + +@pytest.mark.asyncio +async def test_run_colang_turn_v2_honors_configured_instant_actions(): + class Runtime: + def __init__(self): + self.instant_actions = None + + async def process_events(self, events, state, instant_actions, blocking): + self.instant_actions = instant_actions + return [], object() + + runtime = Runtime() + await run_colang_turn( + make_rails("2.x", runtime, instant_actions=["CustomAction"]), + events=[], + state=None, + processing_log=[], + ) + + assert runtime.instant_actions == ["CustomAction"] + + @pytest.mark.asyncio async def test_generate_colang_events_sets_stats_and_passes_processing_log(): class Runtime: diff --git a/tests/rails/llm/test_generation_context.py b/tests/rails/llm/test_generation_context.py new file mode 100644 index 0000000000..9f01a79453 --- /dev/null +++ b/tests/rails/llm/test_generation_context.py @@ -0,0 +1,215 @@ +# 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 + +import pytest + +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.generation import generation_context +from nemoguardrails.rails.llm.generation.generation_context import ( + ensure_explain_info, + explain_info_for_current_context, + start_generation_request_context, + start_generation_stats, +) +from nemoguardrails.rails.llm.options import GenerationOptions +from nemoguardrails.streaming import END_OF_STREAM + + +def test_generation_context_exports(): + assert generation_context.__all__ == [ + "GenerationRequestContext", + "ensure_explain_info", + "explain_info_for_current_context", + "start_generation_request_context", + "start_generation_stats", + ] + + +@pytest.fixture(autouse=True) +def reset_generation_context_vars(): + tokens = [ + explain_info_var.set(None), + generation_options_var.set(None), + llm_stats_var.set(None), + raw_llm_request.set(None), + streaming_handler_var.set(None), + ] + try: + yield + finally: + for var, token in reversed( + [ + (explain_info_var, tokens[0]), + (generation_options_var, tokens[1]), + (llm_stats_var, tokens[2]), + (raw_llm_request, tokens[3]), + (streaming_handler_var, tokens[4]), + ] + ): + var.reset(token) + + +def test_ensure_explain_info_creates_and_reuses_context_value(): + explain_info = ensure_explain_info() + + assert isinstance(explain_info, ExplainInfo) + assert explain_info_var.get() is explain_info + assert ensure_explain_info() is explain_info + + +def test_explain_info_for_current_context_prefers_context_then_fallback(): + fallback = ExplainInfo() + request_explain_info = ExplainInfo() + + assert explain_info_for_current_context(fallback) is fallback + + explain_info_var.set(request_explain_info) + + assert explain_info_for_current_context(fallback) is request_explain_info + + +def test_explain_info_for_current_context_creates_context_when_needed(): + explain_info = explain_info_for_current_context(None) + + assert isinstance(explain_info, ExplainInfo) + assert explain_info_var.get() is explain_info + + +def test_start_generation_stats_sets_stats_and_returns_processing_log(): + llm_stats, processing_log = start_generation_stats() + + assert isinstance(llm_stats, LLMStats) + assert llm_stats_var.get() is llm_stats + assert processing_log == [] + + +@pytest.mark.asyncio +async def test_generation_request_context_restores_previous_bindings_and_closes_stream_once(): + class StreamingHandler: + def __init__(self): + self.chunks = [] + + async def push_chunk(self, chunk): + self.chunks.append(chunk) + + outer_options = GenerationOptions(llm_params={"request_id": "outer"}) + inner_options = GenerationOptions(llm_params={"request_id": "inner"}) + outer_messages = [{"role": "user", "content": "outer"}] + inner_messages = [{"role": "user", "content": "inner"}] + outer_explain_info = ExplainInfo() + outer_stats = LLMStats() + outer_streaming_handler = cast(Any, object()) + inner_streaming_handler = StreamingHandler() + tokens = [ + generation_options_var.set(outer_options), + raw_llm_request.set(outer_messages), + explain_info_var.set(outer_explain_info), + llm_stats_var.set(outer_stats), + streaming_handler_var.set(outer_streaming_handler), + ] + try: + request_context = start_generation_request_context( + gen_options=inner_options, + messages=inner_messages, + streaming_handler=cast(Any, inner_streaming_handler), + ) + + assert request_context.explain_info is outer_explain_info + assert generation_options_var.get() is inner_options + assert raw_llm_request.get() is inner_messages + assert streaming_handler_var.get() is inner_streaming_handler + + request_stats, _ = start_generation_stats() + await request_context.close_streaming_handler() + await request_context.close() + await request_context.close() + + assert inner_streaming_handler.chunks == [END_OF_STREAM] + assert generation_options_var.get() is outer_options + assert raw_llm_request.get() is outer_messages + assert explain_info_var.get() is outer_explain_info + assert llm_stats_var.get() is request_stats + assert streaming_handler_var.get() is outer_streaming_handler + finally: + for var, token in reversed( + [ + (generation_options_var, tokens[0]), + (raw_llm_request, tokens[1]), + (explain_info_var, tokens[2]), + (llm_stats_var, tokens[3]), + (streaming_handler_var, tokens[4]), + ] + ): + var.reset(token) + + +@pytest.mark.asyncio +async def test_generation_request_context_restores_previous_bindings_when_stream_close_fails(): + class FailingStreamingHandler: + async def push_chunk(self, chunk): + assert chunk is END_OF_STREAM + raise RuntimeError("stream close failed") + + outer_options = GenerationOptions(llm_params={"request_id": "outer"}) + inner_options = GenerationOptions(llm_params={"request_id": "inner"}) + outer_messages = [{"role": "user", "content": "outer"}] + inner_messages = [{"role": "user", "content": "inner"}] + outer_explain_info = ExplainInfo() + outer_stats = LLMStats() + outer_streaming_handler = cast(Any, object()) + tokens = [ + generation_options_var.set(outer_options), + raw_llm_request.set(outer_messages), + explain_info_var.set(outer_explain_info), + llm_stats_var.set(outer_stats), + streaming_handler_var.set(outer_streaming_handler), + ] + try: + request_context = start_generation_request_context( + gen_options=inner_options, + messages=inner_messages, + streaming_handler=cast(Any, FailingStreamingHandler()), + ) + + request_stats, _ = start_generation_stats() + with pytest.raises(RuntimeError, match="stream close failed"): + await request_context.close() + + assert generation_options_var.get() is outer_options + assert raw_llm_request.get() is outer_messages + assert explain_info_var.get() is outer_explain_info + assert llm_stats_var.get() is request_stats + assert streaming_handler_var.get() is outer_streaming_handler + finally: + for var, token in reversed( + [ + (generation_options_var, tokens[0]), + (raw_llm_request, tokens[1]), + (explain_info_var, tokens[2]), + (llm_stats_var, tokens[3]), + (streaming_handler_var, tokens[4]), + ] + ): + var.reset(token) diff --git a/tests/rails/llm/test_generation_request.py b/tests/rails/llm/test_generation_request.py new file mode 100644 index 0000000000..b0852fa6e1 --- /dev/null +++ b/tests/rails/llm/test_generation_request.py @@ -0,0 +1,287 @@ +# 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 SimpleNamespace + +import pytest + +from nemoguardrails.exceptions import InvalidStateError +from nemoguardrails.rails.llm import generation as generation_package +from nemoguardrails.rails.llm.config import RailsConfig +from nemoguardrails.rails.llm.generation import generation_request +from nemoguardrails.rails.llm.generation.generation_request import ( + normalize_generation_request, + prepare_generation_request_for_runtime, +) +from nemoguardrails.rails.llm.llmrails import LLMRails +from nemoguardrails.rails.llm.options import GenerationOptions, GenerationRailsOptions + + +def test_generation_package_exports(): + assert generation_package.__all__ == ["generate_standard_async"] + + +def test_generation_request_exports(): + assert generation_request.__all__ == [ + "GenerationRequest", + "PreparedGenerationRequest", + "normalize_generation_request", + "prepare_generation_request_for_runtime", + ] + + +def test_normalize_generation_request_requires_prompt_or_messages(): + with pytest.raises(ValueError, match="Either prompt or messages must be provided."): + normalize_generation_request( + prompt=None, + messages=None, + options=None, + state=None, + ) + + +def test_normalize_generation_request_rejects_prompt_and_messages(): + with pytest.raises(ValueError, match="Only one of prompt or messages can be provided."): + normalize_generation_request( + prompt="hi", + messages=[{"role": "user", "content": "hi"}], + options=None, + state=None, + ) + + +@pytest.mark.asyncio +async def test_generate_async_validates_prompt_and_messages_before_state(): + rails = LLMRails(config=RailsConfig(models=[], colang_version="1.0")) + + with pytest.raises(ValueError, match="Only one of prompt or messages can be provided."): + await rails.generate_async( + prompt="hi", + messages=[{"role": "user", "content": "hi"}], + state={"foo": 1}, + ) + + +def test_normalize_generation_request_converts_prompt_to_user_message(): + request = normalize_generation_request( + prompt="hi", + messages=None, + options=None, + state=None, + ) + + assert request.prompt == "hi" + assert request.messages == [{"role": "user", "content": "hi"}] + assert request.options is None + assert request.state is None + + +def test_normalize_generation_request_preserves_messages_object(): + messages = [{"role": "user", "content": "hi"}] + + request = normalize_generation_request( + prompt=None, + messages=messages, + options=None, + state=None, + ) + + assert request.prompt is None + assert request.messages is messages + assert request.options is None + + +def test_normalize_generation_request_accepts_options_dict_when_not_empty(): + request = normalize_generation_request( + prompt="hi", + messages=None, + options={"output_vars": ["answer"]}, + state=None, + ) + + assert isinstance(request.options, GenerationOptions) + assert request.options.output_vars == ["answer"] + + +def test_normalize_generation_request_preserves_options_object(): + options = GenerationOptions(output_vars=["answer"]) + + request = normalize_generation_request( + prompt="hi", + messages=None, + options=options, + state=None, + ) + + assert request.options is options + + +def test_normalize_generation_request_rejects_empty_options_dict_without_state(): + with pytest.raises(TypeError, match="options must be a dict or GenerationOptions"): + normalize_generation_request( + prompt="hi", + messages=None, + options={}, + state=None, + ) + + +def test_normalize_generation_request_state_forces_generation_options(): + state = {"events": []} + + request = normalize_generation_request( + prompt="hi", + messages=None, + options=None, + state=state, + ) + + assert isinstance(request.options, GenerationOptions) + assert request.state is state + + +def test_normalize_generation_request_accepts_empty_options_dict_with_state(): + request = normalize_generation_request( + prompt="hi", + messages=None, + options={}, + state={"events": []}, + ) + + assert isinstance(request.options, GenerationOptions) + + +def test_normalize_generation_request_deserializes_colang_2_state(monkeypatch): + deserialized = object() + + def fake_json_to_state(state): + assert state == {"flow_states": []} + return deserialized + + monkeypatch.setattr( + "nemoguardrails.rails.llm.generation.generation_request.json_to_state", + fake_json_to_state, + ) + + request = normalize_generation_request( + prompt="hi", + messages=None, + options=None, + state={"version": "2.x", "state": {"flow_states": []}}, + ) + + assert request.state is deserialized + assert isinstance(request.options, GenerationOptions) + + +def test_validate_public_state_preserves_missing_events_guidance(): + config = SimpleNamespace(colang_version="1.0") + + with pytest.raises(InvalidStateError) as exc_info: + generation_request.validate_public_state(config, {"foo": 1}) + + assert str(exc_info.value) == ( + "Invalid Colang 1.0 state format: state must contain an 'events' key. " + "Use an empty dict {} to start a new conversation." + ) + + +def test_prepare_generation_request_for_runtime_preserves_raw_messages_for_context(): + messages = [{"role": "user", "content": "hi"}] + + request = prepare_generation_request_for_runtime( + prompt=None, + messages=messages, + options={"output_vars": ["answer"]}, + state=None, + ) + + assert request.request_messages is messages + assert request.runtime_messages == [ + { + "role": "context", + "content": {"generation_options": GenerationOptions(output_vars=["answer"]).model_dump()}, + }, + {"role": "user", "content": "hi"}, + ] + assert request.runtime_messages is not messages + assert request.needs_llm is True + + +def test_prepare_generation_request_for_runtime_rewrites_assistant_message_when_dialog_disabled(): + messages = [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "draft"}, + ] + + request = prepare_generation_request_for_runtime( + prompt=None, + messages=messages, + options={"rails": {"dialog": False}}, + state=None, + ) + + assert request.request_messages is messages + assert request.runtime_messages == [ + { + "role": "context", + "content": { + "generation_options": GenerationOptions(rails=GenerationRailsOptions(dialog=False)).model_dump(), + "bot_message": "draft", + }, + }, + {"role": "user", "content": "hi"}, + ] + assert messages == [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "draft"}, + ] + assert request.needs_llm is False + + +def test_prepare_generation_request_for_runtime_preserves_assistant_tool_calls(): + assistant_message = { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_safe", + "type": "function", + "function": {"name": "safe_tool", "arguments": '{"param": "safe value"}'}, + } + ], + } + messages = [ + {"role": "user", "content": "Use the requested tool"}, + assistant_message, + ] + + request = prepare_generation_request_for_runtime( + prompt=None, + messages=messages, + options={"rails": {"dialog": False}}, + state=None, + ) + + assert request.runtime_messages == [ + { + "role": "context", + "content": { + "generation_options": GenerationOptions(rails=GenerationRailsOptions(dialog=False)).model_dump(), + }, + }, + messages[0], + assistant_message, + ] diff --git a/tests/rails/llm/test_generation_response.py b/tests/rails/llm/test_generation_response.py new file mode 100644 index 0000000000..319d268200 --- /dev/null +++ b/tests/rails/llm/test_generation_response.py @@ -0,0 +1,187 @@ +# 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.context import llm_response_metadata_var +from nemoguardrails.rails.llm.generation import generation_response +from nemoguardrails.rails.llm.generation.generation_response import ( + GenerationEventMetadata, + generation_event_metadata, + generation_response_from_colang_turn, +) +from nemoguardrails.rails.llm.options import GenerationOptions + + +def _processing_log(): + return [ + {"type": "event", "timestamp": 1.0, "data": {"type": "StartInputRails"}}, + {"type": "event", "timestamp": 2.0, "data": {"type": "InputRailsFinished"}}, + ] + + +def test_generation_event_metadata_extracts_and_clears_context_metadata(): + token = llm_response_metadata_var.set({"provider": "metadata"}) + try: + metadata = generation_event_metadata( + [ + {"type": "BotToolCalls", "tool_calls": [{"id": "call-1"}]}, + {"type": "BotThinking", "content": "reasoning"}, + ] + ) + finally: + llm_response_metadata_var.reset(token) + + assert metadata.tool_calls == [{"id": "call-1"}] + assert metadata.llm_metadata == {"provider": "metadata"} + assert metadata.reasoning_content == "reasoning" + + +def test_generation_response_from_colang_1_turn_populates_response_fields_and_log(): + new_events = [{"type": "StartUtteranceBotAction", "script": "Hello"}] + + response = generation_response_from_colang_turn( + colang_version="1.0", + prompt=None, + new_message={"role": "assistant", "content": "Hello"}, + events=[ + {"type": "ContextUpdate", "data": {"answer": 42}}, + {"type": "UserMessage", "text": "Hi"}, + {"type": "StartUtteranceBotAction", "script": "Hello"}, + ], + new_events=new_events, + processing_log=_processing_log(), + gen_options=GenerationOptions( + output_vars=["answer"], + log=cast(Any, {"internal_events": True, "colang_history": True}), + ), + state={"events": []}, + output_state={"events": ["serialized"]}, + event_metadata=GenerationEventMetadata( + tool_calls=[{"id": "call-1"}], + llm_metadata={"provider": "metadata"}, + reasoning_content="reasoning", + ), + ) + + assert response.response == [{"role": "assistant", "content": "Hello"}] + assert response.tool_calls == [{"id": "call-1"}] + assert response.llm_metadata == {"provider": "metadata"} + assert response.reasoning_content == "reasoning" + assert response.output_data == {"answer": 42} + assert response.state == {"events": ["serialized"]} + assert response.log is not None + assert response.log.internal_events == new_events + assert response.log.colang_history is not None + assert '"Hello"' in response.log.colang_history + + +def test_generation_response_from_colang_1_turn_returns_prompt_content(): + response = generation_response_from_colang_turn( + colang_version="1.0", + prompt="Hi", + new_message={"role": "assistant", "content": "Hello"}, + events=[], + new_events=[], + processing_log=_processing_log(), + gen_options=GenerationOptions(), + state=None, + output_state=None, + event_metadata=GenerationEventMetadata( + tool_calls=None, + llm_metadata=None, + reasoning_content=None, + ), + ) + + assert response.response == "Hello" + + +def test_generation_response_from_colang_1_turn_includes_raw_llm_output(monkeypatch): + raw_response = {"raw": "completion"} + fake_generation_log = SimpleNamespace( + activated_rails=[ + SimpleNamespace( + type="generation", + executed_actions=[ + SimpleNamespace(llm_calls=[SimpleNamespace(raw_response=raw_response)]), + ], + ) + ] + ) + monkeypatch.setattr( + generation_response, + "compute_generation_log", + lambda processing_log: fake_generation_log, + ) + + response = generation_response_from_colang_turn( + colang_version="1.0", + prompt=None, + new_message={"role": "assistant", "content": "Hello"}, + events=[], + new_events=[], + processing_log=[], + gen_options=GenerationOptions(llm_output=True), + state=None, + output_state=None, + event_metadata=GenerationEventMetadata( + tool_calls=None, + llm_metadata=None, + reasoning_content=None, + ), + ) + + assert response.llm_output == raw_response + + +@pytest.mark.parametrize( + ("gen_options", "message"), + [ + ( + GenerationOptions(output_vars=True), + "The `output_vars` option is not supported for Colang 2.0 configurations.", + ), + ( + GenerationOptions(log=cast(Any, {"internal_events": True})), + "The `log` option is not supported for Colang 2.0 configurations.", + ), + ( + GenerationOptions(llm_output=True), + "The `llm_output` option is not supported for Colang 2.0 configurations.", + ), + ], +) +def test_generation_response_from_colang_2_turn_rejects_unsupported_options(gen_options, message): + with pytest.raises(ValueError, match=message): + generation_response_from_colang_turn( + colang_version="2.x", + prompt=None, + new_message={"role": "assistant", "content": "Hello"}, + events=[], + new_events=[], + processing_log=[], + gen_options=gen_options, + state=None, + output_state=None, + event_metadata=GenerationEventMetadata( + tool_calls=None, + llm_metadata=None, + reasoning_content=None, + ), + ) diff --git a/tests/rails/llm/test_generation_tracing.py b/tests/rails/llm/test_generation_tracing.py new file mode 100644 index 0000000000..b69f201ce2 --- /dev/null +++ b/tests/rails/llm/test_generation_tracing.py @@ -0,0 +1,177 @@ +# 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 SimpleNamespace + +import pytest + +from nemoguardrails.rails.llm.generation import generation_tracing +from nemoguardrails.rails.llm.generation.generation_tracing import ( + export_generation_trace, + prepare_generation_tracing, + restore_generation_trace_log, +) +from nemoguardrails.rails.llm.options import ( + ActivatedRail, + GenerationLog, + GenerationLogOptions, + GenerationOptions, + GenerationResponse, +) + + +def test_generation_tracing_exports(): + assert generation_tracing.__all__ == [ + "GenerationTracingState", + "export_generation_trace", + "prepare_generation_tracing", + "restore_generation_trace_log", + ] + + +def test_prepare_generation_tracing_returns_original_options_when_disabled(): + options = GenerationOptions(output_vars=["answer"]) + + state = prepare_generation_tracing( + tracing_enabled=False, + gen_options=options, + ) + + assert state.gen_options is options + assert state.original_log_options is None + + +def test_prepare_generation_tracing_copies_options_and_forces_trace_logs(): + options = GenerationOptions(log=GenerationLogOptions(colang_history=True)) + + state = prepare_generation_tracing( + tracing_enabled=True, + gen_options=options, + ) + + assert state.gen_options is not options + assert state.gen_options is not None + assert state.gen_options.log.activated_rails is True + assert state.gen_options.log.llm_calls is True + assert state.gen_options.log.internal_events is True + assert state.gen_options.log.colang_history is True + assert state.original_log_options == GenerationLogOptions(colang_history=True) + assert options.log.activated_rails is False + assert options.log.llm_calls is False + assert options.log.internal_events is False + + +def test_prepare_generation_tracing_creates_options_when_absent(): + state = prepare_generation_tracing( + tracing_enabled=True, + gen_options=None, + ) + + assert state.gen_options is not None + assert state.gen_options.log.activated_rails is True + assert state.gen_options.log.llm_calls is True + assert state.gen_options.log.internal_events is True + assert state.original_log_options == GenerationLogOptions() + + +def test_restore_generation_trace_log_removes_trace_added_log_when_original_requested_none(): + response = GenerationResponse( + response="hello", + log=GenerationLog( + activated_rails=[ActivatedRail(type="input", name="rail")], + internal_events=[{"type": "event"}], + llm_calls=[], + ), + ) + + restore_generation_trace_log( + response=response, + original_log_options=GenerationLogOptions(), + ) + + assert response.log is None + + +def test_restore_generation_trace_log_keeps_only_originally_requested_log_fields(): + response = GenerationResponse( + response="hello", + log=GenerationLog( + activated_rails=[ActivatedRail(type="input", name="rail")], + internal_events=[{"type": "event"}], + llm_calls=[], + ), + ) + + restore_generation_trace_log( + response=response, + original_log_options=GenerationLogOptions(internal_events=True), + ) + + assert response.log is not None + assert response.log.internal_events == [{"type": "event"}] + assert response.log.activated_rails == [] + assert response.log.llm_calls == [] + + +@pytest.mark.asyncio +async def test_export_generation_trace_uses_tracing_config(monkeypatch): + calls = {} + + class FakeTracer: + def __init__( + self, + *, + input, + response, + adapters, + span_format, + enable_content_capture, + ): + calls["init"] = { + "input": input, + "response": response, + "adapters": adapters, + "span_format": span_format, + "enable_content_capture": enable_content_capture, + } + + async def export_async(self): + calls["exported"] = True + + monkeypatch.setattr("nemoguardrails.tracing.Tracer", FakeTracer) + + response = GenerationResponse(response="hello", log=GenerationLog()) + adapters = [object()] + messages = [{"role": "user", "content": "hi"}] + await export_generation_trace( + tracing_config=SimpleNamespace( + span_format="legacy", + enable_content_capture=True, + ), + log_adapters=adapters, + messages=messages, + response=response, + ) + + assert calls == { + "init": { + "input": messages, + "response": response, + "adapters": adapters, + "span_format": "legacy", + "enable_content_capture": True, + }, + "exported": True, + }