fix(logging): add "Tool" type to message sender labeling (#1412)

Previously, messages of type "tool" were not distinctly labeled in the
LoggingCallbackHandler output, causing them to be grouped under
"System".
This change adds explicit handling for "tool" messages, labeling them as
"Tool" in the logs for improved clarity and debugging.

Author: Pouyanpi

nemoguardrails/logging/callbacks.py

@@ -109,6 +109,8 @@ async def on_chat_model_start(
                     if msg.type == "human"
                     else "Bot"
                     if msg.type == "ai"
+                    else "Tool"
+                    if msg.type == "tool"
                     else "System"
                 )
                 + "[/]"

tests/test_callbacks.py

@@ -13,11 +13,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from unittest.mock import patch
 from uuid import uuid4
 
 import pytest
 from langchain.schema import Generation, LLMResult
-from langchain_core.messages import AIMessage
+from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage
 from langchain_core.outputs import ChatGeneration
 
 from nemoguardrails.context import explain_info_var, llm_call_info_var, llm_stats_var
@@ -168,3 +169,46 @@ async def test_multiple_generations_token_accumulation():
     assert llm_stats.get_stat("total_tokens") == 19
     assert llm_stats.get_stat("total_prompt_tokens") == 12
     assert llm_stats.get_stat("total_completion_tokens") == 7
+
+
+@pytest.mark.asyncio
+async def test_tool_message_labeling_in_logging():
+    """Test that tool messages are labeled as 'Tool' in logging output."""
+    llm_call_info = LLMCallInfo()
+    llm_call_info_var.set(llm_call_info)
+
+    llm_stats = LLMStats()
+    llm_stats_var.set(llm_stats)
+
+    explain_info = ExplainInfo()
+    explain_info_var.set(explain_info)
+
+    handler = LoggingCallbackHandler()
+
+    messages = [
+        HumanMessage(content="Hello"),
+        AIMessage(content="Hi there"),
+        SystemMessage(content="System message"),
+        ToolMessage(content="Tool result", tool_call_id="test_tool_call"),
+    ]
+
+    with patch("nemoguardrails.logging.callbacks.log") as mock_log:
+        await handler.on_chat_model_start(
+            serialized={},
+            messages=[messages],
+            run_id=uuid4(),
+        )
+
+        mock_log.info.assert_called()
+
+        logged_prompt = None
+        for call in mock_log.info.call_args_list:
+            if "Prompt Messages" in str(call):
+                logged_prompt = call[0][1]
+                break
+
+        assert logged_prompt is not None
+        assert "[cyan]User[/]" in logged_prompt
+        assert "[cyan]Bot[/]" in logged_prompt
+        assert "[cyan]System[/]" in logged_prompt
+        assert "[cyan]Tool[/]" in logged_prompt