feat(bot-thinking): improve event logging and add bot thinking display (#1434)

Pouyanpi · tgasser-nv · commit 5882d52cd6e2 · 2025-10-14T15:29:07.000-05:00
add BotThinking event support
diff --git a/nemoguardrails/colang/v1_0/runtime/runtime.py b/nemoguardrails/colang/v1_0/runtime/runtime.py
@@ -159,13 +159,6 @@ async def generate_events(
 
             log.info("Processing event: %s", last_event)
 
-            event_type = last_event["type"]
-            log.info(
-                "Event :: %s %s",
-                event_type,
-                str({k: v for k, v in last_event.items() if k != "type"}),
-            )
-
             # If we need to execute an action, we start doing that.
             if last_event["type"] == "StartInternalSystemAction":
                 next_events = await self._process_start_action(events)
@@ -186,16 +179,23 @@ async def generate_events(
                 if len(next_events) == 0:
                     next_events = [new_event_dict("Listen")]
 
-            # Otherwise, we append the event and continue the processing.
-            events.extend(next_events)
-            new_events.extend(next_events)
-
+            # Log all generated events and add them to processing log
             for event in next_events:
                 if event["type"] != "EventHistoryUpdate":
+                    event_type = event["type"]
+                    log.info(
+                        "Event :: %s %s",
+                        event_type,
+                        str({k: v for k, v in event.items() if k != "type"}),
+                    )
                     processing_log.append(
                         {"type": "event", "timestamp": time(), "data": event}
                     )
 
+            # Append events to the event stream and new_events list
+            events.extend(next_events)
+            new_events.extend(next_events)
+
             # If the next event is a listen, we stop the processing.
             if next_events[-1]["type"] == "Listen":
                 break
diff --git a/nemoguardrails/logging/callbacks.py b/nemoguardrails/logging/callbacks.py
@@ -163,9 +163,27 @@ async def on_llm_end(
             llm_call_info = LLMCallInfo()
         llm_call_info.completion = response.generations[0][0].text
         llm_call_info.finished_at = time()
+
+        completion_text = response.generations[0][0].text
+        reasoning_content = None
+
+        if isinstance(response.generations[0][0], ChatGeneration):
+            chat_gen = response.generations[0][0]
+            if hasattr(chat_gen, "message") and hasattr(
+                chat_gen.message, "additional_kwargs"
+            ):
+                reasoning_content = chat_gen.message.additional_kwargs.get(
+                    "reasoning_content"
+                )
+
+        if reasoning_content:
+            full_completion = f"{reasoning_content}\n---\n{completion_text}"
+        else:
+            full_completion = completion_text
+
         log.info(
             "Completion :: %s",
-            response.generations[0][0].text,
+            full_completion,
             extra={"id": llm_call_info.id, "task": llm_call_info.task},
         )
 
diff --git a/nemoguardrails/logging/verbose.py b/nemoguardrails/logging/verbose.py
@@ -129,6 +129,8 @@ def emit(self, record) -> None:
                                     title = f"[magenta]{event_type[:-8]}[bold]Finished[/][/]"
                             elif event_type.endswith("ActionFailed"):
                                 title = f"[magenta]{event_type[:-6]}[bold]Failed[/][/]"
+                            elif event_type == "BotThinking":
+                                title = f"[yellow bold]{event_type}[/]"
                             else:
                                 title = event_type
                         else:
diff --git a/tests/test_runtime_event_logging.py b/tests/test_runtime_event_logging.py
@@ -0,0 +1,160 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 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
+from unittest.mock import patch
+
+import pytest
+
+from nemoguardrails import RailsConfig
+from tests.utils import TestChat
+
+
+@pytest.mark.asyncio
+async def test_bot_thinking_event_logged_in_runtime(caplog):
+    test_reasoning_trace = "Let me think about this step by step..."
+
+    caplog.set_level(logging.INFO)
+
+    with patch(
+        "nemoguardrails.actions.llm.generation.get_and_clear_reasoning_trace_contextvar"
+    ) as mock_get_reasoning:
+        mock_get_reasoning.return_value = test_reasoning_trace
+
+        config = RailsConfig.from_content(config={"models": [], "passthrough": True})
+        chat = TestChat(config, llm_completions=["The answer is 42"])
+
+        await chat.app.generate_events_async(
+            [{"type": "UserMessage", "text": "What is the answer?"}]
+        )
+
+        bot_thinking_logs = [
+            record
+            for record in caplog.records
+            if "Event :: BotThinking" in record.message
+        ]
+        assert len(bot_thinking_logs) >= 1
+
+
+@pytest.mark.asyncio
+async def test_bot_message_event_logged_in_runtime(caplog):
+    caplog.set_level(logging.INFO)
+
+    config = RailsConfig.from_content(config={"models": [], "passthrough": True})
+    chat = TestChat(config, llm_completions=["The answer is 42"])
+
+    await chat.app.generate_events_async(
+        [{"type": "UserMessage", "text": "What is the answer?"}]
+    )
+
+    bot_message_logs = [
+        record for record in caplog.records if "Event :: BotMessage" in record.message
+    ]
+    assert len(bot_message_logs) >= 1
+
+
+@pytest.mark.asyncio
+async def test_context_update_event_logged_in_runtime(caplog):
+    caplog.set_level(logging.INFO)
+
+    config = RailsConfig.from_content(config={"models": [], "passthrough": True})
+    chat = TestChat(config, llm_completions=["Response"])
+
+    await chat.app.generate_events_async([{"type": "UserMessage", "text": "Hello"}])
+
+    context_update_logs = [
+        record
+        for record in caplog.records
+        if "Event :: ContextUpdate" in record.message
+        or "Event ContextUpdate" in record.message
+    ]
+    assert len(context_update_logs) >= 1
+
+
+@pytest.mark.asyncio
+async def test_all_events_logged_when_multiple_events_generated(caplog):
+    test_reasoning_trace = "Analyzing..."
+
+    caplog.set_level(logging.INFO)
+
+    with patch(
+        "nemoguardrails.actions.llm.generation.get_and_clear_reasoning_trace_contextvar"
+    ) as mock_get_reasoning:
+        mock_get_reasoning.return_value = test_reasoning_trace
+
+        config = RailsConfig.from_content(config={"models": [], "passthrough": True})
+        chat = TestChat(config, llm_completions=["Response"])
+
+        await chat.app.generate_events_async([{"type": "UserMessage", "text": "Test"}])
+
+        bot_thinking_found = any(
+            "Event :: BotThinking" in record.message for record in caplog.records
+        )
+        bot_message_found = any(
+            "Event :: BotMessage" in record.message for record in caplog.records
+        )
+
+        assert bot_thinking_found
+        assert bot_message_found
+
+
+@pytest.mark.asyncio
+async def test_bot_thinking_event_logged_before_bot_message(caplog):
+    test_reasoning_trace = "Step 1: Understand\nStep 2: Respond"
+
+    caplog.set_level(logging.INFO)
+
+    with patch(
+        "nemoguardrails.actions.llm.generation.get_and_clear_reasoning_trace_contextvar"
+    ) as mock_get_reasoning:
+        mock_get_reasoning.return_value = test_reasoning_trace
+
+        config = RailsConfig.from_content(config={"models": [], "passthrough": True})
+        chat = TestChat(config, llm_completions=["Answer"])
+
+        await chat.app.generate_events_async(
+            [{"type": "UserMessage", "text": "Question"}]
+        )
+
+        bot_thinking_idx = None
+        bot_message_idx = None
+
+        for idx, record in enumerate(caplog.records):
+            if "Event :: BotThinking" in record.message and bot_thinking_idx is None:
+                bot_thinking_idx = idx
+            if "Event :: BotMessage" in record.message and bot_message_idx is None:
+                bot_message_idx = idx
+
+        assert bot_thinking_idx is not None
+        assert bot_message_idx is not None
+        assert bot_thinking_idx < bot_message_idx
+
+
+@pytest.mark.asyncio
+async def test_event_history_update_not_logged(caplog):
+    caplog.set_level(logging.INFO)
+
+    config = RailsConfig.from_content(config={"models": [], "passthrough": True})
+    chat = TestChat(config, llm_completions=["Response"])
+
+    await chat.app.generate_events_async([{"type": "UserMessage", "text": "Test"}])
+
+    event_history_update_logs = [
+        record
+        for record in caplog.records
+        if "Event :: EventHistoryUpdate" in record.message
+        or "Event EventHistoryUpdate" in record.message
+    ]
+    assert len(event_history_update_logs) == 0
