Fix case in traverse_dot_path where the value is None when traversing.

testing/invariant/custom_types/trace.py

@@ -5,11 +5,12 @@
 import json
 from typing import Any, Callable, Dict, Generator, List
 
-from invariant.utils.utils import ssl_verification_enabled
 from invariant_sdk.client import Client as InvariantClient
 from invariant_sdk.types.push_traces import PushTracesResponse
 from pydantic import BaseModel
 
+from invariant.utils.utils import ssl_verification_enabled
+
 from .invariant_dict import InvariantDict, InvariantValue
 from .matchers import ContainsImage, Matcher
 
@@ -139,8 +140,9 @@ def match_keyword_filter(
 
 
 def traverse_dot_path(message: dict, path: str) -> Any | None:
-    """Traverse a dictionary using a dot-separated path. If argument is not
-    found, .function will be added as a prefix to the path to search the
+    """Traverse a dictionary using a dot-separated path.
+
+    If argument is not found, "function." will be added as a prefix to the path to search the
     function fields for tool calls.
 
     Args:
@@ -149,15 +151,19 @@ def traverse_dot_path(message: dict, path: str) -> Any | None:
 
     Returns:
         Any: The value at the end of the path, or None if the path does not exist;
-             If the function prefix is added, the second return value will be True, otherwise False.
+             If the "function." prefix is added, the second return value will be True,
+             otherwise False.
     """
     add_function_prefix = False
 
     def _inner(d, _path):
         for k in _path.split("."):
-            if isinstance(d, str) and isinstance(k, str):
-                d = json.loads(d)
-            if k not in d:
+            if isinstance(d, str):
+                try:
+                    d = json.loads(d)  # Attempt to parse string as JSON.
+                except json.JSONDecodeError:
+                    return None
+            if not isinstance(d, dict) or d.get(k) is None:
                 return None
             d = d[k]
         return d

testing/tests/test_trace.py

@@ -1,6 +1,7 @@
 import pytest
 
-from invariant.custom_types.trace import Trace
+
+from invariant.custom_types.trace import Trace, traverse_dot_path
 
 
 @pytest.fixture(name="message_list")
@@ -22,6 +23,14 @@ def fixture_message_list():
     return trace.messages()
 
 
+@pytest.fixture(name="dict_traversal")
+def fixture_for_dict_traversal():
+    return {
+        "user": {"name": "Alice", "details": {"age": 30, "email": "alice@example.com"}},
+        "function": {"tool_calls": {"name": "example_tool", "status": "success"}},
+    }
+
+
 def test_trace_iterable(message_list: Trace):
     trace = message_list
 
@@ -58,3 +67,28 @@ def test_trace_messages_with_index():
         trace.messages(-10)
         trace.messages(-5)
         trace.messages(5)
+
+
+def test_traverse_dot_path_found_match(dict_traversal):
+    value, prefix_added = traverse_dot_path(dict_traversal, "user.details.email")
+    assert value == "alice@example.com"
+    assert not prefix_added
+
+
+def test_traverse_dot_path_no_found_match_with_function_prefix(dict_traversal):
+    value, prefix_added = traverse_dot_path(dict_traversal, "user.details.phone")
+    assert value is None
+    assert prefix_added  # Because function prefix was attempted.
+
+
+def test_traverse_dot_path_found_match_with_function_prefix(dict_traversal):
+    value, prefix_added = traverse_dot_path(dict_traversal, "tool_calls.name")
+    assert value == "example_tool"
+    assert prefix_added  # Because function prefix was attempted.
+
+
+def test_traverse_dot_path_testing_none():
+    dict1 = {"user": None}
+    value, prefix_added = traverse_dot_path(dict1, "user.details")
+    assert value is None
+    assert prefix_added  # Because function prefix was attempted.