Skip to content

Commit b85493c

Browse files
committed
fix(tracing): prevent mutation of user options when tracing is enabled
When tracing is enabled, the system was directly modifying the user's original GenerationOptions object, causing instability when the same options were reused across multiple calls. This fix creates a copy of the options before modification, ensuring the original user options remain unchanged while maintaining full tracing functionality
1 parent 5b8a59c commit b85493c

File tree

2 files changed

+203
-3
lines changed

2 files changed

+203
-3
lines changed

nemoguardrails/rails/llm/llmrails.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -839,6 +839,12 @@ async def generate_async(
839839
if self.config.tracing.enabled:
840840
if options is None:
841841
options = GenerationOptions()
842+
else:
843+
# create a copy of the options to avoid modifying the original
844+
options = GenerationOptions(**options.model_dump())
845+
846+
# enable log options
847+
# it is aggressive, but these are required for tracing
842848
if (
843849
not options.log.activated_rails
844850
or not options.log.llm_calls

tests/test_tracing.py

Lines changed: 197 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,23 +14,25 @@
1414
# limitations under the License.
1515

1616
import asyncio
17-
import os
1817
import unittest
1918
from unittest.mock import AsyncMock, MagicMock, patch
2019

2120
import pytest
2221

23-
from nemoguardrails import LLMRails
22+
from nemoguardrails import LLMRails, RailsConfig
2423
from nemoguardrails.logging.explain import LLMCallInfo
25-
from nemoguardrails.rails.llm.config import RailsConfig, TracingConfig
2624
from nemoguardrails.rails.llm.options import (
2725
ActivatedRail,
2826
ExecutedAction,
2927
GenerationLog,
28+
GenerationLogOptions,
29+
GenerationOptions,
30+
GenerationRailsOptions,
3031
GenerationResponse,
3132
)
3233
from nemoguardrails.tracing.adapters.base import InteractionLogAdapter
3334
from nemoguardrails.tracing.tracer import Tracer, new_uuid
35+
from tests.utils import TestChat
3436

3537

3638
class TestTracer(unittest.TestCase):
@@ -239,5 +241,197 @@ async def test_tracing_enable_no_crash_issue_1093(mockTracer):
239241
assert res.response != None
240242

241243

244+
@pytest.mark.asyncio
245+
async def test_tracing_does_not_mutate_user_options():
246+
"""Test that tracing doesn't modify the user's original GenerationOptions object.
247+
248+
This test verifies the core fix: when tracing is enabled, the user's options
249+
should not be modified. Before the fix, this test would have failed
250+
because the original options object was being mutated.
251+
"""
252+
253+
config = RailsConfig.from_content(
254+
colang_content="""
255+
define user express greeting
256+
"hello"
257+
258+
define flow
259+
user express greeting
260+
bot express greeting
261+
262+
define bot express greeting
263+
"Hello! How can I assist you today?"
264+
""",
265+
config={
266+
"models": [],
267+
"tracing": {"enabled": True, "adapters": [{"name": "FileSystem"}]},
268+
},
269+
)
270+
271+
chat = TestChat(
272+
config,
273+
llm_completions=[
274+
"user express greeting",
275+
"bot express greeting",
276+
"Hello! How can I assist you today?",
277+
],
278+
)
279+
280+
user_options = GenerationOptions(
281+
log=GenerationLogOptions(
282+
activated_rails=False,
283+
llm_calls=False,
284+
internal_events=False,
285+
colang_history=False,
286+
)
287+
)
288+
289+
original_activated_rails = user_options.log.activated_rails
290+
original_llm_calls = user_options.log.llm_calls
291+
original_internal_events = user_options.log.internal_events
292+
original_colang_history = user_options.log.colang_history
293+
294+
# mock file operations to focus on the mutation issue
295+
with patch.object(Tracer, "export_async", return_value=None):
296+
response = await chat.app.generate_async(
297+
messages=[{"role": "user", "content": "hello"}], options=user_options
298+
)
299+
300+
# main fix: no mutation
301+
assert (
302+
user_options.log.activated_rails == original_activated_rails
303+
), "User's original options were modified! This causes instability."
304+
assert (
305+
user_options.log.llm_calls == original_llm_calls
306+
), "User's original options were modified! This causes instability."
307+
assert (
308+
user_options.log.internal_events == original_internal_events
309+
), "User's original options were modified! This causes instability."
310+
assert (
311+
user_options.log.colang_history == original_colang_history
312+
), "User's original options were modified! This causes instability."
313+
314+
# verify that tracing still works
315+
assert response.log is not None, "Tracing should still work correctly"
316+
assert (
317+
response.log.activated_rails is not None
318+
), "Should have activated rails data"
319+
320+
321+
@pytest.mark.asyncio
322+
async def test_tracing_with_none_options():
323+
"""Test that tracing works correctly when no options are provided.
324+
325+
This verifies that the fix doesn't break the case where users don't
326+
provide any options at all.
327+
"""
328+
329+
config = RailsConfig.from_content(
330+
colang_content="""
331+
define user express greeting
332+
"hello"
333+
334+
define flow
335+
user express greeting
336+
bot express greeting
337+
338+
define bot express greeting
339+
"Hello! How can I assist you today?"
340+
""",
341+
config={
342+
"models": [],
343+
"tracing": {"enabled": True, "adapters": [{"name": "FileSystem"}]},
344+
},
345+
)
346+
347+
chat = TestChat(
348+
config,
349+
llm_completions=[
350+
"user express greeting",
351+
"bot express greeting",
352+
"Hello! How can I assist you today?",
353+
],
354+
)
355+
356+
with patch.object(Tracer, "export_async", return_value=None):
357+
response = await chat.app.generate_async(
358+
messages=[{"role": "user", "content": "hello"}], options=None
359+
)
360+
361+
assert response.log is not None
362+
assert response.log.activated_rails is not None
363+
assert response.log.stats is not None
364+
365+
366+
@pytest.mark.asyncio
367+
async def test_tracing_aggressive_override_when_all_disabled():
368+
"""Test that tracing aggressively enables all logging when user disables all options.
369+
370+
When user disables all three tracing related options, tracing still enables
371+
ALL of them to ensure comprehensive logging data.
372+
"""
373+
374+
config = RailsConfig.from_content(
375+
colang_content="""
376+
define user express greeting
377+
"hello"
378+
379+
define flow
380+
user express greeting
381+
bot express greeting
382+
383+
define bot express greeting
384+
"Hello! How can I assist you today?"
385+
""",
386+
config={
387+
"models": [],
388+
"tracing": {"enabled": True, "adapters": [{"name": "FileSystem"}]},
389+
},
390+
)
391+
392+
chat = TestChat(
393+
config,
394+
llm_completions=[
395+
"user express greeting",
396+
"bot express greeting",
397+
"Hello! How can I assist you today?",
398+
],
399+
)
400+
401+
# user explicitly disables ALL tracing related options
402+
user_options = GenerationOptions(
403+
log=GenerationLogOptions(
404+
activated_rails=False,
405+
llm_calls=False,
406+
internal_events=False,
407+
colang_history=True,
408+
)
409+
)
410+
411+
original_activated_rails = user_options.log.activated_rails
412+
original_llm_calls = user_options.log.llm_calls
413+
original_internal_events = user_options.log.internal_events
414+
original_colang_history = user_options.log.colang_history
415+
416+
with patch.object(Tracer, "export_async", return_value=None):
417+
response = await chat.app.generate_async(
418+
messages=[{"role": "user", "content": "hello"}], options=user_options
419+
)
420+
421+
assert user_options.log.activated_rails == original_activated_rails
422+
assert user_options.log.llm_calls == original_llm_calls
423+
assert user_options.log.internal_events == original_internal_events
424+
assert user_options.log.colang_history == original_colang_history
425+
426+
assert response.log is not None
427+
assert response.log.activated_rails is not None
428+
assert response.log.llm_calls is not None
429+
assert response.log.internal_events is not None
430+
431+
assert user_options.log.activated_rails == original_activated_rails
432+
assert user_options.log.llm_calls == original_llm_calls
433+
assert user_options.log.internal_events == original_internal_events
434+
435+
242436
if __name__ == "__main__":
243437
unittest.main()

0 commit comments

Comments
 (0)