fix(logging): correct message type formatting in logs (#1416)

* refactor(logging): replace if-else with dict for type mapping

Author: Pouyanpi

nemoguardrails/logging/callbacks.py

@@ -101,16 +101,17 @@ async def on_chat_model_start(
         if explain_info:
             explain_info.llm_calls.append(llm_call_info)
 
+        type_map = {
+            "human": "User",
+            "ai": "Bot",
+            "tool": "Tool",
+            "system": "System",
+            "developer": "Developer",
+        }
         prompt = "\n" + "\n".join(
             [
                 "[cyan]"
-                + (
-                    "User"
-                    if msg.type == "human"
-                    else "Bot"
-                    if msg.type == "ai"
-                    else "System"
-                )
+                + type_map.get(msg.type, msg.type.title())
                 + "[/]"
                 + "\n"
                 + (msg.content if isinstance(msg.content, str) else "")

pyproject.toml

@@ -152,6 +152,9 @@ yara-python = "^4.5.1"
 opentelemetry-api = "^1.34.1"
 opentelemetry-sdk = "^1.34.1"
 
+# Directories in which to run Pyright type-checking
+[tool.pyright]
+include = ["nemoguardrails/rails/**", "tests/test_callbacks.py"]
 
 [tool.poetry.group.docs]
 optional = true

tests/test_callbacks.py

@@ -17,7 +17,13 @@
 
 import pytest
 from langchain.schema import Generation, LLMResult
-from langchain_core.messages import AIMessage
+from langchain_core.messages import (
+    AIMessage,
+    BaseMessage,
+    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 +174,89 @@ 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
+
+
+@pytest.mark.asyncio
+async def test_unknown_message_type_labeling():
+    """Test that unknown message types display their actual type name."""
+    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()
+
+    class CustomMessage(BaseMessage):
+        def __init__(self, content, msg_type):
+            super().__init__(content=content, type=msg_type)
+
+    messages: list[BaseMessage] = [
+        CustomMessage("Custom message", "custom"),
+        CustomMessage("Function message", "function"),
+    ]
+
+    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]Custom[/]" in logged_prompt
+        assert "[cyan]Function[/]" in logged_prompt