fix(types): harden ChatMessage.from_dict and fix LLMModel.stream protocol

Author: Pouyanpi

nemoguardrails/types.py

@@ -151,6 +151,10 @@ def from_dict(cls, d: Dict[str, Any]) -> "ChatMessage":
                         args_dict = json.loads(raw_args)
                     except json.JSONDecodeError:
                         raise ValueError(f"Tool call arguments are not valid JSON: {raw_args!r}")
+                    if not isinstance(args_dict, dict):
+                        raise ValueError(
+                            f"Tool call arguments must be a JSON object, got {type(args_dict).__name__}: {raw_args!r}"
+                        )
                 else:
                     args_dict = raw_args
 
@@ -165,13 +169,17 @@ def from_dict(cls, d: Dict[str, Any]) -> "ChatMessage":
                     )
                 )
 
+        _standard_keys = {"role", "content", "tool_calls", "tool_call_id", "name", "provider_metadata"}
+        extra = {k: v for k, v in d.items() if k not in _standard_keys}
+        provider_metadata = {**d.get("provider_metadata", {}), **extra}
+
         return cls(
             role=role,
             content=d.get("content"),
             tool_calls=tool_calls,
             tool_call_id=d.get("tool_call_id"),
             name=d.get("name"),
-            provider_metadata=d.get("provider_metadata", {}),
+            provider_metadata=provider_metadata,
         )
 
 
@@ -221,7 +229,7 @@ async def generate(
         **kwargs,
     ) -> "LLMResponse": ...
 
-    async def stream(
+    def stream(
         self,
         prompt: Union[str, List["ChatMessage"]],
         *,

tests/test_types.py

@@ -220,6 +220,20 @@ def test_from_dict_with_malformed_json_arguments(self):
         with pytest.raises(ValueError, match="not valid JSON"):
             ChatMessage.from_dict(d)
 
+    def test_from_dict_with_non_object_json_arguments(self):
+        d = {
+            "role": "assistant",
+            "tool_calls": [
+                {
+                    "id": "tc_1",
+                    "type": "function",
+                    "function": {"name": "search", "arguments": "[]"},
+                }
+            ],
+        }
+        with pytest.raises(ValueError, match="must be a JSON object"):
+            ChatMessage.from_dict(d)
+
     def test_from_dict_with_legacy_flat_tool_calls(self):
         d = {
             "role": "assistant",
@@ -235,6 +249,11 @@ def test_from_dict_captures_provider_metadata(self):
         msg = ChatMessage.from_dict(d)
         assert msg.provider_metadata == {"custom_field": "value", "model": "gpt-4"}
 
+    def test_from_dict_unknown_keys_captured_into_provider_metadata(self):
+        d = {"role": "user", "content": "hi", "unexpected_key": "v"}
+        msg = ChatMessage.from_dict(d)
+        assert msg.provider_metadata["unexpected_key"] == "v"
+
     def test_from_dict_missing_provider_metadata_defaults_to_empty(self):
         msg = ChatMessage.from_dict({"role": "user", "content": "hi"})
         assert msg.provider_metadata == {}