refactor(llmrails): extract generation workflow (4/9)

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"]

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)}

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, []