Respect the provided log options regardless of tracing configuration

Author: RobGeada

nemoguardrails/rails/llm/llmrails.py

@@ -911,6 +911,7 @@ async def generate_async(
             await streaming_handler.push_chunk(END_OF_STREAM)
 
         # IF tracing is enabled we need to set GenerationLog attrs
+        original_log_options = None
         if self.config.tracing.enabled:
             if options is None:
                 options = GenerationOptions()
@@ -921,6 +922,7 @@ async def generate_async(
                 else:
                     # If options is a dict, convert it to GenerationOptions
                     options = GenerationOptions(**options)
+            original_log_options = options.log.model_copy(deep=True)
 
             # enable log options
             # it is aggressive, but these are required for tracing
@@ -1038,6 +1040,26 @@ async def generate_async(
                 )
                 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
+                            )):
+                        del res.log
+                    else:
+                        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.

tests/test_tracing.py

@@ -312,10 +312,7 @@ async def test_tracing_does_not_mutate_user_options():
         ), "User's original options were modified! This causes instability."
 
         # verify that tracing still works
-        assert response.log is not None, "Tracing should still work correctly"
-        assert (
-            response.log.activated_rails is not None
-        ), "Should have activated rails data"
+        assert "log" not in response, "Tracing should still work correctly, without affecting returned log"
 
 
 @pytest.mark.asyncio
@@ -358,17 +355,16 @@ async def test_tracing_with_none_options():
             messages=[{"role": "user", "content": "hello"}], options=None
         )
 
-        assert response.log is not None
-        assert response.log.activated_rails is not None
-        assert response.log.stats is not None
+        assert "log" not in response
 
 
 @pytest.mark.asyncio
 async def test_tracing_aggressive_override_when_all_disabled():
     """Test that tracing aggressively enables all logging when user disables all options.
 
     When user disables all three tracing related options, tracing still enables
-    ALL of them to ensure comprehensive logging data.
+    ALL of them to ensure comprehensive logging data. However, this should not contaminate the
+    returned response object
     """
 
     config = RailsConfig.from_content(
@@ -425,11 +421,10 @@ async def test_tracing_aggressive_override_when_all_disabled():
 
         assert response.log is not None
         assert (
-            response.log.activated_rails is not None
-            and len(response.log.activated_rails) > 0
+            response.log.activated_rails == []
         )
-        assert response.log.llm_calls is not None
-        assert response.log.internal_events is not None
+        assert response.log.llm_calls == []
+        assert response.log.internal_events == []
 
         assert user_options.log.activated_rails == original_activated_rails
         assert user_options.log.llm_calls == original_llm_calls
@@ -439,6 +434,84 @@ async def test_tracing_aggressive_override_when_all_disabled():
         assert user_options.log.internal_events == False
 
 
+@pytest.mark.asyncio
+async def test_tracing_preserves_specific_log_fields():
+    """Test that tracing does not modify the returned log fields
+    """
+
+    config = RailsConfig.from_content(
+        colang_content="""
+        define user express greeting
+            "hello"
+
+        define flow
+            user express greeting
+            bot express greeting
+
+        define bot express greeting
+            "Hello! How can I assist you today?"
+        """,
+        config={
+            "models": [],
+            "tracing": {"enabled": True, "adapters": [{"name": "FileSystem"}]},
+        },
+    )
+
+    chat = TestChat(
+        config,
+        llm_completions=[
+            "user express greeting",
+            "bot express greeting",
+            "Hello! How can I assist you today?",
+        ],
+    )
+
+    # user explicitly disables ALL tracing related options
+    user_options = GenerationOptions(
+        log=GenerationLogOptions(
+            activated_rails=False,
+            llm_calls=True,
+            internal_events=True,
+            colang_history=True,
+        )
+    )
+
+    original_activated_rails = user_options.log.activated_rails
+    original_llm_calls = user_options.log.llm_calls
+    original_internal_events = user_options.log.internal_events
+    original_colang_history = user_options.log.colang_history
+
+    with patch.object(Tracer, "export_async", return_value=None):
+        response = await chat.app.generate_async(
+            messages=[{"role": "user", "content": "hello"}], options=user_options
+        )
+
+        assert user_options.log.activated_rails == original_activated_rails
+        assert user_options.log.llm_calls == original_llm_calls
+        assert user_options.log.internal_events == original_internal_events
+        assert user_options.log.colang_history == original_colang_history
+
+        assert response.log is not None
+        assert (
+            response.log.activated_rails == []
+        )
+        assert (
+                response.log.llm_calls is not None
+                and len(response.log.llm_calls) > 0
+        )
+        assert (
+                response.log.internal_events is not None
+                and len(response.log.internal_events) > 0
+        )
+
+        assert user_options.log.activated_rails == original_activated_rails
+        assert user_options.log.llm_calls == original_llm_calls
+        assert user_options.log.internal_events == original_internal_events
+        assert user_options.log.activated_rails == False
+        assert user_options.log.llm_calls == True
+        assert user_options.log.internal_events == True
+
+
 @pytest.mark.asyncio
 async def test_tracing_aggressive_override_with_dict_options():
     """Test that tracing works correctly when options are passed as a dict.
@@ -502,11 +575,11 @@ async def test_tracing_aggressive_override_with_dict_options():
 
         assert response.log is not None
         assert (
-            response.log.activated_rails is not None
-            and len(response.log.activated_rails) > 0
+            response.log.activated_rails == []
+            and len(response.log.activated_rails) == 0
         )
-        assert response.log.llm_calls is not None
-        assert response.log.internal_events is not None
+        assert response.log.llm_calls == []
+        assert response.log.internal_events == []
 
 
 if __name__ == "__main__":